opensource/sub2api
3
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-09-24 16:05:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(keys): add api key concurrency stats

Browse Source
This commit is contained in:
Bestony@Homelab
2026-07-02 15:58:29 +08:00
parent 7dc7cfce1d
commit 089a7b7fae
20 changed files with 569 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/handler/dto/api_key_mapper_last_used_test.go
+8 -6
View File
@@ -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) {
backend/internal/handler/dto/mappers.go
+26 -25
View File
@@ -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)
backend/internal/handler/dto/types.go
+2
View File
@@ -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"`
backend/internal/handler/gateway_helper.go
+41 -2
View File
@@ -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.
backend/internal/handler/gateway_helper_hotpath_test.go
+69
View File
@@ -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},
backend/internal/handler/openai_gateway_handler.go
+3 -3
View File
@@ -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)
}