From 429eb85e14479fc0552db15b7d41cc56a4cf331d Mon Sep 17 00:00:00 2001
From: Boris Nagaev <bnagaev@gmail.com>
Date: Sat, 3 Aug 2024 23:18:55 -0300
Subject: [PATCH] sweepbatcher: use lnd/clock for timers

Added option: WithClock. Use it for publishDelay timer.
It will be used for testing.
---
 sweepbatcher/sweep_batch.go   |  9 ++++++++-
 sweepbatcher/sweep_batcher.go | 22 ++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/sweepbatcher/sweep_batch.go b/sweepbatcher/sweep_batch.go
index 709c8529..0c308c41 100644
--- a/sweepbatcher/sweep_batch.go
+++ b/sweepbatcher/sweep_batch.go
@@ -24,6 +24,7 @@ import (
 	"github.com/lightninglabs/loop/swap"
 	sweeppkg "github.com/lightninglabs/loop/sweep"
 	"github.com/lightningnetwork/lnd/chainntnfs"
+	"github.com/lightningnetwork/lnd/clock"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/keychain"
 	"github.com/lightningnetwork/lnd/lnrpc/walletrpc"
@@ -134,6 +135,9 @@ type batchConfig struct {
 	// batchConfTarget is the confirmation target of the batch transaction.
 	batchConfTarget int32
 
+	// clock provides methods to work with time and timers.
+	clock clock.Clock
+
 	// batchPublishDelay is the delay between receiving a new block and
 	// publishing the batch transaction.
 	batchPublishDelay time.Duration
@@ -527,6 +531,9 @@ func (b *batch) Run(ctx context.Context) error {
 		return fmt.Errorf("both musig2 signers provided")
 	}
 
+	// Cache clock variable.
+	clock := b.cfg.clock
+
 	blockChan, blockErrChan, err :=
 		b.chainNotifier.RegisterBlockEpochNtfn(runCtx)
 	if err != nil {
@@ -562,7 +569,7 @@ func (b *batch) Run(ctx context.Context) error {
 
 			// Set the timer to publish the batch transaction after
 			// the configured delay.
-			timerChan = time.After(b.cfg.batchPublishDelay)
+			timerChan = clock.TickAfter(b.cfg.batchPublishDelay)
 			b.currentHeight = height
 
 		case <-timerChan:
diff --git a/sweepbatcher/sweep_batcher.go b/sweepbatcher/sweep_batcher.go
index ac5c31e9..a2ee99f3 100644
--- a/sweepbatcher/sweep_batcher.go
+++ b/sweepbatcher/sweep_batcher.go
@@ -16,6 +16,7 @@ import (
 	"github.com/lightninglabs/loop/loopdb"
 	"github.com/lightninglabs/loop/swap"
 	"github.com/lightninglabs/loop/utils"
+	"github.com/lightningnetwork/lnd/clock"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
@@ -253,6 +254,9 @@ type Batcher struct {
 	// exit.
 	wg sync.WaitGroup
 
+	// clock provides methods to work with time and timers.
+	clock clock.Clock
+
 	// customFeeRate provides custom min fee rate per swap. The batch uses
 	// max of the fee rates of its swaps. In this mode confTarget is
 	// ignored and fee bumping by sweepbatcher is disabled.
@@ -267,6 +271,9 @@ type Batcher struct {
 
 // BatcherConfig holds batcher configuration.
 type BatcherConfig struct {
+	// clock provides methods to work with time and timers.
+	clock clock.Clock
+
 	// customFeeRate provides custom min fee rate per swap. The batch uses
 	// max of the fee rates of its swaps. In this mode confTarget is
 	// ignored and fee bumping by sweepbatcher is disabled.
@@ -282,6 +289,14 @@ type BatcherConfig struct {
 // BatcherOption configures batcher behaviour.
 type BatcherOption func(*BatcherConfig)
 
+// WithClock sets the clock used by sweepbatcher and its batches. It is needed
+// to manipulate time in tests.
+func WithClock(clock clock.Clock) BatcherOption {
+	return func(cfg *BatcherConfig) {
+		cfg.clock = clock
+	}
+}
+
 // WithCustomFeeRate instructs sweepbatcher not to fee bump itself and rely on
 // external source of fee rates (FeeRateProvider). To apply a fee rate change,
 // the caller should re-add the sweep by calling AddSweep.
@@ -315,6 +330,11 @@ func NewBatcher(wallet lndclient.WalletKitClient,
 		opt(&cfg)
 	}
 
+	// If WithClock was not provided, use default clock.
+	if cfg.clock == nil {
+		cfg.clock = clock.NewDefaultClock()
+	}
+
 	if cfg.customMuSig2Signer != nil && musig2ServerSigner != nil {
 		panic("customMuSig2Signer must not be used with " +
 			"musig2ServerSigner")
@@ -334,6 +354,7 @@ func NewBatcher(wallet lndclient.WalletKitClient,
 		chainParams:        chainparams,
 		store:              store,
 		sweepStore:         sweepStore,
+		clock:              cfg.clock,
 		customFeeRate:      cfg.customFeeRate,
 		customMuSig2Signer: cfg.customMuSig2Signer,
 	}
@@ -934,6 +955,7 @@ func (b *Batcher) newBatchConfig(maxTimeoutDistance int32) batchConfig {
 		maxTimeoutDistance: maxTimeoutDistance,
 		noBumping:          b.customFeeRate != nil,
 		customMuSig2Signer: b.customMuSig2Signer,
+		clock:              b.clock,
 	}
 }