diff --git a/apps/api/internal/clients/clients_test.go b/apps/api/internal/clients/clients_test.go index 5343aec..2029c71 100644 --- a/apps/api/internal/clients/clients_test.go +++ b/apps/api/internal/clients/clients_test.go @@ -1531,6 +1531,16 @@ func TestGeminiClientImageGenerateBuildsNativeImageBody(t *testing.T) { if data[0].(map[string]any)["b64_json"] != "aW1hZ2U=" { t.Fatalf("unexpected image response: %+v", response.Result) } + if response.Wire == nil || len(response.Wire.RawJSON) != 0 { + t.Fatalf("Gemini image response retained duplicate raw JSON: %+v", response.Wire) + } + candidates, _ := response.Wire.Body["candidates"].([]any) + content, _ := candidates[0].(map[string]any)["content"].(map[string]any) + wireParts, _ := content["parts"].([]any) + inline, _ := wireParts[0].(map[string]any)["inlineData"].(map[string]any) + if inline["data"] != "aW1hZ2U=" { + t.Fatalf("decoded Gemini wire body lost Base64 output: %+v", response.Wire.Body) + } } func TestGeminiClientImageGenerateRejectsPromptBlockedWithoutImage(t *testing.T) { diff --git a/apps/api/internal/clients/gemini.go b/apps/api/internal/clients/gemini.go index 7f7ec09..3f72da2 100644 --- a/apps/api/internal/clients/gemini.go +++ b/apps/api/internal/clients/gemini.go @@ -65,6 +65,13 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error } return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt) } + if wire != nil && geminiImageModelType(request.ModelType) { + // Gemini image responses can contain multi-megabyte Base64 fields. The + // decoded wire body is still used for native synchronous passthrough, so + // retaining the original JSON bytes would keep a second full copy alive + // until the client response finishes. + wire.RawJSON = nil + } if requestID == "" { requestID = requestIDFromResult(output) } diff --git a/apps/api/internal/clients/helpers.go b/apps/api/internal/clients/helpers.go index c7dfd12..c32584f 100644 --- a/apps/api/internal/clients/helpers.go +++ b/apps/api/internal/clients/helpers.go @@ -65,12 +65,14 @@ func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[st Protocol: strings.TrimSpace(protocol), StatusCode: resp.StatusCode, Headers: compatibleResponseHeaders(resp.Header), - RawJSON: append([]byte(nil), raw...), - } - if len(raw) > 0 { - _ = json.Unmarshal(raw, &wire.Body) + // raw is already owned by this response. Reusing it avoids retaining a + // second copy of media responses that may be as large as 128 MiB. + RawJSON: raw, } if readErr != nil { + if len(raw) > 0 { + _ = json.Unmarshal(raw, &wire.Body) + } return nil, wire, &ClientError{ Code: "response_read_error", Message: readErr.Error(), @@ -81,6 +83,9 @@ func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[st } } if tooLarge { + if len(raw) > 0 { + _ = json.Unmarshal(raw, &wire.Body) + } return nil, wire, &ClientError{ Code: "response_too_large", Message: fmt.Sprintf("upstream JSON response exceeds %d bytes", maxBytes), @@ -91,6 +96,9 @@ func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[st } } if resp.StatusCode < 200 || resp.StatusCode >= 300 { + if len(raw) > 0 { + _ = json.Unmarshal(raw, &wire.Body) + } return nil, wire, &ClientError{ Code: statusCodeName(resp.StatusCode), Message: errorMessage(raw, resp.Status), @@ -108,7 +116,10 @@ func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[st if err := json.Unmarshal(raw, &out); err != nil { return nil, wire, &ClientError{Code: "invalid_response", Message: err.Error(), StatusCode: resp.StatusCode, Retryable: false, Wire: wire} } - wire.Body = cloneBody(out) + // Result and wire metadata are read-only until the runner replaces the + // canonical result, so they can share the decoded object instead of keeping + // another full copy of a large Base64 media response. + wire.Body = out return out, wire, nil } diff --git a/apps/api/internal/config/config.go b/apps/api/internal/config/config.go index 1d3dfaa..187015b 100644 --- a/apps/api/internal/config/config.go +++ b/apps/api/internal/config/config.go @@ -59,6 +59,8 @@ type Config struct { BillingEngineMode string ProcessRole string DatabaseMaxConns int + MediaRequestConcurrency int + MediaMaterializationConcurrency int AsyncQueueWorkerEnabled bool AsyncWorkerHardLimit int AsyncWorkerRefreshIntervalSeconds int @@ -117,6 +119,8 @@ func Load() Config { BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")), ProcessRole: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_PROCESS_ROLE", "all"))), DatabaseMaxConns: envInt("AI_GATEWAY_DATABASE_MAX_CONNS", 0), + MediaRequestConcurrency: envIntValidated("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", 16), + MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8), AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true", AsyncWorkerHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048), AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5), @@ -132,6 +136,12 @@ func (c Config) Validate() error { if c.DatabaseMaxConns < 0 || c.DatabaseMaxConns > 1000 { return errors.New("AI_GATEWAY_DATABASE_MAX_CONNS must be between 1 and 1000 when configured") } + if c.MediaMaterializationConcurrency != 0 && (c.MediaMaterializationConcurrency < 1 || c.MediaMaterializationConcurrency > 256) { + return errors.New("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY must be between 1 and 256") + } + if c.MediaRequestConcurrency != 0 && (c.MediaRequestConcurrency < 1 || c.MediaRequestConcurrency > 1024) { + return errors.New("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY must be between 1 and 1024") + } switch strings.ToLower(strings.TrimSpace(c.BillingEngineMode)) { case "", "observe", "enforce", "hold": default: diff --git a/apps/api/internal/config/config_test.go b/apps/api/internal/config/config_test.go index a1ee18f..a3ded19 100644 --- a/apps/api/internal/config/config_test.go +++ b/apps/api/internal/config/config_test.go @@ -131,6 +131,24 @@ func TestValidateProcessRoleAndDatabasePool(t *testing.T) { if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MAX_CONNS") { t.Fatalf("Validate() error = %v, want invalid database max conns", err) } + cfg.DatabaseMaxConns = 16 + cfg.MediaMaterializationConcurrency = 257 + if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_MATERIALIZATION_CONCURRENCY") { + t.Fatalf("Validate() error = %v, want invalid media materialization concurrency", err) + } + t.Setenv("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", "12") + if got := Load().MediaMaterializationConcurrency; got != 12 { + t.Fatalf("media materialization concurrency = %d, want 12", got) + } + cfg = Load() + cfg.MediaRequestConcurrency = 1025 + if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_REQUEST_CONCURRENCY") { + t.Fatalf("Validate() error = %v, want invalid media request concurrency", err) + } + t.Setenv("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", "96") + if got := Load().MediaRequestConcurrency; got != 96 { + t.Fatalf("media request concurrency = %d, want 96", got) + } } func TestValidateTaskHistorySettings(t *testing.T) { diff --git a/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go b/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go index 495c2ee..92f4e6a 100644 --- a/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go +++ b/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go @@ -21,6 +21,173 @@ import ( "github.com/jackc/pgx/v5/pgxpool" ) +func TestAsyncSubmissionReservesConcurrencyOnlyAfterWorkerClaim(t *testing.T) { + databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) + if databaseURL == "" { + t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run async admission timing tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + defer cancel() + applyMigration(t, ctx, databaseURL) + + db, err := store.Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect store: %v", err) + } + defer db.Close() + pool, err := pgxpool.New(ctx, databaseURL) + if err != nil { + t.Fatalf("connect acceptance pool: %v", err) + } + defer pool.Close() + if _, err := pool.Exec(ctx, `UPDATE integration_platforms SET status = 'disabled' WHERE deleted_at IS NULL`); err != nil { + t.Fatalf("isolate acceptance platforms: %v", err) + } + + serverCtx, cancelServer := context.WithCancel(ctx) + defer cancelServer() + server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{ + AppEnv: "test", + HTTPAddr: ":0", + DatabaseURL: databaseURL, + IdentityMode: "hybrid", + JWTSecret: "test-secret", + BillingEngineMode: "observe", + CORSAllowedOrigin: "*", + AsyncQueueWorkerEnabled: true, + AsyncWorkerHardLimit: 1, + AsyncWorkerRefreshIntervalSeconds: 1, + }, db, slog.New(slog.NewTextHandler(io.Discard, nil)))) + defer server.Close() + + adminToken := createAsyncAcceptanceAdmin(t, ctx, pool, server.URL) + if _, err := pool.Exec(ctx, ` +UPDATE gateway_user_groups +SET rate_limit_policy = '{"rules":[{"metric":"concurrent","limit":256,"leaseTtlSeconds":120}]}'::jsonb +WHERE status = 'active'`); err != nil { + t.Fatalf("raise acceptance user-group concurrency: %v", err) + } + suffix := strconv.FormatInt(time.Now().UnixNano(), 10) + model := "worker-claim-admission-" + suffix + platform := createAsyncAcceptancePlatform(t, server.URL, adminToken, "claim-admission-"+suffix, "Claim Admission Simulation", 3) + defer updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 3, "disabled") + modelID := createAsyncAcceptanceModel(t, server.URL, adminToken, platform.ID, model, "video_generate", "inherit", 0, 120) + + waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 1, 15*time.Second) + taskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 3, 5*time.Second) + waitForTaskCount(t, ctx, pool, taskIDs, "running", 1, 10*time.Second) + waitForActiveLeaseCount(t, ctx, pool, modelID, 1, 10*time.Second) + + var admitted, attempts int + if err := pool.QueryRow(ctx, ` +SELECT + (SELECT COUNT(*) + FROM gateway_task_admissions + WHERE task_id = ANY($1::uuid[]) AND status = 'admitted'), + (SELECT COUNT(*) + FROM gateway_task_attempts + WHERE task_id = ANY($1::uuid[]))`, taskIDs).Scan(&admitted, &attempts); err != nil { + t.Fatalf("read queued admission state: %v", err) + } + if admitted != 1 || attempts != 1 { + t.Fatalf("claimed admission state admitted=%d attempts=%d, want 1/1", admitted, attempts) + } + + waitForTaskCount(t, ctx, pool, taskIDs, "succeeded", len(taskIDs), 30*time.Second) + assertAcceptanceAttempts(t, ctx, pool, taskIDs, len(taskIDs)) +} + +func TestAsyncWorkerThousandConcurrentSchedulingAcceptance(t *testing.T) { + databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) + if databaseURL == "" { + t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run thousand-task worker acceptance tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) + defer cancel() + applyMigration(t, ctx, databaseURL) + + apiDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 16) + if err != nil { + t.Fatalf("connect API store: %v", err) + } + defer apiDB.Close() + workerDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 24) + if err != nil { + t.Fatalf("connect Worker store: %v", err) + } + defer workerDB.Close() + pool, err := pgxpool.New(ctx, databaseURL) + if err != nil { + t.Fatalf("connect acceptance pool: %v", err) + } + defer pool.Close() + if _, err := pool.Exec(ctx, `UPDATE integration_platforms SET status = 'disabled' WHERE deleted_at IS NULL`); err != nil { + t.Fatalf("isolate acceptance platforms: %v", err) + } + + serverCtx, cancelServer := context.WithCancel(ctx) + defer cancelServer() + server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{ + AppEnv: "test", + HTTPAddr: ":0", + DatabaseURL: databaseURL, + IdentityMode: "hybrid", + JWTSecret: "test-secret", + BillingEngineMode: "observe", + CORSAllowedOrigin: "*", + ProcessRole: "api", + AsyncQueueWorkerEnabled: true, + AsyncWorkerHardLimit: 1000, + AsyncWorkerRefreshIntervalSeconds: 1, + }, apiDB, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})))) + defer server.Close() + workerServer := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{ + AppEnv: "test", + HTTPAddr: ":0", + DatabaseURL: databaseURL, + IdentityMode: "hybrid", + JWTSecret: "test-secret", + BillingEngineMode: "observe", + CORSAllowedOrigin: "*", + ProcessRole: "worker", + AsyncQueueWorkerEnabled: true, + AsyncWorkerHardLimit: 1000, + AsyncWorkerRefreshIntervalSeconds: 1, + }, workerDB, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})))) + defer workerServer.Close() + + adminToken := createAsyncAcceptanceAdmin(t, ctx, pool, server.URL) + if _, err := pool.Exec(ctx, ` +UPDATE gateway_user_groups +SET rate_limit_policy = '{"rules":[{"metric":"concurrent","limit":1024,"leaseTtlSeconds":120}]}'::jsonb +WHERE status = 'active'`); err != nil { + t.Fatalf("raise acceptance user-group concurrency: %v", err) + } + suffix := strconv.FormatInt(time.Now().UnixNano(), 10) + model := "worker-thousand-" + suffix + platform := createAsyncAcceptancePlatform(t, server.URL, adminToken, "thousand-"+suffix, "Thousand Task Simulation", 1000) + defer updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 1000, "disabled") + modelID := createAsyncAcceptanceModel(t, server.URL, adminToken, platform.ID, model, "video_generate", "override", 1000, 120) + + waitForAsyncWorkerMetric(t, workerServer.URL, "easyai_gateway_async_worker_capacity", 1000, 15*time.Second) + taskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 1000, 45*time.Second) + peakRunning, peakLeases := waitForAcceptanceTasks(t, ctx, pool, taskIDs, modelID, 2*time.Minute) + if peakRunning < 800 || peakLeases < 800 { + t.Fatalf("thousand-task peak running=%d leases=%d, want both >=800", peakRunning, peakLeases) + } + if peakLeases > 1000 { + t.Fatalf("thousand-task lease peak=%d exceeded model concurrent limit 1000", peakLeases) + } + assertAcceptanceAttempts(t, ctx, pool, taskIDs, len(taskIDs)) + t.Logf( + "千级并发证据: tasks=%d duration=%s running_peak=%d lease_peak=%d api_database_max_connections=16 worker_database_max_connections=24", + len(taskIDs), + 45*time.Second, + peakRunning, + peakLeases, + ) +} + func TestAsyncWorkerDynamicConcurrencyAcceptance(t *testing.T) { databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) if databaseURL == "" { @@ -295,6 +462,7 @@ func waitForAcceptanceTasks(t *testing.T, ctx context.Context, pool *pgxpool.Poo t.Helper() deadline := time.Now().Add(timeout) var peakRunning, peakLeases int64 + var lastRunning, lastSucceeded, lastFailed, lastLeases int64 for time.Now().Before(deadline) { var running, succeeded, failed, leases int64 if err := pool.QueryRow(ctx, ` @@ -316,6 +484,10 @@ WHERE scope_type = 'platform_model' } peakRunning = maxInt64(peakRunning, running) peakLeases = maxInt64(peakLeases, leases) + lastRunning = running + lastSucceeded = succeeded + lastFailed = failed + lastLeases = leases if failed > 0 { t.Fatalf("acceptance tasks entered failed/manual status: %d", failed) } @@ -324,7 +496,17 @@ WHERE scope_type = 'platform_model' } time.Sleep(50 * time.Millisecond) } - t.Fatalf("tasks did not finish within %s", timeout) + t.Fatalf( + "tasks did not finish within %s: running=%d succeeded=%d failed=%d leases=%d peak_running=%d peak_leases=%d total=%d", + timeout, + lastRunning, + lastSucceeded, + lastFailed, + lastLeases, + peakRunning, + peakLeases, + len(taskIDs), + ) return 0, 0 } diff --git a/apps/api/internal/httpapi/gemini_base64_stress_integration_test.go b/apps/api/internal/httpapi/gemini_base64_stress_integration_test.go new file mode 100644 index 0000000..17c9063 --- /dev/null +++ b/apps/api/internal/httpapi/gemini_base64_stress_integration_test.go @@ -0,0 +1,753 @@ +package httpapi + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "runtime" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/config" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" + "github.com/jackc/pgx/v5/pgxpool" +) + +func TestGeminiBase64ImageEditThousandSynchronousRequests(t *testing.T) { + if os.Getenv("AI_GATEWAY_RUN_GEMINI_BASE64_STRESS") != "1" { + t.Skip("set AI_GATEWAY_RUN_GEMINI_BASE64_STRESS=1 to run the Gemini Base64 stress acceptance") + } + databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) + if databaseURL == "" { + t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the Gemini Base64 stress acceptance") + } + + const ( + taskCount = 1000 + upstreamConcurrent = 64 + inputVariants = 16 + mediaConcurrent = 8 + ) + imageBytes := stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_IMAGE_BYTES", 256<<10) + upstreamDelay := time.Duration(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_UPSTREAM_DELAY_MS", 4000)) * time.Millisecond + maxHeapGrowth := uint64(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_MAX_HEAP_BYTES", 1<<30)) + clientConnections := stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_CLIENT_CONNECTIONS", 256) + testTimeout := time.Duration(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_TIMEOUT_SECONDS", 1200)) * time.Second + + ctx, cancel := context.WithTimeout(context.Background(), testTimeout) + defer cancel() + applyMigration(t, ctx, databaseURL) + + apiDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 16) + if err != nil { + t.Fatalf("connect API store: %v", err) + } + defer apiDB.Close() + poolConfig, err := pgxpool.ParseConfig(databaseURL) + if err != nil { + t.Fatalf("parse observation pool config: %v", err) + } + poolConfig.MaxConns = 4 + pool, err := pgxpool.NewWithConfig(ctx, poolConfig) + if err != nil { + t.Fatalf("connect observation pool: %v", err) + } + defer pool.Close() + if _, err := pool.Exec(ctx, ` +UPDATE integration_platforms SET status = 'disabled' WHERE deleted_at IS NULL; +UPDATE file_storage_channels SET status = 'disabled' WHERE deleted_at IS NULL; +UPDATE gateway_user_groups +SET rate_limit_policy = '{"rules":[ + {"metric":"concurrent","limit":1024,"leaseTtlSeconds":120}, + {"metric":"queue_size","limit":1000,"maxWaitSeconds":900} +]}'::jsonb +WHERE status = 'active'`); err != nil { + t.Fatalf("isolate Gemini stress acceptance: %v", err) + } + + inputPayloads := make([][]byte, inputVariants) + inputBase64 := make([]string, inputVariants) + validInputs := make(map[string]struct{}, inputVariants) + validInputHashes := make(map[string]struct{}, inputVariants) + for index := range inputPayloads { + inputPayloads[index] = stressImagePayload(imageBytes, byte(index+1)) + inputBase64[index] = base64.StdEncoding.EncodeToString(inputPayloads[index]) + validInputs[inputBase64[index]] = struct{}{} + validInputHashes[stressPayloadSHA256(inputPayloads[index])] = struct{}{} + } + outputPayload := stressImagePayload(imageBytes, 0xf3) + outputBase64 := base64.StdEncoding.EncodeToString(outputPayload) + outputHash := stressPayloadSHA256(outputPayload) + + uploadService := newGeminiStressUploadService(imageBytes, validInputHashes, outputHash) + uploadServer := httptest.NewServer(uploadService) + defer uploadServer.Close() + + var upstreamActive atomic.Int64 + var upstreamPeak atomic.Int64 + var upstreamCalls atomic.Int64 + var invalidInputs atomic.Int64 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalls.Add(1) + active := upstreamActive.Add(1) + defer upstreamActive.Add(-1) + updateAtomicPeak(&upstreamPeak, active) + + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + invalidInputs.Add(1) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + inlineData := geminiStressInlineData(body) + if strings.HasPrefix(inlineData, "data:") { + _, inlineData, _ = strings.Cut(inlineData, ",") + } + if _, ok := validInputs[inlineData]; !ok { + invalidInputs.Add(1) + http.Error(w, "missing or altered Gemini inlineData", http.StatusBadRequest) + return + } + time.Sleep(upstreamDelay) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "candidates": []any{map[string]any{ + "index": 0, + "content": map[string]any{ + "role": "model", + "parts": []any{map[string]any{ + "inlineData": map[string]any{ + "mimeType": "image/png", + "data": outputBase64, + }, + }}, + }, + "finishReason": "STOP", + }}, + "usageMetadata": map[string]any{ + "promptTokenCount": 1, + "candidatesTokenCount": 1, + "totalTokenCount": 2, + }, + }) + })) + defer upstream.Close() + + serverCtx, cancelServer := context.WithCancel(ctx) + defer cancelServer() + storageRoot := t.TempDir() + server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{ + AppEnv: "test", + HTTPAddr: ":0", + DatabaseURL: databaseURL, + DatabaseMaxConns: 16, + IdentityMode: "hybrid", + JWTSecret: "gemini-base64-stress-secret", + BillingEngineMode: "observe", + CORSAllowedOrigin: "*", + ProcessRole: "api", + LocalUploadedStorageDir: filepath.Join(storageRoot, "uploaded"), + LocalGeneratedStorageDir: filepath.Join(storageRoot, "generated"), + MediaRequestConcurrency: 16, + MediaMaterializationConcurrency: mediaConcurrent, + AsyncQueueWorkerEnabled: false, + AsyncWorkerHardLimit: 64, + }, apiDB, slog.New(slog.NewTextHandler(io.Discard, nil)))) + defer server.Close() + + adminToken := createAsyncAcceptanceAdmin(t, ctx, pool, server.URL) + suffix := strconv.FormatInt(time.Now().UnixNano(), 10) + uploadAPIKey := "local-gemini-stress-upload-key" + var storageChannel struct { + ID string `json:"id"` + } + doJSON(t, server.URL, http.MethodPost, "/api/admin/system/file-storage/channels", adminToken, map[string]any{ + "channelKey": "gemini-base64-stress-" + suffix, + "name": "Gemini Base64 Stress Shared Upload", + "provider": "server_main_openapi", + "uploadUrl": uploadServer.URL + "/upload", + "apiKey": uploadAPIKey, + "scenes": []string{ + store.FileStorageSceneUpload, + store.FileStorageSceneRequestAsset, + store.FileStorageSceneImageResult, + }, + "priority": 1, + "status": "enabled", + }, http.StatusCreated, &storageChannel) + uploadService.setAPIKey(uploadAPIKey) + + var gatewayUserID string + if err := pool.QueryRow(ctx, ` +SELECT id::text +FROM gateway_users +WHERE username LIKE 'async_acceptance_%' +ORDER BY created_at DESC +LIMIT 1`).Scan(&gatewayUserID); err != nil { + t.Fatalf("read Gemini stress user: %v", err) + } + doJSON(t, server.URL, http.MethodPatch, "/api/admin/users/"+gatewayUserID+"/wallet", adminToken, map[string]any{ + "currency": "resource", + "balance": 1000000000, + "reason": "seed isolated Gemini Base64 stress wallet", + }, http.StatusOK, &map[string]any{}) + model := "gemini-base64-image-edit-stress-" + suffix + var baseModel struct { + ID string `json:"id"` + } + doJSON(t, server.URL, http.MethodPost, "/api/admin/catalog/base-models", adminToken, map[string]any{ + "providerKey": "gemini", + "canonicalModelKey": "gemini:" + model, + "invocationName": model, + "providerModelName": model, + "modelType": []string{"image_edit"}, + "displayName": model, + "status": "active", + "capabilities": map[string]any{ + "image_edit": map[string]any{ + "support_base64_input": true, + "support_url_input": false, + }, + }, + "baseBillingConfig": map[string]any{}, + }, http.StatusCreated, &baseModel) + + var platform struct { + ID string `json:"id"` + } + doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms", adminToken, map[string]any{ + "provider": "gemini", + "platformKey": "gemini-base64-stress-" + suffix, + "name": "Gemini Base64 Stress", + "baseUrl": upstream.URL, + "authType": "bearer", + "credentials": map[string]any{"apiKey": "local-stress-key"}, + "priority": 1, + "status": "enabled", + "rateLimitPolicy": map[string]any{"rules": []any{ + map[string]any{"metric": "concurrent", "limit": upstreamConcurrent, "leaseTtlSeconds": 120}, + map[string]any{"metric": "queue_size", "limit": 1000, "maxWaitSeconds": 900}, + }}, + }, http.StatusCreated, &platform) + var platformModel struct { + ID string `json:"id"` + } + doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", adminToken, map[string]any{ + "baseModelId": baseModel.ID, + "providerModelName": model, + "modelAlias": model, + "modelType": []string{"image_edit"}, + "rateLimitPolicyMode": "override", + "rateLimitPolicy": map[string]any{"rules": []any{ + map[string]any{"metric": "concurrent", "limit": upstreamConcurrent, "leaseTtlSeconds": 120}, + map[string]any{"metric": "queue_size", "limit": 1000, "maxWaitSeconds": 900}, + }}, + "capabilityOverride": map[string]any{ + "image_edit": map[string]any{ + "support_base64_input": true, + "support_url_input": false, + }, + }, + }, http.StatusCreated, &platformModel) + + requestBodies := make([][]byte, inputVariants) + for index := range requestBodies { + requestBodies[index], err = json.Marshal(map[string]any{ + "contents": []any{map[string]any{ + "role": "user", + "parts": []any{ + map[string]any{"text": "把输入图片改成用于并发验收的蓝色赛博朋克风格"}, + map[string]any{"inlineData": map[string]any{ + "mimeType": "image/png", + "data": inputBase64[index], + }}, + }, + }}, + "generationConfig": map[string]any{ + "responseModalities": []any{"IMAGE"}, + }, + }) + if err != nil { + t.Fatalf("encode stress request %d: %v", index, err) + } + } + + runtime.GC() + var baseline runtime.MemStats + runtime.ReadMemStats(&baseline) + var peakHeap atomic.Uint64 + peakHeap.Store(baseline.HeapAlloc) + var queuePeak atomic.Int64 + var runningPeak atomic.Int64 + samplingDone := make(chan struct{}) + go sampleGeminiStressPressure(ctx, pool, model, &peakHeap, &queuePeak, &runningPeak, samplingDone) + + startedAt := time.Now() + loadClient := &http.Client{ + Timeout: testTimeout, + Transport: &http.Transport{ + MaxIdleConns: clientConnections, + MaxIdleConnsPerHost: clientConnections, + MaxConnsPerHost: clientConnections, + }, + } + var wg sync.WaitGroup + var requestErrors atomic.Int64 + errs := make(chan error, 20) + recordRequestError := func(err error) { + requestErrors.Add(1) + select { + case errs <- err: + default: + } + } + for index := 0; index < taskCount; index++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + server.URL+"/v1beta/models/"+model+":generateContent", + bytes.NewReader(requestBodies[index%inputVariants])) + if err != nil { + recordRequestError(err) + return + } + req.Header.Set("Authorization", "Bearer "+adminToken) + req.Header.Set("Content-Type", "application/json") + resp, err := loadClient.Do(req) + if err != nil { + recordRequestError(fmt.Errorf("request %d: %w", index, err)) + return + } + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + _ = resp.Body.Close() + recordRequestError(fmt.Errorf("request %d status=%d body=%s", index, resp.StatusCode, stressBodyPreview(body))) + return + } + foundOutput, readErr := stressStreamContainsAll(resp.Body, + []byte(outputBase64[:128]), + []byte(outputBase64[len(outputBase64)-128:]), + ) + _ = resp.Body.Close() + if readErr != nil { + recordRequestError(fmt.Errorf("request %d read response: %w", index, readErr)) + return + } + if !foundOutput { + recordRequestError(fmt.Errorf("request %d Gemini response did not contain Base64 image output", index)) + } + }(index) + } + wg.Wait() + close(samplingDone) + close(errs) + for err := range errs { + t.Error(err) + } + if requestErrors.Load() > 20 { + t.Errorf("%d additional request failures omitted", requestErrors.Load()-20) + } + if t.Failed() { + t.FailNow() + } + elapsed := time.Since(startedAt) + + var tasks, succeeded, attempts, duplicateAttemptTasks int + if err := pool.QueryRow(ctx, ` +SELECT + COUNT(*)::int, + COUNT(*) FILTER (WHERE task.status = 'succeeded')::int, + COALESCE(SUM(attempts.attempt_count), 0)::int, + COUNT(*) FILTER (WHERE attempts.attempt_count <> 1)::int +FROM gateway_tasks task +LEFT JOIN LATERAL ( + SELECT COUNT(*)::int AS attempt_count + FROM gateway_task_attempts attempt + 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 { + 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 upstreamCalls.Load() != taskCount || invalidInputs.Load() != 0 { + t.Fatalf("upstream calls=%d invalid_inputs=%d, want %d/0", upstreamCalls.Load(), invalidInputs.Load(), taskCount) + } + if upstreamPeak.Load() > mediaConcurrent || upstreamPeak.Load() < mediaConcurrent/2 { + t.Fatalf( + "upstream peak=%d, want 50%%-%d after media memory admission; queue_peak=%d running_peak=%d heap_peak=%d", + upstreamPeak.Load(), + mediaConcurrent, + queuePeak.Load(), + runningPeak.Load(), + peakHeap.Load()-baseline.HeapAlloc, + ) + } + if queuePeak.Load() == 0 { + t.Fatal("the synchronous distributed admission queue was never observed under the 1000-request burst") + } + requestAssetUploads, generatedUploads, retainedObjects, invalidUploads := uploadService.stats() + if invalidUploads != 0 { + t.Fatalf("shared upload validation failures=%d", invalidUploads) + } + if requestAssetUploads != inputVariants { + t.Fatalf("request asset uploads=%d, want %d deduplicated input variants", requestAssetUploads, inputVariants) + } + if generatedUploads != taskCount { + t.Fatalf("generated result uploads=%d, want %d", generatedUploads, taskCount) + } + if retainedObjects != inputVariants { + t.Fatalf("retained shared input objects=%d, want %d", retainedObjects, inputVariants) + } + uploadedFiles := stressDirectoryFileCount(t, filepath.Join(storageRoot, "uploaded")) + generatedFiles := stressDirectoryFileCount(t, filepath.Join(storageRoot, "generated")) + if uploadedFiles != 0 || generatedFiles != 0 { + t.Fatalf("local stress files uploaded=%d generated=%d, want shared upload only", uploadedFiles, generatedFiles) + } + heapGrowth := peakHeap.Load() - baseline.HeapAlloc + if heapGrowth > maxHeapGrowth { + 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", + taskCount, + imageBytes, + inputVariants, + clientConnections, + upstreamConcurrent, + mediaConcurrent, + elapsed.Round(time.Millisecond), + queuePeak.Load(), + runningPeak.Load(), + upstreamPeak.Load(), + heapGrowth, + requestAssetUploads, + generatedUploads, + uploadedFiles+generatedFiles, + ) +} + +type geminiStressUploadService struct { + expectedBytes int + outputHash string + inputHashes map[string]struct{} + + mu sync.RWMutex + apiKey string + retainedObjects map[string][]byte + + requestAssetUploads atomic.Int64 + generatedUploads atomic.Int64 + invalidUploads atomic.Int64 +} + +func newGeminiStressUploadService(expectedBytes int, inputHashes map[string]struct{}, outputHash string) *geminiStressUploadService { + return &geminiStressUploadService{ + expectedBytes: expectedBytes, + outputHash: outputHash, + inputHashes: inputHashes, + retainedObjects: map[string][]byte{}, + } +} + +func (s *geminiStressUploadService) setAPIKey(apiKey string) { + s.mu.Lock() + s.apiKey = apiKey + s.mu.Unlock() +} + +func (s *geminiStressUploadService) ServeHTTP(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.Path == "/upload": + s.handleUpload(w, r) + case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/objects/"): + s.handleObject(w, r) + default: + http.NotFound(w, r) + } +} + +func (s *geminiStressUploadService) handleUpload(w http.ResponseWriter, r *http.Request) { + s.mu.RLock() + apiKey := s.apiKey + s.mu.RUnlock() + if apiKey == "" || r.Header.Get("Authorization") != "Bearer "+apiKey { + s.invalidUploads.Add(1) + http.Error(w, "invalid upload authorization", http.StatusUnauthorized) + return + } + reader, err := r.MultipartReader() + if err != nil { + s.invalidUploads.Add(1) + http.Error(w, "invalid multipart body", http.StatusBadRequest) + return + } + var ( + filePayload []byte + scene string + ) + for { + part, nextErr := reader.NextPart() + if nextErr == io.EOF { + break + } + if nextErr != nil { + s.invalidUploads.Add(1) + http.Error(w, "read multipart body failed", http.StatusBadRequest) + return + } + switch part.FormName() { + case "file": + filePayload, err = io.ReadAll(io.LimitReader(part, int64(s.expectedBytes)+1)) + case "scene": + var raw []byte + raw, err = io.ReadAll(io.LimitReader(part, 128)) + scene = strings.TrimSpace(string(raw)) + default: + _, err = io.Copy(io.Discard, part) + } + _ = part.Close() + if err != nil { + s.invalidUploads.Add(1) + http.Error(w, "read upload field failed", http.StatusBadRequest) + return + } + } + if len(filePayload) != s.expectedBytes { + s.invalidUploads.Add(1) + http.Error(w, "unexpected upload size", http.StatusBadRequest) + return + } + hash := stressPayloadSHA256(filePayload) + switch scene { + case store.FileStorageSceneRequestAsset: + if _, ok := s.inputHashes[hash]; !ok { + s.invalidUploads.Add(1) + http.Error(w, "unexpected request asset hash", http.StatusBadRequest) + return + } + s.mu.Lock() + if _, exists := s.retainedObjects[hash]; !exists { + s.retainedObjects[hash] = filePayload + } + s.mu.Unlock() + s.requestAssetUploads.Add(1) + case store.FileStorageSceneImageResult: + if hash != s.outputHash { + s.invalidUploads.Add(1) + http.Error(w, "unexpected generated result hash", http.StatusBadRequest) + return + } + s.generatedUploads.Add(1) + default: + s.invalidUploads.Add(1) + http.Error(w, "unexpected upload scene", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "url": "http://" + r.Host + "/objects/" + hash, + }) +} + +func (s *geminiStressUploadService) handleObject(w http.ResponseWriter, r *http.Request) { + hash := strings.TrimPrefix(r.URL.Path, "/objects/") + s.mu.RLock() + payload, ok := s.retainedObjects[hash] + s.mu.RUnlock() + if !ok { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "image/png") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(payload) +} + +func (s *geminiStressUploadService) stats() (requestAssetUploads int, generatedUploads int, retainedObjects int, invalidUploads int) { + s.mu.RLock() + retainedObjects = len(s.retainedObjects) + s.mu.RUnlock() + return int(s.requestAssetUploads.Load()), int(s.generatedUploads.Load()), retainedObjects, int(s.invalidUploads.Load()) +} + +func stressPayloadSHA256(payload []byte) string { + sum := sha256.Sum256(payload) + return hex.EncodeToString(sum[:]) +} + +func stressStreamContainsAll(reader io.Reader, needles ...[]byte) (bool, error) { + maxNeedle := 0 + for _, needle := range needles { + if len(needle) > maxNeedle { + maxNeedle = len(needle) + } + } + if maxNeedle == 0 { + return true, nil + } + found := make([]bool, len(needles)) + buffer := make([]byte, 64<<10+maxNeedle-1) + carried := 0 + for { + n, err := reader.Read(buffer[carried:]) + total := carried + n + window := buffer[:total] + for index, needle := range needles { + if !found[index] && bytes.Contains(window, needle) { + found[index] = true + } + } + if err != nil && err != io.EOF { + return false, err + } + if err == io.EOF { + for _, matched := range found { + if !matched { + return false, nil + } + } + return true, nil + } + carried = min(maxNeedle-1, total) + copy(buffer[:carried], buffer[total-carried:total]) + } +} + +func stressEnvInt(t *testing.T, key string, fallback int) int { + t.Helper() + raw := strings.TrimSpace(os.Getenv(key)) + if raw == "" { + return fallback + } + value, err := strconv.Atoi(raw) + if err != nil || value <= 0 { + t.Fatalf("%s must be a positive integer", key) + } + return value +} + +func stressImagePayload(size int, marker byte) []byte { + payload := make([]byte, size) + copy(payload, []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}) + for index := 8; index < len(payload); index++ { + payload[index] = byte((index*31 + int(marker)) % 251) + } + if len(payload) > 8 { + payload[8] = marker + } + return payload +} + +func geminiStressInlineData(body map[string]any) string { + contents, _ := body["contents"].([]any) + for _, rawContent := range contents { + content, _ := rawContent.(map[string]any) + parts, _ := content["parts"].([]any) + for _, rawPart := range parts { + part, _ := rawPart.(map[string]any) + inline, _ := part["inlineData"].(map[string]any) + if inline == nil { + inline, _ = part["inline_data"].(map[string]any) + } + if data, _ := inline["data"].(string); data != "" { + return data + } + } + } + return "" +} + +func updateAtomicPeak(peak *atomic.Int64, value int64) { + for { + current := peak.Load() + if value <= current || peak.CompareAndSwap(current, value) { + return + } + } +} + +func sampleGeminiStressPressure( + ctx context.Context, + pool *pgxpool.Pool, + model string, + peakHeap *atomic.Uint64, + queuePeak *atomic.Int64, + runningPeak *atomic.Int64, + done <-chan struct{}, +) { + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-done: + return + case <-ticker.C: + var stats runtime.MemStats + runtime.ReadMemStats(&stats) + for { + current := peakHeap.Load() + if stats.HeapAlloc <= current || peakHeap.CompareAndSwap(current, stats.HeapAlloc) { + break + } + } + var waiting, running int64 + err := pool.QueryRow(ctx, ` +SELECT + COUNT(*) FILTER (WHERE admission.status = 'waiting')::bigint, + COUNT(*) FILTER (WHERE task.status = 'running')::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) + if err == nil { + updateAtomicPeak(queuePeak, waiting) + updateAtomicPeak(runningPeak, running) + } + } + } +} + +func stressDirectoryFileCount(t *testing.T, path string) int { + t.Helper() + entries, err := os.ReadDir(path) + if err != nil { + if os.IsNotExist(err) { + return 0 + } + t.Fatalf("read stress directory %s: %v", path, err) + } + count := 0 + for _, entry := range entries { + if !entry.IsDir() { + count++ + } + } + return count +} + +func stressBodyPreview(body []byte) string { + const limit = 512 + if len(body) <= limit { + return string(body) + } + return string(body[:limit]) + "..." +} diff --git a/apps/api/internal/httpapi/gemini_compat.go b/apps/api/internal/httpapi/gemini_compat.go index 98ce2b0..11063cd 100644 --- a/apps/api/internal/httpapi/gemini_compat.go +++ b/apps/api/internal/httpapi/gemini_compat.go @@ -117,6 +117,15 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) { writeGeminiTaskError(http.StatusUnauthorized, "unauthorized", nil, "unauthorized") return } + releaseRequestBody, err := s.acquireMediaRequestBodySlot(r.Context()) + if err != nil { + return + } + defer func() { + if releaseRequestBody != nil { + releaseRequestBody() + } + }() var native map[string]any if err := json.NewDecoder(r.Body).Decode(&native); err != nil { writeGeminiTaskError(http.StatusBadRequest, "invalid json body", nil, "invalid_json_body") @@ -150,6 +159,13 @@ 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. + native = nil + mapping.Body = nil + releaseRequestBody() + releaseRequestBody = nil createInput := store.CreateTaskInput{ Kind: mapping.Kind, Model: mapping.Model, diff --git a/apps/api/internal/httpapi/handlers.go b/apps/api/internal/httpapi/handlers.go index 7e9a173..e55c671 100644 --- a/apps/api/internal/httpapi/handlers.go +++ b/apps/api/internal/httpapi/handlers.go @@ -1171,7 +1171,7 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler { writeTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err)) return } - s.logger.Error("create task failed", "kind", kind, "error_category", "task_create_failed") + s.logger.Error("create task failed", "kind", kind, "error_category", "task_create_failed", "error", err) writeTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed") return } diff --git a/apps/api/internal/httpapi/request_preparation.go b/apps/api/internal/httpapi/request_preparation.go index f0c2a36..c0c2d7b 100644 --- a/apps/api/internal/httpapi/request_preparation.go +++ b/apps/api/internal/httpapi/request_preparation.go @@ -9,9 +9,12 @@ import ( "fmt" "mime" "net/http" + "os" "path/filepath" "sort" + "strconv" "strings" + "sync" "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" @@ -41,6 +44,11 @@ type requestAssetOptions struct { Source string } +type requestAssetLock struct { + mu sync.Mutex + refs int +} + func (s *Server) prepareTaskRequest(ctx context.Context, r *http.Request, user *auth.User, body map[string]any) (preparedTaskRequest, error) { preparedBody, err := s.prepareRequestAssetRefs(ctx, body) if err != nil { @@ -120,13 +128,9 @@ func (s *Server) prepareRequestAssetValue(ctx context.Context, value any, path [ sort.Strings(keys) for _, key := range keys { item := typed[key] - if decoded, ok, err := requestAssetFromValue(key, path, item, typed); err != nil { + if ref, ok, err := s.prepareRequestAssetField(ctx, key, path, item, typed); err != nil { return nil, err } else if ok { - ref, err := s.ensureRequestAsset(ctx, decoded) - if err != nil { - return nil, err - } next[key] = requestAssetWrapper(ref) continue } @@ -152,6 +156,67 @@ func (s *Server) prepareRequestAssetValue(ctx context.Context, value any, path [ } } +func (s *Server) prepareRequestAssetField(ctx context.Context, key string, path []string, value any, siblings map[string]any) (map[string]any, bool, error) { + text, ok := value.(string) + if !ok { + return nil, false, nil + } + raw := strings.TrimSpace(text) + if raw == "" || mediaURLString(raw) { + return nil, false, nil + } + if !strings.HasPrefix(strings.ToLower(raw), "data:") && + !strictRequestBase64Field(key, path) && + !likelyRequestBase64MediaField(key, path, raw) { + return nil, false, nil + } + if err := s.acquireMediaRequestSlot(ctx); err != nil { + return nil, false, err + } + defer s.releaseMediaRequestSlot() + decoded, ok, err := requestAssetFromValue(key, path, value, siblings) + if err != nil || !ok { + return nil, ok, err + } + ref, err := s.ensureRequestAsset(ctx, decoded) + return ref, true, err +} + +func (s *Server) acquireMediaRequestSlot(ctx context.Context) error { + if s.mediaRequestSlots == nil { + return nil + } + select { + case s.mediaRequestSlots <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (s *Server) releaseMediaRequestSlot() { + if s.mediaRequestSlots != nil { + <-s.mediaRequestSlots + } +} + +func (s *Server) acquireMediaRequestBodySlot(ctx context.Context) (func(), error) { + if s.mediaRequestBodySlots == nil { + return func() {}, nil + } + select { + case s.mediaRequestBodySlots <- struct{}{}: + var once sync.Once + return func() { + once.Do(func() { + <-s.mediaRequestBodySlots + }) + }, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + func requestAssetFromValue(key string, path []string, value any, siblings map[string]any) (decodedRequestAsset, bool, error) { text, ok := value.(string) if !ok { @@ -215,6 +280,8 @@ func (s *Server) ensureRequestAssetWithOptions(ctx context.Context, decoded deco if contentType == "" { contentType = "application/octet-stream" } + release := s.acquireRequestAssetLock(sha + "\x00" + contentType + "\x00" + options.UploadScene + "\x00" + strconv.FormatBool(options.RequirePublicURL)) + defer release() now := time.Now() if existing, ok, err := s.store.FindRequestAsset(ctx, sha, contentType); err != nil && !store.IsUndefinedDatabaseObject(err) { return nil, err @@ -286,6 +353,31 @@ func (s *Server) ensureRequestAssetWithOptions(ctx context.Context, decoded deco return requestAssetRef(asset), nil } +func (s *Server) acquireRequestAssetLock(key string) func() { + s.requestAssetLocksMu.Lock() + if s.requestAssetLocks == nil { + s.requestAssetLocks = map[string]*requestAssetLock{} + } + entry := s.requestAssetLocks[key] + if entry == nil { + entry = &requestAssetLock{} + s.requestAssetLocks[key] = entry + } + entry.refs++ + s.requestAssetLocksMu.Unlock() + + entry.mu.Lock() + return func() { + entry.mu.Unlock() + s.requestAssetLocksMu.Lock() + entry.refs-- + if entry.refs == 0 { + delete(s.requestAssetLocks, key) + } + s.requestAssetLocksMu.Unlock() + } +} + func requestConversationKey(r *http.Request, body map[string]any) string { if r != nil { if value := strings.TrimSpace(r.Header.Get("X-EasyAI-Conversation-ID")); value != "" { @@ -397,6 +489,16 @@ func requestAssetStillUsable(asset store.RequestAsset, now time.Time) bool { if asset.ExpiresAt != nil && !asset.ExpiresAt.After(now) { return false } + if strings.EqualFold(strings.TrimSpace(asset.StorageProvider), "local_static") { + localPath := strings.TrimSpace(asset.LocalPath) + info, err := os.Stat(localPath) + if localPath == "" || err != nil || !info.Mode().IsRegular() { + return false + } + if asset.ByteSize > 0 && info.Size() != asset.ByteSize { + return false + } + } return strings.TrimSpace(asset.URL) != "" } diff --git a/apps/api/internal/httpapi/request_preparation_test.go b/apps/api/internal/httpapi/request_preparation_test.go index 25fa05e..638d1a9 100644 --- a/apps/api/internal/httpapi/request_preparation_test.go +++ b/apps/api/internal/httpapi/request_preparation_test.go @@ -16,6 +16,7 @@ import ( "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/config" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" ) func TestRequestAssetFromValueDetectsDataURLAndRawBase64(t *testing.T) { @@ -66,6 +67,33 @@ func TestRequestAssetFromValueDetectsGeminiInlineData(t *testing.T) { } } +func TestRequestAssetStillUsableRequiresExistingLocalFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "asset.png") + if err := os.WriteFile(path, []byte("image"), 0o600); err != nil { + t.Fatalf("write local request asset: %v", err) + } + asset := store.RequestAsset{ + URL: "http://127.0.0.1/static/uploaded/asset.png", + StorageProvider: "local_static", + LocalPath: path, + ByteSize: 5, + } + if !requestAssetStillUsable(asset, time.Now()) { + t.Fatal("existing local request asset was rejected") + } + asset.ByteSize = 6 + if requestAssetStillUsable(asset, time.Now()) { + t.Fatal("truncated local request asset was treated as reusable") + } + asset.ByteSize = 5 + if err := os.Remove(path); err != nil { + t.Fatalf("remove local request asset: %v", err) + } + if requestAssetStillUsable(asset, time.Now()) { + t.Fatal("missing local request asset was treated as reusable") + } +} + func TestCanonicalConversationMessageHashUsesTextAndAssetRefs(t *testing.T) { message := map[string]any{ "role": "user", diff --git a/apps/api/internal/httpapi/server.go b/apps/api/internal/httpapi/server.go index df288ed..ae17df6 100644 --- a/apps/api/internal/httpapi/server.go +++ b/apps/api/internal/httpapi/server.go @@ -32,6 +32,10 @@ type Server struct { runner *runner.Service logger *slog.Logger geminiUploadSessions sync.Map + requestAssetLocksMu sync.Mutex + requestAssetLocks map[string]*requestAssetLock + mediaRequestSlots chan struct{} + mediaRequestBodySlots chan struct{} securityEventReceiver http.Handler securityEventManager *ssfreceiver.ConnectionManager identityRuntime *identityruntime.Manager @@ -68,16 +72,25 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han } func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Store, logger *slog.Logger) http.Handler { + if cfg.MediaMaterializationConcurrency == 0 { + cfg.MediaMaterializationConcurrency = 8 + } + if cfg.MediaRequestConcurrency == 0 { + cfg.MediaRequestConcurrency = 16 + } securityEventMetrics := &ssfreceiver.Metrics{} server := &Server{ - ctx: ctx, - cfg: cfg, - store: db, - oidcUserResolver: db, - auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken), - runner: runner.New(cfg, db, logger, securityEventMetrics), - logger: logger, - billingMetrics: securityEventMetrics, + ctx: ctx, + cfg: cfg, + store: db, + oidcUserResolver: db, + auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken), + runner: runner.New(cfg, db, logger, securityEventMetrics), + logger: logger, + billingMetrics: securityEventMetrics, + requestAssetLocks: map[string]*requestAssetLock{}, + mediaRequestSlots: make(chan struct{}, cfg.MediaMaterializationConcurrency), + mediaRequestBodySlots: make(chan struct{}, cfg.MediaRequestConcurrency), } server.auth.ServerMainInternalKey = cfg.ServerMainInternalKey server.auth.ServerMainInternalSecret = cfg.ServerMainInternalSecret diff --git a/apps/api/internal/runner/admission.go b/apps/api/internal/runner/admission.go index ed54600..662348f 100644 --- a/apps/api/internal/runner/admission.go +++ b/apps/api/internal/runner/admission.go @@ -3,6 +3,7 @@ package runner import ( "context" "errors" + "hash/fnv" "strings" "time" @@ -22,6 +23,17 @@ type taskAdmissionPlan struct { Eligible bool } +type admissionTaskWaiter struct { + wake chan struct{} + waiterID string +} + +const ( + asyncWorkerCapacityScopeKey = "global" + asyncWorkerQueueLimit = 10000 + asyncWorkerMaxWaitSeconds = 24 * 60 * 60 +) + func distributedAdmissionModelType(modelType string) bool { switch strings.ToLower(strings.TrimSpace(modelType)) { case "image_generate", @@ -99,6 +111,12 @@ func (s *Service) buildTaskAdmissionPlan(ctx context.Context, task store.Gateway continue } scopes, groupID := s.admissionScopes(ctx, user, candidate) + if task.AsyncMode { + scopes, err = s.withAsyncWorkerCapacityScope(ctx, scopes) + if err != nil { + return taskAdmissionPlan{}, err + } + } hasConcurrentLimit := false for _, scope := range scopes { if scope.ConcurrentLimit > 0 { @@ -128,10 +146,67 @@ func (s *Service) buildTaskAdmissionPlan(ctx context.Context, task store.Gateway return taskAdmissionPlan{}, store.ErrNoModelCandidate } +func (s *Service) withAsyncWorkerCapacityScope(ctx context.Context, scopes []store.AdmissionScope) ([]store.AdmissionScope, error) { + capacity, err := s.store.ActiveWorkerCapacity(ctx) + if err != nil { + return nil, err + } + if capacity <= 0 { + return scopes, nil + } + out := append([]store.AdmissionScope{}, scopes...) + out = append(out, store.AdmissionScope{ + ScopeType: "worker_capacity", + ScopeKey: asyncWorkerCapacityScopeKey, + ScopeName: "async worker capacity", + ConcurrentLimit: float64(capacity), + Amount: 1, + LeaseTTLSeconds: 120, + QueueLimit: asyncWorkerQueueLimit, + MaxWaitSeconds: asyncWorkerMaxWaitSeconds, + }) + return out, nil +} + func (s *Service) tryTaskAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan, waiterID string) (store.TaskAdmissionResult, error) { return s.tryTaskAdmissionWithAdmittedHook(ctx, task, plan, waiterID, nil) } +func (s *Service) activeAsyncTaskAdmission( + ctx context.Context, + task store.GatewayTask, + plan taskAdmissionPlan, +) (store.TaskAdmissionResult, bool, error) { + if !task.AsyncMode { + return store.TaskAdmissionResult{}, false, nil + } + admission, err := s.store.GetTaskAdmission(ctx, task.ID) + if errors.Is(err, pgx.ErrNoRows) { + return store.TaskAdmissionResult{}, false, nil + } + if err != nil { + return store.TaskAdmissionResult{}, false, err + } + if admission.Status != "admitted" || + admission.PlatformID != plan.Candidate.PlatformID || + admission.PlatformModelID != plan.Candidate.PlatformModelID || + admission.UserGroupID != plan.GroupID { + return store.TaskAdmissionResult{}, false, nil + } + leases, err := s.store.ActiveTaskAdmissionLeases(ctx, task.ID) + if err != nil { + return store.TaskAdmissionResult{}, false, err + } + if len(leases) == 0 { + return store.TaskAdmissionResult{}, false, nil + } + return store.TaskAdmissionResult{ + Admission: admission, + Admitted: true, + Leases: leases, + }, true, nil +} + func (s *Service) tryTaskAdmissionWithAdmittedHook( ctx context.Context, task store.GatewayTask, @@ -139,21 +214,7 @@ func (s *Service) tryTaskAdmissionWithAdmittedHook( waiterID string, onAdmitted func(pgx.Tx) error, ) (store.TaskAdmissionResult, error) { - mode := "sync" - if task.AsyncMode { - mode = "async" - } - input := store.TaskAdmissionInput{ - TaskID: task.ID, - PlatformID: plan.Candidate.PlatformID, - PlatformModelID: plan.Candidate.PlatformModelID, - UserGroupID: plan.GroupID, - QueueKey: plan.Candidate.QueueKey, - Mode: mode, - Priority: task.Priority, - WaiterID: waiterID, - Scopes: plan.Scopes, - } + input := taskAdmissionInput(task, plan, waiterID) current, currentErr := s.store.GetTaskAdmission(ctx, task.ID) if currentErr != nil && !errors.Is(currentErr, pgx.ErrNoRows) { return store.TaskAdmissionResult{}, currentErr @@ -186,14 +247,34 @@ func (s *Service) tryTaskAdmissionWithAdmittedHook( return result, err } +func taskAdmissionInput(task store.GatewayTask, plan taskAdmissionPlan, waiterID string) store.TaskAdmissionInput { + mode := "sync" + if task.AsyncMode { + mode = "async" + } + return store.TaskAdmissionInput{ + TaskID: task.ID, + PlatformID: plan.Candidate.PlatformID, + PlatformModelID: plan.Candidate.PlatformModelID, + UserGroupID: plan.GroupID, + QueueKey: plan.Candidate.QueueKey, + Mode: mode, + Priority: task.Priority, + WaiterID: waiterID, + Scopes: plan.Scopes, + } +} + func (s *Service) waitForSynchronousAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan) (store.TaskAdmissionResult, error) { waiterID := uuid.NewString() return s.waitForTaskAdmission(ctx, task, plan, waiterID) } func (s *Service) waitForTaskAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan, waiterID string) (store.TaskAdmissionResult, error) { - ticker := time.NewTicker(time.Second) - defer ticker.Stop() + taskWake, unregister := s.registerAdmissionTaskWaiter(task.ID, waiterID) + defer unregister() + pollTimer := time.NewTimer(admissionPollInterval(task.ID)) + defer pollTimer.Stop() for { result, err := s.tryTaskAdmission(ctx, task, plan, waiterID) if err != nil || result.Admitted { @@ -204,12 +285,19 @@ func (s *Service) waitForTaskAdmission(ctx context.Context, task store.GatewayTa _ = s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID) s.observeTaskAdmission("cancelled") return store.TaskAdmissionResult{}, ctx.Err() - case <-s.currentAdmissionWake(): - case <-ticker.C: + case <-taskWake: + case <-pollTimer.C: + pollTimer.Reset(admissionPollInterval(task.ID)) } } } +func admissionPollInterval(taskID string) time.Duration { + hasher := fnv.New32a() + _, _ = hasher.Write([]byte(taskID)) + return 30*time.Second + time.Duration(hasher.Sum32()%5000)*time.Millisecond +} + func (s *Service) ensureCandidateAdmission( ctx context.Context, task store.GatewayTask, @@ -218,6 +306,13 @@ func (s *Service) ensureCandidateAdmission( candidate store.RuntimeModelCandidate, ) (store.TaskAdmissionResult, bool, error) { scopes, groupID := s.admissionScopes(ctx, user, candidate) + if task.AsyncMode { + var err error + scopes, err = s.withAsyncWorkerCapacityScope(ctx, scopes) + if err != nil { + return store.TaskAdmissionResult{}, true, err + } + } hasConcurrentLimit := false for _, scope := range scopes { if scope.ConcurrentLimit > 0 { @@ -270,10 +365,12 @@ func (s *Service) observeTaskAdmissionWait(wait time.Duration) { func (s *Service) StartAdmissionNotifier(ctx context.Context) { s.admissionListener.Do(func() { + go s.dispatchWaitingSynchronousAdmissions(ctx) + go s.renewSynchronousAdmissionWaiters(ctx) go func() { for ctx.Err() == nil { - err := s.store.ListenTaskAdmissionNotifications(ctx, func(string) { - s.broadcastAdmissionWake() + err := s.store.ListenTaskAdmissionNotifications(ctx, func(taskID string) { + s.signalAdmissionWake(taskID) }) if ctx.Err() != nil { return @@ -293,17 +390,99 @@ func (s *Service) StartAdmissionNotifier(ctx context.Context) { }) } -func (s *Service) currentAdmissionWake() <-chan struct{} { - s.admissionWakeMu.Lock() - defer s.admissionWakeMu.Unlock() - return s.admissionWake +func (s *Service) dispatchWaitingSynchronousAdmissions(ctx context.Context) { + for { + select { + case <-ctx.Done(): + return + case <-s.admissionWake: + } + taskIDs, err := s.store.ListWaitingTaskAdmissionIDs(ctx, 1) + if err != nil { + if s.logger != nil { + s.logger.Warn("list waiting synchronous admissions failed", "error", err) + } + continue + } + s.admissionWakeMu.Lock() + waiters := make([]chan struct{}, 0, len(taskIDs)) + for _, taskID := range taskIDs { + if waiter := s.admissionTaskWaiters[taskID]; waiter != nil { + waiters = append(waiters, waiter.wake) + } + } + s.admissionWakeMu.Unlock() + for _, taskWake := range waiters { + select { + case taskWake <- struct{}{}: + default: + } + } + } } -func (s *Service) broadcastAdmissionWake() { +func (s *Service) renewSynchronousAdmissionWaiters(ctx context.Context) { + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + s.admissionWakeMu.Lock() + taskIDs := make([]string, 0, len(s.admissionTaskWaiters)) + waiterIDs := make([]string, 0, len(s.admissionTaskWaiters)) + for taskID, waiter := range s.admissionTaskWaiters { + if waiter == nil || strings.TrimSpace(waiter.waiterID) == "" { + continue + } + taskIDs = append(taskIDs, taskID) + waiterIDs = append(waiterIDs, waiter.waiterID) + } + s.admissionWakeMu.Unlock() + if err := s.store.RenewTaskAdmissionWaiters(ctx, taskIDs, waiterIDs); err != nil && s.logger != nil { + s.logger.Warn("renew synchronous admission waiters failed", "waiterCount", len(taskIDs), "error", err) + } + } +} + +func (s *Service) registerAdmissionTaskWaiter(taskID string, waiterID string) (<-chan struct{}, func()) { s.admissionWakeMu.Lock() - close(s.admissionWake) - s.admissionWake = make(chan struct{}) + waiter := &admissionTaskWaiter{wake: make(chan struct{}, 1), waiterID: waiterID} + s.admissionTaskWaiters[taskID] = waiter s.admissionWakeMu.Unlock() + return waiter.wake, func() { + s.admissionWakeMu.Lock() + if s.admissionTaskWaiters[taskID] == waiter { + delete(s.admissionTaskWaiters, taskID) + } + s.admissionWakeMu.Unlock() + } +} + +func (s *Service) signalAdmissionWake(taskID string) { + taskID = strings.TrimSpace(taskID) + if taskID == "" || taskID == "*" { + select { + case s.admissionWake <- struct{}{}: + default: + } + } else { + s.admissionWakeMu.Lock() + waiter := s.admissionTaskWaiters[taskID] + s.admissionWakeMu.Unlock() + if waiter != nil { + select { + case waiter.wake <- struct{}{}: + default: + } + } + } + select { + case s.asyncAdmissionWake <- struct{}{}: + default: + } } func (s *Service) SubmitAsyncTask(ctx context.Context, task store.GatewayTask) error { @@ -335,28 +514,49 @@ func (s *Service) SubmitAsyncTask(ctx context.Context, task store.GatewayTask) e } return nil } - result, err := s.tryTaskAdmissionWithAdmittedHook(ctx, task, plan, "", func(tx pgx.Tx) error { - return s.enqueueAsyncTaskTx(ctx, tx, task.ID, asyncTaskInsertOpts(task)) - }) - if err != nil { + input := taskAdmissionInput(task, plan, "") + current, currentErr := s.store.GetTaskAdmission(ctx, task.ID) + if currentErr != nil && !errors.Is(currentErr, pgx.ErrNoRows) { + return currentErr + } + if currentErr == nil && current.Status == "waiting" && + (current.PlatformID != input.PlatformID || current.PlatformModelID != input.PlatformModelID || current.UserGroupID != input.UserGroupID) { + if _, err := s.store.RebindWaitingTaskAdmission(ctx, input); err != nil { + return err + } + s.observeTaskAdmission("candidate_migrated") + } + // Register FIFO order without creating a River job. A Worker dispatcher + // first reserves both business concurrency and a live execution slot, then + // creates the unique River job in the same transaction. + if _, err := s.store.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil { if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) { return ctx.Err() } _, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error()) return err } - if !result.Admitted { - if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) { - return ctx.Err() - } - return nil - } if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) { return ctx.Err() } return nil } +func (s *Service) dispatchWaitingAsyncTask(ctx context.Context, task store.GatewayTask) error { + user := authUserFromTask(task) + plan, err := s.buildTaskAdmissionPlan(ctx, task, user) + if err != nil { + return err + } + if !plan.Eligible { + return s.EnqueueAsyncTask(ctx, task) + } + _, err = s.tryTaskAdmissionWithAdmittedHook(ctx, task, plan, "", func(tx pgx.Tx) error { + return s.enqueueAsyncTaskTx(ctx, tx, task.ID, asyncTaskInsertOpts(task)) + }) + return err +} + func (s *Service) cancelAsyncSubmissionIfDisconnected(ctx context.Context, taskID string) bool { if ctx.Err() == nil { return false @@ -380,10 +580,10 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) { select { case <-ctx.Done(): return - case <-s.currentAdmissionWake(): + case <-s.asyncAdmissionWake: case <-ticker.C: } - taskIDs, err := s.store.ListWaitingAsyncAdmissionTaskIDs(ctx, 100) + taskIDs, err := s.store.ListWaitingAsyncAdmissionTaskIDs(ctx, 1000) if err != nil { s.logger.Warn("list waiting async admissions failed", "error", err) continue @@ -399,7 +599,7 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) { if task.Status != "queued" || task.RiverJobID > 0 { continue } - if err := s.SubmitAsyncTask(ctx, task); err != nil { + if err := s.dispatchWaitingAsyncTask(ctx, task); err != nil { s.logger.Warn("dispatch waiting async admission failed", "taskID", taskID, "error", err) } } diff --git a/apps/api/internal/runner/proxy.go b/apps/api/internal/runner/proxy.go index 9f74f10..d0ccdc7 100644 --- a/apps/api/internal/runner/proxy.go +++ b/apps/api/internal/runner/proxy.go @@ -21,6 +21,11 @@ type httpClientCache struct { const providerHTTPClientTimeout = 10 * time.Minute +const ( + providerHTTPMaxIdleConnections = 2048 + providerHTTPMaxIdleConnectionsPerHost = 1024 +) + func newHTTPClientCache() *httpClientCache { return &httpClientCache{ none: newHTTPClient(nil), @@ -68,6 +73,12 @@ func (c *httpClientCache) customClient(rawProxy string) (*http.Client, error) { func newHTTPClient(proxy func(*http.Request) (*url.URL, error)) *http.Client { transport := http.DefaultTransport.(*http.Transport).Clone() transport.Proxy = proxy + // Media workers may keep hundreds of slow upstream tasks in flight and + // poll the same provider repeatedly. The standard library only retains two + // idle connections per host, which turns a high-concurrency polling load + // into avoidable TCP/TLS handshakes. + transport.MaxIdleConns = providerHTTPMaxIdleConnections + transport.MaxIdleConnsPerHost = providerHTTPMaxIdleConnectionsPerHost return &http.Client{ Timeout: providerHTTPClientTimeout, Transport: transport, diff --git a/apps/api/internal/runner/proxy_test.go b/apps/api/internal/runner/proxy_test.go index cd6c385..c066d3f 100644 --- a/apps/api/internal/runner/proxy_test.go +++ b/apps/api/internal/runner/proxy_test.go @@ -16,6 +16,20 @@ func TestProviderHTTPClientTimeoutAllowsLongRunningMediaRequests(t *testing.T) { if client.Timeout != 10*time.Minute { t.Fatalf("unexpected provider HTTP timeout: got %s want %s", client.Timeout, 10*time.Minute) } + transport, ok := client.Transport.(*http.Transport) + if !ok { + t.Fatalf("provider transport type = %T, want *http.Transport", client.Transport) + } + if transport.MaxIdleConns != providerHTTPMaxIdleConnections || + transport.MaxIdleConnsPerHost != providerHTTPMaxIdleConnectionsPerHost { + t.Fatalf( + "provider idle connection pool = %d/%d, want %d/%d", + transport.MaxIdleConns, + transport.MaxIdleConnsPerHost, + providerHTTPMaxIdleConnections, + providerHTTPMaxIdleConnectionsPerHost, + ) + } } func TestPlatformProxyModeNoneIgnoresEnvironmentProxy(t *testing.T) { diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index 662475b..0f172d8 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -40,8 +40,11 @@ type Service struct { asyncCapacityLoader func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) admissionWakeMu sync.Mutex admissionWake chan struct{} + asyncAdmissionWake chan struct{} + admissionTaskWaiters map[string]*admissionTaskWaiter admissionListener sync.Once asyncClientFactory func(int) (asyncExecutionClient, error) + mediaResultSlots chan struct{} billingMetrics billingMetricsObserver } @@ -93,6 +96,9 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b if cfg.AsyncWorkerRefreshIntervalSeconds == 0 { cfg.AsyncWorkerRefreshIntervalSeconds = 5 } + if cfg.MediaMaterializationConcurrency == 0 { + cfg.MediaMaterializationConcurrency = 8 + } if cfg.TaskProgressCallbackTimeoutMS == 0 { cfg.TaskProgressCallbackTimeoutMS = 5000 } @@ -139,9 +145,12 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b "universal": clients.UniversalClient{HTTPClient: httpClients.none, ScriptExecutor: scriptExecutor}, "simulation": clients.SimulationClient{}, }, - httpClients: httpClients, - workerInstanceID: asyncWorkerID(), - admissionWake: make(chan struct{}), + httpClients: httpClients, + workerInstanceID: asyncWorkerID(), + admissionWake: make(chan struct{}, 4096), + asyncAdmissionWake: make(chan struct{}, 1), + admissionTaskWaiters: map[string]*admissionTaskWaiter{}, + mediaResultSlots: make(chan struct{}, cfg.MediaMaterializationConcurrency), } if len(observers) > 0 { service.billingMetrics = observers[0] @@ -539,7 +548,11 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, var admissionResult store.TaskAdmissionResult var admissionErr error if task.AsyncMode { - admissionResult, admissionErr = s.tryTaskAdmission(ctx, task, plan, "") + var alreadyAdmitted bool + admissionResult, alreadyAdmitted, admissionErr = s.activeAsyncTaskAdmission(ctx, task, plan) + if admissionErr == nil && !alreadyAdmitted { + admissionResult, admissionErr = s.tryTaskAdmission(ctx, task, plan, "") + } if admissionErr == nil && !admissionResult.Admitted { _ = s.store.ReleaseTaskPreparation(context.WithoutCancel(ctx), task.ID, task.ExecutionToken) return Result{Task: task, Output: task.Result}, &TaskQueuedError{Delay: time.Second} @@ -1272,6 +1285,15 @@ func (s *Service) runCandidate( }) return clients.Response{}, err } + mediaSlotPreacquired := mediaTaskNeedsPreUpstreamMaterializationSlot(candidate.ModelType) + releaseMediaSlot := func() {} + if mediaSlotPreacquired { + releaseMediaSlot, err = s.acquireMediaMaterializationSlot(ctx) + if err != nil { + return clients.Response{}, err + } + defer releaseMediaSlot() + } providerBody, err = s.hydrateProviderRequestAssets(ctx, providerBody, candidate) if err != nil { _ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{ @@ -1465,6 +1487,21 @@ func (s *Service) runCandidate( } return clients.Response{}, err } + if task.AsyncMode { + // Async callers consume the canonical task result through the task API, + // never the original provider wire response. Release the raw JSON before + // decoding/uploading inline media so a large Base64 response is not held + // twice during the most memory-intensive phase. + response.Wire = nil + submissionWire = nil + } + if !mediaSlotPreacquired { + releaseMediaSlot, err = s.acquireMediaResultSlot(ctx, response.Result) + if err != nil { + return clients.Response{}, err + } + defer releaseMediaSlot() + } uploadedResult, err := s.uploadGeneratedAssets(ctx, task.ID, task.Kind, response.Result) if err != nil { metrics := mergeMetrics(taskMetrics(task, user, body, candidate, response, simulated), parameterPreprocessingMetrics(preprocessing), map[string]any{ diff --git a/apps/api/internal/runner/upload.go b/apps/api/internal/runner/upload.go index 526748d..1e3ce60 100644 --- a/apps/api/internal/runner/upload.go +++ b/apps/api/internal/runner/upload.go @@ -20,6 +20,7 @@ import ( "os" "path/filepath" "strings" + "sync" "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/clients" @@ -77,6 +78,40 @@ func defaultGeneratedAssetUploadPolicy() generatedAssetUploadPolicy { } } +func (s *Service) acquireMediaMaterializationSlot(ctx context.Context) (func(), error) { + if s.mediaResultSlots == nil { + return func() {}, nil + } + select { + case s.mediaResultSlots <- struct{}{}: + var once sync.Once + return func() { + once.Do(func() { + <-s.mediaResultSlots + }) + }, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (s *Service) acquireMediaResultSlot(ctx context.Context, result map[string]any) (func(), error) { + if s.mediaResultSlots == nil || + (!TaskResultHasInlineBinary(result) && !generatedRawValueHasInlineMedia(result["raw"], "", nil)) { + return func() {}, nil + } + return s.acquireMediaMaterializationSlot(ctx) +} + +func mediaTaskNeedsPreUpstreamMaterializationSlot(modelType string) bool { + switch canonicalModelType(modelType) { + case "image_generate", "image_edit", "audio_generate", "text_to_speech", "voice_clone", "omni": + return true + default: + return false + } +} + func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, taskKind string, result map[string]any) (map[string]any, error) { data, _ := result["data"].([]any) rawNeedsUpload := generatedRawValueHasInlineMedia(result["raw"], "", nil) diff --git a/apps/api/internal/runner/upload_test.go b/apps/api/internal/runner/upload_test.go index ee97fd4..f731e47 100644 --- a/apps/api/internal/runner/upload_test.go +++ b/apps/api/internal/runner/upload_test.go @@ -53,6 +53,60 @@ func TestGeneratedAssetDecisionUploadsInlineImageBase64(t *testing.T) { } } +func TestMediaResultMaterializationConcurrencyIsBounded(t *testing.T) { + service := New(config.Config{MediaMaterializationConcurrency: 1}, nil, nil) + result := map[string]any{ + "data": []any{map[string]any{ + "b64_json": base64.StdEncoding.EncodeToString([]byte("inline image")), + "mime_type": "image/png", + }}, + } + release, err := service.acquireMediaResultSlot(context.Background(), result) + if err != nil { + t.Fatalf("acquire first media result slot: %v", err) + } + + waitCtx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) + defer cancel() + if _, err := service.acquireMediaResultSlot(waitCtx, result); err == nil { + t.Fatal("second media result materialization bypassed configured concurrency") + } + release() + + release, err = service.acquireMediaResultSlot(context.Background(), result) + if err != nil { + t.Fatalf("acquire released media result slot: %v", err) + } + release() +} + +func TestMediaTaskPreacquiresMaterializationSlotForInlineResultTypes(t *testing.T) { + t.Parallel() + for _, modelType := range []string{ + "image_generate", + "image_edit", + "audio_generate", + "text_to_speech", + "voice_clone", + "omni", + } { + if !mediaTaskNeedsPreUpstreamMaterializationSlot(modelType) { + t.Fatalf("%s should preacquire the media materialization slot", modelType) + } + } + for _, modelType := range []string{ + "text_generate", + "text_embedding", + "text_rerank", + "video_generate", + "image_to_video", + } { + if mediaTaskNeedsPreUpstreamMaterializationSlot(modelType) { + t.Fatalf("%s should not hold a media materialization slot across the upstream call", modelType) + } + } +} + func TestGeneratedAssetDecisionUploadsVectorDocumentWithoutChangingType(t *testing.T) { item := map[string]any{ "type": "file", diff --git a/apps/api/internal/store/admission_queue.go b/apps/api/internal/store/admission_queue.go index 5e8a6b4..f0f6bc0 100644 --- a/apps/api/internal/store/admission_queue.go +++ b/apps/api/internal/store/admission_queue.go @@ -5,6 +5,7 @@ import ( "database/sql" "errors" "fmt" + "math" "sort" "strings" "time" @@ -15,6 +16,7 @@ import ( const ( admissionWaiterLeaseTTL = 15 * time.Second admissionNotifyChannel = "gateway_task_admission" + executionSlotWaitMax = 24 * time.Hour ) var ErrQueueTimeout = errors.New("task admission queue wait timed out") @@ -113,6 +115,111 @@ func (s *Store) TryTaskAdmissionWithAdmittedHook( return s.tryTaskAdmission(ctx, input, onAdmitted) } +// QueueTaskAdmissionWithHook durably registers an asynchronous FIFO waiter and +// runs onQueued in the same transaction. It intentionally does not reserve +// concurrency leases: the Worker dispatcher obtains business and execution +// capacity only when the task reaches the head of every applicable queue. +func (s *Store) QueueTaskAdmissionWithHook( + ctx context.Context, + input TaskAdmissionInput, + onQueued func(pgx.Tx) error, +) (TaskAdmission, error) { + if err := validateTaskAdmissionInput(input); err != nil { + return TaskAdmission{}, err + } + if input.Mode != "async" { + return TaskAdmission{}, errors.New("queued admission registration requires asynchronous mode") + } + tx, err := s.pool.Begin(ctx) + if err != nil { + return TaskAdmission{}, err + } + defer tx.Rollback(ctx) + + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil { + return TaskAdmission{}, err + } + var taskActive bool + if err := tx.QueryRow(ctx, ` +SELECT status IN ('queued', 'running') +FROM gateway_tasks +WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil { + return TaskAdmission{}, err + } + if !taskActive { + return TaskAdmission{}, ErrTaskExecutionFinished + } + + admission, found, err := loadTaskAdmissionTx(ctx, tx, input.TaskID) + if err != nil { + return TaskAdmission{}, err + } + scopes := normalizedAdmissionScopes(input.Scopes) + lockScopes := append([]AdmissionScope{}, scopes...) + if found { + lockScopes = append(lockScopes, + AdmissionScope{ScopeType: "platform_model", ScopeKey: admission.PlatformModelID}, + AdmissionScope{ScopeType: "user_group", ScopeKey: admission.UserGroupID}, + ) + } + for _, scope := range normalizedAdmissionScopes(lockScopes) { + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, admissionLockKey(scope)); err != nil { + return TaskAdmission{}, err + } + } + + if found && admission.Status == "admitted" { + activeLeases, leaseErr := activeTaskAdmissionLeaseCountTx(ctx, tx, input.TaskID) + if leaseErr != nil { + return TaskAdmission{}, leaseErr + } + if activeLeases == 0 { + admission, err = resetTaskAdmissionToWaitingTx(ctx, tx, input) + if err != nil { + return TaskAdmission{}, err + } + } + } + if found && admission.Status == "waiting" && !admission.WaitDeadlineAt.After(time.Now()) { + if err := expireTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { + return TaskAdmission{}, err + } + if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { + return TaskAdmission{}, err + } + if err := tx.Commit(ctx); err != nil { + return TaskAdmission{}, err + } + return TaskAdmission{}, &QueueTimeoutError{TaskID: input.TaskID} + } + if !found { + scopeStates, stateErr := admissionScopeStatesTx(ctx, tx, scopes) + if stateErr != nil { + return TaskAdmission{}, stateErr + } + if admissionErr := validateQueuedAdmissionCapacity(scopeStates); admissionErr != nil { + return TaskAdmission{}, admissionErr + } + admission, err = insertTaskAdmissionTx(ctx, tx, input, admissionDeadline(scopes)) + if err != nil { + return TaskAdmission{}, err + } + found = true + } + if onQueued != nil { + if err := onQueued(tx); err != nil { + return TaskAdmission{}, err + } + } + if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { + return TaskAdmission{}, err + } + if err := tx.Commit(ctx); err != nil { + return TaskAdmission{}, err + } + return admission, nil +} + func (s *Store) tryTaskAdmission( ctx context.Context, input TaskAdmissionInput, @@ -196,7 +303,7 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil { if err := expireTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { return TaskAdmissionResult{}, err } - if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { + if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { return TaskAdmissionResult{}, err } if err := tx.Commit(ctx); err != nil { @@ -276,7 +383,10 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil { if err != nil { return TaskAdmissionResult{}, err } - if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { + // Continue the FIFO chain after one task is admitted. Every API process + // only probes the global head it owns, so free capacity is filled without + // broadcasting an advisory-lock attempt to every waiter. + if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { return TaskAdmissionResult{}, err } if onAdmitted != nil { @@ -312,6 +422,29 @@ func validateNewAdmissionCapacity(states []admissionScopeState) error { return nil } +func validateQueuedAdmissionCapacity(states []admissionScopeState) error { + // Waiting asynchronous tasks have not acquired leases yet. Count them as + // virtual capacity reservations so a burst cannot bypass concurrent and + // queue limits before the dispatcher creates River jobs. + for _, state := range states { + if state.Scope.ConcurrentLimit <= 0 { + continue + } + virtualUsage := state.Active + float64(state.Waiting)*state.Scope.Amount + state.Scope.Amount + if virtualUsage <= state.Scope.ConcurrentLimit { + continue + } + if state.Scope.QueueLimit <= 0 { + return concurrencyAdmissionError(state) + } + virtualQueueDepth := int(math.Ceil(virtualUsage - state.Scope.ConcurrentLimit)) + if virtualQueueDepth > state.Scope.QueueLimit { + return queueFullAdmissionError(state) + } + } + return nil +} + // RebindWaitingTaskAdmission migrates a queued task to a newly selected // candidate without changing its FIFO age or priority. Capacity checks happen // while both the old and new queue locks are held. @@ -469,6 +602,32 @@ WHERE task_id = $1::uuid return nil } +// RenewTaskAdmissionWaiters refreshes all synchronous waiters owned by one API +// process in a single statement. This keeps the 15-second crash-recovery lease +// without running a full advisory-lock admission transaction per waiter. +func (s *Store) RenewTaskAdmissionWaiters(ctx context.Context, taskIDs []string, waiterIDs []string) error { + if len(taskIDs) == 0 { + return nil + } + if len(taskIDs) != len(waiterIDs) { + return errors.New("task admission waiter renewal requires matching task and waiter IDs") + } + _, err := s.pool.Exec(ctx, ` +UPDATE gateway_task_admissions admission +SET waiter_lease_expires_at = now() + $3::interval, + updated_at = now() +FROM unnest($1::uuid[], $2::text[]) AS owned(task_id, waiter_id) +WHERE admission.task_id = owned.task_id + AND admission.waiter_id = owned.waiter_id + AND admission.mode = 'sync' + AND admission.status = 'waiting'`, + taskIDs, + waiterIDs, + admissionWaiterLeaseTTL.String(), + ) + return err +} + func (s *Store) DeleteTaskAdmission(ctx context.Context, taskID string) error { tx, err := s.pool.Begin(ctx) if err != nil { @@ -488,7 +647,7 @@ WHERE task_id = $1::uuid if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil { return err } - if err := notifyTaskAdmissionTx(ctx, tx, taskID); err != nil { + if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { return err } return tx.Commit(ctx) @@ -546,6 +705,52 @@ LIMIT $1`, limit) return taskIDs, rows.Err() } +// ListWaitingTaskAdmissionIDs returns a bounded set of FIFO leaders across +// platform-model and user-group queues. It is used after capacity is released +// so API processes wake only plausible queue heads instead of every +// synchronous waiter. +func (s *Store) ListWaitingTaskAdmissionIDs(ctx context.Context, limit int) ([]string, error) { + if limit <= 0 { + limit = 256 + } + rows, err := s.pool.Query(ctx, ` +WITH ranked AS ( + SELECT task_id, + priority, + enqueued_at, + row_number() OVER ( + PARTITION BY platform_model_id + ORDER BY priority ASC, enqueued_at ASC, task_id ASC + ) AS platform_rank, + row_number() OVER ( + PARTITION BY user_group_id + ORDER BY priority ASC, enqueued_at ASC, task_id ASC + ) AS group_rank + FROM gateway_task_admissions + WHERE status = 'waiting' + AND mode = 'sync' +) +SELECT task_id::text +FROM ranked +WHERE platform_rank <= $1 + AND group_rank <= $1 +ORDER BY priority ASC, enqueued_at ASC, task_id ASC +LIMIT $1`, limit) + if err != nil { + return nil, err + } + defer rows.Close() + taskIDs := make([]string, 0, limit) + for rows.Next() { + var taskID string + if err := rows.Scan(&taskID); err != nil { + return nil, err + } + taskIDs = append(taskIDs, taskID) + } + return taskIDs, rows.Err() +} + func (s *Store) GetTaskAdmission(ctx context.Context, taskID string) (TaskAdmission, error) { admission, found, err := loadTaskAdmission(ctx, s.pool, taskID) if err != nil { @@ -793,6 +998,12 @@ SELECT COUNT(*) FROM gateway_task_admissions WHERE status = 'waiting' AND user_group_id::text = $1`, scope.ScopeKey).Scan(&waiting) return waiting, err + case "worker_capacity": + err := tx.QueryRow(ctx, ` +SELECT COUNT(*) +FROM gateway_task_admissions +WHERE status = 'waiting' AND mode = 'async'`).Scan(&waiting) + return waiting, err default: return 0, nil } @@ -813,7 +1024,10 @@ func admissionDeadline(scopes []AdmissionScope) time.Time { } } if maxWait == 0 { - maxWait = DefaultQueueMaxWaitSeconds + // A task may be registered before a River execution slot is available + // even when business queueing is disabled. Keep that infrastructure wait + // separate from the policy's default 600-second queue timeout. + return time.Now().Add(executionSlotWaitMax) } return time.Now().Add(time.Duration(maxWait) * time.Second) } diff --git a/apps/api/internal/store/admission_queue_integration_test.go b/apps/api/internal/store/admission_queue_integration_test.go index 2d728ce..5c73b98 100644 --- a/apps/api/internal/store/admission_queue_integration_test.go +++ b/apps/api/internal/store/admission_queue_integration_test.go @@ -402,6 +402,59 @@ VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'submitting')`, t.Fatalf("post-submission cancellation changed=%v err=%v, want unchanged", changed, err) } + queuedAtomicTask := createTask(true) + queuedHookFailure := errors.New("synthetic queued hook failure") + _, err = first.QueueTaskAdmissionWithHook(ctx, inputFor(queuedAtomicTask, 100, ""), func(pgx.Tx) error { + return queuedHookFailure + }) + if !errors.Is(err, queuedHookFailure) { + t.Fatalf("queued hook failure = %v, want synthetic failure", err) + } + var queuedAdmissions, queuedLeases int + if err := first.pool.QueryRow(ctx, ` +SELECT + (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid), + (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`, + queuedAtomicTask.ID, + ).Scan(&queuedAdmissions, &queuedLeases); err != nil { + t.Fatalf("read rolled back queued hook state: %v", err) + } + if queuedAdmissions != 0 || queuedLeases != 0 { + t.Fatalf("failed queued hook left admissions=%d leases=%d", queuedAdmissions, queuedLeases) + } + const queuedSyntheticRiverJobID int64 = 987654320 + queuedAdmission, err := first.QueueTaskAdmissionWithHook(ctx, inputFor(queuedAtomicTask, 100, ""), func(tx pgx.Tx) error { + _, hookErr := tx.Exec(ctx, ` +UPDATE gateway_tasks +SET river_job_id = $2 +WHERE id = $1::uuid`, queuedAtomicTask.ID, queuedSyntheticRiverJobID) + return hookErr + }) + if err != nil || queuedAdmission.Status != "waiting" { + t.Fatalf("atomic queued admission=%+v err=%v, want waiting", queuedAdmission, err) + } + var riverJobID int64 + if err := first.pool.QueryRow(ctx, ` +SELECT + river_job_id, + (SELECT COUNT(*) + FROM gateway_concurrency_leases + WHERE task_id = $1::uuid AND released_at IS NULL) +FROM gateway_tasks +WHERE id = $1::uuid`, queuedAtomicTask.ID).Scan(&riverJobID, &queuedLeases); err != nil { + t.Fatalf("read atomic queued River marker: %v", err) + } + if riverJobID != queuedSyntheticRiverJobID || queuedLeases != 0 { + t.Fatalf("queued River marker=%d leases=%d, want %d/0", riverJobID, queuedLeases, queuedSyntheticRiverJobID) + } + result, err = second.TryTaskAdmission(ctx, inputFor(queuedAtomicTask, 100, "")) + if err != nil || !result.Admitted || len(result.Leases) != 1 { + t.Fatalf("worker-time queued admission result=%+v err=%v", result, err) + } + if err := first.DeleteTaskAdmission(ctx, queuedAtomicTask.ID); err != nil { + t.Fatalf("release queued atomic task: %v", err) + } + atomicTask := createTask(true) hookFailure := errors.New("synthetic admitted hook failure") _, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(pgx.Tx) error { @@ -433,7 +486,6 @@ WHERE id = $1::uuid`, atomicTask.ID, syntheticRiverJobID) if err != nil || !result.Admitted { t.Fatalf("atomic admitted hook result=%+v err=%v", result, err) } - var riverJobID int64 if err := first.pool.QueryRow(ctx, ` SELECT river_job_id FROM gateway_tasks diff --git a/apps/api/internal/store/postgres.go b/apps/api/internal/store/postgres.go index 32d36f5..539aa17 100644 --- a/apps/api/internal/store/postgres.go +++ b/apps/api/internal/store/postgres.go @@ -26,6 +26,7 @@ type Store struct { const ( postgresApplicationName = "easyai-ai-gateway" postgresConnectTimeout = 5 * time.Second + postgresPoolWarmLimit = 64 ) func defaultAPIKeyScopes() []string { @@ -88,7 +89,7 @@ func ConnectWithMaxConns(ctx context.Context, databaseURL string, maxConns int) if err != nil { return nil, err } - if err := pool.Ping(ctx); err != nil { + if err := warmPostgresPool(ctx, pool, int(config.MinIdleConns)); err != nil { pool.Close() return nil, err } @@ -104,10 +105,31 @@ func postgresPoolConfig(databaseURL string, maxConns ...int) (*pgxpool.Config, e config.ConnConfig.RuntimeParams["application_name"] = postgresApplicationName if len(maxConns) > 0 && maxConns[0] > 0 { config.MaxConns = int32(maxConns[0]) + config.MinIdleConns = int32(min(maxConns[0], postgresPoolWarmLimit)) } return config, nil } +func warmPostgresPool(ctx context.Context, pool *pgxpool.Pool, count int) error { + if count <= 0 { + return pool.Ping(ctx) + } + connections := make([]*pgxpool.Conn, 0, count) + defer func() { + for _, connection := range connections { + connection.Release() + } + }() + for len(connections) < count { + connection, err := pool.Acquire(ctx) + if err != nil { + return err + } + connections = append(connections, connection) + } + return nil +} + func IsPostgresUnavailable(err error) bool { if err == nil { return false diff --git a/apps/api/internal/store/postgres_config_test.go b/apps/api/internal/store/postgres_config_test.go index 622e055..904431c 100644 --- a/apps/api/internal/store/postgres_config_test.go +++ b/apps/api/internal/store/postgres_config_test.go @@ -27,6 +27,34 @@ func TestPostgresPoolConfigRejectsMalformedURL(t *testing.T) { } } +func TestPostgresPoolConfigWarmsBoundedIdleConnections(t *testing.T) { + config, err := postgresPoolConfig( + "postgresql://gateway:password@localhost:5432/gateway?sslmode=disable", + 24, + ) + if err != nil { + t.Fatalf("parse PostgreSQL pool config: %v", err) + } + if config.MaxConns != 24 || config.MinIdleConns != 24 { + t.Fatalf( + "pool bounds max=%d minIdle=%d, want 24/24", + config.MaxConns, + config.MinIdleConns, + ) + } + + small, err := postgresPoolConfig( + "postgresql://gateway:password@localhost:5432/gateway?sslmode=disable", + 4, + ) + if err != nil { + t.Fatalf("parse small PostgreSQL pool config: %v", err) + } + if small.MinIdleConns != 4 { + t.Fatalf("small pool minIdle=%d, want 4", small.MinIdleConns) + } +} + func TestIsPostgresUnavailableClassifiesConnectivityFailures(t *testing.T) { for _, testCase := range []struct { name string diff --git a/apps/api/internal/store/rate_limit_policy_test.go b/apps/api/internal/store/rate_limit_policy_test.go index 009663b..b6bcef7 100644 --- a/apps/api/internal/store/rate_limit_policy_test.go +++ b/apps/api/internal/store/rate_limit_policy_test.go @@ -1,6 +1,10 @@ package store -import "testing" +import ( + "errors" + "testing" + "time" +) func TestNormalizeAndValidateQueuePolicy(t *testing.T) { t.Run("defaults maximum wait and preserves extensions", func(t *testing.T) { @@ -227,3 +231,64 @@ func TestAsyncWorkerCapacityRespectsUserGroupCeiling(t *testing.T) { }) } } + +func TestAdmissionDeadlineSeparatesExecutionSlotWaitFromPolicyQueueTimeout(t *testing.T) { + now := time.Now() + infrastructureDeadline := admissionDeadline([]AdmissionScope{{ + ConcurrentLimit: 10, + }}) + if wait := infrastructureDeadline.Sub(now); wait < 23*time.Hour || wait > executionSlotWaitMax+time.Second { + t.Fatalf("execution-slot deadline wait=%s, want about %s", wait, executionSlotWaitMax) + } + + defaultPolicyDeadline := admissionDeadline([]AdmissionScope{{ + ConcurrentLimit: 10, + QueueLimit: 20, + }}) + if wait := defaultPolicyDeadline.Sub(now); wait < 599*time.Second || wait > 601*time.Second { + t.Fatalf("default policy deadline wait=%s, want about 600s", wait) + } + + explicitPolicyDeadline := admissionDeadline([]AdmissionScope{{ + ConcurrentLimit: 10, + QueueLimit: 20, + MaxWaitSeconds: 90, + }}) + if wait := explicitPolicyDeadline.Sub(now); wait < 89*time.Second || wait > 91*time.Second { + t.Fatalf("explicit policy deadline wait=%s, want about 90s", wait) + } +} + +func TestQueuedAdmissionCountsFIFOWaitersAsVirtualCapacity(t *testing.T) { + scope := AdmissionScope{ + ScopeType: "platform_model", + ScopeKey: "model-1", + ConcurrentLimit: 1, + Amount: 1, + QueueLimit: 2, + } + for waiting := 0; waiting <= 2; waiting++ { + if err := validateQueuedAdmissionCapacity([]admissionScopeState{{ + Scope: scope, + Waiting: waiting, + }}); err != nil { + t.Fatalf("waiting=%d rejected within concurrent=1 queue=2: %v", waiting, err) + } + } + err := validateQueuedAdmissionCapacity([]admissionScopeState{{ + Scope: scope, + Waiting: 3, + }}) + var limitErr *RateLimitExceededError + if !errors.As(err, &limitErr) || limitErr.Reason != "queue_full" { + t.Fatalf("fourth task error=%v, want queue_full", err) + } + + scope.QueueLimit = 0 + if err := validateQueuedAdmissionCapacity([]admissionScopeState{{ + Scope: scope, + Waiting: 1, + }}); !errors.As(err, &limitErr) || limitErr.Metric != "concurrent" { + t.Fatalf("queue-disabled second task error=%v, want concurrent rate limit", err) + } +} diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index e60e76c..675a193 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -684,7 +684,7 @@ WHERE task_id = $1::uuid if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil { return err } - return notifyTaskAdmissionTx(ctx, tx, taskID) + return notifyTaskAdmissionTx(ctx, tx, "*") }) if err != nil { return GatewayTask{}, false, err @@ -770,7 +770,7 @@ WHERE task_id = $1::uuid if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil { return err } - return notifyTaskAdmissionTx(ctx, tx, taskID) + return notifyTaskAdmissionTx(ctx, tx, "*") }) if err != nil { return GatewayTask{}, false, err diff --git a/apps/api/internal/store/worker_registry.go b/apps/api/internal/store/worker_registry.go index 2f93403..23e4382 100644 --- a/apps/api/internal/store/worker_registry.go +++ b/apps/api/internal/store/worker_registry.go @@ -29,6 +29,18 @@ type WorkerAllocation struct { HeartbeatAt time.Time } +// ActiveWorkerCapacity returns the cluster-wide execution capacity currently +// owned by live Worker instances. +func (s *Store) ActiveWorkerCapacity(ctx context.Context) (int, error) { + var capacity int + err := s.pool.QueryRow(ctx, ` +SELECT COALESCE(SUM(allocated_capacity), 0)::int +FROM gateway_worker_instances +WHERE status = 'active' + AND heartbeat_at > now() - $1::interval`, workerHeartbeatStaleAfter.String()).Scan(&capacity) + return capacity, err +} + func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistrationInput) (WorkerAllocation, error) { input.InstanceID = strings.TrimSpace(input.InstanceID) if input.InstanceID == "" { diff --git a/deploy/kubernetes/production/application.yaml b/deploy/kubernetes/production/application.yaml index 77c0651..d99c3d1 100644 --- a/deploy/kubernetes/production/application.yaml +++ b/deploy/kubernetes/production/application.yaml @@ -54,6 +54,10 @@ spec: value: "false" - name: AI_GATEWAY_DATABASE_MAX_CONNS value: "16" + - name: AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY + value: "16" + - name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY + value: "8" - name: POD_NAMESPACE valueFrom: fieldRef: @@ -179,6 +183,10 @@ spec: value: "false" - name: AI_GATEWAY_DATABASE_MAX_CONNS value: "16" + - name: AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY + value: "16" + - name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY + value: "8" - name: POD_NAMESPACE valueFrom: fieldRef: @@ -304,6 +312,8 @@ spec: value: "true" - name: AI_GATEWAY_DATABASE_MAX_CONNS value: "24" + - name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY + value: "8" - name: POD_NAMESPACE valueFrom: fieldRef: @@ -429,6 +439,8 @@ spec: value: "true" - name: AI_GATEWAY_DATABASE_MAX_CONNS value: "24" + - name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY + value: "8" - name: POD_NAMESPACE valueFrom: fieldRef: