Skip to content

Commit f648b8e

Browse files
author
QTom
committed
feat: increment RPM counter before request forwarding
1 parent 678c3ae commit f648b8e

2 files changed

Lines changed: 23 additions & 0 deletions

File tree

backend/internal/handler/gateway_handler.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
366366
// 账号槽位/等待计数需要在超时或断开时安全回收
367367
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
368368

369+
// RPM 计数递增(调度成功后、Forward 前)
370+
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
371+
if h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID) != nil {
372+
// 失败开放:不阻塞请求
373+
}
374+
}
375+
369376
// 转发请求 - 根据账号平台分流
370377
var result *service.ForwardResult
371378
requestCtx := c.Request.Context()
@@ -549,6 +556,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
549556
// 账号槽位/等待计数需要在超时或断开时安全回收
550557
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
551558

559+
// RPM 计数递增(调度成功后、Forward 前)
560+
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
561+
if h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID) != nil {
562+
// 失败开放:不阻塞请求
563+
}
564+
}
565+
552566
// 转发请求 - 根据账号平台分流
553567
var result *service.ForwardResult
554568
requestCtx := c.Request.Context()

backend/internal/service/gateway_service.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2241,6 +2241,15 @@ func (s *GatewayService) isAccountSchedulableForRPM(ctx context.Context, account
22412241
return true
22422242
}
22432243

2244+
// IncrementAccountRPM increments the RPM counter for the given account.
2245+
func (s *GatewayService) IncrementAccountRPM(ctx context.Context, accountID int64) error {
2246+
if s.rpmCache == nil {
2247+
return nil
2248+
}
2249+
_, err := s.rpmCache.IncrementRPM(ctx, accountID)
2250+
return err
2251+
}
2252+
22442253
// checkAndRegisterSession 检查并注册会话,用于会话数量限制
22452254
// 仅适用于 Anthropic OAuth/SetupToken 账号
22462255
// sessionID: 会话标识符(使用粘性会话的 hash)

0 commit comments

Comments
 (0)