From 82875083883ba30a9f0cb1a9581ca51197c66acc Mon Sep 17 00:00:00 2001 From: 0ceanSlim Date: Thu, 25 Jul 2024 16:43:46 -0400 Subject: [PATCH] refactored some config/limiter orders --- config.example.yml | 36 ++++++++++++------------ main.go | 4 +-- relay/utils/loadConfig.go | 54 ++++++++++++++++++------------------ relay/utils/rateLimiter.go | 56 +++++++++++++++++++------------------- 4 files changed, 76 insertions(+), 74 deletions(-) diff --git a/config.example.yml b/config.example.yml index c404b72..b4bb90d 100644 --- a/config.example.yml +++ b/config.example.yml @@ -6,26 +6,17 @@ server: port: ":8080" # Port for the server to listen on rate_limit: - event_limit: 25 # Global rate limit for events (25 events per second) - event_burst: 50 # Global burst limit for events (allows a burst of 50 events) - ws_limit: 50 # Global rate limit for WebSocket messages (50 messages per second) - ws_burst: 100 # Global burst limit for WebSocket messages (allows a burst of 100 messages) - - kind_limits: # Specific rate limits for different kinds of events - - kind: 0 - limit: 1 # Rate limit for events of kind 0 (1 event per second) - burst: 5 # Burst limit for events of kind 0 (allows a burst of 5 events) - - kind: 1 - limit: 100 # Rate limit for events of kind 1 (100 events per second) - burst: 200 # Burst limit for events of kind 1 (allows a burst of 200 events) - - kind: 3 - limit: 25 # Rate limit for events of kind 3 (25 events per second) - burst: 50 # Burst limit for events of kind 3 (allows a burst of 50 events) + #I'm not sure is global ws_limit is working for ALL websocket messages + ws_limit: 100 # Global rate limit for WebSocket messages (50 messages per second) + ws_burst: 200 # Global burst limit for WebSocket messages (allows a burst of 100 messages) + #event limit doesn't seem to work either + event_limit: 50 # Global rate limit for events (25 events per second) + event_burst: 100 # Global burst limit for events (allows a burst of 50 events) category_limits: # Rate limits based on event categories regular: - limit: 50 # Rate limit for regular events (50 events per second) - burst: 100 # Burst limit for regular events (allows a burst of 100 events) + limit: 25 # Rate limit for regular events (50 events per second) + burst: 50 # Burst limit for regular events (allows a burst of 100 events) replaceable: limit: 10 # Rate limit for replaceable events (10 events per second) burst: 20 # Burst limit for replaceable events (allows a burst of 20 events) @@ -35,3 +26,14 @@ rate_limit: ephemeral: limit: 100 # Rate limit for ephemeral events (100 events per second) burst: 200 # Burst limit for ephemeral events (allows a burst of 200 events) + + kind_limits: # Specific rate limits for different kinds of events + - kind: 0 + limit: 1 # Rate limit for events of kind 0 (1 event per second) + burst: 5 # Burst limit for events of kind 0 (allows a burst of 5 events) + - kind: 1 + limit: 25 # Rate limit for events of kind 1 (100 events per second) + burst: 50 # Burst limit for events of kind 1 (allows a burst of 200 events) + - kind: 3 + limit: 25 # Rate limit for events of kind 3 (25 events per second) + burst: 50 # Burst limit for events of kind 3 (allows a burst of 50 events) diff --git a/main.go b/main.go index d87eed4..101630d 100644 --- a/main.go +++ b/main.go @@ -30,10 +30,10 @@ func main() { // Initialize Rate Limiter rateLimiter := utils.NewRateLimiter( - rate.Limit(config.RateLimit.EventLimit), - config.RateLimit.EventBurst, rate.Limit(config.RateLimit.WsLimit), config.RateLimit.WsBurst, + rate.Limit(config.RateLimit.EventLimit), + config.RateLimit.EventBurst, ) for _, kindLimit := range config.RateLimit.KindLimits { diff --git a/relay/utils/loadConfig.go b/relay/utils/loadConfig.go index 351b9aa..bebb0c5 100644 --- a/relay/utils/loadConfig.go +++ b/relay/utils/loadConfig.go @@ -6,33 +6,6 @@ import ( "gopkg.in/yaml.v2" ) -type RateLimitConfig struct { - WsLimit float64 `yaml:"ws_limit"` - WsBurst int `yaml:"ws_burst"` - EventLimit float64 `yaml:"event_limit"` - EventBurst int `yaml:"event_burst"` - KindLimits []KindLimitConfig `yaml:"kind_limits"` - CategoryLimits map[string]KindLimitConfig `yaml:"category_limits"` -} - -type KindLimitConfig struct { - Kind int `yaml:"kind"` - Limit float64 `yaml:"limit"` - Burst int `yaml:"burst"` -} - -type CategoryLimitConfig struct { - Regular LimitBurst `yaml:"regular"` - Replaceable LimitBurst `yaml:"replaceable"` - ParameterizedReplaceable LimitBurst `yaml:"parameterized_replaceable"` - Ephemeral LimitBurst `yaml:"ephemeral"` -} - -type LimitBurst struct { - Limit float64 `yaml:"limit"` - Burst int `yaml:"burst"` -} - type Config struct { MongoDB struct { URI string `yaml:"uri"` @@ -44,6 +17,33 @@ type Config struct { RateLimit RateLimitConfig `yaml:"rate_limit"` } +type LimitBurst struct { + Limit float64 `yaml:"limit"` + Burst int `yaml:"burst"` +} + +type RateLimitConfig struct { + WsLimit float64 `yaml:"ws_limit"` + WsBurst int `yaml:"ws_burst"` + EventLimit float64 `yaml:"event_limit"` + EventBurst int `yaml:"event_burst"` + CategoryLimits map[string]KindLimitConfig `yaml:"category_limits"` + KindLimits []KindLimitConfig `yaml:"kind_limits"` +} + +type CategoryLimitConfig struct { + Regular LimitBurst `yaml:"regular"` + Replaceable LimitBurst `yaml:"replaceable"` + ParameterizedReplaceable LimitBurst `yaml:"parameterized_replaceable"` + Ephemeral LimitBurst `yaml:"ephemeral"` +} + +type KindLimitConfig struct { + Kind int `yaml:"kind"` + Limit float64 `yaml:"limit"` + Burst int `yaml:"burst"` +} + func LoadConfig(filename string) (*Config, error) { data, err := os.ReadFile(filename) if err != nil { diff --git a/relay/utils/rateLimiter.go b/relay/utils/rateLimiter.go index 169f905..3dfe11e 100644 --- a/relay/utils/rateLimiter.go +++ b/relay/utils/rateLimiter.go @@ -29,6 +29,16 @@ type RateLimiter struct { var rateLimiterInstance *RateLimiter var once sync.Once +func SetRateLimiter(rl *RateLimiter) { + once.Do(func() { + rateLimiterInstance = rl + }) +} + +func GetRateLimiter() *RateLimiter { + return rateLimiterInstance +} + func NewRateLimiter(eventLimit rate.Limit, eventBurst int, wsLimit rate.Limit, wsBurst int) *RateLimiter { return &RateLimiter{ eventLimiter: rate.NewLimiter(eventLimit, eventBurst), @@ -38,24 +48,8 @@ func NewRateLimiter(eventLimit rate.Limit, eventBurst int, wsLimit rate.Limit, w } } -func (rl *RateLimiter) AddKindLimit(kind int, limit rate.Limit, burst int) { - rl.mu.Lock() - defer rl.mu.Unlock() - rl.kindLimiters[kind] = &KindLimiter{ - Limiter: rate.NewLimiter(limit, burst), - Limit: limit, - Burst: burst, - } -} - -func (rl *RateLimiter) AddCategoryLimit(category string, limit rate.Limit, burst int) { - rl.mu.Lock() - defer rl.mu.Unlock() - rl.categoryLimiters[category] = &CategoryLimiter{ - Limiter: rate.NewLimiter(limit, burst), - Limit: limit, - Burst: burst, - } +func (rl *RateLimiter) AllowWs() bool { + return rl.wsLimiter.Allow() } func (rl *RateLimiter) AllowEvent(kind int, category string) bool { @@ -81,16 +75,22 @@ func (rl *RateLimiter) AllowEvent(kind int, category string) bool { return true } -func (rl *RateLimiter) AllowWs() bool { - return rl.wsLimiter.Allow() +func (rl *RateLimiter) AddCategoryLimit(category string, limit rate.Limit, burst int) { + rl.mu.Lock() + defer rl.mu.Unlock() + rl.categoryLimiters[category] = &CategoryLimiter{ + Limiter: rate.NewLimiter(limit, burst), + Limit: limit, + Burst: burst, + } } -func SetRateLimiter(rl *RateLimiter) { - once.Do(func() { - rateLimiterInstance = rl - }) -} - -func GetRateLimiter() *RateLimiter { - return rateLimiterInstance +func (rl *RateLimiter) AddKindLimit(kind int, limit rate.Limit, burst int) { + rl.mu.Lock() + defer rl.mu.Unlock() + rl.kindLimiters[kind] = &KindLimiter{ + Limiter: rate.NewLimiter(limit, burst), + Limit: limit, + Burst: burst, + } }