mirror of
https://github.com/0ceanSlim/grain.git
synced 2024-10-30 01:26:32 +00:00
refactor limit functions
This commit is contained in:
parent
da46e23ae0
commit
b22fbe6aa8
@ -32,7 +32,7 @@ type RateLimiter struct {
|
|||||||
var rateLimiterInstance *RateLimiter
|
var rateLimiterInstance *RateLimiter
|
||||||
var rateOnce sync.Once
|
var rateOnce sync.Once
|
||||||
|
|
||||||
func SetupRateLimiter(cfg *config.ServerConfig) {
|
func SetRateLimit(cfg *config.ServerConfig) {
|
||||||
rateLimiter := NewRateLimiter(
|
rateLimiter := NewRateLimiter(
|
||||||
rate.Limit(cfg.RateLimit.WsLimit),
|
rate.Limit(cfg.RateLimit.WsLimit),
|
||||||
cfg.RateLimit.WsBurst,
|
cfg.RateLimit.WsBurst,
|
||||||
|
@ -12,13 +12,13 @@ type SizeLimiter struct {
|
|||||||
mu sync.RWMutex
|
mu sync.RWMutex
|
||||||
}
|
}
|
||||||
|
|
||||||
func SetupSizeLimiter(cfg *config.ServerConfig) {
|
func SetSizeLimit(cfg *config.ServerConfig) {
|
||||||
sizeLimiter := NewSizeLimiter(cfg.RateLimit.MaxEventSize)
|
sizeLimiter := NewSizeLimiter(cfg.RateLimit.MaxEventSize)
|
||||||
for _, kindSizeLimit := range cfg.RateLimit.KindSizeLimits {
|
for _, kindSizeLimit := range cfg.RateLimit.KindSizeLimits {
|
||||||
sizeLimiter.AddKindSizeLimit(kindSizeLimit.Kind, kindSizeLimit.MaxSize)
|
sizeLimiter.AddKindSizeLimit(kindSizeLimit.Kind, kindSizeLimit.MaxSize)
|
||||||
}
|
}
|
||||||
|
|
||||||
SetSizeLimit(sizeLimiter)
|
SizeLimit(sizeLimiter)
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewSizeLimiter(globalMaxSize int) *SizeLimiter {
|
func NewSizeLimiter(globalMaxSize int) *SizeLimiter {
|
||||||
@ -35,7 +35,7 @@ func GetSizeLimiter() *SizeLimiter {
|
|||||||
return sizeLimiterInstance
|
return sizeLimiterInstance
|
||||||
}
|
}
|
||||||
|
|
||||||
func SetSizeLimit(sl *SizeLimiter) {
|
func SizeLimit(sl *SizeLimiter) {
|
||||||
sizeOnce.Do(func() {
|
sizeOnce.Do(func() {
|
||||||
sizeLimiterInstance = sl
|
sizeLimiterInstance = sl
|
||||||
})
|
})
|
||||||
|
4
main.go
4
main.go
@ -46,8 +46,8 @@ func main() {
|
|||||||
log.Fatal("Error initializing database: ", err)
|
log.Fatal("Error initializing database: ", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
config.SetupRateLimiter(cfg)
|
config.SetRateLimit(cfg)
|
||||||
config.SetupSizeLimiter(cfg)
|
config.SetSizeLimit(cfg)
|
||||||
|
|
||||||
config.ClearTemporaryBans()
|
config.ClearTemporaryBans()
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user