e2037ad756
This commit introduces a major architectural refactoring to improve quota management, centralize logging, and streamline the relay handling logic. Key changes: - **Pre-consume Quota:** Implements a new mechanism to check and reserve user quota *before* making the request to the upstream provider. This ensures more accurate quota deduction and prevents users from exceeding their limits due to concurrent requests. - **Unified Relay Handlers:** Refactors the relay logic to use generic handlers (e.g., `ChatHandler`, `ImageHandler`) instead of provider-specific implementations. This significantly reduces code duplication and simplifies adding new channels. - **Centralized Logger:** A new dedicated `logger` package is introduced, and all system logging calls are migrated to use it, moving this responsibility out of the `common` package. - **Code Reorganization:** DTOs are generalized (e.g., `dalle.go` -> `openai_image.go`) and utility code is moved to more appropriate packages (e.g., `common/http.go` -> `service/http.go`) for better code structure.
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/logger"
|
|
"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 {
|
|
logger.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 }
|
|
}
|