128 lines
3.2 KiB
Go
128 lines
3.2 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/your-username/mev-beta/internal/config"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// 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 fallback endpoints
|
|
for _, endpoint := range cfg.FallbackEndpoints {
|
|
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
|
|
} |