sweeper: implement TriggerSweep

Author: Roasbeef

sweep/mock_test.go

@@ -291,7 +291,7 @@ func (m *MockBumper) Broadcast(req *BumpRequest) <-chan *BumpResult {
 		return nil
 	}
 
-	return args.Get(0).(chan *BumpResult)
+	return args.Get(0).(<-chan *BumpResult)
 }
 
 // MockFeeFunction is a mock implementation of the FeeFunction interface.

sweep/sweeper.go

@@ -326,6 +326,12 @@ type updateResp struct {
 	err        error
 }
 
+// triggerSweepReq is an internal message we'll use to represent an external
+// caller's intent to trigger an immediate sweep of all pending inputs.
+type triggerSweepReq struct {
+	respChan chan int
+}
+
 // UtxoSweeper is responsible for sweeping outputs back into the wallet
 type UtxoSweeper struct {
 	started uint32 // To be used atomically.
@@ -349,6 +355,10 @@ type UtxoSweeper struct {
 	// callers who wish to bump the fee rate of a given input.
 	updateReqs chan *updateReq
 
+	// triggerSweepReqs is a channel that will be sent requests by external
+	// callers who wish to trigger an immediate sweep of all pending inputs.
+	triggerSweepReqs chan *triggerSweepReq
+
 	// inputs is the total set of inputs the UtxoSweeper has been requested
 	// to sweep.
 	inputs InputsMap
@@ -451,6 +461,7 @@ func New(cfg *UtxoSweeperConfig) *UtxoSweeper {
 		spendChan:         make(chan *chainntnfs.SpendDetail),
 		updateReqs:        make(chan *updateReq),
 		pendingSweepsReqs: make(chan *pendingSweepsReq),
+		triggerSweepReqs:  make(chan *triggerSweepReq),
 		quit:              make(chan struct{}),
 		inputs:            make(InputsMap),
 		bumpRespChan:      make(chan *bumpResp, 100),
@@ -711,6 +722,25 @@ func (s *UtxoSweeper) collector() {
 					err)
 			}
 
+		// A new external request has been received to trigger an
+		// immediate sweep of all pending inputs.
+		case req := <-s.triggerSweepReqs:
+			// Update the inputs with the latest height.
+			inputs := s.updateSweeperInputs()
+
+			// Mark all inputs as immediate so they are broadcast
+			// right away. This is necessary for testing where we
+			// want to deterministically trigger sweeps.
+			for _, inp := range inputs {
+				inp.params.Immediate = true
+			}
+
+			// Attempt to sweep any pending inputs.
+			s.sweepPendingInputs(inputs)
+
+			// Send back the number of inputs we attempted to sweep.
+			req.respChan <- len(inputs)
+
 		// A new block comes in, update the bestHeight, perform a check
 		// over all pending inputs and publish sweeping txns if needed.
 		case beat := <-s.BlockbeatChan:
@@ -1203,6 +1233,23 @@ func (s *UtxoSweeper) ListSweeps() ([]chainhash.Hash, error) {
 	return s.cfg.Store.ListSweeps()
 }
 
+// TriggerSweep triggers an immediate attempt to create and broadcast sweep
+// transactions for all pending inputs. This is useful for testing to
+// deterministically control when sweeps are broadcast. This method is
+// thread-safe as it sends a message to the collector goroutine's event loop.
+func (s *UtxoSweeper) TriggerSweep() int {
+	req := &triggerSweepReq{
+		respChan: make(chan int, 1),
+	}
+
+	select {
+	case s.triggerSweepReqs <- req:
+		return <-req.respChan
+	case <-s.quit:
+		return 0
+	}
+}
+
 // mempoolLookup takes an input's outpoint and queries the mempool to see
 // whether it's already been spent in a transaction found in the mempool.
 // Returns the transaction if found.

sweep/sweeper_test.go

@@ -70,6 +70,76 @@ func createMockInput(t *testing.T, s *UtxoSweeper,
 	return inp
 }
 
+// setupMatureMockInput configures a mock input to be mature and ready for
+// sweeping. This centralizes the common pattern of setting up RequiredLockTime,
+// BlocksToMaturity, HeightHint, and OutPoint mocks.
+func setupMatureMockInput(inp *input.MockInput, currentHeight int32,
+	index uint32) {
+
+	inp.On("RequiredLockTime").Return(
+		uint32(currentHeight), false).Maybe()
+	inp.On("BlocksToMaturity").Return(uint32(0)).Maybe()
+	inp.On("HeightHint").Return(uint32(currentHeight)).Maybe()
+	inp.On("OutPoint").Return(wire.OutPoint{Index: index}).Maybe()
+}
+
+// newTestSweeperConfig returns a UtxoSweeperConfig with common test settings.
+// This avoids repeating the GenSweepScript setup across multiple tests.
+func newTestSweeperConfig(wallet Wallet, aggregator UtxoAggregator,
+	publisher Bumper) *UtxoSweeperConfig {
+
+	return &UtxoSweeperConfig{
+		Wallet:     wallet,
+		Aggregator: aggregator,
+		Publisher:  publisher,
+		GenSweepScript: func() fn.Result[lnwallet.AddrWithKey] {
+			return fn.Ok(lnwallet.AddrWithKey{
+				DeliveryAddress: testPubKey.SerializeCompressed(),
+			})
+		},
+		NoDeadlineConfTarget: uint32(DefaultDeadlineDelta),
+	}
+}
+
+// mockInputSet sets up a standard input set mock with common expectations.
+// This centralizes the repetitive pattern of mocking input set methods.
+func mockInputSet(set *MockInputSet, inputs []input.Input) {
+	set.On("Inputs").Return(inputs).Maybe()
+	set.On("NeedWalletInput").Return(false).Once()
+	set.On("DeadlineHeight").Return(int32(testHeight)).Once()
+	set.On("Budget").Return(btcutil.Amount(1)).Once()
+	set.On("StartingFeeRate").Return(
+		fn.None[chainfee.SatPerKWeight]()).Once()
+	set.On("Immediate").Return(true).Once()
+}
+
+// mockAggregatorCluster sets up a standard ClusterInputs mock expectation.
+func mockAggregatorCluster(aggregator *mockUtxoAggregator,
+	sets []InputSet) {
+
+	aggregator.On("ClusterInputs", mock.Anything).Return(sets).Once()
+}
+
+// mockPublisherBroadcast sets up a standard Broadcast mock expectation.
+func mockPublisherBroadcast(publisher *MockBumper) {
+	publisher.On("Broadcast", mock.Anything).Return(
+		make(<-chan *BumpResult)).Once()
+}
+
+// setupMockInputSetForSweep configures a MockInputSet with the standard mocks
+// needed for sweep operations, allowing customization of the immediate flag.
+func setupMockInputSetForSweep(set *MockInputSet, inputs []input.Input,
+	immediate bool) {
+
+	set.On("Inputs").Return(inputs).Maybe()
+	set.On("NeedWalletInput").Return(false).Once()
+	set.On("DeadlineHeight").Return(int32(testHeight)).Once()
+	set.On("Budget").Return(btcutil.Amount(1)).Once()
+	set.On("StartingFeeRate").Return(
+		fn.None[chainfee.SatPerKWeight]()).Once()
+	set.On("Immediate").Return(immediate).Once()
+}
+
 // TestMarkInputsPendingPublish checks that given a list of inputs with
 // different states, only the non-terminal state will be marked as `Published`.
 func TestMarkInputsPendingPublish(t *testing.T) {
@@ -79,7 +149,6 @@ func TestMarkInputsPendingPublish(t *testing.T) {
 
 	// Create a test sweeper.
 	s := New(&UtxoSweeperConfig{})
-
 	// Create a mock input set.
 	set := &MockInputSet{}
 	defer set.AssertExpectations(t)
@@ -1446,3 +1515,147 @@ func TestHandleBumpEventTxUnknownSpendWithRetry(t *testing.T) {
 	// Assert the state of the input is updated.
 	require.Equal(t, PublishFailed, s.inputs[op2].state)
 }
+
+// TestTriggerSweep checks that the `TriggerSweep` method correctly triggers an
+// immediate sweep of all pending inputs and returns the count of inputs
+// attempted.
+func TestTriggerSweep(t *testing.T) {
+	t.Parallel()
+
+	require := require.New(t)
+
+	wallet := &MockWallet{}
+	defer wallet.AssertExpectations(t)
+	aggregator := &mockUtxoAggregator{}
+	defer aggregator.AssertExpectations(t)
+	publisher := &MockBumper{}
+	defer publisher.AssertExpectations(t)
+
+	s := New(newTestSweeperConfig(wallet, aggregator, publisher))
+
+	inp1 := &input.MockInput{}
+	defer inp1.AssertExpectations(t)
+	inp2 := &input.MockInput{}
+	defer inp2.AssertExpectations(t)
+	inp3 := &input.MockInput{}
+	defer inp3.AssertExpectations(t)
+
+	setupMatureMockInput(inp1, s.currentHeight, 1)
+	setupMatureMockInput(inp2, s.currentHeight, 2)
+	setupMatureMockInput(inp3, s.currentHeight, 3)
+
+	input1 := &SweeperInput{state: Init, Input: inp1}
+	input2 := &SweeperInput{state: PublishFailed, Input: inp2}
+	input3 := &SweeperInput{state: PendingPublish, Input: inp3}
+
+	s.inputs = map[wire.OutPoint]*SweeperInput{
+		{Index: 1}: input1,
+		{Index: 2}: input2,
+		{Index: 3}: input3,
+	}
+
+	s.wg.Add(1)
+	go s.collector()
+
+	set := &MockInputSet{}
+	defer set.AssertExpectations(t)
+	setupMockInputSetForSweep(set, []input.Input{inp1, inp2}, true)
+	mockAggregatorCluster(aggregator, []InputSet{set})
+	mockPublisherBroadcast(publisher)
+
+	numInputs := s.TriggerSweep()
+
+	// Only Init and PublishFailed states are sweepable.
+	require.Equal(2, numInputs)
+
+	close(s.quit)
+	s.wg.Wait()
+}
+
+// TestTriggerSweepNoInputs checks that `TriggerSweep` returns 0 when there are
+// no pending inputs to sweep.
+func TestTriggerSweepNoInputs(t *testing.T) {
+	t.Parallel()
+
+	require := require.New(t)
+
+	aggregator := &mockUtxoAggregator{}
+	defer aggregator.AssertExpectations(t)
+
+	s := New(&UtxoSweeperConfig{Aggregator: aggregator})
+	mockAggregatorCluster(aggregator, []InputSet{})
+
+	s.wg.Add(1)
+	go s.collector()
+
+	numInputs := s.TriggerSweep()
+
+	require.Equal(0, numInputs)
+
+	close(s.quit)
+	s.wg.Wait()
+}
+
+// TestTriggerSweepSweeperShutdown checks that `TriggerSweep` returns 0 when
+// the sweeper is shutting down.
+func TestTriggerSweepSweeperShutdown(t *testing.T) {
+	t.Parallel()
+
+	require := require.New(t)
+
+	s := New(&UtxoSweeperConfig{})
+	close(s.quit)
+
+	numInputs := s.TriggerSweep()
+
+	require.Equal(0, numInputs)
+}
+
+// TestTriggerSweepMarksImmediate checks that `TriggerSweep` marks all inputs
+// as immediate before sweeping.
+func TestTriggerSweepMarksImmediate(t *testing.T) {
+	t.Parallel()
+
+	require := require.New(t)
+
+	aggregator := &mockUtxoAggregator{}
+	defer aggregator.AssertExpectations(t)
+	publisher := &MockBumper{}
+	defer publisher.AssertExpectations(t)
+
+	s := New(newTestSweeperConfig(nil, aggregator, publisher))
+
+	inp := &input.MockInput{}
+	defer inp.AssertExpectations(t)
+	setupMatureMockInput(inp, s.currentHeight, 1)
+
+	input1 := &SweeperInput{
+		state: Init,
+		Input: inp,
+		params: Params{
+			Immediate: false,
+		},
+	}
+
+	s.inputs = map[wire.OutPoint]*SweeperInput{
+		{Index: 1}: input1,
+	}
+
+	s.wg.Add(1)
+	go s.collector()
+
+	set := &MockInputSet{}
+	defer set.AssertExpectations(t)
+	setupMockInputSetForSweep(set, []input.Input{inp}, true)
+	mockAggregatorCluster(aggregator, []InputSet{set})
+	mockPublisherBroadcast(publisher)
+
+	numInputs := s.TriggerSweep()
+
+	// Verify TriggerSweep marked the input as immediate.
+	require.True(input1.params.Immediate)
+	require.Equal(1, numInputs)
+
+	close(s.quit)
+	s.wg.Wait()
+}