90 lines
1.9 KiB
Go
90 lines
1.9 KiB
Go
package admission
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
ErrInvalidConfig = errors.New("invalid admission limiter configuration")
|
|
ErrInvalidIdentity = errors.New("invalid admission identity")
|
|
ErrGlobalLimit = errors.New("global admission limit exceeded")
|
|
ErrPerKeyLimit = errors.New("per-key admission limit exceeded")
|
|
ErrUnavailable = errors.New("admission control unavailable")
|
|
)
|
|
|
|
type FixedWindowConfig struct {
|
|
Window time.Duration
|
|
Global int
|
|
PerKey int
|
|
Now func() time.Time
|
|
}
|
|
|
|
// FixedWindow provides one concurrency-safe admission primitive for listener
|
|
// global and stable-client request limits.
|
|
type FixedWindow struct {
|
|
mu sync.Mutex
|
|
|
|
window time.Duration
|
|
global int
|
|
perKey int
|
|
now func() time.Time
|
|
|
|
windowID int64
|
|
globalUsed int
|
|
keyUsed map[string]int
|
|
}
|
|
|
|
func NewFixedWindow(config FixedWindowConfig) (*FixedWindow, error) {
|
|
if config.Window <= 0 || config.Global < 0 || config.PerKey < 0 ||
|
|
(config.Global == 0 && config.PerKey == 0) {
|
|
return nil, ErrInvalidConfig
|
|
}
|
|
if config.Now == nil {
|
|
config.Now = time.Now
|
|
}
|
|
return &FixedWindow{
|
|
window: config.Window,
|
|
global: config.Global,
|
|
perKey: config.PerKey,
|
|
now: config.Now,
|
|
keyUsed: make(map[string]int),
|
|
}, nil
|
|
}
|
|
|
|
func (l *FixedWindow) Admit(ctx context.Context, key string) error {
|
|
if ctx == nil || l == nil || key == "" {
|
|
return ErrInvalidIdentity
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
windowID := l.now().UnixNano() / int64(l.window)
|
|
if windowID != l.windowID {
|
|
l.windowID = windowID
|
|
l.globalUsed = 0
|
|
clear(l.keyUsed)
|
|
}
|
|
if l.global > 0 && l.globalUsed >= l.global {
|
|
return ErrGlobalLimit
|
|
}
|
|
if l.perKey > 0 && l.keyUsed[key] >= l.perKey {
|
|
return ErrPerKeyLimit
|
|
}
|
|
if l.global > 0 {
|
|
l.globalUsed++
|
|
}
|
|
if l.perKey > 0 {
|
|
l.keyUsed[key]++
|
|
}
|
|
return nil
|
|
}
|