Add lock to StuckTxDetector & InMemoryStoreManager

pkg/txm/clientwrappers/dualbroadcast/meta_error_handler_test.go

@@ -34,7 +34,8 @@ func TestMetaErrorHandler(t *testing.T) {
 		setNonce := func(address common.Address, nonce uint64) {}
 		txStoreManager := storage.NewInMemoryStoreManager(lggr, testutils.FixtureChainID)
 		require.NoError(t, txStoreManager.Add(address))
-		txStore := txStoreManager.InMemoryStoreMap[address]
+		txStore, exists := txStoreManager.GetStoreSafe(address)
+		require.True(t, exists)
 		_ = txStore.CreateTransaction(txRequest)
 		tx, err := txStore.UpdateUnstartedTransactionWithNonce(nonce)
 		require.NoError(t, err)
@@ -67,7 +68,8 @@ func TestMetaErrorHandler(t *testing.T) {
 		setNonce := func(address common.Address, nonce uint64) {}
 		txStoreManager := storage.NewInMemoryStoreManager(logger.Test(t), testutils.FixtureChainID)
 		require.NoError(t, txStoreManager.Add(address))
-		txStore := txStoreManager.InMemoryStoreMap[address]
+		txStore, exists := txStoreManager.GetStoreSafe(address)
+		require.True(t, exists)
 		_ = txStore.CreateTransaction(txRequest)
 		tx, err := txStore.UpdateUnstartedTransactionWithNonce(nonce)
 		require.NoError(t, err)
@@ -102,7 +104,8 @@ func TestMetaErrorHandler(t *testing.T) {
 		setNonce := func(address common.Address, nonce uint64) {}
 		txStoreManager := storage.NewInMemoryStoreManager(lggr, testutils.FixtureChainID)
 		require.NoError(t, txStoreManager.Add(address))
-		txStore := txStoreManager.InMemoryStoreMap[address]
+		txStore, exists := txStoreManager.GetStoreSafe(address)
+		require.True(t, exists)
 		_ = txStore.CreateTransaction(txRequest)
 		tx, err := txStore.UpdateUnstartedTransactionWithNonce(nonce)
 		require.NoError(t, err)

pkg/txm/storage/inmemory_store_manager.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math/big"
+	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
 	evmtypes "github.com/ethereum/go-ethereum/core/types"
@@ -16,132 +17,152 @@ import (
 const StoreNotFoundForAddress string = "InMemoryStore for address: %v not found"
 
 type InMemoryStoreManager struct {
+	mu               sync.RWMutex
 	lggr             logger.Logger
 	chainID          *big.Int
-	InMemoryStoreMap map[common.Address]*InMemoryStore
+	inMemoryStoreMap map[common.Address]*InMemoryStore
 }
 
 func NewInMemoryStoreManager(lggr logger.Logger, chainID *big.Int) *InMemoryStoreManager {
 	inMemoryStoreMap := make(map[common.Address]*InMemoryStore)
 	return &InMemoryStoreManager{
 		lggr:             lggr,
 		chainID:          chainID,
-		InMemoryStoreMap: inMemoryStoreMap}
+		inMemoryStoreMap: inMemoryStoreMap}
+}
+
+func (m *InMemoryStoreManager) GetStoreSafe(fromAddress common.Address) (*InMemoryStore, bool) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	store, exists := m.inMemoryStoreMap[fromAddress]
+	return store, exists
 }
 
 func (m *InMemoryStoreManager) AbandonPendingTransactions(_ context.Context, fromAddress common.Address) error {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		store.AbandonPendingTransactions()
 		return nil
 	}
 	return fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) Add(addresses ...common.Address) (err error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	for _, address := range addresses {
-		if _, exists := m.InMemoryStoreMap[address]; exists {
+		if _, exists := m.inMemoryStoreMap[address]; exists {
 			err = errors.Join(err, fmt.Errorf("address %v already exists in store manager", address))
+			continue
 		}
-		m.InMemoryStoreMap[address] = NewInMemoryStore(m.lggr, address, m.chainID)
+		m.inMemoryStoreMap[address] = NewInMemoryStore(m.lggr, address, m.chainID)
 	}
 	return
 }
 
 func (m *InMemoryStoreManager) AppendAttemptToTransaction(_ context.Context, txNonce uint64, fromAddress common.Address, attempt *types.Attempt) (attempts []*types.Attempt, err error) {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.AppendAttemptToTransaction(txNonce, attempt)
 	}
 	return nil, fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) CountUnstartedTransactions(fromAddress common.Address) (int, error) {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.CountUnstartedTransactions(), nil
 	}
 	return 0, fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) CreateEmptyUnconfirmedTransaction(_ context.Context, fromAddress common.Address, nonce uint64, gasLimit uint64) (*types.Transaction, error) {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.CreateEmptyUnconfirmedTransaction(nonce, gasLimit)
 	}
 	return nil, fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) CreateTransaction(_ context.Context, txRequest *types.TxRequest) (*types.Transaction, error) {
-	if store, exists := m.InMemoryStoreMap[txRequest.FromAddress]; exists {
+	if store, exists := m.GetStoreSafe(txRequest.FromAddress); exists {
 		return store.CreateTransaction(txRequest), nil
 	}
 	return nil, fmt.Errorf(StoreNotFoundForAddress, txRequest.FromAddress)
 }
 
 func (m *InMemoryStoreManager) FetchUnconfirmedTransactionAtNonceWithCount(_ context.Context, nonce uint64, fromAddress common.Address) (tx *types.Transaction, count int, err error) {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		tx, count = store.FetchUnconfirmedTransactionAtNonceWithCount(nonce)
 		return
 	}
 	return nil, 0, fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) FetchUnconfirmedTransactions(_ context.Context, fromAddress common.Address) ([]*types.Transaction, error) {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.FetchUnconfirmedTransactions()
 	}
 	return nil, fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) MarkConfirmedAndReorgedTransactions(_ context.Context, nonce uint64, fromAddress common.Address) (confirmedTxs []*types.Transaction, unconfirmedTxIDs []uint64, err error) {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		confirmedTxs, unconfirmedTxIDs, err = store.MarkConfirmedAndReorgedTransactions(nonce)
 		return
 	}
 	return nil, nil, fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) MarkUnconfirmedTransactionPurgeable(_ context.Context, nonce uint64, fromAddress common.Address) error {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.MarkUnconfirmedTransactionPurgeable(nonce)
 	}
 	return fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) UpdateTransactionBroadcast(_ context.Context, txID uint64, nonce uint64, attemptHash common.Hash, fromAddress common.Address) error {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.UpdateTransactionBroadcast(txID, nonce, attemptHash)
 	}
 	return fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) UpdateUnstartedTransactionWithNonce(_ context.Context, fromAddress common.Address, nonce uint64) (*types.Transaction, error) {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.UpdateUnstartedTransactionWithNonce(nonce)
 	}
 	return nil, fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) DeleteAttemptForUnconfirmedTx(_ context.Context, nonce uint64, attempt *types.Attempt, fromAddress common.Address) error {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.DeleteAttemptForUnconfirmedTx(nonce, attempt)
 	}
 	return fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) MarkTxFatal(_ context.Context, tx *types.Transaction, fromAddress common.Address) error {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.MarkTxFatal(tx)
 	}
 	return fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) UpdateSignedAttempt(_ context.Context, txID uint64, attemptID uint64, signedTransaction *evmtypes.Transaction, fromAddress common.Address) error {
-	if store, exists := m.InMemoryStoreMap[fromAddress]; exists {
+	if store, exists := m.GetStoreSafe(fromAddress); exists {
 		return store.UpdateSignedAttempt(txID, attemptID, signedTransaction)
 	}
 	return fmt.Errorf(StoreNotFoundForAddress, fromAddress)
 }
 
 func (m *InMemoryStoreManager) FindTxWithIdempotencyKey(_ context.Context, idempotencyKey string) (*types.Transaction, error) {
-	for _, store := range m.InMemoryStoreMap {
+	m.mu.RLock()
+	stores := make([]*InMemoryStore, 0, len(m.inMemoryStoreMap))
+	for _, store := range m.inMemoryStoreMap {
+		stores = append(stores, store)
+	}
+	m.mu.RUnlock()
+
+	for _, store := range stores {
 		tx := store.FindTxWithIdempotencyKey(idempotencyKey)
 		if tx != nil {
 			return tx, nil

pkg/txm/storage/inmemory_store_manager_test.go

@@ -19,17 +19,24 @@ func TestAdd(t *testing.T) {
 	// Adds a new address
 	err := m.Add(fromAddress)
 	require.NoError(t, err)
-	assert.Len(t, m.InMemoryStoreMap, 1)
+	existingStore, exists := m.GetStoreSafe(fromAddress)
+	require.True(t, exists)
 
 	// Fails if address exists
 	err = m.Add(fromAddress)
 	require.Error(t, err)
+	store, exists := m.GetStoreSafe(fromAddress)
+	require.True(t, exists)
+	assert.Same(t, existingStore, store)
 
 	// Adds multiple addresses
 	fromAddress1 := testutils.NewAddress()
 	fromAddress2 := testutils.NewAddress()
 	addresses := []common.Address{fromAddress1, fromAddress2}
 	err = m.Add(addresses...)
 	require.NoError(t, err)
-	assert.Len(t, m.InMemoryStoreMap, 3)
+	_, exists = m.GetStoreSafe(fromAddress1)
+	require.True(t, exists)
+	_, exists = m.GetStoreSafe(fromAddress2)
+	require.True(t, exists)
 }

pkg/txm/stuck_tx_detector.go

@@ -2,10 +2,7 @@ package txm
 
 import (
 	"context"
-	"encoding/json"
-	"fmt"
-	"io"
-	"net/http"
+	"sync"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -27,6 +24,7 @@ type stuckTxDetector struct {
 	lggr         logger.Logger
 	chainType    chaintype.ChainType
 	config       StuckTxDetectorConfig
+	mu           sync.Mutex
 	lastPurgeMap map[common.Address]time.Time
 }
 
@@ -60,6 +58,9 @@ func (s *stuckTxDetector) DetectStuckTransaction(ctx context.Context, tx *types.
 // so it is more likely to be picked up compared to a transaction that hasn't been broadcasted before. This would avoid slowing down TXM for sebsequent transactions
 // in case the current one is stuck.
 func (s *stuckTxDetector) timeBasedDetection(tx *types.Transaction) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	threshold := (s.config.BlockTime * time.Duration(s.config.StuckTxBlockThreshold))
 	if tx.LastBroadcastAt == nil {
 		if tx.AttemptCount >= maxAttemptsThreshold {
@@ -79,59 +80,3 @@ func (s *stuckTxDetector) timeBasedDetection(tx *types.Transaction) bool {
 	}
 	return false
 }
-
-type APIResponse struct {
-	Status string      `json:"status,omitempty"`
-	Hash   common.Hash `json:"hash,omitempty"`
-}
-
-const (
-	APIStatusPending   = "PENDING"
-	APIStatusIncluded  = "INCLUDED"
-	APIStatusFailed    = "FAILED"
-	APIStatusCancelled = "CANCELLED"
-	APIStatusUnknown   = "UNKNOWN"
-)
-
-// Deprecated: DualBroadcastDetection doesn't provide any significant benefits in terms of speed and time
-// based detection can replace it.
-func (s *stuckTxDetector) DualBroadcastDetection(ctx context.Context, tx *types.Transaction) (bool, error) {
-	for _, attempt := range tx.Attempts {
-		req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.config.DetectionURL+attempt.Hash.String(), nil)
-		if err != nil {
-			return false, fmt.Errorf("failed to make request for txID: %v, attemptHash: %v - %w", tx.ID, attempt.Hash, err)
-		}
-		resp, err := http.DefaultClient.Do(req)
-		if err != nil {
-			return false, fmt.Errorf("failed to get transaction status for txID: %v, attemptHash: %v - %w", tx.ID, attempt.Hash, err)
-		}
-		if resp.StatusCode != http.StatusOK {
-			resp.Body.Close()
-			return false, fmt.Errorf("request %v failed with status: %d", req, resp.StatusCode)
-		}
-		body, err := io.ReadAll(resp.Body)
-		resp.Body.Close()
-		if err != nil {
-			return false, err
-		}
-
-		var apiResponse APIResponse
-		err = json.Unmarshal(body, &apiResponse)
-		if err != nil {
-			return false, fmt.Errorf("failed to unmarshal response for txID: %v, attemptHash: %v - %w: %s", tx.ID, attempt.Hash, err, string(body))
-		}
-		switch apiResponse.Status {
-		case APIStatusPending, APIStatusIncluded:
-			return false, nil
-		case APIStatusFailed, APIStatusCancelled:
-			s.lggr.Debugf("TxID: %v with attempHash: %v was marked as failed/cancelled by the RPC. Transaction is now considered stuck and will be purged.",
-				tx.ID, attempt.Hash)
-			return true, nil
-		case APIStatusUnknown:
-			continue
-		default:
-			continue
-		}
-	}
-	return false, nil
-}

pkg/txm/stuck_tx_detector_test.go

@@ -1,6 +1,7 @@
 package txm
 
 import (
+	"sync"
 	"testing"
 	"time"
 
@@ -95,4 +96,26 @@ func TestTimeBasedDetection(t *testing.T) {
 		assert.True(t, s.timeBasedDetection(tx1))
 		assert.False(t, s.timeBasedDetection(tx2))
 	})
+
+	t.Run("safe to call concurrently for different addresses", func(t *testing.T) {
+		config := StuckTxDetectorConfig{
+			BlockTime:             1 * time.Second,
+			StuckTxBlockThreshold: 10,
+		}
+		s := NewStuckTxDetector(logger.Test(t), "", config)
+		lastBroadcastAt := time.Time{}
+
+		var wg sync.WaitGroup
+		for range 10 {
+			wg.Go(func() {
+				tx := &types.Transaction{
+					ID:              1,
+					LastBroadcastAt: &lastBroadcastAt,
+					FromAddress:     testutils.NewAddress(),
+				}
+				s.timeBasedDetection(tx)
+			})
+		}
+		wg.Wait()
+	})
 }

pkg/txm/txm_test.go

@@ -455,8 +455,8 @@ func TestFlow_ResendTransaction(t *testing.T) {
 	require.NoError(t, tm.BackfillTransactions(t.Context(), address))
 
 	// Set LastBroadcastAt to a time in the past to trigger retry condition
-	txStore := txStoreManager.InMemoryStoreMap[address]
-	require.NotNil(t, txStore)
+	txStore, exists := txStoreManager.GetStoreSafe(address)
+	require.True(t, exists)
 	tx := txStore.UnconfirmedTransactions[initialNonce]
 	require.NotNil(t, tx)
 	pastTime := time.Now().Add(-(config.BlockTime*time.Duration(config.RetryBlockThreshold) + 1*time.Second))