9
9
"github.com/go-redis/redis/v8"
10
10
)
11
11
12
- const redisPrefix = "rate:"
13
-
14
12
type rediser interface {
15
13
Eval (ctx context.Context , script string , keys []string , args ... interface {}) * redis.Cmd
16
14
EvalSha (ctx context.Context , sha1 string , keys []string , args ... interface {}) * redis.Cmd
@@ -73,13 +71,15 @@ func PerHour(rate int) Limit {
73
71
74
72
// Limiter controls how frequently events are allowed to happen.
75
73
type Limiter struct {
76
- rdb rediser
74
+ rdb rediser
75
+ keyPrefix string
77
76
}
78
77
79
78
// NewLimiter returns a new Limiter.
80
- func NewLimiter (rdb rediser ) * Limiter {
79
+ func NewLimiter (rdb rediser , keyPrefix string ) * Limiter {
81
80
return & Limiter {
82
- rdb : rdb ,
81
+ rdb : rdb ,
82
+ keyPrefix : keyPrefix ,
83
83
}
84
84
}
85
85
@@ -96,7 +96,7 @@ func (l Limiter) AllowN(
96
96
n int ,
97
97
) (* Result , error ) {
98
98
values := []interface {}{limit .Burst , limit .Rate , limit .Period .Seconds (), n }
99
- v , err := allowN .Run (ctx , l .rdb , []string {redisPrefix + key }, values ... ).Result ()
99
+ v , err := allowN .Run (ctx , l .rdb , []string {l . keyPrefix + key }, values ... ).Result ()
100
100
if err != nil {
101
101
return nil , err
102
102
}
@@ -132,7 +132,7 @@ func (l Limiter) AllowAtMost(
132
132
n int ,
133
133
) (* Result , error ) {
134
134
values := []interface {}{limit .Burst , limit .Rate , limit .Period .Seconds (), n }
135
- v , err := allowAtMost .Run (ctx , l .rdb , []string {redisPrefix + key }, values ... ).Result ()
135
+ v , err := allowAtMost .Run (ctx , l .rdb , []string {l . keyPrefix + key }, values ... ).Result ()
136
136
if err != nil {
137
137
return nil , err
138
138
}
@@ -161,7 +161,7 @@ func (l Limiter) AllowAtMost(
161
161
162
162
// Reset gets a key and reset all limitations and previous usages
163
163
func (l * Limiter ) Reset (ctx context.Context , key string ) error {
164
- return l .rdb .Del (ctx , redisPrefix + key ).Err ()
164
+ return l .rdb .Del (ctx , l . keyPrefix + key ).Err ()
165
165
}
166
166
167
167
func dur (f float64 ) time.Duration {
0 commit comments