Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

111 строки
2.1 KiB

  1. package limiter
  2. import (
  3. "context"
  4. _ "embed"
  5. "fmt"
  6. "sync"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/go-redis/redis/v8"
  9. )
  10. //go:embed lua/rate_limit.lua
  11. var rateLimitScript string
  12. //go:embed lua/rate_refund.lua
  13. var rateRefundScript string
  14. type RedisLimiter struct {
  15. client *redis.Client
  16. limitScriptSHA string
  17. refundScriptSHA string
  18. }
  19. var (
  20. instance *RedisLimiter
  21. once sync.Once
  22. )
  23. func New(ctx context.Context, r *redis.Client) *RedisLimiter {
  24. once.Do(func() {
  25. // 预加载脚本
  26. limitSHA, err := r.ScriptLoad(ctx, rateLimitScript).Result()
  27. if err != nil {
  28. common.SysLog(fmt.Sprintf("Failed to load rate limit script: %v", err))
  29. }
  30. refundSHA, err := r.ScriptLoad(ctx, rateRefundScript).Result()
  31. if err != nil {
  32. common.SysLog(fmt.Sprintf("Failed to load rate refund script: %v", err))
  33. }
  34. instance = &RedisLimiter{
  35. client: r,
  36. limitScriptSHA: limitSHA,
  37. refundScriptSHA: refundSHA,
  38. }
  39. })
  40. return instance
  41. }
  42. func (rl *RedisLimiter) Allow(ctx context.Context, key string, opts ...Option) (bool, error) {
  43. // 默认配置
  44. config := &Config{
  45. Capacity: 10,
  46. Rate: 1,
  47. Requested: 1,
  48. }
  49. // 应用选项模式
  50. for _, opt := range opts {
  51. opt(config)
  52. }
  53. // 执行限流
  54. result, err := rl.client.EvalSha(
  55. ctx,
  56. rl.limitScriptSHA,
  57. []string{key},
  58. config.Requested,
  59. config.Rate,
  60. config.Capacity,
  61. ).Int()
  62. if err != nil {
  63. return false, fmt.Errorf("rate limit failed: %w", err)
  64. }
  65. return result == 1, nil
  66. }
  67. func (rl *RedisLimiter) Refund(ctx context.Context, key string, requested, capacity int64) error {
  68. _, err := rl.client.EvalSha(
  69. ctx,
  70. rl.refundScriptSHA,
  71. []string{key},
  72. requested,
  73. capacity,
  74. ).Int()
  75. return err
  76. }
  77. // Config 配置选项模式
  78. type Config struct {
  79. Capacity int64
  80. Rate int64
  81. Requested int64
  82. }
  83. type Option func(*Config)
  84. func WithCapacity(c int64) Option {
  85. return func(cfg *Config) { cfg.Capacity = c }
  86. }
  87. func WithRate(r int64) Option {
  88. return func(cfg *Config) { cfg.Rate = r }
  89. }
  90. func WithRequested(n int64) Option {
  91. return func(cfg *Config) { cfg.Requested = n }
  92. }