perf(queue): 限制大媒体任务内存并消除准入惊群

将同步与异步非文本任务的准入唤醒改为按任务和全局队首推进,批量续租等待者,避免大量请求同时争抢 advisory lock 和数据库连接。

对 Base64 请求解析、素材解码、上游媒体执行和结果物化增加分层并发限制,复用请求素材并释放重复 Gemini wire 数据;生产 API 默认入口解析并发 16,媒体物化并发 8。

新增 1000 个同步 Gemini 图像编辑请求的模拟上游压力验收,10 MiB 输入和输出下全部成功,Heap 峰值增长约 2.58 GiB,并验证共享上传哈希、单次 attempt 与本地零落盘。

验证:go test ./... -count=1;go vet ./...;gofmt -l 无输出;kubectl kustomize deploy/kubernetes/production;10 MiB Gemini Base64 千任务压力测试通过。
This commit is contained in:
2026-07-30 13:13:32 +08:00
parent fd3b6bf042
commit ea58d21d03
26 changed files with 1980 additions and 74 deletions
+10
View File
@@ -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) {
+7
View File
@@ -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)
}
+16 -5
View File
@@ -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
}
+10
View File
@@ -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:
+18
View File
@@ -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) {
@@ -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
}
@@ -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]) + "..."
}
@@ -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,
+1 -1
View File
@@ -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
}
@@ -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) != ""
}
@@ -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",
+21 -8
View File
@@ -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
+241 -41
View File
@@ -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)
}
}
+11
View File
@@ -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,
+14
View File
@@ -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) {
+41 -4
View File
@@ -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{
+35
View File
@@ -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)
+54
View File
@@ -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",
+218 -4
View File
@@ -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)
}
@@ -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
+23 -1
View File
@@ -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
@@ -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
@@ -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)
}
}
+2 -2
View File
@@ -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
@@ -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 == "" {
@@ -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: