diff --git a/sweepbatcher/sweep_batcher.go b/sweepbatcher/sweep_batcher.go index 05b27386..3ad1cfda 100644 --- a/sweepbatcher/sweep_batcher.go +++ b/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()) { diff --git a/sweepbatcher/sweep_batcher_test.go b/sweepbatcher/sweep_batcher_test.go index a57541f2..b12ffbf4 100644 --- a/sweepbatcher/sweep_batcher_test.go +++ b/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) {