diff --git a/apps/api/internal/httpapi/gemini_base64_stress_integration_test.go b/apps/api/internal/httpapi/gemini_base64_stress_integration_test.go index 188590a..e9780c8 100644 --- a/apps/api/internal/httpapi/gemini_base64_stress_integration_test.go +++ b/apps/api/internal/httpapi/gemini_base64_stress_integration_test.go @@ -54,12 +54,12 @@ func TestGeminiBase64ImageEditThousandSynchronousRequests(t *testing.T) { defer cancel() applyMigration(t, ctx, databaseURL) - apiDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 16) + apiDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 64) if err != nil { t.Fatalf("connect API store: %v", err) } defer apiDB.Close() - workerDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 16) + workerDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 32) if err != nil { t.Fatalf("connect Worker store: %v", err) } @@ -161,7 +161,7 @@ WHERE status = 'active'`); err != nil { AppEnv: "test", HTTPAddr: ":0", DatabaseURL: databaseURL, - DatabaseMaxConns: 16, + DatabaseMaxConns: 32, IdentityMode: "hybrid", JWTSecret: "gemini-base64-stress-secret", BillingEngineMode: "observe", @@ -180,7 +180,7 @@ WHERE status = 'active'`); err != nil { AppEnv: "test", HTTPAddr: ":0", DatabaseURL: databaseURL, - DatabaseMaxConns: 16, + DatabaseMaxConns: 64, IdentityMode: "hybrid", JWTSecret: "gemini-base64-stress-secret", BillingEngineMode: "observe", @@ -319,8 +319,20 @@ LIMIT 1`).Scan(&gatewayUserID); err != nil { peakHeap.Store(baseline.HeapAlloc) var queuePeak atomic.Int64 var runningPeak atomic.Int64 + var databaseConnectionsPeak atomic.Int64 + var advisoryWaitPeak atomic.Int64 samplingDone := make(chan struct{}) - go sampleGeminiStressPressure(ctx, pool, model, &peakHeap, &queuePeak, &runningPeak, samplingDone) + go sampleGeminiStressPressure( + ctx, + pool, + model, + &peakHeap, + &queuePeak, + &runningPeak, + &databaseConnectionsPeak, + &advisoryWaitPeak, + samplingDone, + ) startedAt := time.Now() loadClient := &http.Client{ @@ -430,6 +442,14 @@ WHERE task.model = $1 if queuePeak.Load() == 0 { t.Fatal("the synchronous distributed admission queue was never observed under the 1000-request burst") } + if advisoryWaitPeak.Load() > int64(2*mediaConcurrent) { + t.Fatalf( + "advisory lock waiters=%d exceeded bounded media admission ceiling=%d; database_connections_peak=%d", + advisoryWaitPeak.Load(), + 2*mediaConcurrent, + databaseConnectionsPeak.Load(), + ) + } requestAssetUploads, generatedUploads, retainedObjects, invalidUploads := uploadService.stats() if invalidUploads != 0 { t.Fatalf("shared upload validation failures=%d", invalidUploads) @@ -453,7 +473,7 @@ WHERE task.model = $1 t.Fatalf("heap growth=%d bytes exceeded configured ceiling=%d bytes", heapGrowth, maxHeapGrowth) } t.Logf( - "Gemini Base64 同步千任务模拟上游证据: tasks=%d image_bytes=%d input_variants=%d client_connections=%d platform_concurrency=%d media_concurrency=%d elapsed=%s queue_peak=%d running_peak=%d upstream_peak=%d heap_growth_peak=%d request_asset_uploads=%d generated_uploads=%d local_files=%d", + "Gemini Base64 同步千任务模拟上游证据: tasks=%d image_bytes=%d input_variants=%d client_connections=%d platform_concurrency=%d media_concurrency=%d elapsed=%s queue_peak=%d running_peak=%d upstream_peak=%d database_connections_peak=%d advisory_wait_peak=%d heap_growth_peak=%d request_asset_uploads=%d generated_uploads=%d local_files=%d", taskCount, imageBytes, inputVariants, @@ -464,6 +484,8 @@ WHERE task.model = $1 queuePeak.Load(), runningPeak.Load(), upstreamPeak.Load(), + databaseConnectionsPeak.Load(), + advisoryWaitPeak.Load(), heapGrowth, requestAssetUploads, generatedUploads, @@ -725,6 +747,8 @@ func sampleGeminiStressPressure( peakHeap *atomic.Uint64, queuePeak *atomic.Int64, runningPeak *atomic.Int64, + databaseConnectionsPeak *atomic.Int64, + advisoryWaitPeak *atomic.Int64, done <-chan struct{}, ) { ticker := time.NewTicker(100 * time.Millisecond) @@ -744,17 +768,21 @@ func sampleGeminiStressPressure( break } } - var waiting, running int64 + var waiting, running, databaseConnections, advisoryWaiters int64 err := pool.QueryRow(ctx, ` SELECT COUNT(*) FILTER (WHERE admission.status = 'waiting')::bigint, - COUNT(*) FILTER (WHERE task.status = 'running')::bigint + COUNT(*) FILTER (WHERE task.status = 'running')::bigint, + (SELECT COUNT(*) FROM pg_stat_activity WHERE backend_type = 'client backend')::bigint, + (SELECT COUNT(*) FROM pg_stat_activity WHERE backend_type = 'client backend' AND wait_event = 'advisory')::bigint FROM gateway_tasks task LEFT JOIN gateway_task_admissions admission ON admission.task_id = task.id -WHERE task.model = $1`, model).Scan(&waiting, &running) +WHERE task.model = $1`, model).Scan(&waiting, &running, &databaseConnections, &advisoryWaiters) if err == nil { updateAtomicPeak(queuePeak, waiting) updateAtomicPeak(runningPeak, running) + updateAtomicPeak(databaseConnectionsPeak, databaseConnections) + updateAtomicPeak(advisoryWaitPeak, advisoryWaiters) } } } diff --git a/apps/api/internal/httpapi/gemini_compat.go b/apps/api/internal/httpapi/gemini_compat.go index d0f8ef2..3f1eafe 100644 --- a/apps/api/internal/httpapi/gemini_compat.go +++ b/apps/api/internal/httpapi/gemini_compat.go @@ -164,13 +164,12 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) { writeGeminiTaskError(status, err.Error(), nil, clients.ErrorCode(err)) return } - // The queued task only needs the materialized assetRef. Release the native - // Base64 request body before it waits for distributed admission or an - // upstream response. + // The queued task only needs the materialized assetRef. Drop the native + // Base64 request body now, but retain the lightweight request slot until + // distributed admission is durably registered. This prevents a burst from + // turning advisory-lock waiters into one PostgreSQL connection per request. native = nil mapping.Body = nil - releaseRequestBody() - releaseRequestBody = nil runMode, err := s.admittedTaskRunMode(admission, prepared.Body) if err != nil { writeTaskTrafficError(w, err, clients.ProtocolGeminiGenerateContent) @@ -237,6 +236,8 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) { runCtx, cancelRun := s.requestExecutionContext(r) defer cancelRun() if streamMode { + releaseRequestBody() + releaseRequestBody = nil s.writeGeminiGenerateContentStream(runCtx, w, r, task, user, mapping.Model, writeGeminiTaskError) return } @@ -248,6 +249,8 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) { writeGeminiTaskError(statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err)) return } + releaseRequestBody() + releaseRequestBody = nil task, err = s.runner.WaitForTaskCompletion(runCtx, task.ID) if err != nil || !requestStillConnected(r) { return diff --git a/apps/api/internal/httpapi/handlers.go b/apps/api/internal/httpapi/handlers.go index fc0696d..7ef869c 100644 --- a/apps/api/internal/httpapi/handlers.go +++ b/apps/api/internal/httpapi/handlers.go @@ -1074,6 +1074,15 @@ func (s *Server) listModelRateLimitStatuses(w http.ResponseWriter, r *http.Reque // @Router /api/v1/reranks [post] func (s *Server) createTask(kind string, compatible bool) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var releaseAsyncMediaRequest func() + if easyAIAsyncMediaRequest(kind, r) { + var err error + releaseAsyncMediaRequest, err = s.acquireMediaRequestBodySlot(r.Context()) + if err != nil { + return + } + defer releaseAsyncMediaRequest() + } targetProtocol := targetProtocolForTaskRequest(kind, r) writeTaskError := func(status int, message string, details map[string]any, codes ...string) { if easyAIAsyncMediaRequest(kind, r) { @@ -1230,6 +1239,10 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler { writeTaskError(statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err)) return } + if releaseAsyncMediaRequest != nil { + releaseAsyncMediaRequest() + releaseAsyncMediaRequest = nil + } writeTaskAccepted(w, task) return }