sweepbatcher: harden AddSweep against ctx closure

Author: hieblmi

sweepbatcher/sweep_batcher.go

@@ -780,16 +780,19 @@ func (b *Batcher) PresignSweepsGroup(ctx context.Context, inputs []Input,
 // times, but the sweeps (including the order of them) must be the same. If
 // notifier is provided, the batcher sends back sweeping results through it.
 func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
-	// If the batcher is shutting down, quit now.
-	select {
-	case <-b.quit:
-		return ErrBatcherShuttingDown
-
-	default:
+	// If the batcher or the caller is shutting down, quit now.
+	err := b.addSweepExitErr(ctx)
+	if err != nil {
+		return err
 	}
 
 	sweeps, err := b.fetchSweeps(ctx, *sweepReq)
 	if err != nil {
+		exitErr := b.addSweepExitErr(ctx)
+		if exitErr != nil {
+			return fmt.Errorf("fetchSweeps failed: %w", exitErr)
+		}
+
 		return fmt.Errorf("fetchSweeps failed: %w", err)
 	}
 
@@ -803,6 +806,12 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 
 	completed, err := b.store.GetSweepStatus(ctx, sweep.outpoint)
 	if err != nil {
+		exitErr := b.addSweepExitErr(ctx)
+		if exitErr != nil {
+			return fmt.Errorf("failed to get the status of "+
+				"sweep %v: %w", sweep.outpoint, exitErr)
+		}
+
 		return fmt.Errorf("failed to get the status of sweep %v: %w",
 			sweep.outpoint, err)
 	}
@@ -816,6 +825,13 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 		// on-chain confirmations to prevent issues caused by reorgs.
 		parentBatch, err = b.store.GetParentBatch(ctx, sweep.outpoint)
 		if err != nil {
+			exitErr := b.addSweepExitErr(ctx)
+			if exitErr != nil {
+				return fmt.Errorf("unable to get parent "+
+					"batch for sweep %x: %w",
+					sweep.swapHash[:6], exitErr)
+			}
+
 			return fmt.Errorf("unable to get parent batch for "+
 				"sweep %x: %w", sweep.swapHash[:6], err)
 		}
@@ -827,6 +843,12 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 
 	minRelayFeeRate, err := b.wallet.MinRelayFee(ctx)
 	if err != nil {
+		exitErr := b.addSweepExitErr(ctx)
+		if exitErr != nil {
+			return fmt.Errorf("failed to get min relay fee: %w",
+				exitErr)
+		}
+
 		return fmt.Errorf("failed to get min relay fee: %w", err)
 	}
 
@@ -861,9 +883,30 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 
 	case <-b.quit:
 		return ErrBatcherShuttingDown
+
+	case <-ctx.Done():
+		err := b.addSweepExitErr(ctx)
+		if err != nil {
+			return err
+		}
+
+		return ctx.Err()
 	}
 }
 
+// addSweepExitErr returns the terminal error to use when AddSweep races with
+// shutdown or caller cancellation.
+func (b *Batcher) addSweepExitErr(ctx context.Context) error {
+	select {
+	case <-b.quit:
+		return ErrBatcherShuttingDown
+
+	default:
+	}
+
+	return ctx.Err()
+}
+
 // testRunInEventLoop runs a function in the event loop blocking until
 // the function returns. For unit tests only!
 func (b *Batcher) testRunInEventLoop(ctx context.Context, handler func()) {

sweepbatcher/sweep_batcher_test.go

@@ -3,6 +3,7 @@ package sweepbatcher
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	"errors"
 	"fmt"
 	"maps"
@@ -3711,6 +3712,47 @@ func (f *sweepFetcherMock) FetchSweep(ctx context.Context, _ lntypes.Hash,
 	return f.store[outpoint], nil
 }
 
+type cancelingSweepFetcher struct {
+	cancel context.CancelFunc
+}
+
+func (f *cancelingSweepFetcher) FetchSweep(context.Context, lntypes.Hash,
+	wire.OutPoint) (*SweepInfo, error) {
+
+	f.cancel()
+
+	return nil, driver.ErrBadConn
+}
+
+// TestAddSweepReturnsContextErrorOnFetchCancellation asserts that AddSweep
+// returns the context cancellation error if sweep fetching fails while the
+// caller context is being canceled.
+func TestAddSweepReturnsContextErrorOnFetchCancellation(t *testing.T) {
+	defer test.Guard(t)()
+
+	lnd := test.NewMockLnd()
+	ctx, cancel := context.WithCancel(context.Background())
+
+	batcher := NewBatcher(
+		lnd.WalletKit, lnd.ChainNotifier, lnd.Signer,
+		testMuSig2SignSweep, testVerifySchnorrSig, lnd.ChainParams,
+		NewStoreMock(), &cancelingSweepFetcher{cancel: cancel},
+	)
+
+	err := batcher.AddSweep(ctx, &SweepRequest{
+		SwapHash: lntypes.Hash{1, 1, 1},
+		Inputs: []Input{{
+			Value: 1111,
+			Outpoint: wire.OutPoint{
+				Hash:  chainhash.Hash{1, 1},
+				Index: 1,
+			},
+		}},
+	})
+	require.ErrorIs(t, err, context.Canceled)
+	require.NotErrorIs(t, err, driver.ErrBadConn)
+}
+
 // testSweepFetcher tests providing custom sweep fetcher to Batcher.
 func testSweepFetcher(t *testing.T, store testStore,
 	batcherStore testBatcherStore) {