easyai-ai-gateway/apps/api/internal/httpapi/gemini_base64_stress_integration_test.go

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]) + "..."
}