mev-beta/pkg/arbitrum/connection.go

package arbitrum

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/ethereum/go-ethereum/ethclient"
	"github.com/fraktal/mev-beta/internal/config"
)

// ConnectionManager manages Arbitrum RPC connections with fallback support
type ConnectionManager struct {
	config             *config.ArbitrumConfig
	primaryClient      *ethclient.Client
	fallbackClients    []*ethclient.Client
	currentClientIndex int
}

// NewConnectionManager creates a new connection manager
func NewConnectionManager(cfg *config.ArbitrumConfig) *ConnectionManager {
	return &ConnectionManager{
		config: cfg,
	}
}

// GetClient returns a connected Ethereum client with automatic fallback
func (cm *ConnectionManager) GetClient(ctx context.Context) (*ethclient.Client, error) {
	// Try primary endpoint first
	if cm.primaryClient == nil {
		primaryEndpoint := cm.getPrimaryEndpoint()
		client, err := cm.connectWithTimeout(ctx, primaryEndpoint)
		if err == nil {
			cm.primaryClient = client
			return client, nil
		}
	} else {
		// Test if primary client is still connected
		if cm.testConnection(ctx, cm.primaryClient) {
			return cm.primaryClient, nil
		}
		// Primary client failed, close it
		cm.primaryClient.Close()
		cm.primaryClient = nil
	}

	// Try fallback endpoints
	fallbackEndpoints := cm.getFallbackEndpoints()
	for i, endpoint := range fallbackEndpoints {
		client, err := cm.connectWithTimeout(ctx, endpoint)
		if err == nil {
			// Store successful fallback client
			if i < len(cm.fallbackClients) {
				if cm.fallbackClients[i] != nil {
					cm.fallbackClients[i].Close()
				}
				cm.fallbackClients[i] = client
			} else {
				cm.fallbackClients = append(cm.fallbackClients, client)
			}
			cm.currentClientIndex = i
			return client, nil
		}
	}

	return nil, fmt.Errorf("all RPC endpoints failed to connect")
}

// getPrimaryEndpoint returns the primary RPC endpoint
func (cm *ConnectionManager) getPrimaryEndpoint() string {
	// Check environment variable first
	if endpoint := os.Getenv("ARBITRUM_RPC_ENDPOINT"); endpoint != "" {
		return endpoint
	}

	// Use config value
	if cm.config != nil && cm.config.RPCEndpoint != "" {
		return cm.config.RPCEndpoint
	}

	// Default fallback
	return "wss://arbitrum-mainnet.core.chainstack.com/f69d14406bc00700da9b936504e1a870"
}

// getFallbackEndpoints returns fallback RPC endpoints
func (cm *ConnectionManager) getFallbackEndpoints() []string {
	var endpoints []string

	// Check environment variable first
	if envEndpoints := os.Getenv("ARBITRUM_FALLBACK_ENDPOINTS"); envEndpoints != "" {
		for _, endpoint := range strings.Split(envEndpoints, ",") {
			if endpoint = strings.TrimSpace(endpoint); endpoint != "" {
				endpoints = append(endpoints, endpoint)
			}
		}
	}

	// Add configured fallback endpoints
	if cm.config != nil {
		for _, endpoint := range cm.config.FallbackEndpoints {
			if endpoint.URL != "" {
				endpoints = append(endpoints, endpoint.URL)
			}
		}
	}

	// Default fallbacks if none configured
	if len(endpoints) == 0 {
		endpoints = []string{
			"https://arb1.arbitrum.io/rpc",
			"https://arbitrum.llamarpc.com",
			"https://arbitrum-one.publicnode.com",
			"https://arbitrum-one.public.blastapi.io",
		}
	}

	return endpoints
}

// connectWithTimeout attempts to connect to an RPC endpoint with timeout
func (cm *ConnectionManager) connectWithTimeout(ctx context.Context, endpoint string) (*ethclient.Client, error) {
	// Create timeout context
	connectCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	// Create client
	client, err := ethclient.DialContext(connectCtx, endpoint)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to %s: %w", endpoint, err)
	}

	// Test connection with a simple call
	if err := cm.testConnection(connectCtx, client); err != nil {
		client.Close()
		return nil, fmt.Errorf("connection test failed for %s: %w", endpoint, err)
	}

	return client, nil
}

// testConnection tests if a client connection is working
func (cm *ConnectionManager) testConnection(ctx context.Context, client *ethclient.Client) error {
	testCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	// Try to get chain ID as a simple connection test
	_, err := client.ChainID(testCtx)
	return err
}

// Close closes all client connections
func (cm *ConnectionManager) Close() {
	if cm.primaryClient != nil {
		cm.primaryClient.Close()
		cm.primaryClient = nil
	}

	for _, client := range cm.fallbackClients {
		if client != nil {
			client.Close()
		}
	}
	cm.fallbackClients = nil
}

// GetClientWithRetry returns a client with automatic retry on failure
func (cm *ConnectionManager) GetClientWithRetry(ctx context.Context, maxRetries int) (*ethclient.Client, error) {
	var lastErr error

	for attempt := 0; attempt < maxRetries; attempt++ {
		client, err := cm.GetClient(ctx)
		if err == nil {
			return client, nil
		}

		lastErr = err

		// Wait before retry (exponential backoff)
		if attempt < maxRetries-1 {
			waitTime := time.Duration(1<<uint(attempt)) * time.Second
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(waitTime):
				// Continue to next attempt
			}
		}
	}

	return nil, fmt.Errorf("failed to connect after %d attempts: %w", maxRetries, lastErr)
}

// GetHealthyClient returns a client that passes health checks
func GetHealthyClient(ctx context.Context) (*ethclient.Client, error) {
	cfg := &config.ArbitrumConfig{} // Use default config
	cm := NewConnectionManager(cfg)
	defer cm.Close()

	return cm.GetClientWithRetry(ctx, 3)
}