sweepbatcher/test: fix races in require.Eventually

The code inside require.Eventually runs in parallel with the event loops
of the batcher and its batches. Accessing fields of the batcher and batches
must be done within an event loop.

To address this, testRunInEventLoop methods were added to the Batcher and batch
types. Unit tests were then rewritten to use this approach when accessing
batcher and batch fields.

Additionally, in many cases, receive operations from RegisterSpendChannel
were moved before require.Eventually. This prevents testRunInEventLoop from
getting stuck in an event loop while blocked on a RegisterSpendChannel send
operation.

Author: starius

sweepbatcher/sweep_batch.go

@@ -215,6 +215,12 @@ type batch struct {
 	// reorgChan is the channel over which reorg notifications are received.
 	reorgChan chan struct{}
 
+	// testReqs is a channel where test requests are received.
+	// This is used only in unit tests! The reason to have this is to
+	// avoid data races in require.Eventually calls running in parallel
+	// to the event loop. See method testRunInEventLoop().
+	testReqs chan *testRequest
+
 	// errChan is the channel over which errors are received.
 	errChan chan error
 
@@ -352,6 +358,7 @@ func NewBatch(cfg batchConfig, bk batchKit) *batch {
 		spendChan:           make(chan *chainntnfs.SpendDetail),
 		confChan:            make(chan *chainntnfs.TxConfirmation, 1),
 		reorgChan:           make(chan struct{}, 1),
+		testReqs:            make(chan *testRequest),
 		errChan:             make(chan error, 1),
 		callEnter:           make(chan struct{}),
 		callLeave:           make(chan struct{}),
@@ -396,6 +403,7 @@ func NewBatchFromDB(cfg batchConfig, bk batchKit) (*batch, error) {
 		spendChan:           make(chan *chainntnfs.SpendDetail),
 		confChan:            make(chan *chainntnfs.TxConfirmation, 1),
 		reorgChan:           make(chan struct{}, 1),
+		testReqs:            make(chan *testRequest),
 		errChan:             make(chan error, 1),
 		callEnter:           make(chan struct{}),
 		callLeave:           make(chan struct{}),
@@ -756,6 +764,10 @@ func (b *batch) Run(ctx context.Context) error {
 				return err
 			}
 
+		case testReq := <-b.testReqs:
+			testReq.handler()
+			close(testReq.quit)
+
 		case err := <-blockErrChan:
 			return err
 
@@ -768,6 +780,36 @@ func (b *batch) Run(ctx context.Context) error {
 	}
 }
 
+// testRunInEventLoop runs a function in the event loop blocking until
+// the function returns. For unit tests only!
+func (b *batch) testRunInEventLoop(ctx context.Context, handler func()) {
+	// If the event loop is finished, run the function.
+	select {
+	case <-b.stopping:
+		handler()
+
+		return
+	default:
+	}
+
+	quit := make(chan struct{})
+	req := &testRequest{
+		handler: handler,
+		quit:    quit,
+	}
+
+	select {
+	case b.testReqs <- req:
+	case <-ctx.Done():
+		return
+	}
+
+	select {
+	case <-quit:
+	case <-ctx.Done():
+	}
+}
+
 // timeout returns minimum timeout as block height among sweeps of the batch.
 // If the batch is empty, return -1.
 func (b *batch) timeout() int32 {

sweepbatcher/sweep_batcher.go

@@ -225,6 +225,16 @@ var (
 	ErrBatcherShuttingDown = errors.New("batcher shutting down")
 )
 
+// testRequest is a function passed to an event loop and a channel used to
+// wait until the function is executed. This is used in unit tests only!
+type testRequest struct {
+	// handler is the function to an event loop.
+	handler func()
+
+	// quit is closed when the handler completes.
+	quit chan struct{}
+}
+
 // Batcher is a system that is responsible for accepting sweep requests and
 // placing them in appropriate batches. It will spin up new batches as needed.
 type Batcher struct {
@@ -234,6 +244,12 @@ type Batcher struct {
 	// sweepReqs is a channel where sweep requests are received.
 	sweepReqs chan SweepRequest
 
+	// testReqs is a channel where test requests are received.
+	// This is used only in unit tests! The reason to have this is to
+	// avoid data races in require.Eventually calls running in parallel
+	// to the event loop. See method testRunInEventLoop().
+	testReqs chan *testRequest
+
 	// errChan is a channel where errors are received.
 	errChan chan error
 
@@ -461,6 +477,7 @@ func NewBatcher(wallet lndclient.WalletKitClient,
 	return &Batcher{
 		batches:             make(map[int32]*batch),
 		sweepReqs:           make(chan SweepRequest),
+		testReqs:            make(chan *testRequest),
 		errChan:             make(chan error, 1),
 		quit:                make(chan struct{}),
 		initDone:            make(chan struct{}),
@@ -528,6 +545,10 @@ func (b *Batcher) Run(ctx context.Context) error {
 				return err
 			}
 
+		case testReq := <-b.testReqs:
+			testReq.handler()
+			close(testReq.quit)
+
 		case err := <-b.errChan:
 			log.Warnf("Batcher received an error: %v.", err)
 			return err
@@ -551,6 +572,36 @@ func (b *Batcher) AddSweep(sweepReq *SweepRequest) error {
 	}
 }
 
+// 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()) {
+	// If the event loop is finished, run the function.
+	select {
+	case <-b.quit:
+		handler()
+
+		return
+	default:
+	}
+
+	quit := make(chan struct{})
+	req := &testRequest{
+		handler: handler,
+		quit:    quit,
+	}
+
+	select {
+	case b.testReqs <- req:
+	case <-ctx.Done():
+		return
+	}
+
+	select {
+	case <-quit:
+	case <-ctx.Done():
+	}
+}
+
 // handleSweep handles a sweep request by either placing it in an existing
 // batch, or by spinning up a new batch for it.
 func (b *Batcher) handleSweep(ctx context.Context, sweep *sweep,