/
keyratelimit.go
100 lines (91 loc) · 2.3 KB
/
keyratelimit.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
package ratelimit
import (
"context"
"fmt"
"time"
)
// Options of MultiLimiter
type Options struct {
Key string // Unique Identifier
IsUnlimited bool
MaxCount uint
Duration time.Duration
}
// Validate given MultiLimiter Options
func (o *Options) Validate() error {
if !o.IsUnlimited {
if o.Key == "" {
return fmt.Errorf("empty keys not allowed")
}
if o.MaxCount == 0 {
return fmt.Errorf("maxcount cannot be zero")
}
if o.Duration == 0 {
return fmt.Errorf("time duration not set")
}
}
return nil
}
// MultiLimiter is wrapper around Limiter than can limit based on a key
type MultiLimiter struct {
limiters map[string]*Limiter
ctx context.Context
}
// Adds new bucket with key
func (m *MultiLimiter) Add(opts *Options) error {
if err := opts.Validate(); err != nil {
return err
}
_, ok := m.limiters[opts.Key]
if ok {
return fmt.Errorf("key already exists")
}
var rlimiter *Limiter
if opts.IsUnlimited {
rlimiter = NewUnlimited(m.ctx)
} else {
rlimiter = New(m.ctx, opts.MaxCount, opts.Duration)
}
m.limiters[opts.Key] = rlimiter
return nil
}
// GetLimit returns current ratelimit of given key
func (m *MultiLimiter) GetLimit(key string) (uint, error) {
limiter, ok := m.limiters[key]
if !ok || limiter == nil {
return 0, fmt.Errorf("key doesnot exist")
}
return limiter.GetLimit(), nil
}
// Take one token from bucket returns error if key not present
func (m *MultiLimiter) Take(key string) error {
limiter, ok := m.limiters[key]
if !ok || limiter == nil {
return fmt.Errorf("key doesnot exist")
}
limiter.Take()
return nil
}
// SleepandReset stops timer removes all tokens and resets with new limit (used for Adaptive Ratelimiting)
func (m *MultiLimiter) SleepandReset(SleepTime time.Duration, opts *Options) error {
if err := opts.Validate(); err != nil {
return err
}
limiter, ok := m.limiters[opts.Key]
if !ok || limiter == nil {
return fmt.Errorf("key doesnot exist")
}
limiter.SleepandReset(SleepTime, opts.MaxCount, opts.Duration)
return nil
}
// NewMultiLimiter : Limits
func NewMultiLimiter(ctx context.Context, opts *Options) (*MultiLimiter, error) {
if err := opts.Validate(); err != nil {
return nil, err
}
multilimiter := &MultiLimiter{
ctx: ctx,
limiters: map[string]*Limiter{},
}
return multilimiter, multilimiter.Add(opts)
}