sweepbatcher: harden AddSweep against ctx closure

sweepbatcher/sweep_batcher.go

@@ -780,16 +780,21 @@ 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.shutdownOrCancelErrIfAny(ctx)
+	if err != nil {
+		return err
 	}
 
 	sweeps, err := b.fetchSweeps(ctx, *sweepReq)
 	if err != nil {
+		if exitErr := b.shutdownOrCancelErrIfAny(ctx); exitErr != nil {
+			infof("fetchSweeps failed during shutdown, returning "+
+				"%v instead of %v.", exitErr, err)
+
+			return exitErr
+		}
+
 		return fmt.Errorf("fetchSweeps failed: %w", err)
 	}
 
@@ -803,6 +808,14 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 
 	completed, err := b.store.GetSweepStatus(ctx, sweep.outpoint)
 	if err != nil {
+		if exitErr := b.shutdownOrCancelErrIfAny(ctx); exitErr != nil {
+			infof("GetSweepStatus failed for sweep %v during "+
+				"shutdown, returning %v instead of %v.",
+				sweep.outpoint, exitErr, err)
+
+			return exitErr
+		}
+
 		return fmt.Errorf("failed to get the status of sweep %v: %w",
 			sweep.outpoint, err)
 	}
@@ -816,6 +829,14 @@ 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 {
+			if exitErr := b.shutdownOrCancelErrIfAny(ctx); exitErr != nil {
+				infof("GetParentBatch failed for sweep %v "+
+					"during shutdown, returning %v instead "+
+					"of %v.", sweep.outpoint, exitErr, err)
+
+				return exitErr
+			}
+
 			return fmt.Errorf("unable to get parent batch for "+
 				"sweep %x: %w", sweep.swapHash[:6], err)
 		}
@@ -827,6 +848,13 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 
 	minRelayFeeRate, err := b.wallet.MinRelayFee(ctx)
 	if err != nil {
+		if exitErr := b.shutdownOrCancelErrIfAny(ctx); exitErr != nil {
+			infof("MinRelayFee failed during shutdown, returning "+
+				"%v instead of %v.", exitErr, err)
+
+			return exitErr
+		}
+
 		return fmt.Errorf("failed to get min relay fee: %w", err)
 	}
 
@@ -839,6 +867,15 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 			b.chainParams,
 		)
 		if err != nil {
+			if exitErr := b.shutdownOrCancelErrIfAny(ctx); exitErr != nil {
+				infof("ensurePresigned failed for primary "+
+					"sweep %v during shutdown, returning %v "+
+					"instead of %v.", sweep.outpoint,
+					exitErr, err)
+
+				return exitErr
+			}
+
 			return fmt.Errorf("inputs with primarySweep %v were "+
 				"not presigned (call PresignSweepsGroup "+
 				"first): %w", sweep.outpoint, err)
@@ -861,9 +898,26 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 
 	case <-b.quit:
 		return ErrBatcherShuttingDown
+
+	case <-ctx.Done():
+		return b.shutdownOrCancelErrIfAny(ctx)
 	}
 }
 
+// shutdownOrCancelErrIfAny returns the terminal error to use when caller-facing
+// batcher methods race with shutdown or caller cancellation. It returns nil if
+// the operation should continue.
+func (b *Batcher) shutdownOrCancelErrIfAny(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,60 @@ func (f *sweepFetcherMock) FetchSweep(ctx context.Context, _ lntypes.Hash,
 	return f.store[outpoint], nil
 }
 
+// cancelingSweepFetcher cancels its caller context while returning a backend
+// fetch error.
+type cancelingSweepFetcher struct {
+	cancel context.CancelFunc
+}
+
+func (f *cancelingSweepFetcher) FetchSweep(context.Context, lntypes.Hash,
+	wire.OutPoint) (*SweepInfo, error) {
+
+	// Simulate the caller canceling while the backend returns a
+	// driver-level error.
+	f.cancel()
+
+	return nil, driver.ErrBadConn
+}
+
+// testAddSweepReturnsContextErrorOnFetchCancellation asserts that AddSweep
+// returns context.Canceled instead of a backend error when sweep fetching races
+// with caller cancellation.
+func testAddSweepReturnsContextErrorOnFetchCancellation(t *testing.T,
+	_ testStore, batcherStore testBatcherStore) {
+
+	defer test.Guard(t)()
+
+	lnd := test.NewMockLnd()
+	ctx, cancel := context.WithCancel(t.Context())
+
+	batcher := NewBatcher(
+		lnd.WalletKit, lnd.ChainNotifier, lnd.Signer,
+		testMuSig2SignSweep, testVerifySchnorrSig, lnd.ChainParams,
+		batcherStore, &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)
+}
+
+// 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) {
+	runTests(t, testAddSweepReturnsContextErrorOnFetchCancellation)
+}
+
 // testSweepFetcher tests providing custom sweep fetcher to Batcher.
 func testSweepFetcher(t *testing.T, store testStore,
 	batcherStore testBatcherStore) {