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:
@@ -1531,6 +1531,16 @@ func TestGeminiClientImageGenerateBuildsNativeImageBody(t *testing.T) {
|
|||||||
if data[0].(map[string]any)["b64_json"] != "aW1hZ2U=" {
|
if data[0].(map[string]any)["b64_json"] != "aW1hZ2U=" {
|
||||||
t.Fatalf("unexpected image response: %+v", response.Result)
|
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) {
|
func TestGeminiClientImageGenerateRejectsPromptBlockedWithoutImage(t *testing.T) {
|
||||||
|
|||||||
@@ -65,6 +65,13 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error
|
|||||||
}
|
}
|
||||||
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
|
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 == "" {
|
if requestID == "" {
|
||||||
requestID = requestIDFromResult(output)
|
requestID = requestIDFromResult(output)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -65,12 +65,14 @@ func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[st
|
|||||||
Protocol: strings.TrimSpace(protocol),
|
Protocol: strings.TrimSpace(protocol),
|
||||||
StatusCode: resp.StatusCode,
|
StatusCode: resp.StatusCode,
|
||||||
Headers: compatibleResponseHeaders(resp.Header),
|
Headers: compatibleResponseHeaders(resp.Header),
|
||||||
RawJSON: append([]byte(nil), raw...),
|
// 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.
|
||||||
if len(raw) > 0 {
|
RawJSON: raw,
|
||||||
_ = json.Unmarshal(raw, &wire.Body)
|
|
||||||
}
|
}
|
||||||
if readErr != nil {
|
if readErr != nil {
|
||||||
|
if len(raw) > 0 {
|
||||||
|
_ = json.Unmarshal(raw, &wire.Body)
|
||||||
|
}
|
||||||
return nil, wire, &ClientError{
|
return nil, wire, &ClientError{
|
||||||
Code: "response_read_error",
|
Code: "response_read_error",
|
||||||
Message: readErr.Error(),
|
Message: readErr.Error(),
|
||||||
@@ -81,6 +83,9 @@ func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[st
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
if tooLarge {
|
if tooLarge {
|
||||||
|
if len(raw) > 0 {
|
||||||
|
_ = json.Unmarshal(raw, &wire.Body)
|
||||||
|
}
|
||||||
return nil, wire, &ClientError{
|
return nil, wire, &ClientError{
|
||||||
Code: "response_too_large",
|
Code: "response_too_large",
|
||||||
Message: fmt.Sprintf("upstream JSON response exceeds %d bytes", maxBytes),
|
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 resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
if len(raw) > 0 {
|
||||||
|
_ = json.Unmarshal(raw, &wire.Body)
|
||||||
|
}
|
||||||
return nil, wire, &ClientError{
|
return nil, wire, &ClientError{
|
||||||
Code: statusCodeName(resp.StatusCode),
|
Code: statusCodeName(resp.StatusCode),
|
||||||
Message: errorMessage(raw, resp.Status),
|
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 {
|
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}
|
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
|
return out, wire, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -59,6 +59,8 @@ type Config struct {
|
|||||||
BillingEngineMode string
|
BillingEngineMode string
|
||||||
ProcessRole string
|
ProcessRole string
|
||||||
DatabaseMaxConns int
|
DatabaseMaxConns int
|
||||||
|
MediaRequestConcurrency int
|
||||||
|
MediaMaterializationConcurrency int
|
||||||
AsyncQueueWorkerEnabled bool
|
AsyncQueueWorkerEnabled bool
|
||||||
AsyncWorkerHardLimit int
|
AsyncWorkerHardLimit int
|
||||||
AsyncWorkerRefreshIntervalSeconds int
|
AsyncWorkerRefreshIntervalSeconds int
|
||||||
@@ -117,6 +119,8 @@ func Load() Config {
|
|||||||
BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")),
|
BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")),
|
||||||
ProcessRole: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_PROCESS_ROLE", "all"))),
|
ProcessRole: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_PROCESS_ROLE", "all"))),
|
||||||
DatabaseMaxConns: envInt("AI_GATEWAY_DATABASE_MAX_CONNS", 0),
|
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",
|
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
|
||||||
AsyncWorkerHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
|
AsyncWorkerHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
|
||||||
AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5),
|
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 {
|
if c.DatabaseMaxConns < 0 || c.DatabaseMaxConns > 1000 {
|
||||||
return errors.New("AI_GATEWAY_DATABASE_MAX_CONNS must be between 1 and 1000 when configured")
|
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)) {
|
switch strings.ToLower(strings.TrimSpace(c.BillingEngineMode)) {
|
||||||
case "", "observe", "enforce", "hold":
|
case "", "observe", "enforce", "hold":
|
||||||
default:
|
default:
|
||||||
|
|||||||
@@ -131,6 +131,24 @@ func TestValidateProcessRoleAndDatabasePool(t *testing.T) {
|
|||||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MAX_CONNS") {
|
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MAX_CONNS") {
|
||||||
t.Fatalf("Validate() error = %v, want invalid database max conns", err)
|
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) {
|
func TestValidateTaskHistorySettings(t *testing.T) {
|
||||||
|
|||||||
@@ -21,6 +21,173 @@ import (
|
|||||||
"github.com/jackc/pgx/v5/pgxpool"
|
"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) {
|
func TestAsyncWorkerDynamicConcurrencyAcceptance(t *testing.T) {
|
||||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||||
if databaseURL == "" {
|
if databaseURL == "" {
|
||||||
@@ -295,6 +462,7 @@ func waitForAcceptanceTasks(t *testing.T, ctx context.Context, pool *pgxpool.Poo
|
|||||||
t.Helper()
|
t.Helper()
|
||||||
deadline := time.Now().Add(timeout)
|
deadline := time.Now().Add(timeout)
|
||||||
var peakRunning, peakLeases int64
|
var peakRunning, peakLeases int64
|
||||||
|
var lastRunning, lastSucceeded, lastFailed, lastLeases int64
|
||||||
for time.Now().Before(deadline) {
|
for time.Now().Before(deadline) {
|
||||||
var running, succeeded, failed, leases int64
|
var running, succeeded, failed, leases int64
|
||||||
if err := pool.QueryRow(ctx, `
|
if err := pool.QueryRow(ctx, `
|
||||||
@@ -316,6 +484,10 @@ WHERE scope_type = 'platform_model'
|
|||||||
}
|
}
|
||||||
peakRunning = maxInt64(peakRunning, running)
|
peakRunning = maxInt64(peakRunning, running)
|
||||||
peakLeases = maxInt64(peakLeases, leases)
|
peakLeases = maxInt64(peakLeases, leases)
|
||||||
|
lastRunning = running
|
||||||
|
lastSucceeded = succeeded
|
||||||
|
lastFailed = failed
|
||||||
|
lastLeases = leases
|
||||||
if failed > 0 {
|
if failed > 0 {
|
||||||
t.Fatalf("acceptance tasks entered failed/manual status: %d", failed)
|
t.Fatalf("acceptance tasks entered failed/manual status: %d", failed)
|
||||||
}
|
}
|
||||||
@@ -324,7 +496,17 @@ WHERE scope_type = 'platform_model'
|
|||||||
}
|
}
|
||||||
time.Sleep(50 * time.Millisecond)
|
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
|
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")
|
writeGeminiTaskError(http.StatusUnauthorized, "unauthorized", nil, "unauthorized")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
releaseRequestBody, err := s.acquireMediaRequestBodySlot(r.Context())
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer func() {
|
||||||
|
if releaseRequestBody != nil {
|
||||||
|
releaseRequestBody()
|
||||||
|
}
|
||||||
|
}()
|
||||||
var native map[string]any
|
var native map[string]any
|
||||||
if err := json.NewDecoder(r.Body).Decode(&native); err != nil {
|
if err := json.NewDecoder(r.Body).Decode(&native); err != nil {
|
||||||
writeGeminiTaskError(http.StatusBadRequest, "invalid json body", nil, "invalid_json_body")
|
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))
|
writeGeminiTaskError(status, err.Error(), nil, clients.ErrorCode(err))
|
||||||
return
|
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{
|
createInput := store.CreateTaskInput{
|
||||||
Kind: mapping.Kind,
|
Kind: mapping.Kind,
|
||||||
Model: mapping.Model,
|
Model: mapping.Model,
|
||||||
|
|||||||
@@ -1171,7 +1171,7 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
|||||||
writeTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
|
writeTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
|
||||||
return
|
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")
|
writeTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,9 +9,12 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"mime"
|
"mime"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"sort"
|
"sort"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||||
@@ -41,6 +44,11 @@ type requestAssetOptions struct {
|
|||||||
Source string
|
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) {
|
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)
|
preparedBody, err := s.prepareRequestAssetRefs(ctx, body)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -120,13 +128,9 @@ func (s *Server) prepareRequestAssetValue(ctx context.Context, value any, path [
|
|||||||
sort.Strings(keys)
|
sort.Strings(keys)
|
||||||
for _, key := range keys {
|
for _, key := range keys {
|
||||||
item := typed[key]
|
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
|
return nil, err
|
||||||
} else if ok {
|
} else if ok {
|
||||||
ref, err := s.ensureRequestAsset(ctx, decoded)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
next[key] = requestAssetWrapper(ref)
|
next[key] = requestAssetWrapper(ref)
|
||||||
continue
|
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) {
|
func requestAssetFromValue(key string, path []string, value any, siblings map[string]any) (decodedRequestAsset, bool, error) {
|
||||||
text, ok := value.(string)
|
text, ok := value.(string)
|
||||||
if !ok {
|
if !ok {
|
||||||
@@ -215,6 +280,8 @@ func (s *Server) ensureRequestAssetWithOptions(ctx context.Context, decoded deco
|
|||||||
if contentType == "" {
|
if contentType == "" {
|
||||||
contentType = "application/octet-stream"
|
contentType = "application/octet-stream"
|
||||||
}
|
}
|
||||||
|
release := s.acquireRequestAssetLock(sha + "\x00" + contentType + "\x00" + options.UploadScene + "\x00" + strconv.FormatBool(options.RequirePublicURL))
|
||||||
|
defer release()
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
if existing, ok, err := s.store.FindRequestAsset(ctx, sha, contentType); err != nil && !store.IsUndefinedDatabaseObject(err) {
|
if existing, ok, err := s.store.FindRequestAsset(ctx, sha, contentType); err != nil && !store.IsUndefinedDatabaseObject(err) {
|
||||||
return nil, err
|
return nil, err
|
||||||
@@ -286,6 +353,31 @@ func (s *Server) ensureRequestAssetWithOptions(ctx context.Context, decoded deco
|
|||||||
return requestAssetRef(asset), nil
|
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 {
|
func requestConversationKey(r *http.Request, body map[string]any) string {
|
||||||
if r != nil {
|
if r != nil {
|
||||||
if value := strings.TrimSpace(r.Header.Get("X-EasyAI-Conversation-ID")); value != "" {
|
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) {
|
if asset.ExpiresAt != nil && !asset.ExpiresAt.After(now) {
|
||||||
return false
|
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) != ""
|
return strings.TrimSpace(asset.URL) != ""
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
"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) {
|
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) {
|
func TestCanonicalConversationMessageHashUsesTextAndAssetRefs(t *testing.T) {
|
||||||
message := map[string]any{
|
message := map[string]any{
|
||||||
"role": "user",
|
"role": "user",
|
||||||
|
|||||||
@@ -32,6 +32,10 @@ type Server struct {
|
|||||||
runner *runner.Service
|
runner *runner.Service
|
||||||
logger *slog.Logger
|
logger *slog.Logger
|
||||||
geminiUploadSessions sync.Map
|
geminiUploadSessions sync.Map
|
||||||
|
requestAssetLocksMu sync.Mutex
|
||||||
|
requestAssetLocks map[string]*requestAssetLock
|
||||||
|
mediaRequestSlots chan struct{}
|
||||||
|
mediaRequestBodySlots chan struct{}
|
||||||
securityEventReceiver http.Handler
|
securityEventReceiver http.Handler
|
||||||
securityEventManager *ssfreceiver.ConnectionManager
|
securityEventManager *ssfreceiver.ConnectionManager
|
||||||
identityRuntime *identityruntime.Manager
|
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 {
|
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{}
|
securityEventMetrics := &ssfreceiver.Metrics{}
|
||||||
server := &Server{
|
server := &Server{
|
||||||
ctx: ctx,
|
ctx: ctx,
|
||||||
cfg: cfg,
|
cfg: cfg,
|
||||||
store: db,
|
store: db,
|
||||||
oidcUserResolver: db,
|
oidcUserResolver: db,
|
||||||
auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken),
|
auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken),
|
||||||
runner: runner.New(cfg, db, logger, securityEventMetrics),
|
runner: runner.New(cfg, db, logger, securityEventMetrics),
|
||||||
logger: logger,
|
logger: logger,
|
||||||
billingMetrics: securityEventMetrics,
|
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.ServerMainInternalKey = cfg.ServerMainInternalKey
|
||||||
server.auth.ServerMainInternalSecret = cfg.ServerMainInternalSecret
|
server.auth.ServerMainInternalSecret = cfg.ServerMainInternalSecret
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package runner
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
|
"hash/fnv"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -22,6 +23,17 @@ type taskAdmissionPlan struct {
|
|||||||
Eligible bool
|
Eligible bool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type admissionTaskWaiter struct {
|
||||||
|
wake chan struct{}
|
||||||
|
waiterID string
|
||||||
|
}
|
||||||
|
|
||||||
|
const (
|
||||||
|
asyncWorkerCapacityScopeKey = "global"
|
||||||
|
asyncWorkerQueueLimit = 10000
|
||||||
|
asyncWorkerMaxWaitSeconds = 24 * 60 * 60
|
||||||
|
)
|
||||||
|
|
||||||
func distributedAdmissionModelType(modelType string) bool {
|
func distributedAdmissionModelType(modelType string) bool {
|
||||||
switch strings.ToLower(strings.TrimSpace(modelType)) {
|
switch strings.ToLower(strings.TrimSpace(modelType)) {
|
||||||
case "image_generate",
|
case "image_generate",
|
||||||
@@ -99,6 +111,12 @@ func (s *Service) buildTaskAdmissionPlan(ctx context.Context, task store.Gateway
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
scopes, groupID := s.admissionScopes(ctx, user, candidate)
|
scopes, groupID := s.admissionScopes(ctx, user, candidate)
|
||||||
|
if task.AsyncMode {
|
||||||
|
scopes, err = s.withAsyncWorkerCapacityScope(ctx, scopes)
|
||||||
|
if err != nil {
|
||||||
|
return taskAdmissionPlan{}, err
|
||||||
|
}
|
||||||
|
}
|
||||||
hasConcurrentLimit := false
|
hasConcurrentLimit := false
|
||||||
for _, scope := range scopes {
|
for _, scope := range scopes {
|
||||||
if scope.ConcurrentLimit > 0 {
|
if scope.ConcurrentLimit > 0 {
|
||||||
@@ -128,10 +146,67 @@ func (s *Service) buildTaskAdmissionPlan(ctx context.Context, task store.Gateway
|
|||||||
return taskAdmissionPlan{}, store.ErrNoModelCandidate
|
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) {
|
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)
|
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(
|
func (s *Service) tryTaskAdmissionWithAdmittedHook(
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
task store.GatewayTask,
|
task store.GatewayTask,
|
||||||
@@ -139,21 +214,7 @@ func (s *Service) tryTaskAdmissionWithAdmittedHook(
|
|||||||
waiterID string,
|
waiterID string,
|
||||||
onAdmitted func(pgx.Tx) error,
|
onAdmitted func(pgx.Tx) error,
|
||||||
) (store.TaskAdmissionResult, error) {
|
) (store.TaskAdmissionResult, error) {
|
||||||
mode := "sync"
|
input := taskAdmissionInput(task, plan, waiterID)
|
||||||
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,
|
|
||||||
}
|
|
||||||
current, currentErr := s.store.GetTaskAdmission(ctx, task.ID)
|
current, currentErr := s.store.GetTaskAdmission(ctx, task.ID)
|
||||||
if currentErr != nil && !errors.Is(currentErr, pgx.ErrNoRows) {
|
if currentErr != nil && !errors.Is(currentErr, pgx.ErrNoRows) {
|
||||||
return store.TaskAdmissionResult{}, currentErr
|
return store.TaskAdmissionResult{}, currentErr
|
||||||
@@ -186,14 +247,34 @@ func (s *Service) tryTaskAdmissionWithAdmittedHook(
|
|||||||
return result, err
|
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) {
|
func (s *Service) waitForSynchronousAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan) (store.TaskAdmissionResult, error) {
|
||||||
waiterID := uuid.NewString()
|
waiterID := uuid.NewString()
|
||||||
return s.waitForTaskAdmission(ctx, task, plan, waiterID)
|
return s.waitForTaskAdmission(ctx, task, plan, waiterID)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *Service) waitForTaskAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan, waiterID string) (store.TaskAdmissionResult, error) {
|
func (s *Service) waitForTaskAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan, waiterID string) (store.TaskAdmissionResult, error) {
|
||||||
ticker := time.NewTicker(time.Second)
|
taskWake, unregister := s.registerAdmissionTaskWaiter(task.ID, waiterID)
|
||||||
defer ticker.Stop()
|
defer unregister()
|
||||||
|
pollTimer := time.NewTimer(admissionPollInterval(task.ID))
|
||||||
|
defer pollTimer.Stop()
|
||||||
for {
|
for {
|
||||||
result, err := s.tryTaskAdmission(ctx, task, plan, waiterID)
|
result, err := s.tryTaskAdmission(ctx, task, plan, waiterID)
|
||||||
if err != nil || result.Admitted {
|
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.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID)
|
||||||
s.observeTaskAdmission("cancelled")
|
s.observeTaskAdmission("cancelled")
|
||||||
return store.TaskAdmissionResult{}, ctx.Err()
|
return store.TaskAdmissionResult{}, ctx.Err()
|
||||||
case <-s.currentAdmissionWake():
|
case <-taskWake:
|
||||||
case <-ticker.C:
|
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(
|
func (s *Service) ensureCandidateAdmission(
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
task store.GatewayTask,
|
task store.GatewayTask,
|
||||||
@@ -218,6 +306,13 @@ func (s *Service) ensureCandidateAdmission(
|
|||||||
candidate store.RuntimeModelCandidate,
|
candidate store.RuntimeModelCandidate,
|
||||||
) (store.TaskAdmissionResult, bool, error) {
|
) (store.TaskAdmissionResult, bool, error) {
|
||||||
scopes, groupID := s.admissionScopes(ctx, user, candidate)
|
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
|
hasConcurrentLimit := false
|
||||||
for _, scope := range scopes {
|
for _, scope := range scopes {
|
||||||
if scope.ConcurrentLimit > 0 {
|
if scope.ConcurrentLimit > 0 {
|
||||||
@@ -270,10 +365,12 @@ func (s *Service) observeTaskAdmissionWait(wait time.Duration) {
|
|||||||
|
|
||||||
func (s *Service) StartAdmissionNotifier(ctx context.Context) {
|
func (s *Service) StartAdmissionNotifier(ctx context.Context) {
|
||||||
s.admissionListener.Do(func() {
|
s.admissionListener.Do(func() {
|
||||||
|
go s.dispatchWaitingSynchronousAdmissions(ctx)
|
||||||
|
go s.renewSynchronousAdmissionWaiters(ctx)
|
||||||
go func() {
|
go func() {
|
||||||
for ctx.Err() == nil {
|
for ctx.Err() == nil {
|
||||||
err := s.store.ListenTaskAdmissionNotifications(ctx, func(string) {
|
err := s.store.ListenTaskAdmissionNotifications(ctx, func(taskID string) {
|
||||||
s.broadcastAdmissionWake()
|
s.signalAdmissionWake(taskID)
|
||||||
})
|
})
|
||||||
if ctx.Err() != nil {
|
if ctx.Err() != nil {
|
||||||
return
|
return
|
||||||
@@ -293,17 +390,99 @@ func (s *Service) StartAdmissionNotifier(ctx context.Context) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *Service) currentAdmissionWake() <-chan struct{} {
|
func (s *Service) dispatchWaitingSynchronousAdmissions(ctx context.Context) {
|
||||||
s.admissionWakeMu.Lock()
|
for {
|
||||||
defer s.admissionWakeMu.Unlock()
|
select {
|
||||||
return s.admissionWake
|
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()
|
s.admissionWakeMu.Lock()
|
||||||
close(s.admissionWake)
|
waiter := &admissionTaskWaiter{wake: make(chan struct{}, 1), waiterID: waiterID}
|
||||||
s.admissionWake = make(chan struct{})
|
s.admissionTaskWaiters[taskID] = waiter
|
||||||
s.admissionWakeMu.Unlock()
|
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 {
|
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
|
return nil
|
||||||
}
|
}
|
||||||
result, err := s.tryTaskAdmissionWithAdmittedHook(ctx, task, plan, "", func(tx pgx.Tx) error {
|
input := taskAdmissionInput(task, plan, "")
|
||||||
return s.enqueueAsyncTaskTx(ctx, tx, task.ID, asyncTaskInsertOpts(task))
|
current, currentErr := s.store.GetTaskAdmission(ctx, task.ID)
|
||||||
})
|
if currentErr != nil && !errors.Is(currentErr, pgx.ErrNoRows) {
|
||||||
if err != nil {
|
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) {
|
if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
|
||||||
return ctx.Err()
|
return ctx.Err()
|
||||||
}
|
}
|
||||||
_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error())
|
_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error())
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if !result.Admitted {
|
|
||||||
if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
|
|
||||||
return ctx.Err()
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
|
if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
|
||||||
return ctx.Err()
|
return ctx.Err()
|
||||||
}
|
}
|
||||||
return nil
|
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 {
|
func (s *Service) cancelAsyncSubmissionIfDisconnected(ctx context.Context, taskID string) bool {
|
||||||
if ctx.Err() == nil {
|
if ctx.Err() == nil {
|
||||||
return false
|
return false
|
||||||
@@ -380,10 +580,10 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) {
|
|||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
return
|
return
|
||||||
case <-s.currentAdmissionWake():
|
case <-s.asyncAdmissionWake:
|
||||||
case <-ticker.C:
|
case <-ticker.C:
|
||||||
}
|
}
|
||||||
taskIDs, err := s.store.ListWaitingAsyncAdmissionTaskIDs(ctx, 100)
|
taskIDs, err := s.store.ListWaitingAsyncAdmissionTaskIDs(ctx, 1000)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
s.logger.Warn("list waiting async admissions failed", "error", err)
|
s.logger.Warn("list waiting async admissions failed", "error", err)
|
||||||
continue
|
continue
|
||||||
@@ -399,7 +599,7 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) {
|
|||||||
if task.Status != "queued" || task.RiverJobID > 0 {
|
if task.Status != "queued" || task.RiverJobID > 0 {
|
||||||
continue
|
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)
|
s.logger.Warn("dispatch waiting async admission failed", "taskID", taskID, "error", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -21,6 +21,11 @@ type httpClientCache struct {
|
|||||||
|
|
||||||
const providerHTTPClientTimeout = 10 * time.Minute
|
const providerHTTPClientTimeout = 10 * time.Minute
|
||||||
|
|
||||||
|
const (
|
||||||
|
providerHTTPMaxIdleConnections = 2048
|
||||||
|
providerHTTPMaxIdleConnectionsPerHost = 1024
|
||||||
|
)
|
||||||
|
|
||||||
func newHTTPClientCache() *httpClientCache {
|
func newHTTPClientCache() *httpClientCache {
|
||||||
return &httpClientCache{
|
return &httpClientCache{
|
||||||
none: newHTTPClient(nil),
|
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 {
|
func newHTTPClient(proxy func(*http.Request) (*url.URL, error)) *http.Client {
|
||||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||||
transport.Proxy = proxy
|
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{
|
return &http.Client{
|
||||||
Timeout: providerHTTPClientTimeout,
|
Timeout: providerHTTPClientTimeout,
|
||||||
Transport: transport,
|
Transport: transport,
|
||||||
|
|||||||
@@ -16,6 +16,20 @@ func TestProviderHTTPClientTimeoutAllowsLongRunningMediaRequests(t *testing.T) {
|
|||||||
if client.Timeout != 10*time.Minute {
|
if client.Timeout != 10*time.Minute {
|
||||||
t.Fatalf("unexpected provider HTTP timeout: got %s want %s", 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) {
|
func TestPlatformProxyModeNoneIgnoresEnvironmentProxy(t *testing.T) {
|
||||||
|
|||||||
@@ -40,8 +40,11 @@ type Service struct {
|
|||||||
asyncCapacityLoader func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error)
|
asyncCapacityLoader func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error)
|
||||||
admissionWakeMu sync.Mutex
|
admissionWakeMu sync.Mutex
|
||||||
admissionWake chan struct{}
|
admissionWake chan struct{}
|
||||||
|
asyncAdmissionWake chan struct{}
|
||||||
|
admissionTaskWaiters map[string]*admissionTaskWaiter
|
||||||
admissionListener sync.Once
|
admissionListener sync.Once
|
||||||
asyncClientFactory func(int) (asyncExecutionClient, error)
|
asyncClientFactory func(int) (asyncExecutionClient, error)
|
||||||
|
mediaResultSlots chan struct{}
|
||||||
billingMetrics billingMetricsObserver
|
billingMetrics billingMetricsObserver
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -93,6 +96,9 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
|
|||||||
if cfg.AsyncWorkerRefreshIntervalSeconds == 0 {
|
if cfg.AsyncWorkerRefreshIntervalSeconds == 0 {
|
||||||
cfg.AsyncWorkerRefreshIntervalSeconds = 5
|
cfg.AsyncWorkerRefreshIntervalSeconds = 5
|
||||||
}
|
}
|
||||||
|
if cfg.MediaMaterializationConcurrency == 0 {
|
||||||
|
cfg.MediaMaterializationConcurrency = 8
|
||||||
|
}
|
||||||
if cfg.TaskProgressCallbackTimeoutMS == 0 {
|
if cfg.TaskProgressCallbackTimeoutMS == 0 {
|
||||||
cfg.TaskProgressCallbackTimeoutMS = 5000
|
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},
|
"universal": clients.UniversalClient{HTTPClient: httpClients.none, ScriptExecutor: scriptExecutor},
|
||||||
"simulation": clients.SimulationClient{},
|
"simulation": clients.SimulationClient{},
|
||||||
},
|
},
|
||||||
httpClients: httpClients,
|
httpClients: httpClients,
|
||||||
workerInstanceID: asyncWorkerID(),
|
workerInstanceID: asyncWorkerID(),
|
||||||
admissionWake: make(chan struct{}),
|
admissionWake: make(chan struct{}, 4096),
|
||||||
|
asyncAdmissionWake: make(chan struct{}, 1),
|
||||||
|
admissionTaskWaiters: map[string]*admissionTaskWaiter{},
|
||||||
|
mediaResultSlots: make(chan struct{}, cfg.MediaMaterializationConcurrency),
|
||||||
}
|
}
|
||||||
if len(observers) > 0 {
|
if len(observers) > 0 {
|
||||||
service.billingMetrics = 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 admissionResult store.TaskAdmissionResult
|
||||||
var admissionErr error
|
var admissionErr error
|
||||||
if task.AsyncMode {
|
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 {
|
if admissionErr == nil && !admissionResult.Admitted {
|
||||||
_ = s.store.ReleaseTaskPreparation(context.WithoutCancel(ctx), task.ID, task.ExecutionToken)
|
_ = s.store.ReleaseTaskPreparation(context.WithoutCancel(ctx), task.ID, task.ExecutionToken)
|
||||||
return Result{Task: task, Output: task.Result}, &TaskQueuedError{Delay: time.Second}
|
return Result{Task: task, Output: task.Result}, &TaskQueuedError{Delay: time.Second}
|
||||||
@@ -1272,6 +1285,15 @@ func (s *Service) runCandidate(
|
|||||||
})
|
})
|
||||||
return clients.Response{}, err
|
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)
|
providerBody, err = s.hydrateProviderRequestAssets(ctx, providerBody, candidate)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
|
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
|
||||||
@@ -1465,6 +1487,21 @@ func (s *Service) runCandidate(
|
|||||||
}
|
}
|
||||||
return clients.Response{}, err
|
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)
|
uploadedResult, err := s.uploadGeneratedAssets(ctx, task.ID, task.Kind, response.Result)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
metrics := mergeMetrics(taskMetrics(task, user, body, candidate, response, simulated), parameterPreprocessingMetrics(preprocessing), map[string]any{
|
metrics := mergeMetrics(taskMetrics(task, user, body, candidate, response, simulated), parameterPreprocessingMetrics(preprocessing), map[string]any{
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ import (
|
|||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
"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) {
|
func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, taskKind string, result map[string]any) (map[string]any, error) {
|
||||||
data, _ := result["data"].([]any)
|
data, _ := result["data"].([]any)
|
||||||
rawNeedsUpload := generatedRawValueHasInlineMedia(result["raw"], "", nil)
|
rawNeedsUpload := generatedRawValueHasInlineMedia(result["raw"], "", nil)
|
||||||
|
|||||||
@@ -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) {
|
func TestGeneratedAssetDecisionUploadsVectorDocumentWithoutChangingType(t *testing.T) {
|
||||||
item := map[string]any{
|
item := map[string]any{
|
||||||
"type": "file",
|
"type": "file",
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import (
|
|||||||
"database/sql"
|
"database/sql"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"math"
|
||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
@@ -15,6 +16,7 @@ import (
|
|||||||
const (
|
const (
|
||||||
admissionWaiterLeaseTTL = 15 * time.Second
|
admissionWaiterLeaseTTL = 15 * time.Second
|
||||||
admissionNotifyChannel = "gateway_task_admission"
|
admissionNotifyChannel = "gateway_task_admission"
|
||||||
|
executionSlotWaitMax = 24 * time.Hour
|
||||||
)
|
)
|
||||||
|
|
||||||
var ErrQueueTimeout = errors.New("task admission queue wait timed out")
|
var ErrQueueTimeout = errors.New("task admission queue wait timed out")
|
||||||
@@ -113,6 +115,111 @@ func (s *Store) TryTaskAdmissionWithAdmittedHook(
|
|||||||
return s.tryTaskAdmission(ctx, input, onAdmitted)
|
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(
|
func (s *Store) tryTaskAdmission(
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
input TaskAdmissionInput,
|
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 {
|
if err := expireTaskAdmissionTx(ctx, tx, input.TaskID); err != nil {
|
||||||
return TaskAdmissionResult{}, err
|
return TaskAdmissionResult{}, err
|
||||||
}
|
}
|
||||||
if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); err != nil {
|
if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil {
|
||||||
return TaskAdmissionResult{}, err
|
return TaskAdmissionResult{}, err
|
||||||
}
|
}
|
||||||
if err := tx.Commit(ctx); err != nil {
|
if err := tx.Commit(ctx); err != nil {
|
||||||
@@ -276,7 +383,10 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return TaskAdmissionResult{}, err
|
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
|
return TaskAdmissionResult{}, err
|
||||||
}
|
}
|
||||||
if onAdmitted != nil {
|
if onAdmitted != nil {
|
||||||
@@ -312,6 +422,29 @@ func validateNewAdmissionCapacity(states []admissionScopeState) error {
|
|||||||
return nil
|
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
|
// RebindWaitingTaskAdmission migrates a queued task to a newly selected
|
||||||
// candidate without changing its FIFO age or priority. Capacity checks happen
|
// candidate without changing its FIFO age or priority. Capacity checks happen
|
||||||
// while both the old and new queue locks are held.
|
// while both the old and new queue locks are held.
|
||||||
@@ -469,6 +602,32 @@ WHERE task_id = $1::uuid
|
|||||||
return nil
|
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 {
|
func (s *Store) DeleteTaskAdmission(ctx context.Context, taskID string) error {
|
||||||
tx, err := s.pool.Begin(ctx)
|
tx, err := s.pool.Begin(ctx)
|
||||||
if err != nil {
|
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 {
|
if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if err := notifyTaskAdmissionTx(ctx, tx, taskID); err != nil {
|
if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
return tx.Commit(ctx)
|
return tx.Commit(ctx)
|
||||||
@@ -546,6 +705,52 @@ LIMIT $1`, limit)
|
|||||||
return taskIDs, rows.Err()
|
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) {
|
func (s *Store) GetTaskAdmission(ctx context.Context, taskID string) (TaskAdmission, error) {
|
||||||
admission, found, err := loadTaskAdmission(ctx, s.pool, taskID)
|
admission, found, err := loadTaskAdmission(ctx, s.pool, taskID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -793,6 +998,12 @@ SELECT COUNT(*)
|
|||||||
FROM gateway_task_admissions
|
FROM gateway_task_admissions
|
||||||
WHERE status = 'waiting' AND user_group_id::text = $1`, scope.ScopeKey).Scan(&waiting)
|
WHERE status = 'waiting' AND user_group_id::text = $1`, scope.ScopeKey).Scan(&waiting)
|
||||||
return waiting, err
|
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:
|
default:
|
||||||
return 0, nil
|
return 0, nil
|
||||||
}
|
}
|
||||||
@@ -813,7 +1024,10 @@ func admissionDeadline(scopes []AdmissionScope) time.Time {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
if maxWait == 0 {
|
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)
|
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)
|
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)
|
atomicTask := createTask(true)
|
||||||
hookFailure := errors.New("synthetic admitted hook failure")
|
hookFailure := errors.New("synthetic admitted hook failure")
|
||||||
_, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(pgx.Tx) error {
|
_, 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 {
|
if err != nil || !result.Admitted {
|
||||||
t.Fatalf("atomic admitted hook result=%+v err=%v", result, err)
|
t.Fatalf("atomic admitted hook result=%+v err=%v", result, err)
|
||||||
}
|
}
|
||||||
var riverJobID int64
|
|
||||||
if err := first.pool.QueryRow(ctx, `
|
if err := first.pool.QueryRow(ctx, `
|
||||||
SELECT river_job_id
|
SELECT river_job_id
|
||||||
FROM gateway_tasks
|
FROM gateway_tasks
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ type Store struct {
|
|||||||
const (
|
const (
|
||||||
postgresApplicationName = "easyai-ai-gateway"
|
postgresApplicationName = "easyai-ai-gateway"
|
||||||
postgresConnectTimeout = 5 * time.Second
|
postgresConnectTimeout = 5 * time.Second
|
||||||
|
postgresPoolWarmLimit = 64
|
||||||
)
|
)
|
||||||
|
|
||||||
func defaultAPIKeyScopes() []string {
|
func defaultAPIKeyScopes() []string {
|
||||||
@@ -88,7 +89,7 @@ func ConnectWithMaxConns(ctx context.Context, databaseURL string, maxConns int)
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if err := pool.Ping(ctx); err != nil {
|
if err := warmPostgresPool(ctx, pool, int(config.MinIdleConns)); err != nil {
|
||||||
pool.Close()
|
pool.Close()
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
@@ -104,10 +105,31 @@ func postgresPoolConfig(databaseURL string, maxConns ...int) (*pgxpool.Config, e
|
|||||||
config.ConnConfig.RuntimeParams["application_name"] = postgresApplicationName
|
config.ConnConfig.RuntimeParams["application_name"] = postgresApplicationName
|
||||||
if len(maxConns) > 0 && maxConns[0] > 0 {
|
if len(maxConns) > 0 && maxConns[0] > 0 {
|
||||||
config.MaxConns = int32(maxConns[0])
|
config.MaxConns = int32(maxConns[0])
|
||||||
|
config.MinIdleConns = int32(min(maxConns[0], postgresPoolWarmLimit))
|
||||||
}
|
}
|
||||||
return config, nil
|
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 {
|
func IsPostgresUnavailable(err error) bool {
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return false
|
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) {
|
func TestIsPostgresUnavailableClassifiesConnectivityFailures(t *testing.T) {
|
||||||
for _, testCase := range []struct {
|
for _, testCase := range []struct {
|
||||||
name string
|
name string
|
||||||
|
|||||||
@@ -1,6 +1,10 @@
|
|||||||
package store
|
package store
|
||||||
|
|
||||||
import "testing"
|
import (
|
||||||
|
"errors"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
func TestNormalizeAndValidateQueuePolicy(t *testing.T) {
|
func TestNormalizeAndValidateQueuePolicy(t *testing.T) {
|
||||||
t.Run("defaults maximum wait and preserves extensions", func(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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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 {
|
if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
return notifyTaskAdmissionTx(ctx, tx, taskID)
|
return notifyTaskAdmissionTx(ctx, tx, "*")
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return GatewayTask{}, false, err
|
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 {
|
if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
return notifyTaskAdmissionTx(ctx, tx, taskID)
|
return notifyTaskAdmissionTx(ctx, tx, "*")
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return GatewayTask{}, false, err
|
return GatewayTask{}, false, err
|
||||||
|
|||||||
@@ -29,6 +29,18 @@ type WorkerAllocation struct {
|
|||||||
HeartbeatAt time.Time
|
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) {
|
func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistrationInput) (WorkerAllocation, error) {
|
||||||
input.InstanceID = strings.TrimSpace(input.InstanceID)
|
input.InstanceID = strings.TrimSpace(input.InstanceID)
|
||||||
if input.InstanceID == "" {
|
if input.InstanceID == "" {
|
||||||
|
|||||||
@@ -54,6 +54,10 @@ spec:
|
|||||||
value: "false"
|
value: "false"
|
||||||
- name: AI_GATEWAY_DATABASE_MAX_CONNS
|
- name: AI_GATEWAY_DATABASE_MAX_CONNS
|
||||||
value: "16"
|
value: "16"
|
||||||
|
- name: AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY
|
||||||
|
value: "16"
|
||||||
|
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
|
||||||
|
value: "8"
|
||||||
- name: POD_NAMESPACE
|
- name: POD_NAMESPACE
|
||||||
valueFrom:
|
valueFrom:
|
||||||
fieldRef:
|
fieldRef:
|
||||||
@@ -179,6 +183,10 @@ spec:
|
|||||||
value: "false"
|
value: "false"
|
||||||
- name: AI_GATEWAY_DATABASE_MAX_CONNS
|
- name: AI_GATEWAY_DATABASE_MAX_CONNS
|
||||||
value: "16"
|
value: "16"
|
||||||
|
- name: AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY
|
||||||
|
value: "16"
|
||||||
|
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
|
||||||
|
value: "8"
|
||||||
- name: POD_NAMESPACE
|
- name: POD_NAMESPACE
|
||||||
valueFrom:
|
valueFrom:
|
||||||
fieldRef:
|
fieldRef:
|
||||||
@@ -304,6 +312,8 @@ spec:
|
|||||||
value: "true"
|
value: "true"
|
||||||
- name: AI_GATEWAY_DATABASE_MAX_CONNS
|
- name: AI_GATEWAY_DATABASE_MAX_CONNS
|
||||||
value: "24"
|
value: "24"
|
||||||
|
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
|
||||||
|
value: "8"
|
||||||
- name: POD_NAMESPACE
|
- name: POD_NAMESPACE
|
||||||
valueFrom:
|
valueFrom:
|
||||||
fieldRef:
|
fieldRef:
|
||||||
@@ -429,6 +439,8 @@ spec:
|
|||||||
value: "true"
|
value: "true"
|
||||||
- name: AI_GATEWAY_DATABASE_MAX_CONNS
|
- name: AI_GATEWAY_DATABASE_MAX_CONNS
|
||||||
value: "24"
|
value: "24"
|
||||||
|
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
|
||||||
|
value: "8"
|
||||||
- name: POD_NAMESPACE
|
- name: POD_NAMESPACE
|
||||||
valueFrom:
|
valueFrom:
|
||||||
fieldRef:
|
fieldRef:
|
||||||
|
|||||||
Reference in New Issue
Block a user