grain/config/sizeLimiter.go

61 lines
1.1 KiB
Go
Raw Normal View History

2024-07-30 15:27:38 +00:00
package config
2024-07-26 20:46:01 +00:00
import (
"sync"
)
type SizeLimiter struct {
globalMaxSize int
kindSizeLimits map[int]int
mu sync.RWMutex
}
func NewSizeLimiter(globalMaxSize int) *SizeLimiter {
return &SizeLimiter{
globalMaxSize: globalMaxSize,
kindSizeLimits: make(map[int]int),
}
}
var sizeLimiterInstance *SizeLimiter
var sizeOnce sync.Once
func GetSizeLimiter() *SizeLimiter {
return sizeLimiterInstance
}
func SetSizeLimiter(sl *SizeLimiter) {
sizeOnce.Do(func() {
sizeLimiterInstance = sl
})
}
func (sl *SizeLimiter) SetGlobalMaxSize(maxSize int) {
sl.mu.Lock()
defer sl.mu.Unlock()
sl.globalMaxSize = maxSize
}
func (sl *SizeLimiter) AddKindSizeLimit(kind int, maxSize int) {
sl.mu.Lock()
defer sl.mu.Unlock()
sl.kindSizeLimits[kind] = maxSize
}
func (sl *SizeLimiter) AllowSize(kind int, size int) (bool, string) {
sl.mu.RLock()
defer sl.mu.RUnlock()
if size > sl.globalMaxSize {
return false, "Global event size limit exceeded"
}
if maxSize, exists := sl.kindSizeLimits[kind]; exists {
if size > maxSize {
return false, "Event size limit exceeded for kind"
}
}
return true, ""
}