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