init: 1.0.0
This commit is contained in:
Chang lue Tsen
@@ -0,0 +1,133 @@
|
||||
package limit
|
||||
|
||||
import (
|
||||
"context"
|
||||
_ "embed"
|
||||
"errors"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
const (
|
||||
// Unknown means not initialized state.
|
||||
Unknown = iota
|
||||
// Allowed means allowed state.
|
||||
Allowed
|
||||
// HitQuota means this request exactly hit the quota.
|
||||
HitQuota
|
||||
// OverQuota means passed the quota.
|
||||
OverQuota
|
||||
|
||||
internalOverQuota = 0
|
||||
internalAllowed = 1
|
||||
internalHitQuota = 2
|
||||
)
|
||||
|
||||
var (
|
||||
// ErrUnknownCode is an error that represents unknown status code.
|
||||
ErrUnknownCode = errors.New("unknown status code")
|
||||
|
||||
//go:embed periodscript.lua
|
||||
periodLuaScript string
|
||||
periodScript = redis.NewScript(periodLuaScript)
|
||||
)
|
||||
|
||||
type (
|
||||
// PeriodOption defines the method to customize a PeriodLimit.
|
||||
PeriodOption func(l *PeriodLimit)
|
||||
|
||||
// A PeriodLimit is used to limit requests during a period of time.
|
||||
PeriodLimit struct {
|
||||
period int
|
||||
quota int
|
||||
limitStore *redis.Client
|
||||
keyPrefix string
|
||||
align bool
|
||||
}
|
||||
)
|
||||
|
||||
// NewPeriodLimit returns a PeriodLimit with given parameters.
|
||||
func NewPeriodLimit(period, quota int, limitStore *redis.Client, keyPrefix string,
|
||||
opts ...PeriodOption) *PeriodLimit {
|
||||
limiter := &PeriodLimit{
|
||||
period: period,
|
||||
quota: quota,
|
||||
limitStore: limitStore,
|
||||
keyPrefix: keyPrefix,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(limiter)
|
||||
}
|
||||
|
||||
return limiter
|
||||
}
|
||||
|
||||
// Take requests a permit, it returns the permit state.
|
||||
func (h *PeriodLimit) Take(key string) (int, error) {
|
||||
return h.TakeCtx(context.Background(), key)
|
||||
}
|
||||
|
||||
// TakeCtx requests a permit with context, it returns the permit state.
|
||||
func (h *PeriodLimit) TakeCtx(ctx context.Context, key string) (int, error) {
|
||||
resp, err := periodScript.Run(ctx, h.limitStore, []string{h.keyPrefix + key}, []string{
|
||||
strconv.Itoa(h.quota),
|
||||
strconv.Itoa(h.calcExpireSeconds()),
|
||||
}).Result()
|
||||
|
||||
if err != nil {
|
||||
return Unknown, err
|
||||
}
|
||||
|
||||
code, ok := resp.(int64)
|
||||
if !ok {
|
||||
return Unknown, ErrUnknownCode
|
||||
}
|
||||
|
||||
switch code {
|
||||
case internalOverQuota:
|
||||
return OverQuota, nil
|
||||
case internalAllowed:
|
||||
return Allowed, nil
|
||||
case internalHitQuota:
|
||||
return HitQuota, nil
|
||||
default:
|
||||
return Unknown, ErrUnknownCode
|
||||
}
|
||||
}
|
||||
|
||||
func (h *PeriodLimit) calcExpireSeconds() int {
|
||||
if h.align {
|
||||
now := time.Now()
|
||||
_, offset := now.Zone()
|
||||
unix := now.Unix() + int64(offset)
|
||||
return h.period - int(unix%int64(h.period))
|
||||
}
|
||||
|
||||
return h.period
|
||||
}
|
||||
|
||||
// Align returns a func to customize a PeriodLimit with alignment.
|
||||
// For example, if we want to limit end users with 5 sms verification messages every day,
|
||||
// we need to align with the local timezone and the start of the day.
|
||||
func Align() PeriodOption {
|
||||
return func(l *PeriodLimit) {
|
||||
l.align = true
|
||||
}
|
||||
}
|
||||
|
||||
// ParsePermitState returns the permit state by the code.
|
||||
func (h *PeriodLimit) ParsePermitState(code int) bool {
|
||||
switch code {
|
||||
case Allowed:
|
||||
return true
|
||||
case HitQuota:
|
||||
return true
|
||||
case OverQuota:
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package limit
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestPeriodLimit_Take(t *testing.T) {
|
||||
testPeriodLimit(t)
|
||||
}
|
||||
|
||||
func TestPeriodLimit_TakeWithAlign(t *testing.T) {
|
||||
testPeriodLimit(t, Align())
|
||||
}
|
||||
|
||||
func TestPeriodLimit_RedisUnavailable(t *testing.T) {
|
||||
//t.Skipf("skip this test because it's not stable")
|
||||
const (
|
||||
seconds = 1
|
||||
quota = 5
|
||||
)
|
||||
rds := redis.NewClient(&redis.Options{
|
||||
Addr: "localhost:12345",
|
||||
})
|
||||
|
||||
l := NewPeriodLimit(seconds, quota, rds, "periodlimit:")
|
||||
val, err := l.Take("first")
|
||||
assert.NotNil(t, err)
|
||||
assert.Equal(t, 0, val)
|
||||
}
|
||||
|
||||
func testPeriodLimit(t *testing.T, opts ...PeriodOption) {
|
||||
store, _ := CreateRedisWithClean(t)
|
||||
const (
|
||||
seconds = 1
|
||||
total = 100
|
||||
quota = 5
|
||||
)
|
||||
l := NewPeriodLimit(seconds, quota, store, "periodlimit", opts...)
|
||||
var allowed, hitQuota, overQuota int
|
||||
for i := 0; i < total; i++ {
|
||||
val, err := l.Take("first")
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
switch val {
|
||||
case Allowed:
|
||||
allowed++
|
||||
case HitQuota:
|
||||
hitQuota++
|
||||
case OverQuota:
|
||||
overQuota++
|
||||
default:
|
||||
t.Error("unknown status")
|
||||
}
|
||||
}
|
||||
assert.Equal(t, quota-1, allowed)
|
||||
assert.Equal(t, 1, hitQuota)
|
||||
assert.Equal(t, total-quota, overQuota)
|
||||
}
|
||||
|
||||
func TestQuotaFull(t *testing.T) {
|
||||
rds, _ := CreateRedisWithClean(t)
|
||||
l := NewPeriodLimit(1, 1, rds, "periodlimit")
|
||||
val, err := l.Take("first")
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, HitQuota, val)
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
-- to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
|
||||
local limit = tonumber(ARGV[1])
|
||||
local window = tonumber(ARGV[2])
|
||||
local current = redis.call("INCRBY", KEYS[1], 1)
|
||||
if current == 1 then
|
||||
redis.call("expire", KEYS[1], window)
|
||||
end
|
||||
if current < limit then
|
||||
return 1
|
||||
elseif current == limit then
|
||||
return 2
|
||||
else
|
||||
return 0
|
||||
end
|
||||
@@ -0,0 +1,160 @@
|
||||
package limit
|
||||
|
||||
import (
|
||||
"context"
|
||||
_ "embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/perfect-panel/ppanel-server/pkg/errorx"
|
||||
"github.com/perfect-panel/ppanel-server/pkg/logger"
|
||||
"github.com/redis/go-redis/v9"
|
||||
xrate "golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
const (
|
||||
tokenFormat = "{%s}.tokens"
|
||||
timestampFormat = "{%s}.ts"
|
||||
pingInterval = time.Millisecond * 100
|
||||
)
|
||||
|
||||
var (
|
||||
//go:embed tokenscript.lua
|
||||
tokenLuaScript string
|
||||
tokenScript = redis.NewScript(tokenLuaScript)
|
||||
)
|
||||
|
||||
// A TokenLimiter controls how frequently events are allowed to happen with in one second.
|
||||
type TokenLimiter struct {
|
||||
rate int
|
||||
burst int
|
||||
store *redis.Client
|
||||
tokenKey string
|
||||
timestampKey string
|
||||
rescueLock sync.Mutex
|
||||
redisAlive uint32
|
||||
monitorStarted bool
|
||||
rescueLimiter *xrate.Limiter
|
||||
}
|
||||
|
||||
// NewTokenLimiter returns a new TokenLimiter that allows events up to rate and permits
|
||||
// bursts of at most burst tokens.
|
||||
func NewTokenLimiter(rate, burst int, store *redis.Client, key string) *TokenLimiter {
|
||||
tokenKey := fmt.Sprintf(tokenFormat, key)
|
||||
timestampKey := fmt.Sprintf(timestampFormat, key)
|
||||
|
||||
return &TokenLimiter{
|
||||
rate: rate,
|
||||
burst: burst,
|
||||
store: store,
|
||||
tokenKey: tokenKey,
|
||||
timestampKey: timestampKey,
|
||||
redisAlive: 1,
|
||||
rescueLimiter: xrate.NewLimiter(xrate.Every(time.Second/time.Duration(rate)), burst),
|
||||
}
|
||||
}
|
||||
|
||||
// Allow is shorthand for AllowN(time.Now(), 1).
|
||||
func (lim *TokenLimiter) Allow() bool {
|
||||
return lim.AllowN(time.Now(), 1)
|
||||
}
|
||||
|
||||
// AllowCtx is shorthand for AllowNCtx(ctx,time.Now(), 1) with incoming context.
|
||||
func (lim *TokenLimiter) AllowCtx(ctx context.Context) bool {
|
||||
return lim.AllowNCtx(ctx, time.Now(), 1)
|
||||
}
|
||||
|
||||
// AllowN reports whether n events may happen at time now.
|
||||
// Use this method if you intend to drop / skip events that exceed the rate.
|
||||
// Otherwise, use Reserve or Wait.
|
||||
func (lim *TokenLimiter) AllowN(now time.Time, n int) bool {
|
||||
return lim.reserveN(context.Background(), now, n)
|
||||
}
|
||||
|
||||
// AllowNCtx reports whether n events may happen at time now with incoming context.
|
||||
// Use this method if you intend to drop / skip events that exceed the rate.
|
||||
// Otherwise, use Reserve or Wait.
|
||||
func (lim *TokenLimiter) AllowNCtx(ctx context.Context, now time.Time, n int) bool {
|
||||
return lim.reserveN(ctx, now, n)
|
||||
}
|
||||
|
||||
func (lim *TokenLimiter) reserveN(ctx context.Context, now time.Time, n int) bool {
|
||||
if atomic.LoadUint32(&lim.redisAlive) == 0 {
|
||||
return lim.rescueLimiter.AllowN(now, n)
|
||||
}
|
||||
|
||||
resp, err := tokenScript.Run(ctx, lim.store, []string{lim.tokenKey, lim.timestampKey}, []string{
|
||||
strconv.Itoa(lim.rate),
|
||||
strconv.Itoa(lim.burst),
|
||||
strconv.FormatInt(now.Unix(), 10),
|
||||
strconv.Itoa(n),
|
||||
}).Result()
|
||||
// redis allowed == false
|
||||
// Lua boolean false -> r Nil bulk reply
|
||||
if errors.Is(err, redis.Nil) {
|
||||
return false
|
||||
}
|
||||
if errorx.In(err, context.DeadlineExceeded, context.Canceled) {
|
||||
logger.WithContext(ctx).Error("fail to use rate limiter", logger.Field("error", err.Error()))
|
||||
return false
|
||||
}
|
||||
if err != nil {
|
||||
//log.Errorf(ctx, "fail to eval redis script: %v, use in-process limiter for rescue", err)
|
||||
lim.startMonitor()
|
||||
return lim.rescueLimiter.AllowN(now, n)
|
||||
}
|
||||
|
||||
code, ok := resp.(int64)
|
||||
if !ok {
|
||||
logger.Error("fail to eval redis script, use in-process limiter for rescue", logger.Field("response", resp))
|
||||
lim.startMonitor()
|
||||
return lim.rescueLimiter.AllowN(now, n)
|
||||
}
|
||||
|
||||
// redis allowed == true
|
||||
// Lua boolean true -> r integer reply with value of 1
|
||||
return code == 1
|
||||
}
|
||||
|
||||
func (lim *TokenLimiter) startMonitor() {
|
||||
lim.rescueLock.Lock()
|
||||
defer lim.rescueLock.Unlock()
|
||||
|
||||
if lim.monitorStarted {
|
||||
return
|
||||
}
|
||||
|
||||
lim.monitorStarted = true
|
||||
atomic.StoreUint32(&lim.redisAlive, 0)
|
||||
|
||||
go lim.waitForRedis()
|
||||
}
|
||||
|
||||
func (lim *TokenLimiter) waitForRedis() {
|
||||
ticker := time.NewTicker(pingInterval)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
lim.rescueLock.Lock()
|
||||
lim.monitorStarted = false
|
||||
lim.rescueLock.Unlock()
|
||||
}()
|
||||
|
||||
for range ticker.C {
|
||||
if lim.StorePingCtx() {
|
||||
atomic.StoreUint32(&lim.redisAlive, 1)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (lim *TokenLimiter) StorePingCtx() bool {
|
||||
v, err := lim.store.Ping(context.Background()).Result()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return v == "PONG"
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package limit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestTokenLimit_WithCtx(t *testing.T) {
|
||||
const (
|
||||
total = 100
|
||||
rate = 5
|
||||
burst = 10
|
||||
)
|
||||
store, _ := CreateRedisWithClean(t)
|
||||
l := NewTokenLimiter(rate, burst, store, "tokenlimit")
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
ok := l.AllowCtx(ctx)
|
||||
assert.True(t, ok)
|
||||
|
||||
cancel()
|
||||
for i := 0; i < total; i++ {
|
||||
ok := l.AllowCtx(ctx)
|
||||
assert.False(t, ok)
|
||||
assert.False(t, l.monitorStarted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenLimit_Take(t *testing.T) {
|
||||
store, _ := CreateRedisWithClean(t)
|
||||
|
||||
const (
|
||||
total = 100
|
||||
rate = 5
|
||||
burst = 10
|
||||
)
|
||||
l := NewTokenLimiter(rate, burst, store, "tokenlimit")
|
||||
var allowed int
|
||||
for i := 0; i < total; i++ {
|
||||
time.Sleep(time.Second / time.Duration(total))
|
||||
if l.Allow() {
|
||||
allowed++
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, allowed >= burst+rate)
|
||||
}
|
||||
|
||||
func TestTokenLimit_TakeBurst(t *testing.T) {
|
||||
store, _ := CreateRedisWithClean(t)
|
||||
|
||||
const (
|
||||
total = 100
|
||||
rate = 5
|
||||
burst = 10
|
||||
)
|
||||
l := NewTokenLimiter(rate, burst, store, "tokenlimit")
|
||||
var allowed int
|
||||
for i := 0; i < total; i++ {
|
||||
if l.Allow() {
|
||||
allowed++
|
||||
}
|
||||
}
|
||||
|
||||
assert.True(t, allowed >= burst)
|
||||
}
|
||||
|
||||
// CreateRedisWithClean returns an in process redis.Redis and a clean function.
|
||||
func CreateRedisWithClean(t *testing.T) (r *redis.Client, clean func()) {
|
||||
mr := miniredis.RunT(t)
|
||||
return redis.NewClient(&redis.Options{
|
||||
Addr: mr.Addr(),
|
||||
}), mr.Close
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
-- to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
|
||||
-- KEYS[1] as tokens_key
|
||||
-- KEYS[2] as timestamp_key
|
||||
local rate = tonumber(ARGV[1])
|
||||
local capacity = tonumber(ARGV[2])
|
||||
local now = tonumber(ARGV[3])
|
||||
local requested = tonumber(ARGV[4])
|
||||
local fill_time = capacity/rate
|
||||
local ttl = math.floor(fill_time*2)
|
||||
local last_tokens = tonumber(redis.call("get", KEYS[1]))
|
||||
if last_tokens == nil then
|
||||
last_tokens = capacity
|
||||
end
|
||||
|
||||
local last_refreshed = tonumber(redis.call("get", KEYS[2]))
|
||||
if last_refreshed == nil then
|
||||
last_refreshed = 0
|
||||
end
|
||||
|
||||
local delta = math.max(0, now-last_refreshed)
|
||||
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
|
||||
local allowed = filled_tokens >= requested
|
||||
local new_tokens = filled_tokens
|
||||
if allowed then
|
||||
new_tokens = filled_tokens - requested
|
||||
end
|
||||
|
||||
redis.call("setex", KEYS[1], ttl, new_tokens)
|
||||
redis.call("setex", KEYS[2], ttl, now)
|
||||
|
||||
return allowed
|
||||
Reference in New Issue
Block a user