fb44c8bf03
包含Go API项目的所有源代码、配置文件、Docker配置、文档和前端资源 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
90 lines
1.5 KiB
Go
90 lines
1.5 KiB
Go
package limiter
|
|
|
|
import (
|
|
"context"
|
|
_ "embed"
|
|
"fmt"
|
|
"github.com/go-redis/redis/v8"
|
|
"one-api/common"
|
|
"sync"
|
|
)
|
|
|
|
//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 }
|
|
}
|