refactored some config/limiter orders

This commit is contained in:
0ceanSlim 2024-07-25 16:43:46 -04:00
parent 7365ef5c11
commit 8287508388
4 changed files with 76 additions and 74 deletions

View File

@ -6,26 +6,17 @@ server:
port: ":8080" # Port for the server to listen on port: ":8080" # Port for the server to listen on
rate_limit: rate_limit:
event_limit: 25 # Global rate limit for events (25 events per second) #I'm not sure is global ws_limit is working for ALL websocket messages
event_burst: 50 # Global burst limit for events (allows a burst of 50 events) ws_limit: 100 # Global rate limit for WebSocket messages (50 messages per second)
ws_limit: 50 # 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)
ws_burst: 100 # 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)
kind_limits: # Specific rate limits for different kinds of events event_burst: 100 # Global burst limit for events (allows a burst of 50 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)
category_limits: # Rate limits based on event categories category_limits: # Rate limits based on event categories
regular: regular:
limit: 50 # Rate limit for regular events (50 events per second) limit: 25 # Rate limit for regular events (50 events per second)
burst: 100 # Burst limit for regular events (allows a burst of 100 events) burst: 50 # Burst limit for regular events (allows a burst of 100 events)
replaceable: replaceable:
limit: 10 # Rate limit for replaceable events (10 events per second) limit: 10 # Rate limit for replaceable events (10 events per second)
burst: 20 # Burst limit for replaceable events (allows a burst of 20 events) burst: 20 # Burst limit for replaceable events (allows a burst of 20 events)
@ -35,3 +26,14 @@ rate_limit:
ephemeral: ephemeral:
limit: 100 # Rate limit for ephemeral events (100 events per second) limit: 100 # Rate limit for ephemeral events (100 events per second)
burst: 200 # Burst limit for ephemeral events (allows a burst of 200 events) 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)

View File

@ -30,10 +30,10 @@ func main() {
// Initialize Rate Limiter // Initialize Rate Limiter
rateLimiter := utils.NewRateLimiter( rateLimiter := utils.NewRateLimiter(
rate.Limit(config.RateLimit.EventLimit),
config.RateLimit.EventBurst,
rate.Limit(config.RateLimit.WsLimit), rate.Limit(config.RateLimit.WsLimit),
config.RateLimit.WsBurst, config.RateLimit.WsBurst,
rate.Limit(config.RateLimit.EventLimit),
config.RateLimit.EventBurst,
) )
for _, kindLimit := range config.RateLimit.KindLimits { for _, kindLimit := range config.RateLimit.KindLimits {

View File

@ -6,33 +6,6 @@ import (
"gopkg.in/yaml.v2" "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 { type Config struct {
MongoDB struct { MongoDB struct {
URI string `yaml:"uri"` URI string `yaml:"uri"`
@ -44,6 +17,33 @@ type Config struct {
RateLimit RateLimitConfig `yaml:"rate_limit"` 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) { func LoadConfig(filename string) (*Config, error) {
data, err := os.ReadFile(filename) data, err := os.ReadFile(filename)
if err != nil { if err != nil {

View File

@ -29,6 +29,16 @@ type RateLimiter struct {
var rateLimiterInstance *RateLimiter var rateLimiterInstance *RateLimiter
var once sync.Once 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 { func NewRateLimiter(eventLimit rate.Limit, eventBurst int, wsLimit rate.Limit, wsBurst int) *RateLimiter {
return &RateLimiter{ return &RateLimiter{
eventLimiter: rate.NewLimiter(eventLimit, eventBurst), 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) { func (rl *RateLimiter) AllowWs() bool {
rl.mu.Lock() return rl.wsLimiter.Allow()
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) AllowEvent(kind int, category string) bool { func (rl *RateLimiter) AllowEvent(kind int, category string) bool {
@ -81,16 +75,22 @@ func (rl *RateLimiter) AllowEvent(kind int, category string) bool {
return true return true
} }
func (rl *RateLimiter) AllowWs() bool { func (rl *RateLimiter) AddCategoryLimit(category string, limit rate.Limit, burst int) {
return rl.wsLimiter.Allow() rl.mu.Lock()
defer rl.mu.Unlock()
rl.categoryLimiters[category] = &CategoryLimiter{
Limiter: rate.NewLimiter(limit, burst),
Limit: limit,
Burst: burst,
}
} }
func SetRateLimiter(rl *RateLimiter) { func (rl *RateLimiter) AddKindLimit(kind int, limit rate.Limit, burst int) {
once.Do(func() { rl.mu.Lock()
rateLimiterInstance = rl defer rl.mu.Unlock()
}) rl.kindLimiters[kind] = &KindLimiter{
Limiter: rate.NewLimiter(limit, burst),
Limit: limit,
Burst: burst,
} }
func GetRateLimiter() *RateLimiter {
return rateLimiterInstance
} }