From d0c7fd8aac0fd7476a975e8a621181655054f443 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yong2452@gmail.com>
Date: Fri, 25 Oct 2024 15:32:30 +0800
Subject: [PATCH] sweep: add new interface method `Immediate`

This prepares the following commit where we now let the fee bumpr
decides whether to broadcast immediately or not.
---
 sweep/fee_bumper.go   |  4 ++++
 sweep/mock_test.go    |  7 +++++++
 sweep/sweeper.go      |  1 +
 sweep/sweeper_test.go |  2 ++
 sweep/tx_input_set.go | 22 ++++++++++++++++++++++
 5 files changed, 36 insertions(+)

diff --git a/sweep/fee_bumper.go b/sweep/fee_bumper.go
index bfe550886..1dff9bea8 100644
--- a/sweep/fee_bumper.go
+++ b/sweep/fee_bumper.go
@@ -143,6 +143,10 @@ type BumpRequest struct {
 	// ExtraTxOut tracks if this bump request has an optional set of extra
 	// outputs to add to the transaction.
 	ExtraTxOut fn.Option[SweepOutput]
+
+	// Immediate is used to specify that the tx should be broadcast
+	// immediately.
+	Immediate bool
 }
 
 // MaxFeeRateAllowed returns the maximum fee rate allowed for the given
diff --git a/sweep/mock_test.go b/sweep/mock_test.go
index eeeb28396..d42b0320d 100644
--- a/sweep/mock_test.go
+++ b/sweep/mock_test.go
@@ -268,6 +268,13 @@ func (m *MockInputSet) StartingFeeRate() fn.Option[chainfee.SatPerKWeight] {
 	return args.Get(0).(fn.Option[chainfee.SatPerKWeight])
 }
 
+// Immediate returns whether the inputs should be swept immediately.
+func (m *MockInputSet) Immediate() bool {
+	args := m.Called()
+
+	return args.Bool(0)
+}
+
 // MockBumper is a mock implementation of the interface Bumper.
 type MockBumper struct {
 	mock.Mock
diff --git a/sweep/sweeper.go b/sweep/sweeper.go
index a2a48beb2..6245d7941 100644
--- a/sweep/sweeper.go
+++ b/sweep/sweeper.go
@@ -827,6 +827,7 @@ func (s *UtxoSweeper) sweep(set InputSet) error {
 		DeliveryAddress: sweepAddr,
 		MaxFeeRate:      s.cfg.MaxFeeRate.FeePerKWeight(),
 		StartingFeeRate: set.StartingFeeRate(),
+		Immediate:       set.Immediate(),
 		// TODO(yy): pass the strategy here.
 	}
 
diff --git a/sweep/sweeper_test.go b/sweep/sweeper_test.go
index 7d99ba93b..2b527f5c8 100644
--- a/sweep/sweeper_test.go
+++ b/sweep/sweeper_test.go
@@ -704,11 +704,13 @@ func TestSweepPendingInputs(t *testing.T) {
 	setNeedWallet.On("Budget").Return(btcutil.Amount(1)).Once()
 	setNeedWallet.On("StartingFeeRate").Return(
 		fn.None[chainfee.SatPerKWeight]()).Once()
+	setNeedWallet.On("Immediate").Return(false).Once()
 	normalSet.On("Inputs").Return(nil).Maybe()
 	normalSet.On("DeadlineHeight").Return(testHeight).Once()
 	normalSet.On("Budget").Return(btcutil.Amount(1)).Once()
 	normalSet.On("StartingFeeRate").Return(
 		fn.None[chainfee.SatPerKWeight]()).Once()
+	normalSet.On("Immediate").Return(false).Once()
 
 	// Make pending inputs for testing. We don't need real values here as
 	// the returned clusters are mocked.
diff --git a/sweep/tx_input_set.go b/sweep/tx_input_set.go
index adae7cf13..b80d52b0e 100644
--- a/sweep/tx_input_set.go
+++ b/sweep/tx_input_set.go
@@ -64,6 +64,13 @@ type InputSet interface {
 	// StartingFeeRate returns the max starting fee rate found in the
 	// inputs.
 	StartingFeeRate() fn.Option[chainfee.SatPerKWeight]
+
+	// Immediate returns a boolean to indicate whether the tx made from
+	// this input set should be published immediately.
+	//
+	// TODO(yy): create a new method `Params` to combine the informational
+	// methods DeadlineHeight, Budget, StartingFeeRate and Immediate.
+	Immediate() bool
 }
 
 // createWalletTxInput converts a wallet utxo into an object that can be added
@@ -414,3 +421,18 @@ func (b *BudgetInputSet) StartingFeeRate() fn.Option[chainfee.SatPerKWeight] {
 
 	return startingFeeRate
 }
+
+// Immediate returns whether the inputs should be swept immediately.
+//
+// NOTE: part of the InputSet interface.
+func (b *BudgetInputSet) Immediate() bool {
+	for _, inp := range b.inputs {
+		// As long as one of the inputs is immediate, the whole set is
+		// immediate.
+		if inp.params.Immediate {
+			return true
+		}
+	}
+
+	return false
+}