fix(admission): 限制媒体任务准入锁等待并发

媒体请求从认证、任务创建到异步准入登记共用可配置的请求并发令牌，完成登记后立即释放。这样仍会尽早丢弃 Base64 正文，但不会让高并发请求各占一个 PostgreSQL 连接等待 advisory lock。\n\n扩展 Gemini 双角色压力测试，使用生产同档 API/Worker 连接池并记录数据库连接与 advisory lock 峰值。\n\n验证：go test ./... -count=1；go vet ./...；pnpm openapi；迁移安全检查；256 个 256 KiB 异步 Gemini 请求全部通过，连接峰值 45、锁等待峰值 14。
2026-07-31 06:05:54 +08:00
parent 4af86b22ee
commit 4fa44f4c41
3 changed files with 58 additions and 14 deletions
@@ -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)
 			}
 		}
 	}
@@ -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
@@ -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
 		}