43 lines
1.4 KiB
Go
43 lines
1.4 KiB
Go
package stocklock
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
// 滑动窗口：ZSET 里存每次请求的时间戳（毫秒），先删掉窗口外的，再数窗口内的
|
|
var slidingScript = redis.NewScript(`
|
|
local key, now, window, limit, member = KEYS[1], tonumber(ARGV[1]), tonumber(ARGV[2]), tonumber(ARGV[3]), ARGV[4]
|
|
redis.call("ZREMRANGEBYSCORE", key, 0, now - window)
|
|
if redis.call("ZCARD", key) >= limit then
|
|
return 0
|
|
end
|
|
redis.call("ZADD", key, now, member)
|
|
redis.call("PEXPIRE", key, window)
|
|
return 1`)
|
|
|
|
type SlidingWindow struct {
|
|
rdb *redis.Client
|
|
limit int
|
|
window time.Duration
|
|
seq func() string // 生成唯一成员，避免同一毫秒的请求互相覆盖
|
|
}
|
|
|
|
func NewSlidingWindow(rdb *redis.Client, limit int, window time.Duration) *SlidingWindow {
|
|
var n atomic.Uint64 // 并发安全的序号；多实例部署时成员里再加上实例标识或用 UUID
|
|
return &SlidingWindow{rdb: rdb, limit: limit, window: window, seq: func() string {
|
|
return fmt.Sprintf("%d-%d", time.Now().UnixNano(), n.Add(1))
|
|
}}
|
|
}
|
|
|
|
// Allow 判断 key（如 "rl:user:42"）在当前窗口内是否还能通过
|
|
func (s *SlidingWindow) Allow(ctx context.Context, key string, now time.Time) (bool, error) {
|
|
n, err := slidingScript.Run(ctx, s.rdb, []string{key},
|
|
now.UnixMilli(), s.window.Milliseconds(), s.limit, s.seq()).Int()
|
|
return n == 1, err
|
|
}
|