fix: 修复大规模数据性能与流式状态清理

2026-08-15 18:09:54 +08:00
parent 072ed8f168
commit ddebe3c4a2
15 changed files with 439 additions and 97 deletions
@@ -17,6 +17,7 @@ import (

 type App struct {
 	mu                sync.RWMutex
+	priceMu           sync.Mutex
 	config            Config
 	usage             *collection.Service
 	store             *repository.SQLiteUsageRepository
@@ -230,6 +231,10 @@ func (a *App) handleRequestComplete(raw []byte) ([]byte, error) {
 	if a.usage == nil {
 		return nil, fmt.Errorf("用量数据库尚未初始化")
 	}
+	pendingManagedKeyID := ""
+	if pending, ok := a.pending.LoadAndDelete(completion.RequestID); ok {
+		pendingManagedKeyID, _ = pending.(string)
+	}
 	if err := a.store.CompleteBillingAdmission(context.Background(), completion.RequestID, completion.CompletedAt); err != nil {
 		return nil, fmt.Errorf("释放并发占用: %w", err)
 	}
@@ -238,8 +243,8 @@ func (a *App) handleRequestComplete(raw []byte) ([]byte, error) {
 	if key, keyErr := a.store.ManagedKeyByReference(context.Background(), managedKeyID); keyErr == nil {
 		managedKeyID = key.ID
 	}
-	if pending, ok := a.pending.LoadAndDelete(completion.RequestID); managedKeyID == "" && ok {
-		managedKeyID, _ = pending.(string)
+	if managedKeyID == "" {
+		managedKeyID = pendingManagedKeyID
 	}
 	if err := a.usage.ObserveRequest(context.Background(), collection.RequestRecord{
 		ManagedKeyID: managedKeyID, RequestID: completion.RequestID, TraceID: completion.TraceID,
@@ -186,6 +186,101 @@ func TestConcurrentUsage(t *testing.T) {
 	}
 }

+func TestSSELifecycleCleanupDoesNotAccumulatePendingState(t *testing.T) {
+	app := NewApp()
+	if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
+		t.Fatal(err)
+	}
+	defer app.Shutdown()
+	if _, err := app.store.UpdateBilling(context.Background(), "key_default", managedaccess.BillingSettings{
+		QuotaMicros: 1_000_000_000, ResetPeriod: managedaccess.ResetNone, MaxConcurrency: 64,
+	}, time.Now().UTC()); err != nil {
+		t.Fatal(err)
+	}
+	scope := managedaccess.CallerScope("key_default")
+	outcomes := [...]RequestCompletionOutcome{
+		RequestCompletionSucceeded, RequestCompletionFailed, RequestCompletionCanceled, RequestCompletionRejected,
+	}
+	for index := range 128 {
+		requestID := fmt.Sprintf("sse-%d", index)
+		beforeRaw, _ := json.Marshal(RequestInterceptRequest{
+			RequestID: requestID, Model: "deepseek-v4-flash", RequestedModel: "deepseek-v4-flash",
+			Stream: true, Metadata: map[string]any{callerScopeMetadata: scope},
+		})
+		beforeResponse, err := app.HandleMethod(MethodRequestBefore, beforeRaw)
+		if err != nil {
+			t.Fatal(err)
+		}
+		var beforeEnvelope Envelope
+		_ = json.Unmarshal(beforeResponse, &beforeEnvelope)
+		var beforeResult RequestInterceptResponse
+		_ = json.Unmarshal(beforeEnvelope.Result, &beforeResult)
+		if beforeResult.Terminate {
+			t.Fatalf("SSE request %s was denied: %+v", requestID, beforeResult)
+		}
+		completionRaw, _ := json.Marshal(RequestCompletion{
+			RequestID: requestID, Model: "deepseek-v4-flash", Stream: true,
+			StartedAt: time.Now().Add(-time.Second), CompletedAt: time.Now(),
+			Outcome: outcomes[index%len(outcomes)], Metadata: map[string]any{callerScopeMetadata: scope},
+		})
+		if _, err := app.HandleMethod(MethodRequestComplete, completionRaw); err != nil {
+			t.Fatal(err)
+		}
+	}
+	pending := 0
+	app.pending.Range(func(_, _ any) bool { pending++; return true })
+	if pending != 0 {
+		t.Fatalf("pending SSE requests = %d", pending)
+	}
+	state, err := app.store.BillingState(context.Background(), "key_default", time.Now().UTC())
+	if err != nil || state.ActiveRequests != 0 {
+		t.Fatalf("SSE billing state=%+v err=%v", state, err)
+	}
+}
+
+func TestFailedSSECompletionStillClearsPendingState(t *testing.T) {
+	app := NewApp()
+	if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := app.store.UpdateBilling(context.Background(), "key_default", managedaccess.BillingSettings{
+		QuotaMicros: 1_000_000, ResetPeriod: managedaccess.ResetNone, MaxConcurrency: 4,
+	}, time.Now().UTC()); err != nil {
+		t.Fatal(err)
+	}
+	scope := managedaccess.CallerScope("key_default")
+	beforeRaw, _ := json.Marshal(RequestInterceptRequest{
+		RequestID: "sse-database-failure", Model: "deepseek-v4-flash", Stream: true,
+		Metadata: map[string]any{callerScopeMetadata: scope},
+	})
+	beforeResponse, err := app.HandleMethod(MethodRequestBefore, beforeRaw)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var beforeEnvelope Envelope
+	_ = json.Unmarshal(beforeResponse, &beforeEnvelope)
+	var beforeResult RequestInterceptResponse
+	_ = json.Unmarshal(beforeEnvelope.Result, &beforeResult)
+	if beforeResult.Terminate {
+		t.Fatalf("SSE request was denied: %+v", beforeResult)
+	}
+	if err := app.store.Close(); err != nil {
+		t.Fatal(err)
+	}
+	completionRaw, _ := json.Marshal(RequestCompletion{
+		RequestID: "sse-database-failure", Stream: true, Outcome: RequestCompletionCanceled,
+		CompletedAt: time.Now(), Metadata: map[string]any{callerScopeMetadata: scope},
+	})
+	if _, err := app.HandleMethod(MethodRequestComplete, completionRaw); err == nil {
+		t.Fatal("completion unexpectedly succeeded against a closed database")
+	}
+	pending := 0
+	app.pending.Range(func(_, _ any) bool { pending++; return true })
+	if pending != 0 {
+		t.Fatalf("pending requests after failed completion = %d", pending)
+	}
+}
+
 func TestUnknownMethodReturnsErrorEnvelope(t *testing.T) {
 	raw, err := NewApp().HandleMethod("missing", nil)
 	if err != nil {
@@ -63,16 +63,33 @@ func (a *App) putPrice(body []byte) ManagementResponse {
 	if err != nil {
 		return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": err.Error()}})
 	}
+	a.priceMu.Lock()
+	defer a.priceMu.Unlock()
 	a.mu.Lock()
-	defer a.mu.Unlock()
-	if a.store == nil {
+	store := a.store
+	if store == nil {
+		a.mu.Unlock()
 		return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}})
 	}
-	if err := a.store.UpsertPrice(context.Background(), policy); err != nil {
+	if err := store.UpsertPrice(context.Background(), policy); err != nil {
+		a.mu.Unlock()
 		return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
 	}
 	a.prices[normalizeModelName(policy.Model)] = policy
-	if _, err := a.store.BackfillMissingCosts(context.Background(), policy); err != nil {
+	a.mu.Unlock()
+	a.mu.RLock()
+	currentStore := a.store == store
+	a.mu.RUnlock()
+	if !currentStore {
+		return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_reconfigured", "message": "价格已经保存，请在重新配置后重试历史补算"}})
+	}
+	if _, err := store.BackfillMissingCosts(context.Background(), policy); err != nil {
+		a.mu.RLock()
+		reconfigured := a.store != store
+		a.mu.RUnlock()
+		if reconfigured {
+			return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_reconfigured", "message": "价格已经保存，重新配置中止了历史补算"}})
+		}
 		return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
 	}
 	return jsonManagementResponse(http.StatusOK, policyToDTO(policy))
@@ -87,14 +104,19 @@ func (a *App) deletePrice(body []byte) ManagementResponse {
 	if request.Model == "" {
 		return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": "model 不能为空"}})
 	}
+	a.priceMu.Lock()
+	defer a.priceMu.Unlock()
 	a.mu.Lock()
-	defer a.mu.Unlock()
-	if a.store == nil {
+	store := a.store
+	if store == nil {
+		a.mu.Unlock()
 		return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}})
 	}
-	if err := a.store.DeletePrice(context.Background(), request.Model); err != nil {
+	if err := store.DeletePrice(context.Background(), request.Model); err != nil {
+		a.mu.Unlock()
 		return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
 	}
+	defer a.mu.Unlock()
 	delete(a.prices, normalizeModelName(request.Model))
 	return jsonManagementResponse(http.StatusOK, map[string]any{"deleted": request.Model})
 }