Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 8752327

Browse files
committed
support customized key prefix
1 parent d327cbf commit 8752327

File tree

2 files changed

+40
-8
lines changed

2 files changed

+40
-8
lines changed

‎rate.go‎

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -73,14 +73,34 @@ func PerHour(rate int) Limit {
7373

7474
// Limiter controls how frequently events are allowed to happen.
7575
type Limiter struct {
76-
rdb rediser
76+
rdb rediser
77+
keyPrefix string
7778
}
7879

80+
type Option func(*Limiter)
81+
7982
// NewLimiter returns a new Limiter.
80-
func NewLimiter(rdb rediser) *Limiter {
81-
return &Limiter{
83+
func NewLimiter(rdb rediser, options...Option) *Limiter {
84+
limiter:= &Limiter{
8285
rdb: rdb,
8386
}
87+
88+
for _, opt := range options {
89+
opt(limiter)
90+
}
91+
92+
if limiter.keyPrefix == "" {
93+
limiter.keyPrefix = redisPrefix
94+
}
95+
96+
return limiter
97+
}
98+
99+
// WithKeyPrefix is a functional option to set the redis key prefix.
100+
func WithKeyPrefix(keyPrefix string) Option {
101+
return func(l *Limiter) {
102+
l.keyPrefix = keyPrefix
103+
}
84104
}
85105

86106
// Allow is a shortcut for AllowN(ctx, key, limit, 1).
@@ -96,7 +116,7 @@ func (l Limiter) AllowN(
96116
n int,
97117
) (*Result, error) {
98118
values := []interface{}{limit.Burst, limit.Rate, limit.Period.Seconds(), n}
99-
v, err := allowN.Run(ctx, l.rdb, []string{redisPrefix + key}, values...).Result()
119+
v, err := allowN.Run(ctx, l.rdb, []string{l.keyPrefix + key}, values...).Result()
100120
if err != nil {
101121
return nil, err
102122
}
@@ -132,7 +152,7 @@ func (l Limiter) AllowAtMost(
132152
n int,
133153
) (*Result, error) {
134154
values := []interface{}{limit.Burst, limit.Rate, limit.Period.Seconds(), n}
135-
v, err := allowAtMost.Run(ctx, l.rdb, []string{redisPrefix + key}, values...).Result()
155+
v, err := allowAtMost.Run(ctx, l.rdb, []string{l.keyPrefix + key}, values...).Result()
136156
if err != nil {
137157
return nil, err
138158
}
@@ -161,7 +181,7 @@ func (l Limiter) AllowAtMost(
161181

162182
// Reset gets a key and reset all limitations and previous usages
163183
func (l *Limiter) Reset(ctx context.Context, key string) error {
164-
return l.rdb.Del(ctx, redisPrefix+key).Err()
184+
return l.rdb.Del(ctx, l.keyPrefix+key).Err()
165185
}
166186

167187
func dur(f float64) time.Duration {

‎rate_test.go‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,22 @@ func rateLimiter() *redis_rate.Limiter {
2121
return redis_rate.NewLimiter(ring)
2222
}
2323

24+
func TestAllow_WithKeyPrefix(t *testing.T) {
25+
ring := redis.NewRing(&redis.RingOptions{
26+
Addrs: map[string]string{"server0": ":6379"},
27+
})
28+
if err := ring.FlushDB(context.TODO()).Err(); err != nil {
29+
panic(err)
30+
}
31+
testAllow(t, redis_rate.NewLimiter(ring, redis_rate.WithKeyPrefix("redis_rate:")))
32+
}
33+
2434
func TestAllow(t *testing.T) {
25-
ctx := context.Background()
35+
testAllow(t, rateLimiter())
36+
}
2637

27-
l := rateLimiter()
38+
func testAllow(t *testing.T, l *redis_rate.Limiter) {
39+
ctx := context.Background()
2840

2941
limit := redis_rate.PerSecond(10)
3042
require.Equal(t, limit.String(), "10 req/s (burst 10)")

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /