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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user