package arbitrum import ( "context" "fmt" "time" "github.com/ethereum/go-ethereum/ethclient" "github.com/fraktal/mev-beta/internal/config" "github.com/fraktal/mev-beta/internal/logger" ) // RoundRobinClient wraps a client and tracks round-robin usage type RoundRobinClient struct { manager *RPCManager ctx context.Context logger *logger.Logger lastIdx int readCalls int64 writeCalls int64 } // NewRoundRobinClient creates a new round-robin client wrapper func NewRoundRobinClient(manager *RPCManager, ctx context.Context, logger *logger.Logger) *RoundRobinClient { return &RoundRobinClient{ manager: manager, ctx: ctx, logger: logger, lastIdx: -1, } } // GetClientForRead returns the next RPC client for a read operation using round-robin func (rr *RoundRobinClient) GetClientForRead() (*ethclient.Client, error) { if rr.manager == nil { return nil, fmt.Errorf("RPC manager not initialized") } client, idx, err := rr.manager.GetNextClient(rr.ctx) if err != nil { return nil, err } rr.lastIdx = idx rr.readCalls++ if client == nil || client.Client == nil { return nil, fmt.Errorf("client at index %d is nil", idx) } return client.Client, nil } // RecordReadSuccess records a successful read operation func (rr *RoundRobinClient) RecordReadSuccess(responseTime time.Duration) { if rr.lastIdx >= 0 && rr.manager != nil { rr.manager.RecordSuccess(rr.lastIdx, responseTime) } } // RecordReadFailure records a failed read operation func (rr *RoundRobinClient) RecordReadFailure() { if rr.lastIdx >= 0 && rr.manager != nil { rr.manager.RecordFailure(rr.lastIdx) } } // GetClientForWrite returns a read-optimized client for write operations // Uses least-failure strategy to prefer stable endpoints func (rr *RoundRobinClient) GetClientForWrite() (*ethclient.Client, error) { if rr.manager == nil { return nil, fmt.Errorf("RPC manager not initialized") } // Temporarily switch to least-failures policy for writes currentPolicy := rr.manager.rotationPolicy rr.manager.SetRotationPolicy(LeastFailures) defer rr.manager.SetRotationPolicy(currentPolicy) client, idx, err := rr.manager.GetNextClient(rr.ctx) if err != nil { return nil, err } rr.lastIdx = idx rr.writeCalls++ if client == nil || client.Client == nil { return nil, fmt.Errorf("client at index %d is nil", idx) } return client.Client, nil } // RecordWriteSuccess records a successful write operation func (rr *RoundRobinClient) RecordWriteSuccess(responseTime time.Duration) { if rr.lastIdx >= 0 && rr.manager != nil { rr.manager.RecordSuccess(rr.lastIdx, responseTime) } } // RecordWriteFailure records a failed write operation func (rr *RoundRobinClient) RecordWriteFailure() { if rr.lastIdx >= 0 && rr.manager != nil { rr.manager.RecordFailure(rr.lastIdx) } } // GetLoadBalancingStats returns statistics about load distribution func (rr *RoundRobinClient) GetLoadBalancingStats() map[string]interface{} { if rr.manager == nil { return map[string]interface{}{ "error": "RPC manager not initialized", } } return map[string]interface{}{ "total_reads": rr.readCalls, "total_writes": rr.writeCalls, "rpc_stats": rr.manager.GetStats(), } } // InitializeRPCRoundRobin sets up round-robin RPC management for a connection manager func InitializeRPCRoundRobin(cm *ConnectionManager, endpoints []string) error { if cm == nil { return fmt.Errorf("connection manager is nil") } if len(endpoints) == 0 { cm.logger.Warn("⚠️ No additional endpoints provided for round-robin initialization") return nil } // Connect to each endpoint and add to RPC manager connectedCount := 0 for _, endpoint := range endpoints { client, err := cm.connectWithTimeout(context.Background(), endpoint) if err != nil { cm.logger.Warn(fmt.Sprintf("Failed to connect to endpoint %s: %v", endpoint, err)) continue } if err := cm.rpcManager.AddEndpoint(client, endpoint); err != nil { cm.logger.Warn(fmt.Sprintf("Failed to add endpoint to RPC manager: %v", err)) client.Client.Close() continue } connectedCount++ } cm.logger.Info(fmt.Sprintf("✅ Initialized round-robin with %d/%d endpoints", connectedCount, len(endpoints))) // Perform initial health check healthCheckCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() if err := cm.rpcManager.HealthCheckAll(healthCheckCtx); err != nil { cm.logger.Warn(fmt.Sprintf("⚠️ Health check failed: %v", err)) } return nil } // ConfigureRPCLoadBalancing configures load balancing strategy for RPC endpoints func ConfigureRPCLoadBalancing(cm *ConnectionManager, strategy RotationPolicy) error { if cm == nil { return fmt.Errorf("connection manager is nil") } cm.SetRPCRotationPolicy(strategy) cm.logger.Info(fmt.Sprintf("📊 RPC load balancing configured with strategy: %s", strategy)) return nil } // GetConnectionManagerWithRoundRobin creates a connection manager with round-robin already set up func GetConnectionManagerWithRoundRobin(cfg *config.ArbitrumConfig, logger *logger.Logger, endpoints []string) (*ConnectionManager, error) { if cfg == nil { return nil, fmt.Errorf("config must not be nil") } cm := NewConnectionManager(cfg, logger) // Initialize with provided endpoints if len(endpoints) > 0 { if err := InitializeRPCRoundRobin(cm, endpoints); err != nil { logger.Warn(fmt.Sprintf("Failed to initialize round-robin: %v", err)) } } // Set health-aware strategy by default cm.SetRPCRotationPolicy(HealthAware) return cm, nil }