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:
@@ -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]) + "..."
|
||||
}
|
||||
Reference in New Issue
Block a user