chii
eb972dfa93
单二进制三角色 + Dock 适配器 + Swarm stack 达到可部署态；mgr1 实测订阅经 central-proxy bootstrap，探活 alive=41/52。 Co-authored-by: Cursor <cursoragent@cursor.com>
60 lines
1.6 KiB
Go
60 lines
1.6 KiB
Go
// ratelimit.go：每 key 进程内令牌桶（rpm，design-arch §2.3）。
|
||
// gateway 无状态多副本部署时为「每副本」限流，接受该近似（无 Redis 约束）。
|
||
package gateway
|
||
|
||
import (
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// tokenBucket 单 key 令牌桶。
|
||
type tokenBucket struct {
|
||
tokens float64
|
||
last time.Time
|
||
rpm int
|
||
capacity float64
|
||
}
|
||
|
||
// RateLimiter 每 key rpm 令牌桶（进程内）。
|
||
type RateLimiter struct {
|
||
mu sync.Mutex
|
||
buckets map[int64]*tokenBucket
|
||
now func() time.Time
|
||
}
|
||
|
||
// NewRateLimiter 构造。
|
||
func NewRateLimiter() *RateLimiter {
|
||
return &RateLimiter{buckets: map[int64]*tokenBucket{}, now: time.Now}
|
||
}
|
||
|
||
// Allow 判定 key 是否可放行；返回 (是否放行, 剩余令牌, 重置秒数)。
|
||
// rpm<=0 视为不限流（放行）。
|
||
func (r *RateLimiter) Allow(keyID int64, rpm int) (bool, int, int) {
|
||
if rpm <= 0 {
|
||
return true, 0, 0
|
||
}
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
now := r.now()
|
||
b, ok := r.buckets[keyID]
|
||
if !ok || b.rpm != rpm {
|
||
b = &tokenBucket{tokens: float64(rpm), rpm: rpm, capacity: float64(rpm), last: now}
|
||
r.buckets[keyID] = b
|
||
}
|
||
// 补充令牌：速率 rpm/60 每秒
|
||
elapsed := now.Sub(b.last).Seconds()
|
||
b.last = now
|
||
b.tokens += elapsed * float64(rpm) / 60.0
|
||
if b.tokens > b.capacity {
|
||
b.tokens = b.capacity
|
||
}
|
||
if b.tokens < 1 {
|
||
// 需要等多少秒才凑满 1 个令牌
|
||
need := (1 - b.tokens) / (float64(rpm) / 60.0)
|
||
return false, 0, int(need) + 1
|
||
}
|
||
b.tokens--
|
||
remaining := int(b.tokens)
|
||
reset := int((b.capacity - b.tokens) / (float64(rpm) / 60.0))
|
||
return true, remaining, reset
|
||
}
|