sweepbatcher: fix race in store_mock

This commit is contained in:
Boris Nagaev 2025-02-25 00:20:59 -03:00
parent e86ccb9bfc
commit f0f64f83c3
No known key found for this signature in database

View file

@ -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,7 +121,11 @@ 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
}
@ -112,6 +133,9 @@ func (s *StoreMock) UpsertSweep(ctx context.Context, sweep *dbSweep) error {
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")