/
strategy.go
136 lines (109 loc) · 2.64 KB
/
strategy.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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
package rate
import (
"encoding/json"
"time"
"github.com/pkg/errors"
"golang.org/x/time/rate"
)
type LimitAlgoType string
const (
// rate limit algorithms, only `fixed_window` and `token bucket` are supported for now.
LimitAlgoFixedWindow LimitAlgoType = "fixed_window"
LimitAlgoTokenBucket LimitAlgoType = "token_bucket"
)
type LimitType int
const (
// rate limit types, only `ip` or `key` are supported for now.
LimitTypeByKey LimitType = iota
LimitTypeByIp
)
const (
// pre-defined default strategy name
DefaultStrategy = "default"
)
// Strategy rate limit strategy
type Strategy struct {
ID uint32 // strategy ID
Name string // strategy name
LimitOptions map[string]interface{} // resource => limit option
}
func NewStrategy(id uint32, name string) *Strategy {
return &Strategy{
ID: id,
Name: name,
LimitOptions: make(map[string]interface{}),
}
}
// UnmarshalJSON implements `json.Unmarshaler`
func (s *Strategy) UnmarshalJSON(data []byte) error {
tmpRules := make(map[string]*LimitRule)
if err := json.Unmarshal(data, &tmpRules); err != nil {
return errors.WithMessage(err, "malformed json format")
}
for resource, rule := range tmpRules {
s.LimitOptions[resource] = rule.Option
}
return nil
}
// FixedWindowOption limit option for fixed window
type FixedWindowOption struct {
Interval time.Duration
Quota int
}
// UnmarshalJSON implements `json.Unmarshaler`
func (fwo *FixedWindowOption) UnmarshalJSON(data []byte) error {
var tmp struct {
Interval string
Quota int
}
if err := json.Unmarshal(data, &tmp); err != nil {
return err
}
fwo.Quota = tmp.Quota
interval, err := time.ParseDuration(tmp.Interval)
if err == nil {
fwo.Interval = interval
}
return err
}
// FixedWindowOption limit option for token bucket
type TokenBucketOption struct {
Rate rate.Limit
Burst int
}
func NewTokenBucketOption(r, b int) TokenBucketOption {
return TokenBucketOption{
Rate: rate.Limit(r),
Burst: b,
}
}
// LimitRule resource limit rule
type LimitRule struct {
Algo LimitAlgoType
Option interface{}
}
func (r *LimitRule) UnmarshalJSON(data []byte) (err error) {
var tmp struct {
Algo LimitAlgoType
Option json.RawMessage
}
if err := json.Unmarshal(data, &tmp); err != nil {
return err
}
r.Algo = tmp.Algo
switch tmp.Algo {
case LimitAlgoFixedWindow:
var fwopt FixedWindowOption
if err = json.Unmarshal(tmp.Option, &fwopt); err == nil {
r.Option = fwopt
}
case LimitAlgoTokenBucket:
var tbopt TokenBucketOption
if err = json.Unmarshal(tmp.Option, &tbopt); err == nil {
r.Option = tbopt
}
default:
return errors.New("invalid rate limit algorithm")
}
return err
}