/
mem-ratelimter.go
63 lines (53 loc) · 1.48 KB
/
mem-ratelimter.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
package ratelimiter
import (
"context"
"time"
"github.com/Raj63/go-sdk/logger"
"github.com/patrickmn/go-cache"
"golang.org/x/time/rate"
)
// MemRatelimiter limiter
type MemRatelimiter struct {
*rate.Limiter
*cache.Cache
Expire time.Duration
logger *logger.Logger
}
var (
// MemRatelimiterCacheExpiration MemRatelimiter key
MemRatelimiterCacheExpiration = time.Minute * 60
// MemRatelimiterCacheCleanInterval MemRatelimiter
MemRatelimiterCacheCleanInterval = time.Minute * 60
)
// NewMemRatelimiter mem limiter
func NewMemRatelimiter(logger *logger.Logger) *MemRatelimiter {
// mem cache
memCache := cache.New(MemRatelimiterCacheExpiration, MemRatelimiterCacheCleanInterval)
return &MemRatelimiter{
Cache: memCache,
logger: logger,
}
}
// Allow time/rate & token bucket key
// tokenFillInterval Token
// bucketSize Token
func (r *MemRatelimiter) Allow(ctx context.Context, key string, tokenFillInterval time.Duration, bucketSize int) bool {
if tokenFillInterval.Seconds() <= 0 || bucketSize <= 0 {
return false
}
tokenRate := rate.Every(tokenFillInterval)
limiterI, exists := r.Cache.Get(key)
if !exists {
limiter := rate.NewLimiter(tokenRate, bucketSize)
limiter.Allow()
r.Cache.Set(key, limiter, MemRatelimiterCacheExpiration)
return true
}
if limiter, ok := limiterI.(*rate.Limiter); ok {
isAllow := limiter.Allow()
r.Cache.Set(key, limiter, MemRatelimiterCacheExpiration)
return isAllow
}
r.logger.Errorf("MemRatelimiter assert limiter error")
return true
}