diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index a6fb5266aa..a412563a6c 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -67,7 +67,11 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { serviceUserPlatformQuotaRepository := repository.NewUserPlatformQuotaServiceAdapter(userPlatformQuotaRepository) billingCacheService := service.ProvideBillingCacheService(billingCache, userRepository, userSubscriptionRepository, apiKeyRepository, userRPMCache, userGroupRateRepository, configConfig, serviceUserPlatformQuotaRepository) apiKeyCache := repository.NewAPIKeyCache(redisClient) - apiKeyService := service.ProvideAPIKeyService(apiKeyRepository, userRepository, groupRepository, userSubscriptionRepository, userGroupRateRepository, apiKeyCache, configConfig, billingCacheService) + concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig) + schedulerCache := repository.ProvideSchedulerCache(redisClient, configConfig) + accountRepository := repository.NewAccountRepository(client, db, schedulerCache) + concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig) + apiKeyService := service.ProvideAPIKeyService(apiKeyRepository, userRepository, groupRepository, userSubscriptionRepository, userGroupRateRepository, apiKeyCache, configConfig, billingCacheService, concurrencyService) apiKeyAuthCacheInvalidator := service.ProvideAPIKeyAuthCacheInvalidator(apiKeyService) promoService := service.NewPromoService(promoCodeRepository, userRepository, billingCacheService, client, apiKeyAuthCacheInvalidator) subscriptionService := service.NewSubscriptionService(groupRepository, userSubscriptionRepository, billingCacheService, client, configConfig) @@ -92,10 +96,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { usageLogRepository := repository.NewUsageLogRepository(client, db) usageService := service.NewUsageService(usageLogRepository, userRepository, client, apiKeyAuthCacheInvalidator) opsRepository := repository.NewOpsRepository(db) - schedulerCache := repository.ProvideSchedulerCache(redisClient, configConfig) - accountRepository := repository.NewAccountRepository(client, db, schedulerCache) - concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig) - concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig) usageBillingRepository := repository.NewUsageBillingRepository(client, db) gatewayCache := repository.NewGatewayCache(redisClient) schedulerOutboxRepository := repository.NewSchedulerOutboxRepository(db) diff --git a/backend/internal/handler/dto/api_key_mapper_last_used_test.go b/backend/internal/handler/dto/api_key_mapper_last_used_test.go index 99644ced7f..d63baba91a 100644 --- a/backend/internal/handler/dto/api_key_mapper_last_used_test.go +++ b/backend/internal/handler/dto/api_key_mapper_last_used_test.go @@ -11,18 +11,20 @@ import ( func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) { lastUsed := time.Now().UTC().Truncate(time.Second) src := &service.APIKey{ - ID: 1, - UserID: 2, - Key: "sk-map-last-used", - Name: "Mapper", - Status: service.StatusActive, - LastUsedAt: &lastUsed, + ID: 1, + UserID: 2, + Key: "sk-map-last-used", + Name: "Mapper", + Status: service.StatusActive, + LastUsedAt: &lastUsed, + CurrentConcurrency: 3, } out := APIKeyFromService(src) require.NotNil(t, out) require.NotNil(t, out.LastUsedAt) require.WithinDuration(t, lastUsed, *out.LastUsedAt, time.Second) + require.Equal(t, 3, out.CurrentConcurrency) } func TestAPIKeyFromService_MapsNilLastUsedAt(t *testing.T) { diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 72679de37b..02a0a38d7b 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -79,31 +79,32 @@ func APIKeyFromService(k *service.APIKey) *APIKey { return nil } out := &APIKey{ - ID: k.ID, - UserID: k.UserID, - Key: k.Key, - Name: k.Name, - GroupID: k.GroupID, - Status: k.Status, - IPWhitelist: k.IPWhitelist, - IPBlacklist: k.IPBlacklist, - LastUsedAt: k.LastUsedAt, - Quota: k.Quota, - QuotaUsed: k.QuotaUsed, - ExpiresAt: k.ExpiresAt, - CreatedAt: k.CreatedAt, - UpdatedAt: k.UpdatedAt, - RateLimit5h: k.RateLimit5h, - RateLimit1d: k.RateLimit1d, - RateLimit7d: k.RateLimit7d, - Usage5h: k.EffectiveUsage5h(), - Usage1d: k.EffectiveUsage1d(), - Usage7d: k.EffectiveUsage7d(), - Window5hStart: k.Window5hStart, - Window1dStart: k.Window1dStart, - Window7dStart: k.Window7dStart, - User: UserFromServiceShallow(k.User), - Group: GroupFromServiceShallow(k.Group), + ID: k.ID, + UserID: k.UserID, + Key: k.Key, + Name: k.Name, + GroupID: k.GroupID, + Status: k.Status, + IPWhitelist: k.IPWhitelist, + IPBlacklist: k.IPBlacklist, + LastUsedAt: k.LastUsedAt, + Quota: k.Quota, + QuotaUsed: k.QuotaUsed, + ExpiresAt: k.ExpiresAt, + CreatedAt: k.CreatedAt, + UpdatedAt: k.UpdatedAt, + CurrentConcurrency: k.CurrentConcurrency, + RateLimit5h: k.RateLimit5h, + RateLimit1d: k.RateLimit1d, + RateLimit7d: k.RateLimit7d, + Usage5h: k.EffectiveUsage5h(), + Usage1d: k.EffectiveUsage1d(), + Usage7d: k.EffectiveUsage7d(), + Window5hStart: k.Window5hStart, + Window1dStart: k.Window1dStart, + Window7dStart: k.Window7dStart, + User: UserFromServiceShallow(k.User), + Group: GroupFromServiceShallow(k.Group), } if k.Window5hStart != nil && !service.IsWindowExpired(k.Window5hStart, service.RateLimitWindow5h) { t := k.Window5hStart.Add(service.RateLimitWindow5h) diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index cef84465ef..06dce2beaa 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -63,6 +63,8 @@ type APIKey struct { ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = never expires) CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` + // CurrentConcurrency is the real-time active request count for this API key. + CurrentConcurrency int `json:"current_concurrency"` // Rate limit fields RateLimit5h float64 `json:"rate_limit_5h"` diff --git a/backend/internal/handler/gateway_helper.go b/backend/internal/handler/gateway_helper.go index b948ac8fc7..48110da93f 100644 --- a/backend/internal/handler/gateway_helper.go +++ b/backend/internal/handler/gateway_helper.go @@ -10,6 +10,7 @@ import ( "sync" "time" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" @@ -211,6 +212,14 @@ func (h *ConcurrencyHelper) TryAcquireUserSlot(ctx context.Context, userID int64 return result.ReleaseFunc, true, nil } +func (h *ConcurrencyHelper) TryAcquireUserSlotForAPIKey(ctx context.Context, userID int64, maxConcurrency int, apiKeyID int64) (func(), bool, error) { + releaseFunc, acquired, err := h.TryAcquireUserSlot(ctx, userID, maxConcurrency) + if err != nil || !acquired { + return releaseFunc, acquired, err + } + return h.withAPIKeySlot(ctx, apiKeyID, releaseFunc), true, nil +} + // TryAcquireAccountSlot 尝试立即获取账号并发槽位。 // 返回值: (releaseFunc, acquired, error) func (h *ConcurrencyHelper) TryAcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (func(), bool, error) { @@ -241,7 +250,7 @@ func (h *ConcurrencyHelper) acquireUserSlotWithWaitTimeout(c *gin.Context, userI } if acquired { - return releaseFunc, nil + return h.withAPIKeySlotFromGin(c, releaseFunc), nil } queueLimit := service.CalculateMaxWait(maxConcurrency) - maxConcurrency @@ -258,7 +267,37 @@ func (h *ConcurrencyHelper) acquireUserSlotWithWaitTimeout(c *gin.Context, userI defer h.DecrementWaitCount(ctx, userID) // Need to wait - handle streaming ping if needed - return h.waitForSlotWithPingTimeout(c, "user", userID, maxConcurrency, timeout, isStream, streamStarted, false) + releaseFunc, err = h.waitForSlotWithPingTimeout(c, "user", userID, maxConcurrency, timeout, isStream, streamStarted, false) + if err != nil { + return nil, err + } + return h.withAPIKeySlotFromGin(c, releaseFunc), nil +} + +func (h *ConcurrencyHelper) withAPIKeySlotFromGin(c *gin.Context, releaseFunc func()) func() { + if c == nil { + return releaseFunc + } + apiKey, ok := middleware2.GetAPIKeyFromContext(c) + if !ok || apiKey == nil { + return releaseFunc + } + return h.withAPIKeySlot(c.Request.Context(), apiKey.ID, releaseFunc) +} + +func (h *ConcurrencyHelper) withAPIKeySlot(ctx context.Context, apiKeyID int64, releaseFunc func()) func() { + if h == nil || h.concurrencyService == nil || apiKeyID <= 0 { + return releaseFunc + } + apiKeyReleaseFunc := h.concurrencyService.TrackAPIKeySlot(ctx, apiKeyID) + return func() { + if releaseFunc != nil { + releaseFunc() + } + if apiKeyReleaseFunc != nil { + apiKeyReleaseFunc() + } + } } // AcquireAccountSlotWithWait acquires an account concurrency slot, waiting if necessary. diff --git a/backend/internal/handler/gateway_helper_hotpath_test.go b/backend/internal/handler/gateway_helper_hotpath_test.go index 65dc849683..2f0c3261c7 100644 --- a/backend/internal/handler/gateway_helper_hotpath_test.go +++ b/backend/internal/handler/gateway_helper_hotpath_test.go @@ -9,6 +9,7 @@ import ( "testing" "time" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -29,6 +30,9 @@ type helperConcurrencyCacheStub struct { waitDecrementCalls int waitMaxWait int waitIncrementHook func() + apiKeyTrackCalls int + apiKeyReleaseCalls int + apiKeyTrackIDs []int64 } func (s *helperConcurrencyCacheStub) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) { @@ -97,6 +101,29 @@ func (s *helperConcurrencyCacheStub) GetUserConcurrency(ctx context.Context, use return 0, nil } +func (s *helperConcurrencyCacheStub) TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.apiKeyTrackCalls++ + s.apiKeyTrackIDs = append(s.apiKeyTrackIDs, apiKeyID) + return nil +} + +func (s *helperConcurrencyCacheStub) ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.apiKeyReleaseCalls++ + return nil +} + +func (s *helperConcurrencyCacheStub) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) { + out := make(map[int64]int, len(apiKeyIDs)) + for _, apiKeyID := range apiKeyIDs { + out[apiKeyID] = 0 + } + return out, nil +} + func (s *helperConcurrencyCacheStub) IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error) { s.mu.Lock() s.waitIncrementCalls++ @@ -266,6 +293,48 @@ func TestAcquireUserSlotWithWait_ImmediateAcquireSkipsWaitQueue(t *testing.T) { require.Equal(t, 1, cache.userReleaseCalls) } +func TestAcquireUserSlotWithWait_TracksAPIKeySlot(t *testing.T) { + cache := &helperConcurrencyCacheStub{ + userSeq: []bool{true}, + } + concurrency := service.NewConcurrencyService(cache) + helper := NewConcurrencyHelper(concurrency, SSEPingFormatNone, 5*time.Millisecond) + c, _ := newHelperTestContext(http.MethodPost, "/v1/messages") + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ID: 77}) + streamStarted := false + + release, err := helper.acquireUserSlotWithWaitTimeout(c, 202, 3, time.Second, false, &streamStarted) + require.NoError(t, err) + require.NotNil(t, release) + require.Equal(t, 1, cache.apiKeyTrackCalls) + require.Equal(t, []int64{77}, cache.apiKeyTrackIDs) + + release() + + require.Equal(t, 1, cache.userReleaseCalls) + require.Equal(t, 1, cache.apiKeyReleaseCalls) +} + +func TestTryAcquireUserSlotForAPIKey_TracksAPIKeySlot(t *testing.T) { + cache := &helperConcurrencyCacheStub{ + userSeq: []bool{true}, + } + concurrency := service.NewConcurrencyService(cache) + helper := NewConcurrencyHelper(concurrency, SSEPingFormatNone, 5*time.Millisecond) + + release, acquired, err := helper.TryAcquireUserSlotForAPIKey(context.Background(), 202, 3, 77) + require.NoError(t, err) + require.True(t, acquired) + require.NotNil(t, release) + require.Equal(t, 1, cache.apiKeyTrackCalls) + require.Equal(t, []int64{77}, cache.apiKeyTrackIDs) + + release() + + require.Equal(t, 1, cache.userReleaseCalls) + require.Equal(t, 1, cache.apiKeyReleaseCalls) +} + func TestAcquireUserSlotWithWait_WaitSuccessDecrementsBeforeReturn(t *testing.T) { cache := &helperConcurrencyCacheStub{ userSeq: []bool{false, true}, diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index ac346fea3f..995c29f936 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -1318,7 +1318,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // 必须尽早注册，确保任何 early return 都能释放已获取的并发槽位。 defer releaseTurnSlots() - userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlot(ctx, subject.UserID, subject.Concurrency) + userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlotForAPIKey(ctx, subject.UserID, subject.Concurrency, apiKey.ID) if err != nil { reqLog.Warn("openai.websocket_user_slot_acquire_failed", zap.Error(err)) closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to acquire user concurrency slot") @@ -1333,7 +1333,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { if currentUserRelease != nil { return true } - userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlot(ctx, subject.UserID, subject.Concurrency) + userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlotForAPIKey(ctx, subject.UserID, subject.Concurrency, apiKey.ID) if err != nil { reqLog.Warn("openai.websocket_user_slot_reacquire_failed", zap.Error(err)) closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to acquire user concurrency slot") @@ -1484,7 +1484,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // 防御式清理：避免异常路径下旧槽位覆盖导致泄漏。 releaseTurnSlots() // 非首轮 turn 需要重新抢占并发槽位，避免长连接空闲占槽。 - userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlot(ctx, subject.UserID, subject.Concurrency) + userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlotForAPIKey(ctx, subject.UserID, subject.Concurrency, apiKey.ID) if err != nil { return service.NewOpenAIWSClientCloseError(coderws.StatusInternalError, "failed to acquire user concurrency slot", err) } diff --git a/backend/internal/repository/concurrency_cache.go b/backend/internal/repository/concurrency_cache.go index 5e6f10062f..eb827309a3 100644 --- a/backend/internal/repository/concurrency_cache.go +++ b/backend/internal/repository/concurrency_cache.go @@ -27,6 +27,8 @@ const ( accountSlotKeyPrefix = "concurrency:account:" // 格式: concurrency:user:{userID} userSlotKeyPrefix = "concurrency:user:" + // 格式: concurrency:api_key:{apiKeyID} + apiKeySlotKeyPrefix = "concurrency:api_key:" // 等待队列计数器格式: concurrency:wait:{userID} waitQueueKeyPrefix = "concurrency:wait:" // 账号级等待队列计数器格式: wait:account:{accountID} @@ -99,6 +101,28 @@ var ( return redis.call('ZCARD', key) `) + // trackSlotScript 记录 stats-only 槽位，不做并发上限判断。 + // KEYS[1] = 有序集合键 + // ARGV[1] = TTL（秒） + // ARGV[2] = requestID + trackSlotScript = redis.NewScript(` + -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME') + -- replicates correctly. No-op on Redis 5.0+ (effects replication is default). + redis.replicate_commands() + local key = KEYS[1] + local ttl = tonumber(ARGV[1]) + local requestID = ARGV[2] + + local timeResult = redis.call('TIME') + local now = tonumber(timeResult[1]) + local expireBefore = now - ttl + + redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore) + redis.call('ZADD', key, now, requestID) + redis.call('EXPIRE', key, ttl) + return 1 + `) + // incrementWaitScript - refreshes TTL on each increment to keep queue depth accurate // KEYS[1] = wait queue key // ARGV[1] = maxWait @@ -231,6 +255,10 @@ func userSlotKey(userID int64) string { return fmt.Sprintf("%s%d", userSlotKeyPrefix, userID) } +func apiKeySlotKey(apiKeyID int64) string { + return fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID) +} + func waitQueueKey(userID int64) string { return fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID) } @@ -330,6 +358,54 @@ func (c *concurrencyCache) GetUserConcurrency(ctx context.Context, userID int64) return result, nil } +func (c *concurrencyCache) TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error { + key := apiKeySlotKey(apiKeyID) + _, err := trackSlotScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds, requestID).Result() + return err +} + +func (c *concurrencyCache) ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error { + key := apiKeySlotKey(apiKeyID) + return c.rdb.ZRem(ctx, key, requestID).Err() +} + +func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) { + if len(apiKeyIDs) == 0 { + return map[int64]int{}, nil + } + + now, err := c.rdb.Time(ctx).Result() + if err != nil { + return nil, fmt.Errorf("redis TIME: %w", err) + } + cutoffTime := now.Unix() - int64(c.slotTTLSeconds) + + pipe := c.rdb.Pipeline() + type apiKeyCmd struct { + apiKeyID int64 + zcardCmd *redis.IntCmd + } + cmds := make([]apiKeyCmd, 0, len(apiKeyIDs)) + for _, apiKeyID := range apiKeyIDs { + slotKey := apiKeySlotKeyPrefix + strconv.FormatInt(apiKeyID, 10) + pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10)) + cmds = append(cmds, apiKeyCmd{ + apiKeyID: apiKeyID, + zcardCmd: pipe.ZCard(ctx, slotKey), + }) + } + + if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) { + return nil, fmt.Errorf("pipeline exec: %w", err) + } + + result := make(map[int64]int, len(apiKeyIDs)) + for _, cmd := range cmds { + result[cmd.apiKeyID] = int(cmd.zcardCmd.Val()) + } + return result, nil +} + // Wait queue operations func (c *concurrencyCache) IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error) { @@ -509,7 +585,7 @@ func (c *concurrencyCache) CleanupStaleProcessSlots(ctx context.Context, activeR } // 1. 清理有序集合中非当前进程前缀的成员 - slotPatterns := []string{accountSlotKeyPrefix + "*", userSlotKeyPrefix + "*"} + slotPatterns := []string{accountSlotKeyPrefix + "*", userSlotKeyPrefix + "*", apiKeySlotKeyPrefix + "*"} for _, pattern := range slotPatterns { if err := c.cleanupSlotsByPattern(ctx, pattern, activeRequestPrefix); err != nil { return err diff --git a/backend/internal/repository/concurrency_cache_integration_test.go b/backend/internal/repository/concurrency_cache_integration_test.go index 5da94fc258..bcc5249c84 100644 --- a/backend/internal/repository/concurrency_cache_integration_test.go +++ b/backend/internal/repository/concurrency_cache_integration_test.go @@ -3,6 +3,7 @@ package repository import ( + "context" "errors" "fmt" "testing" @@ -34,6 +35,18 @@ func (s *ConcurrencyCacheSuite) SetupTest() { s.cache = NewConcurrencyCache(s.rdb, testSlotTTLMinutes, int(testSlotTTL.Seconds())) } +type apiKeyConcurrencyCacheForTest interface { + TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error + ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error + GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) +} + +func (s *ConcurrencyCacheSuite) apiKeyConcurrencyCache() apiKeyConcurrencyCacheForTest { + cache, ok := s.cache.(apiKeyConcurrencyCacheForTest) + require.True(s.T(), ok) + return cache +} + func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() { accountID := int64(10) reqID1, reqID2, reqID3 := "req1", "req2", "req3" @@ -160,6 +173,34 @@ func (s *ConcurrencyCacheSuite) TestUserSlot_TTL() { s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL) } +func (s *ConcurrencyCacheSuite) TestAPIKeySlot_TrackReleaseAndBatchCount() { + cache := s.apiKeyConcurrencyCache() + apiKeyID := int64(300) + emptyAPIKeyID := int64(301) + slotKey := fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID) + + require.NoError(s.T(), cache.TrackAPIKeySlot(s.ctx, apiKeyID, "req1")) + require.NoError(s.T(), cache.TrackAPIKeySlot(s.ctx, apiKeyID, "req2")) + + counts, err := cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID, emptyAPIKeyID}) + require.NoError(s.T(), err) + require.Equal(s.T(), map[int64]int{apiKeyID: 2, emptyAPIKeyID: 0}, counts) + + ttl, err := s.rdb.TTL(s.ctx, slotKey).Result() + require.NoError(s.T(), err, "TTL") + s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL) + + require.NoError(s.T(), cache.ReleaseAPIKeySlot(s.ctx, apiKeyID, "req1")) + counts, err = cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID}) + require.NoError(s.T(), err) + require.Equal(s.T(), 1, counts[apiKeyID]) + + require.NoError(s.T(), cache.ReleaseAPIKeySlot(s.ctx, apiKeyID, "req2")) + counts, err = cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID}) + require.NoError(s.T(), err) + require.Equal(s.T(), 0, counts[apiKeyID]) +} + func (s *ConcurrencyCacheSuite) TestWaitQueue_IncrementAndDecrement() { userID := int64(20) waitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID) @@ -254,8 +295,10 @@ func (s *ConcurrencyCacheSuite) TestAccountWaitQueue_IncrementAndDecrement() { func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() { accountID := int64(901) userID := int64(902) + apiKeyID := int64(903) accountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID) userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, userID) + apiKeyKey := fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID) userWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID) accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID) @@ -268,6 +311,10 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() { redis.Z{Score: float64(now), Member: "oldproc-2"}, redis.Z{Score: float64(now), Member: "keep-2"}, ).Err()) + require.NoError(s.T(), s.rdb.ZAdd(s.ctx, apiKeyKey, + redis.Z{Score: float64(now), Member: "oldproc-3"}, + redis.Z{Score: float64(now), Member: "keep-3"}, + ).Err()) require.NoError(s.T(), s.rdb.Set(s.ctx, userWaitKey, 3, time.Minute).Err()) require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 2, time.Minute).Err()) @@ -281,6 +328,10 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() { require.NoError(s.T(), err) require.Equal(s.T(), []string{"keep-2"}, userMembers) + apiKeyMembers, err := s.rdb.ZRange(s.ctx, apiKeyKey, 0, -1).Result() + require.NoError(s.T(), err) + require.Equal(s.T(), []string{"keep-3"}, apiKeyMembers) + _, err = s.rdb.Get(s.ctx, userWaitKey).Result() require.True(s.T(), errors.Is(err, redis.Nil)) diff --git a/backend/internal/service/api_key.go b/backend/internal/service/api_key.go index ec20b0a9bf..dfc3ec1c5a 100644 --- a/backend/internal/service/api_key.go +++ b/backend/internal/service/api_key.go @@ -44,6 +44,7 @@ type APIKey struct { UpdatedAt time.Time User *User Group *Group + CurrentConcurrency int // Quota fields Quota float64 // Quota limit in USD (0 = unlimited) diff --git a/backend/internal/service/api_key_service.go b/backend/internal/service/api_key_service.go index de9b908dc7..8903be65ee 100644 --- a/backend/internal/service/api_key_service.go +++ b/backend/internal/service/api_key_service.go @@ -203,6 +203,7 @@ type APIKeyService struct { userGroupRateRepo UserGroupRateRepository cache APIKeyCache rateLimitCacheInvalid RateLimitCacheInvalidator // optional: invalidate Redis rate limit cache + concurrencyService *ConcurrencyService cfg *config.Config authCacheL1 *ristretto.Cache authCfg apiKeyAuthCacheConfig @@ -240,6 +241,10 @@ func (s *APIKeyService) SetRateLimitCacheInvalidator(inv RateLimitCacheInvalidat s.rateLimitCacheInvalid = inv } +func (s *APIKeyService) SetConcurrencyService(concurrencyService *ConcurrencyService) { + s.concurrencyService = concurrencyService +} + func (s *APIKeyService) compileAPIKeyIPRules(apiKey *APIKey) { if apiKey == nil { return @@ -436,9 +441,40 @@ func (s *APIKeyService) List(ctx context.Context, userID int64, params paginatio if err != nil { return nil, nil, fmt.Errorf("list api keys: %w", err) } + s.fillCurrentConcurrency(ctx, keys) return keys, pagination, nil } +func (s *APIKeyService) fillCurrentConcurrency(ctx context.Context, keys []APIKey) { + if s == nil || s.concurrencyService == nil || len(keys) == 0 { + return + } + ids := make([]int64, 0, len(keys)) + for i := range keys { + if keys[i].ID > 0 { + ids = append(ids, keys[i].ID) + } + } + counts, err := s.concurrencyService.GetAPIKeyConcurrencyBatch(ctx, ids) + if err != nil { + return + } + for i := range keys { + keys[i].CurrentConcurrency = counts[keys[i].ID] + } +} + +func (s *APIKeyService) currentConcurrencyForAPIKey(ctx context.Context, apiKeyID int64) int { + if s == nil || s.concurrencyService == nil || apiKeyID <= 0 { + return 0 + } + counts, err := s.concurrencyService.GetAPIKeyConcurrencyBatch(ctx, []int64{apiKeyID}) + if err != nil { + return 0 + } + return counts[apiKeyID] +} + func (s *APIKeyService) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) { if len(apiKeyIDs) == 0 { return []int64{}, nil @@ -458,6 +494,9 @@ func (s *APIKeyService) GetByID(ctx context.Context, id int64) (*APIKey, error) return nil, fmt.Errorf("get api key: %w", err) } s.compileAPIKeyIPRules(apiKey) + if apiKey != nil { + apiKey.CurrentConcurrency = s.currentConcurrencyForAPIKey(ctx, apiKey.ID) + } return apiKey, nil } diff --git a/backend/internal/service/api_key_service_delete_test.go b/backend/internal/service/api_key_service_delete_test.go index 8664c03bd7..25ad1edb15 100644 --- a/backend/internal/service/api_key_service_delete_test.go +++ b/backend/internal/service/api_key_service_delete_test.go @@ -300,6 +300,40 @@ func TestApiKeyService_Delete_NotFound(t *testing.T) { require.Empty(t, cache.deleteAuthKeys) } +func TestAPIKeyService_List_FillsCurrentConcurrency(t *testing.T) { + repo := &apiKeyRepoStub{ + allowListByUserID: true, + listByUserIDKeys: []APIKey{ + {ID: 10, UserID: 7, Key: "sk-10", Name: "key-10"}, + {ID: 11, UserID: 7, Key: "sk-11", Name: "key-11"}, + }, + } + concurrency := NewConcurrencyService(&stubConcurrencyCacheForTest{ + apiKeyConcurrency: map[int64]int{10: 2, 11: 0}, + }) + svc := &APIKeyService{apiKeyRepo: repo, concurrencyService: concurrency} + + keys, _, err := svc.List(context.Background(), 7, pagination.PaginationParams{Page: 1, PageSize: 20}, APIKeyListFilters{}) + require.NoError(t, err) + require.Len(t, keys, 2) + require.Equal(t, 2, keys[0].CurrentConcurrency) + require.Equal(t, 0, keys[1].CurrentConcurrency) +} + +func TestAPIKeyService_GetByID_FillsCurrentConcurrency(t *testing.T) { + repo := &apiKeyRepoStub{ + apiKey: &APIKey{ID: 10, UserID: 7, Key: "sk-10", Name: "key-10"}, + } + concurrency := NewConcurrencyService(&stubConcurrencyCacheForTest{ + apiKeyConcurrency: map[int64]int{10: 4}, + }) + svc := &APIKeyService{apiKeyRepo: repo, concurrencyService: concurrency} + + key, err := svc.GetByID(context.Background(), 10) + require.NoError(t, err) + require.Equal(t, 4, key.CurrentConcurrency) +} + // TestApiKeyService_Delete_DeleteFails 测试删除操作失败时的错误处理。 // 预期行为： // - GetKeyAndOwnerID 返回正确的所有者 ID diff --git a/backend/internal/service/concurrency_service.go b/backend/internal/service/concurrency_service.go index 712fc1a749..49b34b2457 100644 --- a/backend/internal/service/concurrency_service.go +++ b/backend/internal/service/concurrency_service.go @@ -52,6 +52,12 @@ type ConcurrencyCache interface { CleanupStaleProcessSlots(ctx context.Context, activeRequestPrefix string) error } +type APIKeyConcurrencyCache interface { + TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error + ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error + GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) +} + var ( requestIDPrefix = initRequestIDPrefix() requestIDCounter atomic.Uint64 @@ -89,6 +95,8 @@ const ( defaultAccountLoadBatchCacheTTL = 200 * time.Millisecond accountLoadBatchFetchTimeout = 3 * time.Second maxAccountLoadBatchCacheEntries = 256 + apiKeyConcurrencyFetchTimeout = 3 * time.Second + apiKeySlotTrackTimeout = 2 * time.Second ) // ConcurrencyService 管理账号和用户的并发限制。 @@ -237,6 +245,77 @@ func (s *ConcurrencyService) AcquireUserSlot(ctx context.Context, userID int64, }, nil } +// TrackAPIKeySlot records one active request slot for an API key without +// applying key-level concurrency limits. It is fail-open: Redis errors are +// logged and return a no-op release function. +func (s *ConcurrencyService) TrackAPIKeySlot(ctx context.Context, apiKeyID int64) func() { + if s == nil || s.cache == nil || apiKeyID <= 0 { + return func() {} + } + cache, ok := s.cache.(APIKeyConcurrencyCache) + if !ok { + return func() {} + } + + requestID := generateRequestID() + baseCtx := context.Background() + if ctx != nil { + baseCtx = context.WithoutCancel(ctx) + } + trackCtx, cancel := context.WithTimeout(baseCtx, apiKeySlotTrackTimeout) + err := cache.TrackAPIKeySlot(trackCtx, apiKeyID, requestID) + cancel() + if err != nil { + logger.LegacyPrintf("service.concurrency", "Warning: failed to track api key slot for %d (req=%s): %v", apiKeyID, requestID, err) + return func() {} + } + + return func() { + bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := cache.ReleaseAPIKeySlot(bgCtx, apiKeyID, requestID); err != nil { + logger.LegacyPrintf("service.concurrency", "Warning: failed to release api key slot for %d (req=%s): %v", apiKeyID, requestID, err) + } + } +} + +// GetAPIKeyConcurrencyBatch gets real-time active request counts for API keys. +// Stats are best-effort: missing Redis support or Redis errors return zeroes. +func (s *ConcurrencyService) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) { + result := zeroAPIKeyConcurrencyMap(apiKeyIDs) + if len(apiKeyIDs) == 0 { + return result, nil + } + if s == nil || s.cache == nil { + return result, nil + } + cache, ok := s.cache.(APIKeyConcurrencyCache) + if !ok { + return result, nil + } + + redisCtx, cancel := context.WithTimeout(context.Background(), apiKeyConcurrencyFetchTimeout) + defer cancel() + + counts, err := cache.GetAPIKeyConcurrencyBatch(redisCtx, apiKeyIDs) + if err != nil { + logger.LegacyPrintf("service.concurrency", "Warning: get api key concurrency batch failed: %v", err) + return result, nil + } + for _, apiKeyID := range apiKeyIDs { + result[apiKeyID] = counts[apiKeyID] + } + return result, nil +} + +func zeroAPIKeyConcurrencyMap(apiKeyIDs []int64) map[int64]int { + result := make(map[int64]int, len(apiKeyIDs)) + for _, apiKeyID := range apiKeyIDs { + result[apiKeyID] = 0 + } + return result +} + // ============================================ // Wait Queue Count Methods // ============================================ diff --git a/backend/internal/service/concurrency_service_test.go b/backend/internal/service/concurrency_service_test.go index 7d5f501dc5..504ec1ee48 100644 --- a/backend/internal/service/concurrency_service_test.go +++ b/backend/internal/service/concurrency_service_test.go @@ -16,25 +16,33 @@ import ( // stubConcurrencyCacheForTest 用于并发服务单元测试的缓存桩 type stubConcurrencyCacheForTest struct { - acquireResult bool - acquireErr error - releaseErr error - concurrency int - concurrencyErr error - waitAllowed bool - waitErr error - waitCount int - waitCountErr error - loadBatch map[int64]*AccountLoadInfo - loadBatchErr error - usersLoadBatch map[int64]*UserLoadInfo - usersLoadErr error - cleanupErr error + acquireResult bool + acquireErr error + releaseErr error + concurrency int + concurrencyErr error + waitAllowed bool + waitErr error + waitCount int + waitCountErr error + loadBatch map[int64]*AccountLoadInfo + loadBatchErr error + usersLoadBatch map[int64]*UserLoadInfo + usersLoadErr error + cleanupErr error + apiKeyTrackErr error + apiKeyReleaseErr error + apiKeyConcurrency map[int64]int + apiKeyConcurrencyErr error // 记录调用 - releasedAccountIDs []int64 - releasedRequestIDs []string - loadBatchCalls atomic.Int64 + releasedAccountIDs []int64 + releasedRequestIDs []string + loadBatchCalls atomic.Int64 + trackedAPIKeyIDs []int64 + trackedAPIKeyRequestIDs []string + releasedAPIKeyIDs []int64 + releasedAPIKeyRequestIDs []string } var _ ConcurrencyCache = (*stubConcurrencyCacheForTest)(nil) @@ -78,6 +86,26 @@ func (c *stubConcurrencyCacheForTest) ReleaseUserSlot(_ context.Context, _ int64 func (c *stubConcurrencyCacheForTest) GetUserConcurrency(_ context.Context, _ int64) (int, error) { return c.concurrency, c.concurrencyErr } +func (c *stubConcurrencyCacheForTest) TrackAPIKeySlot(_ context.Context, apiKeyID int64, requestID string) error { + c.trackedAPIKeyIDs = append(c.trackedAPIKeyIDs, apiKeyID) + c.trackedAPIKeyRequestIDs = append(c.trackedAPIKeyRequestIDs, requestID) + return c.apiKeyTrackErr +} +func (c *stubConcurrencyCacheForTest) ReleaseAPIKeySlot(_ context.Context, apiKeyID int64, requestID string) error { + c.releasedAPIKeyIDs = append(c.releasedAPIKeyIDs, apiKeyID) + c.releasedAPIKeyRequestIDs = append(c.releasedAPIKeyRequestIDs, requestID) + return c.apiKeyReleaseErr +} +func (c *stubConcurrencyCacheForTest) GetAPIKeyConcurrencyBatch(_ context.Context, apiKeyIDs []int64) (map[int64]int, error) { + if c.apiKeyConcurrencyErr != nil { + return nil, c.apiKeyConcurrencyErr + } + result := make(map[int64]int, len(apiKeyIDs)) + for _, apiKeyID := range apiKeyIDs { + result[apiKeyID] = c.apiKeyConcurrency[apiKeyID] + } + return result, nil +} func (c *stubConcurrencyCacheForTest) IncrementWaitCount(_ context.Context, _ int64, _ int) (bool, error) { return c.waitAllowed, c.waitErr } @@ -197,6 +225,62 @@ func TestAcquireUserSlot_UnlimitedConcurrency(t *testing.T) { require.True(t, result.Acquired) } +func TestTrackAPIKeySlot_ReleaseDecrements(t *testing.T) { + cache := &stubConcurrencyCacheForTest{} + svc := NewConcurrencyService(cache) + + release := svc.TrackAPIKeySlot(context.Background(), 88) + require.NotNil(t, release) + require.Equal(t, []int64{88}, cache.trackedAPIKeyIDs) + require.Len(t, cache.trackedAPIKeyRequestIDs, 1) + require.NotEmpty(t, cache.trackedAPIKeyRequestIDs[0]) + + release() + + require.Equal(t, []int64{88}, cache.releasedAPIKeyIDs) + require.Equal(t, cache.trackedAPIKeyRequestIDs, cache.releasedAPIKeyRequestIDs) +} + +func TestTrackAPIKeySlot_FailOpen(t *testing.T) { + cache := &stubConcurrencyCacheForTest{apiKeyTrackErr: errors.New("redis down")} + svc := NewConcurrencyService(cache) + + release := svc.TrackAPIKeySlot(context.Background(), 88) + require.NotNil(t, release) + require.Equal(t, []int64{88}, cache.trackedAPIKeyIDs) + + require.NotPanics(t, release) + require.Empty(t, cache.releasedAPIKeyIDs) +} + +func TestGetAPIKeyConcurrencyBatch_Fallbacks(t *testing.T) { + t.Run("nil cache returns zeroes", func(t *testing.T) { + svc := &ConcurrencyService{cache: nil} + + counts, err := svc.GetAPIKeyConcurrencyBatch(context.Background(), []int64{1, 2}) + require.NoError(t, err) + require.Equal(t, map[int64]int{1: 0, 2: 0}, counts) + }) + + t.Run("redis error returns zeroes", func(t *testing.T) { + cache := &stubConcurrencyCacheForTest{apiKeyConcurrencyErr: errors.New("redis down")} + svc := NewConcurrencyService(cache) + + counts, err := svc.GetAPIKeyConcurrencyBatch(context.Background(), []int64{1, 2}) + require.NoError(t, err) + require.Equal(t, map[int64]int{1: 0, 2: 0}, counts) + }) + + t.Run("success returns counts", func(t *testing.T) { + cache := &stubConcurrencyCacheForTest{apiKeyConcurrency: map[int64]int{1: 3, 2: 0}} + svc := NewConcurrencyService(cache) + + counts, err := svc.GetAPIKeyConcurrencyBatch(context.Background(), []int64{1, 2}) + require.NoError(t, err) + require.Equal(t, map[int64]int{1: 3, 2: 0}, counts) + }) +} + func TestGenerateRequestID_UsesStablePrefixAndMonotonicCounter(t *testing.T) { id1 := generateRequestID() id2 := generateRequestID() diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 7278a4e0c5..908783eb88 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -537,9 +537,11 @@ func ProvideAPIKeyService( cache APIKeyCache, cfg *config.Config, billingCacheService *BillingCacheService, + concurrencyService *ConcurrencyService, ) *APIKeyService { svc := NewAPIKeyService(apiKeyRepo, userRepo, groupRepo, userSubRepo, userGroupRateRepo, cache, cfg) svc.SetRateLimitCacheInvalidator(billingCacheService) + svc.SetConcurrencyService(concurrencyService) return svc } diff --git a/frontend/src/i18n/locales/en.ts b/frontend/src/i18n/locales/en.ts index 3469d3fff2..5a2e1d1a12 100644 --- a/frontend/src/i18n/locales/en.ts +++ b/frontend/src/i18n/locales/en.ts @@ -743,6 +743,7 @@ export default { deleteConfirmMessage: "Are you sure you want to delete '{name}'? This action cannot be undone.", apiKey: 'API Key', group: 'Group', + currentConcurrency: 'Current Concurrency', noGroup: 'No group', searchGroup: 'Search groups...', noGroupFound: 'No groups found', diff --git a/frontend/src/i18n/locales/zh.ts b/frontend/src/i18n/locales/zh.ts index 3fa1e1371d..a6cb37bf95 100644 --- a/frontend/src/i18n/locales/zh.ts +++ b/frontend/src/i18n/locales/zh.ts @@ -742,6 +742,7 @@ export default { deleteConfirmMessage: "确定要删除 '{name}' 吗？此操作无法撤销。", apiKey: 'API 密钥', group: '分组', + currentConcurrency: '当前并发', noGroup: '无分组', searchGroup: '搜索分组...', noGroupFound: '未找到匹配的分组', diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 9e185b0fc2..8b74679e68 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -577,6 +577,7 @@ export interface ApiKey { expires_at: string | null // Expiration time (null = never expires) created_at: string updated_at: string + current_concurrency: number group?: Group rate_limit_5h: number rate_limit_1d: number diff --git a/frontend/src/views/user/KeysView.vue b/frontend/src/views/user/KeysView.vue index 598e88ed96..a704e0ae2f 100644 --- a/frontend/src/views/user/KeysView.vue +++ b/frontend/src/views/user/KeysView.vue @@ -166,6 +166,19 @@ + +