fix(worker): 将 Gemini 同步请求交由异步队列执行

将非流式 Gemini generateContent 改为 River Worker 执行,API 通过批量状态查询等待完成,避免高并发同步请求耗尽 API 数据库连接池。\n\n生成图片在持久化时保存带哈希的资产引用,响应恢复时下载并校验后重建 Base64,数据库不保存媒体原文。请求体物化增加前置内存门禁,并扩展双 API/Worker PostgreSQL 压力测试。\n\n验证:go test ./... -count=1;go vet ./...;pnpm openapi;迁移安全检查;64 与 256 请求双角色异步 Gemini 压力测试。
This commit is contained in:
2026-07-31 05:49:45 +08:00
parent 352e23e099
commit 4af86b22ee
12 changed files with 516 additions and 48 deletions
@@ -43,7 +43,7 @@ func TestGeminiBase64ImageEditThousandSynchronousRequests(t *testing.T) {
mediaConcurrent = 8
)
baselineProfile := acceptanceworkload.GeminiBaseline
taskCount := baselineProfile.Requests
taskCount := stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_REQUESTS", baselineProfile.Requests)
imageBytes := stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_IMAGE_BYTES", baselineProfile.InputBytes)
upstreamDelay := time.Duration(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_UPSTREAM_DELAY_MS", int(baselineProfile.DelayMin/time.Millisecond))) * time.Millisecond
maxHeapGrowth := uint64(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_MAX_HEAP_BYTES", 1<<30))
@@ -59,6 +59,11 @@ func TestGeminiBase64ImageEditThousandSynchronousRequests(t *testing.T) {
t.Fatalf("connect API store: %v", err)
}
defer apiDB.Close()
workerDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 16)
if err != nil {
t.Fatalf("connect Worker store: %v", err)
}
defer workerDB.Close()
poolConfig, err := pgxpool.ParseConfig(databaseURL)
if err != nil {
t.Fatalf("parse observation pool config: %v", err)
@@ -152,6 +157,25 @@ WHERE status = 'active'`); err != nil {
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
storageRoot := t.TempDir()
workerConfig := config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
DatabaseMaxConns: 16,
IdentityMode: "hybrid",
JWTSecret: "gemini-base64-stress-secret",
BillingEngineMode: "observe",
ProcessRole: "worker",
LocalUploadedStorageDir: filepath.Join(storageRoot, "worker-uploaded"),
LocalGeneratedStorageDir: filepath.Join(storageRoot, "worker-generated"),
MediaRequestConcurrency: 16,
MediaMaterializationConcurrency: mediaConcurrent,
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: upstreamConcurrent,
AsyncWorkerInstanceHardLimit: upstreamConcurrent,
AsyncWorkerRefreshIntervalSeconds: 1,
}
_ = NewServerWithContext(serverCtx, workerConfig, workerDB, slog.New(slog.NewTextHandler(io.Discard, nil)))
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
@@ -369,10 +393,11 @@ LIMIT 1`).Scan(&gatewayUserID); err != nil {
}
elapsed := time.Since(startedAt)
var tasks, succeeded, attempts, duplicateAttemptTasks int
var tasks, asyncTasks, succeeded, attempts, duplicateAttemptTasks int
if err := pool.QueryRow(ctx, `
SELECT
COUNT(*)::int,
COUNT(*) FILTER (WHERE task.async_mode)::int,
COUNT(*) FILTER (WHERE task.status = 'succeeded')::int,
COALESCE(SUM(attempts.attempt_count), 0)::int,
COUNT(*) FILTER (WHERE attempts.attempt_count <> 1)::int
@@ -383,11 +408,11 @@ LEFT JOIN LATERAL (
WHERE attempt.task_id = task.id
) attempts ON true
WHERE task.model = $1
AND task.created_at >= $2`, model, startedAt).Scan(&tasks, &succeeded, &attempts, &duplicateAttemptTasks); err != nil {
AND task.created_at >= $2`, model, startedAt).Scan(&tasks, &asyncTasks, &succeeded, &attempts, &duplicateAttemptTasks); err != nil {
t.Fatalf("read Gemini stress task results: %v", err)
}
if tasks != taskCount || succeeded != taskCount || attempts != taskCount || duplicateAttemptTasks != 0 {
t.Fatalf("task results tasks=%d succeeded=%d attempts=%d duplicate_attempt_tasks=%d", tasks, succeeded, attempts, duplicateAttemptTasks)
if tasks != taskCount || asyncTasks != taskCount || succeeded != taskCount || attempts != taskCount || duplicateAttemptTasks != 0 {
t.Fatalf("task results tasks=%d async=%d succeeded=%d attempts=%d duplicate_attempt_tasks=%d", tasks, asyncTasks, succeeded, attempts, duplicateAttemptTasks)
}
if upstreamCalls.Load() != int64(taskCount) || invalidInputs.Load() != 0 {
t.Fatalf("upstream calls=%d invalid_inputs=%d, want %d/0", upstreamCalls.Load(), invalidInputs.Load(), taskCount)
@@ -454,6 +479,7 @@ type geminiStressUploadService struct {
mu sync.RWMutex
apiKey string
retainedObjects map[string][]byte
generatedObject []byte
requestAssetUploads atomic.Int64
generatedUploads atomic.Int64
@@ -557,6 +583,9 @@ func (s *geminiStressUploadService) handleUpload(w http.ResponseWriter, r *http.
http.Error(w, "unexpected generated result hash", http.StatusBadRequest)
return
}
s.mu.Lock()
s.generatedObject = filePayload
s.mu.Unlock()
s.generatedUploads.Add(1)
default:
s.invalidUploads.Add(1)
@@ -573,6 +602,9 @@ func (s *geminiStressUploadService) handleObject(w http.ResponseWriter, r *http.
hash := strings.TrimPrefix(r.URL.Path, "/objects/")
s.mu.RLock()
payload, ok := s.retainedObjects[hash]
if !ok && hash == s.outputHash && len(s.generatedObject) > 0 {
payload, ok = s.generatedObject, true
}
s.mu.RUnlock()
if !ok {
http.NotFound(w, r)