mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
feat(keys): add api key concurrency stats
This commit is contained in:
Bestony@Homelab
@@ -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)
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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},
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
// ============================================
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -742,6 +742,7 @@ export default {
|
||||
deleteConfirmMessage: "确定要删除 '{name}' 吗？此操作无法撤销。",
|
||||
apiKey: 'API 密钥',
|
||||
group: '分组',
|
||||
currentConcurrency: '当前并发',
|
||||
noGroup: '无分组',
|
||||
searchGroup: '搜索分组...',
|
||||
noGroupFound: '未找到匹配的分组',
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -166,6 +166,19 @@
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<template #cell-current_concurrency="{ value }">
|
||||
<span
|
||||
:class="[
|
||||
'inline-flex min-w-8 items-center justify-center rounded px-2 py-1 text-sm font-semibold tabular-nums',
|
||||
(value ?? 0) > 0
|
||||
? 'bg-emerald-50 text-emerald-700 ring-1 ring-emerald-200 dark:bg-emerald-900/25 dark:text-emerald-300 dark:ring-emerald-800'
|
||||
: 'bg-gray-100 text-gray-500 dark:bg-dark-700 dark:text-dark-400'
|
||||
]"
|
||||
>
|
||||
{{ value ?? 0 }}
|
||||
</span>
|
||||
</template>
|
||||
|
||||
<template #cell-usage="{ row }">
|
||||
<div class="text-sm">
|
||||
<div class="flex items-center gap-1.5">
|
||||
@@ -1135,6 +1148,7 @@ const allColumns = computed<Column[]>(() => [
|
||||
{ key: 'name', label: t('common.name'), sortable: true },
|
||||
{ key: 'key', label: t('keys.apiKey'), sortable: false },
|
||||
{ key: 'group', label: t('keys.group'), sortable: false },
|
||||
{ key: 'current_concurrency', label: t('keys.currentConcurrency'), sortable: false },
|
||||
{ key: 'usage', label: t('keys.usage'), sortable: false },
|
||||
{ key: 'rate_limit', label: t('keys.rateLimitColumn'), sortable: false },
|
||||
{ key: 'expires_at', label: t('keys.expiresAt'), sortable: true },
|
||||
|
||||
@@ -42,6 +42,7 @@ const messages: Record<string, string> = {
|
||||
'keys.created': 'Created',
|
||||
'keys.expiresAt': 'Expires',
|
||||
'keys.group': 'Group',
|
||||
'keys.currentConcurrency': 'Current Concurrency',
|
||||
'keys.lastUsedAt': 'Last Used',
|
||||
'keys.rateLimitColumn': 'Rate Limit',
|
||||
'keys.searchPlaceholder': 'Search name or key...',
|
||||
@@ -117,6 +118,7 @@ const createApiKey = (): ApiKey => ({
|
||||
expires_at: null,
|
||||
created_at: '2026-06-27T00:00:00Z',
|
||||
updated_at: '2026-06-27T00:00:00Z',
|
||||
current_concurrency: 3,
|
||||
rate_limit_5h: 0,
|
||||
rate_limit_1d: 0,
|
||||
rate_limit_7d: 0,
|
||||
@@ -154,6 +156,9 @@ const DataTableStub = {
|
||||
<div data-test="columns">{{ columns.map((col) => col.key).join(',') }}</div>
|
||||
<div v-for="row in data" :key="row.id">
|
||||
<slot name="cell-name" :value="row.name" :row="row" />
|
||||
<div data-test="current-concurrency">
|
||||
<slot name="cell-current_concurrency" :value="row.current_concurrency" :row="row" />
|
||||
</div>
|
||||
</div>
|
||||
<slot name="empty" />
|
||||
</div>
|
||||
@@ -251,6 +256,7 @@ describe('user KeysView column settings', () => {
|
||||
'name',
|
||||
'key',
|
||||
'group',
|
||||
'current_concurrency',
|
||||
'usage',
|
||||
'expires_at',
|
||||
'status',
|
||||
@@ -282,6 +288,7 @@ describe('user KeysView column settings', () => {
|
||||
expect(visibleColumnKeys(wrapper)).toEqual([
|
||||
'name',
|
||||
'key',
|
||||
'current_concurrency',
|
||||
'usage',
|
||||
'rate_limit',
|
||||
'expires_at',
|
||||
@@ -299,8 +306,15 @@ describe('user KeysView column settings', () => {
|
||||
|
||||
const columnMenuText = wrapper.text()
|
||||
expect(columnMenuText).toContain('API Key')
|
||||
expect(columnMenuText).toContain('Current Concurrency')
|
||||
expect(columnMenuText).toContain('Rate Limit')
|
||||
expect(columnMenuText).not.toContain('Name')
|
||||
expect(columnMenuText).not.toContain('Actions')
|
||||
})
|
||||
|
||||
it('renders the current concurrency value', async () => {
|
||||
const wrapper = await mountView()
|
||||
|
||||
expect(wrapper.get('[data-test="current-concurrency"]').text()).toBe('3')
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user