mev-beta/internal/ratelimit/manager.go

package ratelimit

import (
	"context"
	"fmt"
	"sync"

	"golang.org/x/time/rate"

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

// LimiterManager manages rate limiters for multiple endpoints
type LimiterManager struct {
	limiters map[string]*EndpointLimiter
	mu       sync.RWMutex
}

// EndpointLimiter represents a rate limiter for a specific endpoint
type EndpointLimiter struct {
	URL     string
	Limiter *rate.Limiter
	Config  config.RateLimitConfig
}

// NewLimiterManager creates a new LimiterManager
func NewLimiterManager(cfg *config.ArbitrumConfig) *LimiterManager {
	lm := &LimiterManager{
		limiters: make(map[string]*EndpointLimiter),
	}

	// Create limiter for primary endpoint
	limiter := createLimiter(cfg.RateLimit)
	lm.limiters[cfg.RPCEndpoint] = &EndpointLimiter{
		URL:     cfg.RPCEndpoint,
		Limiter: limiter,
		Config:  cfg.RateLimit,
	}

	// Create limiters for reading endpoints
	for _, endpoint := range cfg.ReadingEndpoints {
		limiter := createLimiter(endpoint.RateLimit)
		lm.limiters[endpoint.URL] = &EndpointLimiter{
			URL:     endpoint.URL,
			Limiter: limiter,
			Config:  endpoint.RateLimit,
		}
	}

	// Create limiters for execution endpoints
	for _, endpoint := range cfg.ExecutionEndpoints {
		limiter := createLimiter(endpoint.RateLimit)
		lm.limiters[endpoint.URL] = &EndpointLimiter{
			URL:     endpoint.URL,
			Limiter: limiter,
			Config:  endpoint.RateLimit,
		}
	}

	return lm
}

// createLimiter creates a rate limiter based on the configuration
func createLimiter(cfg config.RateLimitConfig) *rate.Limiter {
	// Create a rate limiter with the specified rate and burst
	r := rate.Limit(cfg.RequestsPerSecond)
	return rate.NewLimiter(r, cfg.Burst)
}

// WaitForLimit waits for the rate limiter to allow a request
func (lm *LimiterManager) WaitForLimit(ctx context.Context, endpointURL string) error {
	lm.mu.RLock()
	limiter, exists := lm.limiters[endpointURL]
	lm.mu.RUnlock()

	if !exists {
		return fmt.Errorf("no rate limiter found for endpoint: %s", endpointURL)
	}

	// Wait for permission to make a request
	return limiter.Limiter.Wait(ctx)
}

// TryWaitForLimit tries to wait for the rate limiter to allow a request without blocking
func (lm *LimiterManager) TryWaitForLimit(ctx context.Context, endpointURL string) error {
	lm.mu.RLock()
	limiter, exists := lm.limiters[endpointURL]
	lm.mu.RUnlock()

	if !exists {
		return fmt.Errorf("no rate limiter found for endpoint: %s", endpointURL)
	}

	// Try to wait for permission to make a request without blocking
	if !limiter.Limiter.Allow() {
		return fmt.Errorf("rate limit exceeded for endpoint: %s", endpointURL)
	}

	return nil
}

// GetLimiter returns the rate limiter for a specific endpoint
func (lm *LimiterManager) GetLimiter(endpointURL string) (*rate.Limiter, error) {
	lm.mu.RLock()
	limiter, exists := lm.limiters[endpointURL]
	lm.mu.RUnlock()

	if !exists {
		return nil, fmt.Errorf("no rate limiter found for endpoint: %s", endpointURL)
	}

	return limiter.Limiter, nil
}

// UpdateLimiter updates the rate limiter for an endpoint
func (lm *LimiterManager) UpdateLimiter(endpointURL string, cfg config.RateLimitConfig) {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	limiter := createLimiter(cfg)
	lm.limiters[endpointURL] = &EndpointLimiter{
		URL:     endpointURL,
		Limiter: limiter,
		Config:  cfg,
	}
}

// GetEndpoints returns all endpoint URLs
func (lm *LimiterManager) GetEndpoints() []string {
	lm.mu.RLock()
	defer lm.mu.RUnlock()

	endpoints := make([]string, 0, len(lm.limiters))
	for url := range lm.limiters {
		endpoints = append(endpoints, url)
	}

	return endpoints
}