weichen
29f14ac5f2
Details (original message): Add multi-instance groundwork: blobstore seam for artifacts, PG LISTEN/NOTIFY event fanout, and per-tenant rate limiting with native Redis backend
83 lines
2.0 KiB
Go
83 lines
2.0 KiB
Go
// Package ratelimit 提供按租户的限流（阶段 3）：
|
|
// - Memory：进程内固定窗口（默认，单实例档，零外部服务）；
|
|
// - Redis：多实例共享固定窗口（原生 Redis，设置 runtime.redisAddr 启用）。
|
|
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
type Limiter interface {
|
|
// Allow 判定 key（租户 owner）在窗口内第 N 次请求是否放行。
|
|
Allow(ctx context.Context, key string, limit int, window time.Duration) bool
|
|
}
|
|
|
|
// ---- 进程内固定窗口 ----
|
|
|
|
type memWindow struct {
|
|
n int
|
|
reset time.Time
|
|
}
|
|
|
|
type Memory struct {
|
|
mu sync.Mutex
|
|
windows map[string]*memWindow
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewMemory() *Memory {
|
|
return &Memory{windows: map[string]*memWindow{}, now: time.Now}
|
|
}
|
|
|
|
func (m *Memory) Allow(_ context.Context, key string, limit int, window time.Duration) bool {
|
|
if limit <= 0 {
|
|
return true
|
|
}
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
now := m.now()
|
|
w, ok := m.windows[key]
|
|
if !ok || !now.Before(w.reset) {
|
|
m.windows[key] = &memWindow{n: 1, reset: now.Add(window)}
|
|
return true
|
|
}
|
|
if w.n >= limit {
|
|
return false
|
|
}
|
|
w.n++
|
|
return true
|
|
}
|
|
|
|
// ---- Redis 固定窗口（INCR + EXPIRE；首击设窗）----
|
|
|
|
type Redis struct {
|
|
rdb *redis.Client
|
|
}
|
|
|
|
func NewRedis(addr, password string, db int) *Redis {
|
|
return &Redis{rdb: redis.NewClient(&redis.Options{Addr: addr, Password: password, DB: db})}
|
|
}
|
|
|
|
func (r *Redis) Allow(ctx context.Context, key string, limit int, window time.Duration) bool {
|
|
if limit <= 0 {
|
|
return true
|
|
}
|
|
rkey := "rl:" + key + ":" + time.Now().UTC().Truncate(window).Format("150405")
|
|
n, err := r.rdb.Incr(ctx, rkey).Result()
|
|
if err != nil {
|
|
return true // Redis 故障不阻断请求（限流是柔性行为，放行并记日志于调用方）
|
|
}
|
|
if n == 1 {
|
|
r.rdb.Expire(ctx, rkey, window+time.Second)
|
|
}
|
|
return n <= int64(limit)
|
|
}
|
|
|
|
func (r *Redis) Ping(ctx context.Context) error { return r.rdb.Ping(ctx).Err() }
|
|
|
|
func (r *Redis) Close() error { return r.rdb.Close() }
|