李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
313 lines
9.0 KiB
Go
313 lines
9.0 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
cryptorand "crypto/rand"
|
||
"encoding/hex"
|
||
"fmt"
|
||
"math"
|
||
"math/rand/v2"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||
"github.com/tidwall/gjson"
|
||
)
|
||
|
||
// UserMsgQueueCache 用户消息串行队列 Redis 缓存接口
|
||
type UserMsgQueueCache interface {
|
||
// AcquireLock 尝试获取账号级串行锁
|
||
AcquireLock(ctx context.Context, accountID int64, requestID string, lockTtlMs int) (acquired bool, err error)
|
||
// ReleaseLock 释放锁并记录完成时间
|
||
ReleaseLock(ctx context.Context, accountID int64, requestID string) (released bool, err error)
|
||
// GetLastCompletedMs 获取上次完成时间（毫秒时间戳，Redis TIME 源）
|
||
GetLastCompletedMs(ctx context.Context, accountID int64) (int64, error)
|
||
// GetCurrentTimeMs 获取 Redis 服务器当前时间（毫秒），与 ReleaseLock 记录的时间源一致
|
||
GetCurrentTimeMs(ctx context.Context) (int64, error)
|
||
// ReconcileExpiredLockCandidates 处理锁索引中的到期候选，按真实 PTTL 清理或刷新索引
|
||
ReconcileExpiredLockCandidates(ctx context.Context, maxCount int) (cleaned int, err error)
|
||
}
|
||
|
||
// QueueLockResult 锁获取结果
|
||
type QueueLockResult struct {
|
||
Acquired bool
|
||
RequestID string
|
||
}
|
||
|
||
// UserMessageQueueService 用户消息串行队列服务
|
||
// 对真实用户消息实施账号级串行化 + RPM 自适应延迟
|
||
type UserMessageQueueService struct {
|
||
cache UserMsgQueueCache
|
||
rpmCache RPMCache
|
||
cfg *config.UserMessageQueueConfig
|
||
stopCh chan struct{} // graceful shutdown
|
||
stopOnce sync.Once // 确保 Stop() 并发安全
|
||
}
|
||
|
||
// NewUserMessageQueueService 创建用户消息串行队列服务
|
||
func NewUserMessageQueueService(cache UserMsgQueueCache, rpmCache RPMCache, cfg *config.UserMessageQueueConfig) *UserMessageQueueService {
|
||
return &UserMessageQueueService{
|
||
cache: cache,
|
||
rpmCache: rpmCache,
|
||
cfg: cfg,
|
||
stopCh: make(chan struct{}),
|
||
}
|
||
}
|
||
|
||
// IsRealUserMessage 检测是否为真实用户消息（非 tool_result）
|
||
// 与 claude-relay-service 的检测逻辑一致：
|
||
// 1. messages 非空
|
||
// 2. 最后一条消息 role == "user"
|
||
// 3. 最后一条消息 content（如果是数组）中不含 type:"tool_result" / "tool_use_result"
|
||
func IsRealUserMessage(parsed *ParsedRequest) bool {
|
||
if parsed == nil {
|
||
return false
|
||
}
|
||
messagesRaw := parsed.MessagesRaw()
|
||
if len(messagesRaw) == 0 {
|
||
return false
|
||
}
|
||
|
||
messages := gjson.ParseBytes(messagesRaw)
|
||
if !messages.IsArray() {
|
||
return false
|
||
}
|
||
lastMsg := gjson.Result{}
|
||
messages.ForEach(func(_, msg gjson.Result) bool {
|
||
lastMsg = msg
|
||
return true
|
||
})
|
||
if !lastMsg.Exists() || !lastMsg.IsObject() {
|
||
return false
|
||
}
|
||
if lastMsg.Get("role").String() != "user" {
|
||
return false
|
||
}
|
||
|
||
content := lastMsg.Get("content")
|
||
if !content.Exists() {
|
||
return true
|
||
}
|
||
if !content.IsArray() {
|
||
return true
|
||
}
|
||
|
||
isReal := true
|
||
content.ForEach(func(_, item gjson.Result) bool {
|
||
itemType := item.Get("type").String()
|
||
if itemType == "tool_result" || itemType == "tool_use_result" {
|
||
isReal = false
|
||
return false
|
||
}
|
||
return true
|
||
})
|
||
return isReal
|
||
}
|
||
|
||
// TryAcquire 尝试立即获取串行锁
|
||
func (s *UserMessageQueueService) TryAcquire(ctx context.Context, accountID int64) (*QueueLockResult, error) {
|
||
if s.cache == nil {
|
||
return &QueueLockResult{Acquired: true}, nil // fail-open
|
||
}
|
||
|
||
requestID := generateUMQRequestID()
|
||
lockTTL := s.cfg.LockTTLMs
|
||
if lockTTL <= 0 {
|
||
lockTTL = 120000
|
||
}
|
||
|
||
acquired, err := s.cache.AcquireLock(ctx, accountID, requestID, lockTTL)
|
||
if err != nil {
|
||
logger.LegacyPrintf("service.umq", "AcquireLock failed for account %d: %v", accountID, err)
|
||
return &QueueLockResult{Acquired: true}, nil // fail-open
|
||
}
|
||
|
||
return &QueueLockResult{
|
||
Acquired: acquired,
|
||
RequestID: requestID,
|
||
}, nil
|
||
}
|
||
|
||
// Release 释放串行锁
|
||
func (s *UserMessageQueueService) Release(ctx context.Context, accountID int64, requestID string) error {
|
||
if s.cache == nil || requestID == "" {
|
||
return nil
|
||
}
|
||
released, err := s.cache.ReleaseLock(ctx, accountID, requestID)
|
||
if err != nil {
|
||
logger.LegacyPrintf("service.umq", "ReleaseLock failed for account %d: %v", accountID, err)
|
||
return err
|
||
}
|
||
if !released {
|
||
logger.LegacyPrintf("service.umq", "ReleaseLock no-op for account %d (requestID mismatch or expired)", accountID)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// EnforceDelay 根据 RPM 负载执行自适应延迟
|
||
// 使用 Redis TIME 确保与 releaseLockScript 记录的时间源一致
|
||
func (s *UserMessageQueueService) EnforceDelay(ctx context.Context, accountID int64, baseRPM int) error {
|
||
if s.cache == nil {
|
||
return nil
|
||
}
|
||
|
||
// 先检查历史记录：没有历史则无需延迟，避免不必要的 RPM 查询
|
||
lastMs, err := s.cache.GetLastCompletedMs(ctx, accountID)
|
||
if err != nil {
|
||
logger.LegacyPrintf("service.umq", "GetLastCompletedMs failed for account %d: %v", accountID, err)
|
||
return nil // fail-open
|
||
}
|
||
if lastMs == 0 {
|
||
return nil // 没有历史记录，无需延迟
|
||
}
|
||
|
||
delay := s.CalculateRPMAwareDelay(ctx, accountID, baseRPM)
|
||
if delay <= 0 {
|
||
return nil
|
||
}
|
||
|
||
// 获取 Redis 当前时间（与 lastMs 同源，避免时钟偏差）
|
||
nowMs, err := s.cache.GetCurrentTimeMs(ctx)
|
||
if err != nil {
|
||
logger.LegacyPrintf("service.umq", "GetCurrentTimeMs failed: %v", err)
|
||
return nil // fail-open
|
||
}
|
||
|
||
elapsed := time.Duration(nowMs-lastMs) * time.Millisecond
|
||
if elapsed < 0 {
|
||
// 时钟异常（Redis 故障转移等），fail-open
|
||
return nil
|
||
}
|
||
remaining := delay - elapsed
|
||
if remaining <= 0 {
|
||
return nil
|
||
}
|
||
|
||
// 执行延迟
|
||
timer := time.NewTimer(remaining)
|
||
defer timer.Stop()
|
||
select {
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
case <-timer.C:
|
||
return nil
|
||
}
|
||
}
|
||
|
||
// CalculateRPMAwareDelay 根据当前 RPM 负载计算自适应延迟
|
||
// ratio = currentRPM / baseRPM
|
||
// ratio < 0.5 → MinDelay
|
||
// 0.5 ≤ ratio < 0.8 → 线性插值 MinDelay..MaxDelay
|
||
// ratio ≥ 0.8 → MaxDelay
|
||
// 返回值包含 ±15% 随机抖动（anti-detection + 避免惊群效应）
|
||
func (s *UserMessageQueueService) CalculateRPMAwareDelay(ctx context.Context, accountID int64, baseRPM int) time.Duration {
|
||
minDelay := time.Duration(s.cfg.MinDelayMs) * time.Millisecond
|
||
maxDelay := time.Duration(s.cfg.MaxDelayMs) * time.Millisecond
|
||
|
||
if minDelay <= 0 {
|
||
minDelay = 200 * time.Millisecond
|
||
}
|
||
if maxDelay <= 0 {
|
||
maxDelay = 2000 * time.Millisecond
|
||
}
|
||
// 防止配置错误：minDelay > maxDelay 时交换
|
||
if minDelay > maxDelay {
|
||
minDelay, maxDelay = maxDelay, minDelay
|
||
}
|
||
|
||
var baseDelay time.Duration
|
||
|
||
if baseRPM <= 0 || s.rpmCache == nil {
|
||
baseDelay = minDelay
|
||
} else {
|
||
currentRPM, err := s.rpmCache.GetRPM(ctx, accountID)
|
||
if err != nil {
|
||
logger.LegacyPrintf("service.umq", "GetRPM failed for account %d: %v", accountID, err)
|
||
baseDelay = minDelay // fail-open
|
||
} else {
|
||
ratio := float64(currentRPM) / float64(baseRPM)
|
||
if ratio < 0.5 {
|
||
baseDelay = minDelay
|
||
} else if ratio >= 0.8 {
|
||
baseDelay = maxDelay
|
||
} else {
|
||
// 线性插值: 0.5 → minDelay, 0.8 → maxDelay
|
||
t := (ratio - 0.5) / 0.3
|
||
interpolated := float64(minDelay) + t*(float64(maxDelay)-float64(minDelay))
|
||
baseDelay = time.Duration(math.Round(interpolated))
|
||
}
|
||
}
|
||
}
|
||
|
||
// ±15% 随机抖动
|
||
return applyJitter(baseDelay, 0.15)
|
||
}
|
||
|
||
// StartCleanupWorker 启动孤儿锁清理 worker。
|
||
// worker 只处理锁索引中的到期候选，真正删除前由 cache 层再次校验锁 PTTL。
|
||
func (s *UserMessageQueueService) StartCleanupWorker(interval time.Duration) {
|
||
if s == nil || s.cache == nil || interval <= 0 {
|
||
return
|
||
}
|
||
|
||
runCleanup := func() {
|
||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
defer cancel()
|
||
|
||
// 每轮限制处理数量，避免清理任务在大量过期候选时长时间占用 Redis。
|
||
cleaned, err := s.cache.ReconcileExpiredLockCandidates(ctx, 1000)
|
||
if err != nil {
|
||
logger.LegacyPrintf("service.umq", "Cleanup reconcile failed: %v", err)
|
||
return
|
||
}
|
||
|
||
if cleaned > 0 {
|
||
logger.LegacyPrintf("service.umq", "Cleanup completed: released %d orphaned locks", cleaned)
|
||
}
|
||
}
|
||
|
||
go func() {
|
||
ticker := time.NewTicker(interval)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-s.stopCh:
|
||
return
|
||
case <-ticker.C:
|
||
runCleanup()
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
// Stop 停止后台 cleanup worker
|
||
func (s *UserMessageQueueService) Stop() {
|
||
if s != nil && s.stopCh != nil {
|
||
s.stopOnce.Do(func() {
|
||
close(s.stopCh)
|
||
})
|
||
}
|
||
}
|
||
|
||
// applyJitter 对延迟值施加 ±jitterPct 的随机抖动
|
||
// 使用 math/rand/v2（Go 1.22+ 自动使用 crypto/rand 种子），与 nextBackoff 一致
|
||
// 例如 applyJitter(200ms, 0.15) 返回 170ms ~ 230ms
|
||
func applyJitter(d time.Duration, jitterPct float64) time.Duration {
|
||
if d <= 0 || jitterPct <= 0 {
|
||
return d
|
||
}
|
||
// [-jitterPct, +jitterPct]
|
||
jitter := (rand.Float64()*2 - 1) * jitterPct
|
||
return time.Duration(float64(d) * (1 + jitter))
|
||
}
|
||
|
||
// generateUMQRequestID 生成唯一请求 ID（与 generateRequestID 一致的 fallback 模式）
|
||
func generateUMQRequestID() string {
|
||
b := make([]byte, 16)
|
||
if _, err := cryptorand.Read(b); err != nil {
|
||
return fmt.Sprintf("%x", time.Now().UnixNano())
|
||
}
|
||
return hex.EncodeToString(b)
|
||
}
|