mirror of
https://github.com/0ceanSlim/grain.git
synced 2024-10-29 17:16:31 +00:00
refactored some config/limiter orders
This commit is contained in:
parent
7365ef5c11
commit
8287508388
@ -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)
|
||||
|
4
main.go
4
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 {
|
||||
|
@ -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 {
|
||||
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user