mev-beta/pkg/arbitrage/nonce_manager.go

package arbitrage

import (
	"context"
	"fmt"
	"sync"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/ethclient"
)

// NonceManager provides thread-safe nonce management for transaction submission
// Prevents nonce collisions when submitting multiple transactions rapidly
type NonceManager struct {
	mu sync.Mutex

	client  *ethclient.Client
	account common.Address

	// Track the last nonce we've assigned
	lastNonce uint64

	// Track pending nonces to avoid reuse
	pending map[uint64]bool

	// Initialized flag
	initialized bool
}

// NewNonceManager creates a new nonce manager for the given account
func NewNonceManager(client *ethclient.Client, account common.Address) *NonceManager {
	return &NonceManager{
		client:      client,
		account:     account,
		pending:     make(map[uint64]bool),
		initialized: false,
	}
}

// GetNextNonce returns the next available nonce for transaction submission
// This method is thread-safe and prevents nonce collisions
func (nm *NonceManager) GetNextNonce(ctx context.Context) (uint64, error) {
	nm.mu.Lock()
	defer nm.mu.Unlock()

	// Get current pending nonce from network
	currentNonce, err := nm.client.PendingNonceAt(ctx, nm.account)
	if err != nil {
		return 0, fmt.Errorf("failed to get pending nonce: %w", err)
	}

	// First time initialization
	if !nm.initialized {
		// On first call, hand back the network's pending nonce so we don't
		// skip a slot and create gaps that block execution.
		nm.lastNonce = currentNonce
		nm.initialized = true
		nm.pending[currentNonce] = true
		return currentNonce, nil
	}

	// Determine next nonce to use
	var nextNonce uint64

	// If network nonce is higher than our last assigned, use network nonce
	// This handles cases where transactions confirmed between calls
	if currentNonce > nm.lastNonce {
		nextNonce = currentNonce
		nm.lastNonce = currentNonce
		// Clear pending nonces below current (they've been mined)
		nm.clearPendingBefore(currentNonce)
	} else {
		// Otherwise increment our last nonce
		nextNonce = nm.lastNonce + 1
		nm.lastNonce = nextNonce
	}

	// Mark this nonce as pending
	nm.pending[nextNonce] = true

	return nextNonce, nil
}

// MarkConfirmed marks a nonce as confirmed (mined in a block)
// This allows the nonce manager to clean up its pending tracking
func (nm *NonceManager) MarkConfirmed(nonce uint64) {
	nm.mu.Lock()
	defer nm.mu.Unlock()

	delete(nm.pending, nonce)
}

// MarkFailed marks a nonce as failed (transaction rejected)
// This allows the nonce to be potentially reused
func (nm *NonceManager) MarkFailed(nonce uint64) {
	nm.mu.Lock()
	defer nm.mu.Unlock()

	delete(nm.pending, nonce)

	// Reset lastNonce if this was the last one we assigned
	if nonce == nm.lastNonce && nonce > 0 {
		nm.lastNonce = nonce - 1
	}
}

// GetPendingCount returns the number of pending nonces
func (nm *NonceManager) GetPendingCount() int {
	nm.mu.Lock()
	defer nm.mu.Unlock()

	return len(nm.pending)
}

// Reset resets the nonce manager state
// Should be called if you want to re-sync with network state
func (nm *NonceManager) Reset() {
	nm.mu.Lock()
	defer nm.mu.Unlock()

	nm.pending = make(map[uint64]bool)
	nm.initialized = false
	nm.lastNonce = 0
}

// clearPendingBefore removes pending nonces below the given threshold
// (internal method, mutex must be held by caller)
func (nm *NonceManager) clearPendingBefore(threshold uint64) {
	for nonce := range nm.pending {
		if nonce < threshold {
			delete(nm.pending, nonce)
		}
	}
}

// GetCurrentNonce returns the last assigned nonce without incrementing
func (nm *NonceManager) GetCurrentNonce() uint64 {
	nm.mu.Lock()
	defer nm.mu.Unlock()

	return nm.lastNonce
}

// IsPending checks if a nonce is currently pending
func (nm *NonceManager) IsPending(nonce uint64) bool {
	nm.mu.Lock()
	defer nm.mu.Unlock()

	return nm.pending[nonce]
}