Fix: update gas coin reservation (#458)

* move coin reservation to preparePTBTransaction in TXM and make it atomic over all coins

* update gas coin manager and associated tests

Author: faisal-chainlink

relayer/txm/confirmer.go

@@ -104,7 +104,7 @@ func handleSuccess(txm *SuiTxm, tx SuiTx) error {
 	}
 	txm.lggr.Infow("Transaction finalized", "transactionID", tx.TransactionID)
 
-	if err := txm.coinManager.ReleaseCoins(tx.TransactionID); err != nil {
+	if err := tx.CoinManager.ReleaseCoins(tx.TransactionID); err != nil {
 		// This error is not critical, can be safely ignored as the coins will auto-release after the default TTL
 		txm.lggr.Debugw("Failed to release coins", "transactionID", tx.TransactionID, "error", err)
 	}
@@ -117,8 +117,8 @@ func handleTransactionError(ctx context.Context, txm *SuiTxm, tx SuiTx, result *
 
 	txError := suierrors.ParseSuiErrorMessage(result.Error)
 
-	// Check if the error is a locked object error, mark the coin as reserved if it is not already
-	// to avoid other transactions from using it
+	// Check if the error is a locked object error, and exclude the coin so other
+	// transactions (and this transaction's own coin-refresh retry) do not re-select it.
 	if objectID, version, ok := suierrors.ExtractLockedObjectRef(result.Error); ok {
 		txm.lggr.Infow("Detected locked coin at confirmation time",
 			"txID", tx.TransactionID,
@@ -127,28 +127,20 @@ func handleTransactionError(ctx context.Context, txm *SuiTxm, tx SuiTx, result *
 		)
 
 		coinID, err := transaction.ConvertSuiAddressStringToBytes(models.SuiAddress(objectID))
-		if err == nil && !txm.coinManager.IsCoinReserved(*coinID) {
-			// Coin lock duration
-			expiry := DefaultLockedCoinTTL
-
-			// The coin is not recorded is not marked as reserved, mark it as reserved
-			err = txm.coinManager.TryReserveCoins(ctx, tx.TransactionID, []transaction.SuiObjectRef{
-				{
-					ObjectId: *coinID,
-					Version:  0,
-					Digest:   nil,
-				},
-			}, &expiry)
-
-			if err != nil {
-				// This is not a critical error, so we continue
-				txm.lggr.Debugw(
-					"Failed to mark locked coin as reserved",
-					"transactionID", tx.TransactionID,
-					"objectID", objectID,
-					"error", err,
-				)
-			}
+		if err != nil {
+			txm.lggr.Debugw(
+				"Failed to convert locked coin object ID",
+				"transactionID", tx.TransactionID,
+				"objectID", objectID,
+				"error", err,
+			)
+		} else {
+			// Use a standalone, long-lived exclusion keyed by the coin itself (not by
+			// txID). The upcoming coin-refresh retry calls ReleaseCoins(txID), which
+			// would otherwise wipe a txID-scoped reservation for this coin; the
+			// standalone exclusion survives that release and keeps the locked coin out
+			// of re-selection until its TTL expires.
+			tx.CoinManager.ReserveLockedCoin(*coinID, DefaultLockedCoinTTL)
 		}
 	}
 
@@ -188,6 +180,16 @@ func handleGasBumpRetry(ctx context.Context, txm *SuiTxm, tx SuiTx, txError *sui
 		return err
 	}
 
+	// Release the current reservation before the transaction is rebuilt. UpdateTransactionGas
+	// re-selects and re-reserves coins via preparePTBTransaction; without releasing first, the
+	// rebuild would exclude this transaction's own still-reserved coins and either pick a
+	// different set (orphaning the old reservation until its TTL) or fail if no other coins are
+	// free. Releasing lets the rebuild re-select the same coins with the bumped budget.
+	if err := tx.CoinManager.ReleaseCoins(tx.TransactionID); err != nil {
+		// Not critical - coins auto-release after TTL.
+		txm.lggr.Debugw("Failed to release coins before gas bump rebuild", "transactionID", tx.TransactionID, "error", err)
+	}
+
 	if err := txm.transactionRepository.UpdateTransactionGas(ctx, txm.keystoreService, txm.suiGateway, tx.TransactionID, &updatedGas); err != nil {
 		txm.lggr.Errorw("Failed to update transaction gas", "transactionID", tx.TransactionID, "error", err)
 		return err
@@ -206,7 +208,7 @@ func handleCoinRefreshRetry(ctx context.Context, txm *SuiTxm, tx SuiTx, txError
 	txm.lggr.Infow("Coin refresh strategy - refreshing coins for locked coin error", "transactionID", tx.TransactionID)
 
 	// Release the old coins that are locked
-	if err := txm.coinManager.ReleaseCoins(tx.TransactionID); err != nil {
+	if err := tx.CoinManager.ReleaseCoins(tx.TransactionID); err != nil {
 		// This is not critical - coins will auto-release after TTL
 		txm.lggr.Debugw("Failed to release old coins", "transactionID", tx.TransactionID, "error", err)
 	}
@@ -275,7 +277,7 @@ func markTransactionFailed(txm *SuiTxm, tx SuiTx, txError *suierrors.SuiError) e
 
 	txm.lggr.Infow("Transaction failed", "transactionID", tx.TransactionID)
 
-	if err := txm.coinManager.ReleaseCoins(tx.TransactionID); err != nil {
+	if err := tx.CoinManager.ReleaseCoins(tx.TransactionID); err != nil {
 		// This error is not critical, can be safely ignored as the coins will auto-release after the default TTL
 		txm.lggr.Debugw("Failed to release coins", "transactionID", tx.TransactionID, "error", err)
 	}

relayer/txm/gas_coin_manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/hex"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/block-vision/sui-go-sdk/models"
@@ -21,15 +22,25 @@ const (
 
 type GasCoinManager interface {
 	TryReserveCoins(ctx context.Context, txID string, paymentCoins []transaction.SuiObjectRef, expiry *time.Duration) error
+	ReserveLockedCoin(coinID models.SuiAddressBytes, ttl time.Duration)
 	ReleaseCoins(txID string) error
 	IsCoinReserved(coinID models.SuiAddressBytes) bool
 }
 
+// lockedCoinKeyPrefix namespaces the cache entries used to exclude coins that the
+// chain reported as locked. These entries are deliberately kept separate from the
+// txID-associated reservations so that ReleaseCoins (which only unwinds a txID's
+// reservation) can never remove a locked-coin exclusion before its TTL expires.
+const lockedCoinKeyPrefix = "locked:"
+
 // SuiGasCoinManager is the concrete implementation of GasCoinManager.
 type SuiGasCoinManager struct {
 	lggr       logger.Logger
 	client     client.SuiPTBClient
 	coinsCache *cache.Cache
+	// mu guards the whole check-then-reserve sequence so that reservations are
+	// atomic and all-or-nothing across the full coin set.
+	mu sync.Mutex
 }
 
 // NewGasCoinManager creates a new SuiGasCoinManager.
@@ -48,29 +59,62 @@ func (m *SuiGasCoinManager) TryReserveCoins(
 	coinIDs []transaction.SuiObjectRef,
 	expiry *time.Duration,
 ) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	// First pass: verify the entire set is available before writing anything.
+	// Holding the lock across the check and the writes closes the check-then-set
+	// race between concurrent callers, and validating upfront makes the
+	// reservation all-or-nothing (no partial reservations leaked on failure).
 	for _, coin := range coinIDs {
-		if m.IsCoinReserved(coin.ObjectId) {
+		if m.isCoinReserved(coin.ObjectId) {
 			return fmt.Errorf("coin %s is already reserved", hex.EncodeToString(coin.ObjectId[:]))
 		}
+	}
 
-		coinID := hex.EncodeToString(coin.ObjectId[:])
-		expiresAt := DefaultAllocationTimeout
-
-		if expiry != nil {
-			expiresAt = *expiry
-		}
+	expiresAt := DefaultAllocationTimeout
+	if expiry != nil {
+		expiresAt = *expiry
+	}
 
+	// Second pass: reserve the full set now that every coin is known to be free.
+	for _, coin := range coinIDs {
+		coinID := hex.EncodeToString(coin.ObjectId[:])
 		m.coinsCache.Set(coinID, true, expiresAt)
 	}
 
-	m.coinsCache.Set(txID, coinIDs, DefaultAllocationTimeout)
+	// Track the reserved set under the txID using the same TTL so ReleaseCoins can
+	// always unwind the individual coin entries for the reservation's lifetime.
+	m.coinsCache.Set(txID, coinIDs, expiresAt)
 
 	return nil
 }
 
+// ReserveLockedCoin marks a single coin object as excluded for the given TTL,
+// independently of any transaction ID. Unlike TryReserveCoins, the exclusion is
+// not associated with a txID and therefore is never removed by ReleaseCoins; it
+// persists until the TTL expires. This is used to exclude a coin that the chain
+// reported as locked so it is not re-selected while the on-chain lock persists,
+// even if that same coin was previously reserved as a payment coin (whose
+// reservation is released during the coin-refresh retry).
+func (m *SuiGasCoinManager) ReserveLockedCoin(coinID models.SuiAddressBytes, ttl time.Duration) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.coinsCache.Set(lockedCoinCacheKey(coinID), true, ttl)
+}
+
+// lockedCoinCacheKey returns the namespaced cache key for a locked-coin exclusion.
+func lockedCoinCacheKey(coinID models.SuiAddressBytes) string {
+	return lockedCoinKeyPrefix + hex.EncodeToString(coinID[:])
+}
+
 // ReleaseCoins only releases reservations stored under a txID key (txID -> []SuiObjectRef).
 // It does not work with coinID keys (coinID -> bool) and cannot unlock those entries directly.
 func (m *SuiGasCoinManager) ReleaseCoins(txID string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	coinIDs, ok := m.coinsCache.Get(txID)
 	if !ok {
 		return fmt.Errorf("no coins reserved for transaction %s", txID)
@@ -86,7 +130,27 @@ func (m *SuiGasCoinManager) ReleaseCoins(txID string) error {
 }
 
 func (m *SuiGasCoinManager) IsCoinReserved(coinID models.SuiAddressBytes) bool {
-	coinIDStr := hex.EncodeToString(coinID[:])
-	isReserved, found := m.coinsCache.Get(coinIDStr)
-	return found && isReserved.(bool)
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.isCoinReserved(coinID)
+}
+
+// isCoinReserved is the lock-free implementation of IsCoinReserved. Callers must
+// already hold m.mu. A coin is considered reserved if it has either a txID-scoped
+// reservation (TryReserveCoins) or a standalone locked-coin exclusion
+// (ReserveLockedCoin).
+func (m *SuiGasCoinManager) isCoinReserved(coinID models.SuiAddressBytes) bool {
+	if isReserved, found := m.coinsCache.Get(hex.EncodeToString(coinID[:])); found {
+		if reserved, ok := isReserved.(bool); ok && reserved {
+			return true
+		}
+	}
+
+	if locked, found := m.coinsCache.Get(lockedCoinCacheKey(coinID)); found {
+		if isLocked, ok := locked.(bool); ok && isLocked {
+			return true
+		}
+	}
+
+	return false
 }

relayer/txm/gas_coin_manager_test.go

@@ -5,6 +5,8 @@ package txm_test
 import (
 	"context"
 	"fmt"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -18,6 +20,14 @@ import (
 	"github.com/smartcontractkit/chainlink-sui/relayer/txm"
 )
 
+// mustCoinRef builds a SuiObjectRef from a 0x-prefixed 32-byte address string.
+func mustCoinRef(t *testing.T, addr string) transaction.SuiObjectRef {
+	t.Helper()
+	b, err := transaction.ConvertSuiAddressStringToBytes(models.SuiAddress(addr))
+	require.NoError(t, err)
+	return transaction.SuiObjectRef{ObjectId: *b, Version: 1, Digest: nil}
+}
+
 func TestSuiGasCoinManager_TryReserveCoins(t *testing.T) {
 	lggr := logger.Test(t)
 	mockClient := &testutils.FakeSuiPTBClient{}
@@ -137,3 +147,137 @@ func TestSuiGasCoinManager_TryReserveCoins(t *testing.T) {
 		}, 45*time.Second, 10*time.Second)
 	})
 }
+
+// TestSuiGasCoinManager_AtomicReservation exercises the atomic, all-or-nothing
+// behaviour of TryReserveCoins that guards against two transactions reserving the
+// same gas coin.
+func TestSuiGasCoinManager_AtomicReservation(t *testing.T) {
+	lggr := logger.Test(t)
+	mockClient := &testutils.FakeSuiPTBClient{}
+	ctx := context.Background()
+
+	t.Run("concurrent reservations of the same coin: exactly one wins", func(t *testing.T) {
+		gcm := txm.NewGasCoinManager(lggr, mockClient)
+		coin := mustCoinRef(t, fmt.Sprintf("0x%064x", 1))
+
+		const numGoroutines = 50
+		var wg sync.WaitGroup
+		var successes atomic.Int32
+		start := make(chan struct{})
+
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func(i int) {
+				defer wg.Done()
+				<-start // release all goroutines at once to maximize contention
+				txID := fmt.Sprintf("tx-%d", i)
+				if err := gcm.TryReserveCoins(ctx, txID, []transaction.SuiObjectRef{coin}, nil); err == nil {
+					successes.Add(1)
+				}
+			}(i)
+		}
+
+		close(start)
+		wg.Wait()
+
+		assert.Equal(t, int32(1), successes.Load(), "exactly one caller should reserve the coin")
+		assert.True(t, gcm.IsCoinReserved(coin.ObjectId))
+	})
+
+	t.Run("reservation is all-or-nothing on partial overlap", func(t *testing.T) {
+		gcm := txm.NewGasCoinManager(lggr, mockClient)
+		coinA := mustCoinRef(t, fmt.Sprintf("0x%064x", 0xAA))
+		coinB := mustCoinRef(t, fmt.Sprintf("0x%064x", 0xBB))
+
+		// Reserve coinA under tx1.
+		require.NoError(t, gcm.TryReserveCoins(ctx, "tx1", []transaction.SuiObjectRef{coinA}, nil))
+
+		// tx2 requests [coinB, coinA]; coinA is already taken, so the whole call must
+		// fail without leaving coinB reserved.
+		err := gcm.TryReserveCoins(ctx, "tx2", []transaction.SuiObjectRef{coinB, coinA}, nil)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "is already reserved")
+		assert.False(t, gcm.IsCoinReserved(coinB.ObjectId), "coinB must not be leaked when the reservation fails")
+
+		// The failed reservation left no txID mapping to release.
+		assert.Error(t, gcm.ReleaseCoins("tx2"))
+
+		// coinB is still free and can be claimed by another transaction.
+		assert.NoError(t, gcm.TryReserveCoins(ctx, "tx3", []transaction.SuiObjectRef{coinB}, nil))
+	})
+
+	t.Run("concurrent disjoint reservations all succeed", func(t *testing.T) {
+		gcm := txm.NewGasCoinManager(lggr, mockClient)
+
+		const numGoroutines = 50
+		// Pre-build coin refs so the goroutines don't call require/t inside them.
+		coins := make([]transaction.SuiObjectRef, numGoroutines)
+		for i := range coins {
+			coins[i] = mustCoinRef(t, fmt.Sprintf("0x%064x", 0x1000+i))
+		}
+
+		var wg sync.WaitGroup
+		var successes atomic.Int32
+		start := make(chan struct{})
+
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func(i int) {
+				defer wg.Done()
+				<-start
+				if err := gcm.TryReserveCoins(ctx, fmt.Sprintf("tx-%d", i), []transaction.SuiObjectRef{coins[i]}, nil); err == nil {
+					successes.Add(1)
+				}
+			}(i)
+		}
+
+		close(start)
+		wg.Wait()
+
+		assert.Equal(t, int32(numGoroutines), successes.Load(), "all disjoint reservations should succeed")
+	})
+}
+
+// TestSuiGasCoinManager_ReserveLockedCoin verifies that a locked-coin exclusion is
+// independent of any txID reservation and survives ReleaseCoins, which is what keeps
+// a chain-locked coin out of re-selection during the coin-refresh retry.
+func TestSuiGasCoinManager_ReserveLockedCoin(t *testing.T) {
+	lggr := logger.Test(t)
+	mockClient := &testutils.FakeSuiPTBClient{}
+	ctx := context.Background()
+
+	t.Run("exclusion survives ReleaseCoins of the payment reservation", func(t *testing.T) {
+		gcm := txm.NewGasCoinManager(lggr, mockClient)
+		coin := mustCoinRef(t, fmt.Sprintf("0x%064x", 0xCC))
+
+		// The coin was reserved as a normal payment coin under the transaction's ID.
+		require.NoError(t, gcm.TryReserveCoins(ctx, "tx-locked", []transaction.SuiObjectRef{coin}, nil))
+		assert.True(t, gcm.IsCoinReserved(coin.ObjectId))
+
+		// The chain reports it as locked; add a standalone exclusion (as handleTransactionError does).
+		gcm.ReserveLockedCoin(coin.ObjectId, txm.DefaultLockedCoinTTL)
+
+		// The coin-refresh retry releases the payment reservation for this txID.
+		require.NoError(t, gcm.ReleaseCoins("tx-locked"))
+
+		// The standalone exclusion must survive, so the coin is still excluded and
+		// cannot be re-selected or re-reserved by any transaction.
+		assert.True(t, gcm.IsCoinReserved(coin.ObjectId), "locked coin exclusion must survive ReleaseCoins")
+		err := gcm.TryReserveCoins(ctx, "tx-other", []transaction.SuiObjectRef{coin}, nil)
+		assert.Error(t, err, "a locked coin must not be reservable by another transaction")
+	})
+
+	t.Run("exclusion is created even when the coin was never a payment reservation", func(t *testing.T) {
+		gcm := txm.NewGasCoinManager(lggr, mockClient)
+		coin := mustCoinRef(t, fmt.Sprintf("0x%064x", 0xDD))
+
+		assert.False(t, gcm.IsCoinReserved(coin.ObjectId))
+		gcm.ReserveLockedCoin(coin.ObjectId, txm.DefaultLockedCoinTTL)
+		assert.True(t, gcm.IsCoinReserved(coin.ObjectId))
+
+		// There is no txID mapping for a locked-only exclusion, so ReleaseCoins keyed
+		// by the coin's own hex finds nothing and the exclusion remains.
+		assert.Error(t, gcm.ReleaseCoins(fmt.Sprintf("0x%064x", 0xDD)))
+		assert.True(t, gcm.IsCoinReserved(coin.ObjectId))
+	})
+}