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。
This commit is contained in:
@@ -54,12 +54,12 @@ func TestGeminiBase64ImageEditThousandSynchronousRequests(t *testing.T) {
|
|||||||
defer cancel()
|
defer cancel()
|
||||||
applyMigration(t, ctx, databaseURL)
|
applyMigration(t, ctx, databaseURL)
|
||||||
|
|
||||||
apiDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 16)
|
apiDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 64)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("connect API store: %v", err)
|
t.Fatalf("connect API store: %v", err)
|
||||||
}
|
}
|
||||||
defer apiDB.Close()
|
defer apiDB.Close()
|
||||||
workerDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 16)
|
workerDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 32)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("connect Worker store: %v", err)
|
t.Fatalf("connect Worker store: %v", err)
|
||||||
}
|
}
|
||||||
@@ -161,7 +161,7 @@ WHERE status = 'active'`); err != nil {
|
|||||||
AppEnv: "test",
|
AppEnv: "test",
|
||||||
HTTPAddr: ":0",
|
HTTPAddr: ":0",
|
||||||
DatabaseURL: databaseURL,
|
DatabaseURL: databaseURL,
|
||||||
DatabaseMaxConns: 16,
|
DatabaseMaxConns: 32,
|
||||||
IdentityMode: "hybrid",
|
IdentityMode: "hybrid",
|
||||||
JWTSecret: "gemini-base64-stress-secret",
|
JWTSecret: "gemini-base64-stress-secret",
|
||||||
BillingEngineMode: "observe",
|
BillingEngineMode: "observe",
|
||||||
@@ -180,7 +180,7 @@ WHERE status = 'active'`); err != nil {
|
|||||||
AppEnv: "test",
|
AppEnv: "test",
|
||||||
HTTPAddr: ":0",
|
HTTPAddr: ":0",
|
||||||
DatabaseURL: databaseURL,
|
DatabaseURL: databaseURL,
|
||||||
DatabaseMaxConns: 16,
|
DatabaseMaxConns: 64,
|
||||||
IdentityMode: "hybrid",
|
IdentityMode: "hybrid",
|
||||||
JWTSecret: "gemini-base64-stress-secret",
|
JWTSecret: "gemini-base64-stress-secret",
|
||||||
BillingEngineMode: "observe",
|
BillingEngineMode: "observe",
|
||||||
@@ -319,8 +319,20 @@ LIMIT 1`).Scan(&gatewayUserID); err != nil {
|
|||||||
peakHeap.Store(baseline.HeapAlloc)
|
peakHeap.Store(baseline.HeapAlloc)
|
||||||
var queuePeak atomic.Int64
|
var queuePeak atomic.Int64
|
||||||
var runningPeak atomic.Int64
|
var runningPeak atomic.Int64
|
||||||
|
var databaseConnectionsPeak atomic.Int64
|
||||||
|
var advisoryWaitPeak atomic.Int64
|
||||||
samplingDone := make(chan struct{})
|
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()
|
startedAt := time.Now()
|
||||||
loadClient := &http.Client{
|
loadClient := &http.Client{
|
||||||
@@ -430,6 +442,14 @@ WHERE task.model = $1
|
|||||||
if queuePeak.Load() == 0 {
|
if queuePeak.Load() == 0 {
|
||||||
t.Fatal("the synchronous distributed admission queue was never observed under the 1000-request burst")
|
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()
|
requestAssetUploads, generatedUploads, retainedObjects, invalidUploads := uploadService.stats()
|
||||||
if invalidUploads != 0 {
|
if invalidUploads != 0 {
|
||||||
t.Fatalf("shared upload validation failures=%d", invalidUploads)
|
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.Fatalf("heap growth=%d bytes exceeded configured ceiling=%d bytes", heapGrowth, maxHeapGrowth)
|
||||||
}
|
}
|
||||||
t.Logf(
|
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,
|
taskCount,
|
||||||
imageBytes,
|
imageBytes,
|
||||||
inputVariants,
|
inputVariants,
|
||||||
@@ -464,6 +484,8 @@ WHERE task.model = $1
|
|||||||
queuePeak.Load(),
|
queuePeak.Load(),
|
||||||
runningPeak.Load(),
|
runningPeak.Load(),
|
||||||
upstreamPeak.Load(),
|
upstreamPeak.Load(),
|
||||||
|
databaseConnectionsPeak.Load(),
|
||||||
|
advisoryWaitPeak.Load(),
|
||||||
heapGrowth,
|
heapGrowth,
|
||||||
requestAssetUploads,
|
requestAssetUploads,
|
||||||
generatedUploads,
|
generatedUploads,
|
||||||
@@ -725,6 +747,8 @@ func sampleGeminiStressPressure(
|
|||||||
peakHeap *atomic.Uint64,
|
peakHeap *atomic.Uint64,
|
||||||
queuePeak *atomic.Int64,
|
queuePeak *atomic.Int64,
|
||||||
runningPeak *atomic.Int64,
|
runningPeak *atomic.Int64,
|
||||||
|
databaseConnectionsPeak *atomic.Int64,
|
||||||
|
advisoryWaitPeak *atomic.Int64,
|
||||||
done <-chan struct{},
|
done <-chan struct{},
|
||||||
) {
|
) {
|
||||||
ticker := time.NewTicker(100 * time.Millisecond)
|
ticker := time.NewTicker(100 * time.Millisecond)
|
||||||
@@ -744,17 +768,21 @@ func sampleGeminiStressPressure(
|
|||||||
break
|
break
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
var waiting, running int64
|
var waiting, running, databaseConnections, advisoryWaiters int64
|
||||||
err := pool.QueryRow(ctx, `
|
err := pool.QueryRow(ctx, `
|
||||||
SELECT
|
SELECT
|
||||||
COUNT(*) FILTER (WHERE admission.status = 'waiting')::bigint,
|
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
|
FROM gateway_tasks task
|
||||||
LEFT JOIN gateway_task_admissions admission ON admission.task_id = task.id
|
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 {
|
if err == nil {
|
||||||
updateAtomicPeak(queuePeak, waiting)
|
updateAtomicPeak(queuePeak, waiting)
|
||||||
updateAtomicPeak(runningPeak, running)
|
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))
|
writeGeminiTaskError(status, err.Error(), nil, clients.ErrorCode(err))
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
// The queued task only needs the materialized assetRef. Release the native
|
// The queued task only needs the materialized assetRef. Drop the native
|
||||||
// Base64 request body before it waits for distributed admission or an
|
// Base64 request body now, but retain the lightweight request slot until
|
||||||
// upstream response.
|
// distributed admission is durably registered. This prevents a burst from
|
||||||
|
// turning advisory-lock waiters into one PostgreSQL connection per request.
|
||||||
native = nil
|
native = nil
|
||||||
mapping.Body = nil
|
mapping.Body = nil
|
||||||
releaseRequestBody()
|
|
||||||
releaseRequestBody = nil
|
|
||||||
runMode, err := s.admittedTaskRunMode(admission, prepared.Body)
|
runMode, err := s.admittedTaskRunMode(admission, prepared.Body)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
writeTaskTrafficError(w, err, clients.ProtocolGeminiGenerateContent)
|
writeTaskTrafficError(w, err, clients.ProtocolGeminiGenerateContent)
|
||||||
@@ -237,6 +236,8 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
|||||||
runCtx, cancelRun := s.requestExecutionContext(r)
|
runCtx, cancelRun := s.requestExecutionContext(r)
|
||||||
defer cancelRun()
|
defer cancelRun()
|
||||||
if streamMode {
|
if streamMode {
|
||||||
|
releaseRequestBody()
|
||||||
|
releaseRequestBody = nil
|
||||||
s.writeGeminiGenerateContentStream(runCtx, w, r, task, user, mapping.Model, writeGeminiTaskError)
|
s.writeGeminiGenerateContentStream(runCtx, w, r, task, user, mapping.Model, writeGeminiTaskError)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@@ -248,6 +249,8 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
|||||||
writeGeminiTaskError(statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err))
|
writeGeminiTaskError(statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err))
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
releaseRequestBody()
|
||||||
|
releaseRequestBody = nil
|
||||||
task, err = s.runner.WaitForTaskCompletion(runCtx, task.ID)
|
task, err = s.runner.WaitForTaskCompletion(runCtx, task.ID)
|
||||||
if err != nil || !requestStillConnected(r) {
|
if err != nil || !requestStillConnected(r) {
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -1074,6 +1074,15 @@ func (s *Server) listModelRateLimitStatuses(w http.ResponseWriter, r *http.Reque
|
|||||||
// @Router /api/v1/reranks [post]
|
// @Router /api/v1/reranks [post]
|
||||||
func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
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)
|
targetProtocol := targetProtocolForTaskRequest(kind, r)
|
||||||
writeTaskError := func(status int, message string, details map[string]any, codes ...string) {
|
writeTaskError := func(status int, message string, details map[string]any, codes ...string) {
|
||||||
if easyAIAsyncMediaRequest(kind, r) {
|
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))
|
writeTaskError(statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err))
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
if releaseAsyncMediaRequest != nil {
|
||||||
|
releaseAsyncMediaRequest()
|
||||||
|
releaseAsyncMediaRequest = nil
|
||||||
|
}
|
||||||
writeTaskAccepted(w, task)
|
writeTaskAccepted(w, task)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user