File size: 1,606 Bytes
4674012 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | package limiter
import (
"context"
_ "embed"
"fmt"
"sync"
"github.com/QuantumNous/new-api/common"
"github.com/go-redis/redis/v8"
)
//go:embed lua/rate_limit.lua
var rateLimitScript string
type RedisLimiter struct {
client *redis.Client
limitScriptSHA string
}
var (
instance *RedisLimiter
once sync.Once
)
func New(ctx context.Context, r *redis.Client) *RedisLimiter {
once.Do(func() {
// 预加载脚本
limitSHA, err := r.ScriptLoad(ctx, rateLimitScript).Result()
if err != nil {
common.SysLog(fmt.Sprintf("Failed to load rate limit script: %v", err))
}
instance = &RedisLimiter{
client: r,
limitScriptSHA: limitSHA,
}
})
return instance
}
func (rl *RedisLimiter) Allow(ctx context.Context, key string, opts ...Option) (bool, error) {
// 默认配置
config := &Config{
Capacity: 10,
Rate: 1,
Requested: 1,
}
// 应用选项模式
for _, opt := range opts {
opt(config)
}
// 执行限流
result, err := rl.client.EvalSha(
ctx,
rl.limitScriptSHA,
[]string{key},
config.Requested,
config.Rate,
config.Capacity,
).Int()
if err != nil {
return false, fmt.Errorf("rate limit failed: %w", err)
}
return result == 1, nil
}
// Config 配置选项模式
type Config struct {
Capacity int64
Rate int64
Requested int64
}
type Option func(*Config)
func WithCapacity(c int64) Option {
return func(cfg *Config) { cfg.Capacity = c }
}
func WithRate(r int64) Option {
return func(cfg *Config) { cfg.Rate = r }
}
func WithRequested(n int64) Option {
return func(cfg *Config) { cfg.Requested = n }
}
|