diff --git a/sweepbatcher/sweep_batch.go b/sweepbatcher/sweep_batch.go index 0374745a..c837064d 100644 --- a/sweepbatcher/sweep_batch.go +++ b/sweepbatcher/sweep_batch.go @@ -193,8 +193,12 @@ type batch struct { // main event loop. callLeave chan struct{} - // stopped signals that the batch has stopped. - stopped chan struct{} + // stopping signals that the batch is stopping. + stopping chan struct{} + + // finished signals that the batch has stopped and all child goroutines + // have finished. + finished chan struct{} // quit is owned by the parent batcher and signals that the batch must // stop. @@ -273,7 +277,10 @@ func (b *batch) scheduleNextCall() (func(), error) { case <-b.quit: return func() {}, ErrBatcherShuttingDown - case <-b.stopped: + case <-b.stopping: + return func() {}, ErrBatchShuttingDown + + case <-b.finished: return func() {}, ErrBatchShuttingDown } @@ -297,7 +304,8 @@ func NewBatch(cfg batchConfig, bk batchKit) *batch { errChan: make(chan error, 1), callEnter: make(chan struct{}), callLeave: make(chan struct{}), - stopped: make(chan struct{}), + stopping: make(chan struct{}), + finished: make(chan struct{}), quit: bk.quit, batchTxid: bk.batchTxid, wallet: bk.wallet, @@ -340,7 +348,8 @@ func NewBatchFromDB(cfg batchConfig, bk batchKit) (*batch, error) { errChan: make(chan error, 1), callEnter: make(chan struct{}), callLeave: make(chan struct{}), - stopped: make(chan struct{}), + stopping: make(chan struct{}), + finished: make(chan struct{}), quit: bk.quit, batchTxid: bk.batchTxid, batchPkScript: bk.batchPkScript, @@ -460,7 +469,7 @@ func (b *batch) sweepExists(hash lntypes.Hash) bool { // Wait waits for the batch to gracefully stop. func (b *batch) Wait() { b.log.Infof("Stopping") - b.wg.Wait() + <-b.finished } // Run is the batch's main event loop. @@ -468,8 +477,12 @@ func (b *batch) Run(ctx context.Context) error { runCtx, cancel := context.WithCancel(ctx) defer func() { cancel() - close(b.stopped) + close(b.stopping) + + // Make sure not to call b.wg.Wait from any other place to avoid + // race condition between b.wg.Add(1) and b.wg.Wait(). b.wg.Wait() + close(b.finished) }() if b.muSig2SignSweep == nil { diff --git a/sweepbatcher/sweep_batcher_test.go b/sweepbatcher/sweep_batcher_test.go index e5dc2f61..d211e88b 100644 --- a/sweepbatcher/sweep_batcher_test.go +++ b/sweepbatcher/sweep_batcher_test.go @@ -1538,3 +1538,99 @@ func TestSweepFetcher(t *testing.T) { // Make sure the batcher exited without an error. checkBatcherError(t, runErr) } + +// TestSweepBatcherCloseDuringAdding tests that sweep batcher works correctly +// if it is closed (stops running) during AddSweep call. +func TestSweepBatcherCloseDuringAdding(t *testing.T) { + defer test.Guard(t)() + + lnd := test.NewMockLnd() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + store := loopdb.NewStoreMock(t) + sweepStore, err := NewSweepFetcherFromSwapStore(store, lnd.ChainParams) + require.NoError(t, err) + + batcherStore := NewStoreMock() + + batcher := NewBatcher(lnd.WalletKit, lnd.ChainNotifier, lnd.Signer, + testMuSig2SignSweep, nil, lnd.ChainParams, batcherStore, + sweepStore) + go func() { + err := batcher.Run(ctx) + checkBatcherError(t, err) + }() + + // Add many swaps. + for i := byte(1); i < 255; i++ { + swapHash := lntypes.Hash{i, i, i} + + // Create a swap contract. + swap := &loopdb.LoopOutContract{ + SwapContract: loopdb.SwapContract{ + CltvExpiry: 111, + AmountRequested: 111, + }, + + SwapInvoice: swapInvoice, + } + + err = store.CreateLoopOut(ctx, swapHash, swap) + require.NoError(t, err) + store.AssertLoopOutStored() + } + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // Add many sweeps. + for i := byte(1); i < 255; i++ { + // Create a sweep request. + sweepReq := SweepRequest{ + SwapHash: lntypes.Hash{i, i, i}, + Value: 111, + Outpoint: wire.OutPoint{ + Hash: chainhash.Hash{i, i}, + Index: 1, + }, + Notifier: &dummyNotifier, + } + + // Deliver sweep request to batcher. + err := batcher.AddSweep(&sweepReq) + if err == ErrBatcherShuttingDown { + break + } + require.NoError(t, err) + } + }() + + wg.Add(1) + go func() { + defer wg.Done() + // Close sweepbatcher during addings. + time.Sleep(1 * time.Millisecond) + cancel() + }() + + // We don't know how many spend notification registrations will be + // issued, so accept them while waiting for two goroutines to stop. + quit := make(chan struct{}) + registrationChan := make(chan struct{}) + go func() { + defer close(registrationChan) + for { + select { + case <-lnd.RegisterSpendChannel: + case <-quit: + return + } + } + }() + + wg.Wait() + close(quit) + <-registrationChan +}