sweepbatcher: fix race in store_mock

Author: starius

sweepbatcher/store_mock.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"sort"
+	"sync"
 
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/lightningnetwork/lnd/lntypes"
@@ -13,6 +14,7 @@ import (
 type StoreMock struct {
 	batches map[int32]dbBatch
 	sweeps  map[lntypes.Hash]dbSweep
+	mu      sync.Mutex
 }
 
 // NewStoreMock instantiates a new mock store.
@@ -28,6 +30,9 @@ func NewStoreMock() *StoreMock {
 func (s *StoreMock) FetchUnconfirmedSweepBatches(ctx context.Context) (
 	[]*dbBatch, error) {
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	result := []*dbBatch{}
 	for _, batch := range s.batches {
 		batch := batch
@@ -44,6 +49,9 @@ func (s *StoreMock) FetchUnconfirmedSweepBatches(ctx context.Context) (
 func (s *StoreMock) InsertSweepBatch(ctx context.Context,
 	batch *dbBatch) (int32, error) {
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	var id int32
 
 	if len(s.batches) == 0 {
@@ -66,12 +74,18 @@ func (s *StoreMock) DropBatch(ctx context.Context, id int32) error {
 func (s *StoreMock) UpdateSweepBatch(ctx context.Context,
 	batch *dbBatch) error {
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	s.batches[batch.ID] = *batch
 	return nil
 }
 
 // ConfirmBatch confirms a batch.
 func (s *StoreMock) ConfirmBatch(ctx context.Context, id int32) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	batch, ok := s.batches[id]
 	if !ok {
 		return errors.New("batch not found")
@@ -87,6 +101,9 @@ func (s *StoreMock) ConfirmBatch(ctx context.Context, id int32) error {
 func (s *StoreMock) FetchBatchSweeps(ctx context.Context,
 	id int32) ([]*dbSweep, error) {
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	result := []*dbSweep{}
 	for _, sweep := range s.sweeps {
 		sweep := sweep
@@ -104,14 +121,21 @@ func (s *StoreMock) FetchBatchSweeps(ctx context.Context,
 
 // UpsertSweep inserts a sweep into the database, or updates an existing sweep.
 func (s *StoreMock) UpsertSweep(ctx context.Context, sweep *dbSweep) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	s.sweeps[sweep.SwapHash] = *sweep
+
 	return nil
 }
 
 // GetSweepStatus returns the status of a sweep.
 func (s *StoreMock) GetSweepStatus(ctx context.Context,
 	swapHash lntypes.Hash) (bool, error) {
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	sweep, ok := s.sweeps[swapHash]
 	if !ok {
 		return false, nil
@@ -127,6 +151,9 @@ func (s *StoreMock) Close() error {
 
 // AssertSweepStored asserts that a sweep is stored.
 func (s *StoreMock) AssertSweepStored(id lntypes.Hash) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	_, ok := s.sweeps[id]
 	return ok
 }
@@ -135,6 +162,9 @@ func (s *StoreMock) AssertSweepStored(id lntypes.Hash) bool {
 func (s *StoreMock) GetParentBatch(ctx context.Context, swapHash lntypes.Hash) (
 	*dbBatch, error) {
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	for _, sweep := range s.sweeps {
 		if sweep.SwapHash == swapHash {
 			batch, ok := s.batches[sweep.BatchID]
@@ -153,6 +183,9 @@ func (s *StoreMock) GetParentBatch(ctx context.Context, swapHash lntypes.Hash) (
 func (s *StoreMock) TotalSweptAmount(ctx context.Context, batchID int32) (
 	btcutil.Amount, error) {
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	batch, ok := s.batches[batchID]
 	if !ok {
 		return 0, errors.New("batch not found")