wangjia
3c687e5e0b
新包 internal/ratelimit：内存/redis(GCRA+Lua) 双实现 + 出错逐调用降级内存 （fail-open 到内存不 fail-closed）。REDIS_ADDR 空=内存模式，行为与既往一致； 配置后跨重启保状态、支持多实例。miniredis 全覆盖测试，零真实外部依赖。 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
84 lines
2.3 KiB
Go
84 lines
2.3 KiB
Go
// Package ratelimit 提供限流/登录失败锁/日配额计数的统一状态存储抽象。
|
||
//
|
||
// 两个实现：内存（默认，单实例）与 Redis（REDIS_ADDR 非空时启用，跨重启保状态、
|
||
// 支持未来多实例共享）。Redis 出错时逐调用降级到内存（fail-open 到内存而非
|
||
// fail-closed，防护不中断，见 fallback.go）。
|
||
package ratelimit
|
||
|
||
import (
|
||
"context"
|
||
"log"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
"golang.org/x/time/rate"
|
||
)
|
||
|
||
// Result 一次限流判定。RetryAfter 供 429 的 Retry-After 头。
|
||
type Result struct {
|
||
Allowed bool
|
||
RetryAfter time.Duration
|
||
}
|
||
|
||
// Limiter 速率限流（令牌桶/GCRA 语义：速率 r，突发 burst）。
|
||
type Limiter interface {
|
||
Allow(key string) Result
|
||
}
|
||
|
||
// Counter 累加计数（日配额等）：首次自增时设 TTL，返回累加后的值。
|
||
type Counter interface {
|
||
Incr(key string, ttl time.Duration) (int64, error)
|
||
}
|
||
|
||
// FailLocker 登录失败锁：max 次失败锁定 lockFor（max<=0 不锁）；
|
||
// 锁定即清零计数；成功登录 Reset。
|
||
type FailLocker interface {
|
||
Locked(key string) bool
|
||
RecordFailure(key string, max int, lockFor time.Duration)
|
||
Reset(key string)
|
||
}
|
||
|
||
// Store 聚合工厂。
|
||
type Store interface {
|
||
NewLimiter(r rate.Limit, burst int) Limiter
|
||
Counter() Counter
|
||
FailLocker() FailLocker
|
||
}
|
||
|
||
var (
|
||
defMu sync.Mutex
|
||
def Store
|
||
)
|
||
|
||
// Init 由 main.go 在 config.Load 之后、router.Setup 之前调用。
|
||
// addr 为空 → 内存模式；Redis Ping 失败 → 打警告并落内存（启动不因 Redis 挂而失败）。
|
||
func Init(addr string) {
|
||
defMu.Lock()
|
||
defer defMu.Unlock()
|
||
if addr == "" {
|
||
def = newMemoryStore()
|
||
return
|
||
}
|
||
rdb := redis.NewClient(&redis.Options{Addr: addr})
|
||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||
defer cancel()
|
||
if err := rdb.Ping(ctx).Err(); err != nil {
|
||
log.Printf("[ratelimit] redis %s 不可用（%v），降级为内存限流", addr, err)
|
||
def = newMemoryStore()
|
||
return
|
||
}
|
||
log.Printf("[ratelimit] 限流/登录锁状态外置 redis %s", addr)
|
||
def = newFallbackStore(rdb)
|
||
}
|
||
|
||
// Default 返回全局 Store；未 Init 时惰性落内存（测试零配置即用内存路径）。
|
||
func Default() Store {
|
||
defMu.Lock()
|
||
defer defMu.Unlock()
|
||
if def == nil {
|
||
def = newMemoryStore()
|
||
}
|
||
return def
|
||
}
|