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
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)

View File

@ -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 {

View File

@ -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 {

View File

@ -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,
}
}