Files
easyai-ai-gateway/apps/api/cmd/acceptance-load/main.go
T
wangbo be6ce7f78a fix(acceptance): 修复线上双站压测失真
将线上压力生成器拆分为宁波、香港宿主机上的 K3s 外进程,按全局序号分片并保守合并站点报告,避免操作者本机上行带宽成为容量瓶颈。\n\n为每次负载执行生成独立幂等键和 Run 级图片变体,阻止三轮验收重放旧任务或复用历史媒体缓存,并新增直接 OSS 物化一致性门禁。\n\n验证:Go 全量测试、go vet、迁移安全检查、bash -n、ShellCheck、manual-release-test 和双宿主机 linux/amd64 启动冒烟通过。
2026-08-01 01:02:07 +08:00

1271 lines
39 KiB
Go

package main
import (
"bufio"
"bytes"
"context"
"crypto/sha256"
"crypto/tls"
"crypto/x509"
"encoding/base64"
"encoding/binary"
"encoding/hex"
"encoding/json"
"errors"
"flag"
"fmt"
"hash/crc32"
"io"
"net/http"
"net/url"
"os"
"os/signal"
"regexp"
"sort"
"strconv"
"strings"
"sync"
"syscall"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceworkload"
)
const (
runHeader = "X-EasyAI-Acceptance-Run"
tokenHeader = "X-EasyAI-Acceptance-Token"
upstreamHeader = "X-EasyAI-Acceptance-Upstream"
)
var reportURLPattern = regexp.MustCompile(`(?i)(?:https?|postgres(?:ql)?):\/\/[^\s"'<>]+`)
type options struct {
gateways []string
gatewayTLSName string
gatewayCAFile string
emulatorURL string
apiKeys []string
runID string
runToken string
geminiModel string
videoModel string
profile string
reportPath string
realImageURLs []string
timeout time.Duration
mixedDuration time.Duration
imageRate float64
videoRate float64
shardIndex int
shardCount int
executionID string
}
type report struct {
SchemaVersion string `json:"schemaVersion"`
RunID string `json:"runId"`
Profile string `json:"profile"`
StartedAt time.Time `json:"startedAt"`
FinishedAt time.Time `json:"finishedAt"`
Passed bool `json:"passed"`
Phases []phaseReport `json:"phases"`
Failure string `json:"failure,omitempty"`
FailureOperation string `json:"failureOperation,omitempty"`
SecretSafe bool `json:"secretSafe"`
}
type phaseReport struct {
Name string `json:"name"`
Requests int `json:"requests"`
Completed int `json:"completed"`
Failed int `json:"failed"`
DurationMS int64 `json:"durationMs"`
SubmissionDurationMS int64 `json:"submissionDurationMs,omitempty"`
P50MS float64 `json:"p50Ms"`
P95MS float64 `json:"p95Ms"`
P99MS float64 `json:"p99Ms"`
DecodedOutputBytes int64 `json:"decodedOutputBytes,omitempty"`
OutputSHA256 string `json:"outputSha256,omitempty"`
UniqueTaskIDs int `json:"uniqueTaskIds,omitempty"`
UniqueImageCombos int `json:"uniqueImageCombinations,omitempty"`
ForcedConversionTasks int `json:"forcedConversionTasks,omitempty"`
Throttled int `json:"throttled,omitempty"`
Unexpected5xx int `json:"unexpected5xx,omitempty"`
OfferedRatePerSecond float64 `json:"offeredRatePerSecond,omitempty"`
}
type phaseResult struct {
report phaseReport
latencies []time.Duration
err error
}
type httpStatusError struct {
Status int
Body string
}
type operationError struct {
Operation string
Err error
}
func (e *operationError) Error() string { return e.Operation + ": " + e.Err.Error() }
func (e *operationError) Unwrap() error { return e.Err }
func withOperation(operation string, err error) error {
if err == nil {
return nil
}
var existing *operationError
if errors.As(err, &existing) {
return err
}
return &operationError{Operation: operation, Err: err}
}
func (e *httpStatusError) Error() string {
return fmt.Sprintf("HTTP %d: %s", e.Status, e.Body)
}
func main() {
opts, err := parseOptions()
if err != nil {
fmt.Fprintln(os.Stderr, "acceptance load configuration error:", err)
os.Exit(2)
}
signalContext, stopSignals := signal.NotifyContext(
context.Background(),
os.Interrupt,
syscall.SIGTERM,
)
defer stopSignals()
ctx, cancel := context.WithTimeout(signalContext, opts.timeout)
defer cancel()
result := report{
SchemaVersion: "acceptance-load-report/v1",
RunID: opts.runID, Profile: opts.profile, StartedAt: time.Now().UTC(), SecretSafe: true,
}
phaseResults, runErr := run(ctx, opts)
for _, phase := range phaseResults {
phase.report.P50MS = percentileMilliseconds(phase.latencies, 0.50)
phase.report.P95MS = percentileMilliseconds(phase.latencies, 0.95)
phase.report.P99MS = percentileMilliseconds(phase.latencies, 0.99)
result.Phases = append(result.Phases, phase.report)
if runErr == nil && phase.err != nil {
runErr = phase.err
}
}
result.FinishedAt = time.Now().UTC()
result.Passed = runErr == nil
if runErr != nil {
result.Failure = redactError(runErr.Error(), opts)
var operationErr *operationError
if errors.As(runErr, &operationErr) {
result.FailureOperation = operationErr.Operation
}
}
payload, _ := json.MarshalIndent(result, "", " ")
fmt.Println(string(payload))
if opts.reportPath != "" {
if err := writeReportExclusive(opts.reportPath, payload); err != nil {
fmt.Fprintln(os.Stderr, "write acceptance report failed:", err)
os.Exit(1)
}
}
if runErr != nil {
os.Exit(1)
}
}
func writeReportExclusive(path string, payload []byte) error {
reportFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
if err != nil {
return err
}
if _, err := reportFile.Write(append(payload, '\n')); err != nil {
_ = reportFile.Close()
return err
}
return reportFile.Close()
}
func parseOptions() (options, error) {
var profile string
var reportPath string
var timeout time.Duration
var mixedDuration time.Duration
var imageRate float64
var videoRate float64
var shardIndex int
var shardCount int
var executionID string
flag.StringVar(&profile, "profile", env("AI_GATEWAY_ACCEPTANCE_PROFILE", "simulated-all"), "simulated-all, Gemini profile, video profile, or real-canary")
flag.StringVar(&reportPath, "report", env("AI_GATEWAY_ACCEPTANCE_REPORT", ""), "secret-safe JSON report path")
flag.DurationVar(&timeout, "timeout", 45*time.Minute, "overall timeout")
flag.DurationVar(&mixedDuration, "duration", envDuration("AI_GATEWAY_ACCEPTANCE_MIXED_DURATION", 10*time.Minute), "mixed workload duration")
flag.Float64Var(&imageRate, "image-rate", envFloat("AI_GATEWAY_ACCEPTANCE_IMAGE_RATE", 0), "mixed image requests per second")
flag.Float64Var(&videoRate, "video-rate", envFloat("AI_GATEWAY_ACCEPTANCE_VIDEO_RATE", 0), "mixed video requests per second")
flag.IntVar(&shardIndex, "shard-index", 0, "zero-based distributed load shard index")
flag.IntVar(&shardCount, "shard-count", 1, "distributed load shard count")
flag.StringVar(&executionID, "execution-id", "", "unique workload execution identifier")
flag.Parse()
opts := options{
gateways: splitCSV(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAYS")),
gatewayTLSName: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME")),
gatewayCAFile: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAY_CA_FILE")),
emulatorURL: strings.TrimRight(strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_EMULATOR_URL")), "/"),
apiKeys: splitCSV(env("AI_GATEWAY_ACCEPTANCE_API_KEYS", os.Getenv("AI_GATEWAY_ACCEPTANCE_API_KEY"))),
runID: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_RUN_ID")),
runToken: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_RUN_TOKEN")),
geminiModel: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_GEMINI_MODEL")),
videoModel: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_VIDEO_MODEL")),
profile: strings.ToLower(strings.TrimSpace(profile)),
reportPath: strings.TrimSpace(reportPath),
realImageURLs: splitCSV(os.Getenv("AI_GATEWAY_ACCEPTANCE_REAL_IMAGE_URLS")),
timeout: timeout,
mixedDuration: mixedDuration,
imageRate: imageRate,
videoRate: videoRate,
shardIndex: shardIndex,
shardCount: shardCount,
executionID: strings.ToLower(strings.TrimSpace(executionID)),
}
if opts.executionID == "" {
opts.executionID = opts.profile
}
if matched, _ := regexp.MatchString(`^[a-z0-9][a-z0-9-]{0,127}$`, opts.executionID); !matched {
return options{}, errors.New("execution-id must contain only lowercase letters, numbers, and hyphens")
}
if opts.shardCount < 1 || opts.shardCount > 32 || opts.shardIndex < 0 || opts.shardIndex >= opts.shardCount {
return options{}, fmt.Errorf("invalid load shard %d/%d", opts.shardIndex, opts.shardCount)
}
if len(opts.gateways) < 1 || len(opts.gateways) > 2 || (opts.shardCount == 1 && len(opts.gateways) != 2) {
return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAYS must contain two API base URLs, or one URL for a distributed shard")
}
for index := range opts.gateways {
opts.gateways[index] = strings.TrimRight(opts.gateways[index], "/")
if _, err := url.ParseRequestURI(opts.gateways[index]); err != nil {
return options{}, fmt.Errorf("gateway %d URL is invalid", index+1)
}
}
if opts.gatewayTLSName != "" &&
(strings.ContainsAny(opts.gatewayTLSName, "/:@") || !strings.Contains(opts.gatewayTLSName, ".")) {
return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME must be a DNS hostname")
}
if len(opts.apiKeys) == 0 || opts.runID == "" || opts.runToken == "" {
return options{}, errors.New("acceptance API key, Run ID, and Run Token environment variables are required")
}
if opts.geminiModel == "" || opts.videoModel == "" {
return options{}, errors.New("Gemini and video model environment variables are required")
}
if opts.timeout <= 0 {
return options{}, errors.New("timeout must be positive")
}
switch opts.profile {
case "simulated-smoke", "simulated-all", "gemini-baseline", "gemini-large", "gemini-peak", "video-throughput", "video-recovery",
"mixed-soak", "mixed-overload":
if opts.emulatorURL == "" {
return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_EMULATOR_URL is required for simulated profiles")
}
if (opts.profile == "mixed-soak" || opts.profile == "mixed-overload") &&
(opts.mixedDuration <= 0 || opts.imageRate < 0 || opts.videoRate < 0 ||
opts.imageRate+opts.videoRate <= 0) {
return options{}, errors.New("mixed profiles require a positive duration and offered rate")
}
case "real-canary":
if len(opts.realImageURLs) < 3 {
return options{}, errors.New("real-canary requires at least three public AI_GATEWAY_ACCEPTANCE_REAL_IMAGE_URLS")
}
default:
return options{}, fmt.Errorf("unsupported profile %q", opts.profile)
}
return opts, nil
}
func (o options) shardRequestCount(total int) int {
if total <= o.shardIndex {
return 0
}
return (total-1-o.shardIndex)/o.shardCount + 1
}
func (o options) logicalRequestIndex(localIndex int) int {
return o.shardIndex + localIndex*o.shardCount
}
func run(ctx context.Context, opts options) ([]phaseResult, error) {
var tlsConfig *tls.Config
if opts.gatewayTLSName != "" || opts.gatewayCAFile != "" {
rootCAs, err := acceptanceRootCAs(opts.gatewayCAFile)
if err != nil {
return nil, err
}
tlsConfig = &tls.Config{
MinVersion: tls.VersionTLS12,
ServerName: opts.gatewayTLSName,
RootCAs: rootCAs,
}
}
client := &http.Client{
Timeout: opts.timeout,
Transport: &http.Transport{
MaxIdleConns: 2048,
MaxIdleConnsPerHost: 1024,
MaxConnsPerHost: 1024,
ForceAttemptHTTP2: true,
TLSClientConfig: tlsConfig,
},
}
results := make([]phaseResult, 0, 5)
runPhase := func(name string) error {
var result phaseResult
if profile, ok := acceptanceworkload.GeminiProfileByName(name); ok {
result = runGemini(
ctx, client, opts, name, opts.shardRequestCount(profile.Requests),
profile.InputBytes, profile.OutputBytes, false, opts.shardIndex, opts.shardCount,
)
} else {
switch name {
case "smoke-gemini":
result = runGemini(ctx, client, opts, name, opts.shardRequestCount(4), 256<<10, 256<<10, false, opts.shardIndex, opts.shardCount)
case "smoke-video":
result = runVideo(ctx, client, opts, name, opts.shardRequestCount(10), false, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount)
case "video-throughput":
result = runVideo(ctx, client, opts, name, opts.shardRequestCount(1200), false, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount)
case "video-recovery":
result = runVideo(ctx, client, opts, name, opts.shardRequestCount(96), true, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount)
case "mixed-soak", "mixed-overload":
result = runMixed(ctx, client, opts, name == "mixed-overload")
case "real-gemini-canary":
result = runGemini(ctx, client, opts, name, 1, 256<<10, 0, true, 0, 1)
case "real-video-canary":
result = runVideo(ctx, client, opts, name, 1, false, opts.realImageURLs, true, 0, 1)
default:
return fmt.Errorf("unknown phase %q", name)
}
}
results = append(results, result)
return result.err
}
phases := []string{opts.profile}
if opts.profile == "simulated-smoke" {
phases = []string{"smoke-gemini", "smoke-video"}
} else if opts.profile == "simulated-all" {
phases = []string{"gemini-baseline", "gemini-large", "gemini-peak", "video-throughput", "video-recovery"}
} else if opts.profile == "real-canary" {
phases = []string{"real-gemini-canary", "real-video-canary"}
}
for _, phase := range phases {
if err := runPhase(phase); err != nil {
return results, err
}
}
return results, nil
}
func runGemini(
ctx context.Context,
client *http.Client,
opts options,
name string,
requestCount int,
inputBytes int,
expectedOutputBytes int,
realUpstream bool,
requestOffset int,
requestStride int,
) phaseResult {
if requestStride < 1 {
requestStride = 1
}
startedAt := time.Now()
expectedHash := ""
if expectedOutputBytes > 0 {
sum := sha256.Sum256(paddedPNG(expectedOutputBytes))
expectedHash = hex.EncodeToString(sum[:])
}
latencies := make([]time.Duration, requestCount)
var decodedBytes int64
var completed int
var failures int
var firstErr error
var mu sync.Mutex
slots := make(chan struct{}, requestCount)
var wg sync.WaitGroup
for index := 0; index < requestCount; index++ {
index := index
wg.Add(1)
go func() {
defer wg.Done()
logicalIndex := requestOffset + index*requestStride
select {
case slots <- struct{}{}:
case <-ctx.Done():
return
}
defer func() { <-slots }()
requestStarted := time.Now()
endpoint := opts.gateways[logicalIndex%len(opts.gateways)] + "/v1beta/models/" + url.PathEscape(opts.geminiModel) + ":generateContent"
requestBody := streamGeminiRequestBody(paddedPNGVariant(inputBytes, geminiInputVariant(opts.runID+"/"+opts.executionID, logicalIndex)))
req, err := http.NewRequestWithContext(
ctx,
http.MethodPost,
endpoint,
requestBody,
)
if err != nil {
_ = requestBody.Close()
}
if err == nil {
opts.setHeaders(req, logicalIndex, realUpstream)
req.Header.Set("Content-Type", "application/json")
var response *http.Response
response, err = client.Do(req)
if err == nil {
if response.StatusCode != http.StatusOK {
err = withOperation("gemini_submit", responseStatusError(response))
} else {
var size int64
var outputHash string
size, outputHash, err = streamGeminiImageHash(response.Body)
err = withOperation("gemini_response", err)
_ = response.Body.Close()
if err == nil && expectedOutputBytes > 0 && (size != int64(expectedOutputBytes) || outputHash != expectedHash) {
err = withOperation("gemini_response", fmt.Errorf("Gemini output mismatch: bytes=%d hash=%s", size, outputHash))
}
if err == nil {
mu.Lock()
decodedBytes += size
if expectedHash == "" {
expectedHash = outputHash
}
mu.Unlock()
}
}
}
}
err = withOperation("gemini_submit", err)
latency := time.Since(requestStarted)
mu.Lock()
latencies[index] = latency
if err != nil {
failures++
if firstErr == nil {
firstErr = err
}
} else {
completed++
}
mu.Unlock()
}()
}
wg.Wait()
elapsed := time.Since(startedAt)
result := phaseReport{
Name: name, Requests: requestCount, Completed: completed, Failed: failures,
DurationMS: elapsed.Milliseconds(), DecodedOutputBytes: decodedBytes, OutputSHA256: expectedHash,
}
if firstErr == nil && name == "gemini-baseline" && elapsed > 8*time.Minute {
firstErr = fmt.Errorf("Gemini baseline exceeded 8 minutes: %s", elapsed)
}
if firstErr == nil && completed != requestCount {
firstErr = fmt.Errorf("Gemini completed %d of %d requests", completed, requestCount)
}
return phaseResult{report: result, latencies: latencies, err: firstErr}
}
func streamGeminiRequestBody(input []byte) io.ReadCloser {
reader, writer := io.Pipe()
go func() {
closeWithError := func(err error) {
_ = writer.CloseWithError(err)
}
if _, err := io.WriteString(writer, `{"contents":[{"role":"user","parts":[{"text":"将参考图编辑为蓝色赛博朋克风格,保留主体结构"},{"inlineData":{"mimeType":"image/png","data":"`); err != nil {
closeWithError(err)
return
}
encoder := base64.NewEncoder(base64.StdEncoding, writer)
if _, err := encoder.Write(input); err != nil {
closeWithError(err)
return
}
if err := encoder.Close(); err != nil {
closeWithError(err)
return
}
if _, err := io.WriteString(writer, `"}}]}],"generationConfig":{"responseModalities":["IMAGE"]}}`); err != nil {
closeWithError(err)
return
}
_ = writer.Close()
}()
return reader
}
func runVideo(
ctx context.Context,
client *http.Client,
opts options,
name string,
requestCount int,
longRun bool,
imageURLs []string,
realUpstream bool,
requestOffset int,
requestStride int,
) phaseResult {
if requestStride < 1 {
requestStride = 1
}
startedAt := time.Now()
combinationCount := 128
if realUpstream {
combinationCount = 1
}
combinations := videoCombinations(imageURLs, combinationCount)
var taskIDs = make([]string, requestCount)
latencies := make([]time.Duration, requestCount)
startedTasks := make([]time.Time, requestCount)
var completed int
var failures int
var forcedConversions int
var firstErr error
var mu sync.Mutex
var wg sync.WaitGroup
submitStartedAt := time.Now()
for index := 0; index < requestCount; index++ {
index := index
wg.Add(1)
go func() {
defer wg.Done()
logicalIndex := requestOffset + index*requestStride
requestStarted := time.Now()
startedTasks[index] = requestStarted
imageCount := acceptanceworkload.VideoImageCount(logicalIndex)
combo := append([]string(nil), combinations[logicalIndex%len(combinations)]...)
if len(combo) > imageCount {
combo = combo[:imageCount]
}
if !realUpstream && logicalIndex%4 == 0 {
combo[0] = imageURLs[12+(logicalIndex/4)%4]
}
content := make([]any, 0, len(combo)+1)
prompt := "多参考图生成连续运镜视频,保持人物、服装和场景一致"
if longRun {
prompt += " acceptance-long-recovery"
}
if !realUpstream && logicalIndex%4 == 0 {
prompt += " acceptance-force-conversion"
}
content = append(content, map[string]any{"type": "text", "text": prompt})
for imageIndex, imageURL := range combo {
role := "reference_image"
if logicalIndex%5 == 0 && imageIndex == 0 {
role = "first_frame"
}
if logicalIndex%5 == 0 && len(combo) > 3 && imageIndex == len(combo)-1 {
role = "last_frame"
}
content = append(content, map[string]any{
"type": "image_url", "role": role,
"image_url": map[string]any{"url": imageURL},
})
}
body, _ := json.Marshal(map[string]any{
"model": opts.videoModel, "modelType": "omni_video", "content": content, "seed": logicalIndex + 1,
"duration": 5, "ratio": "16:9", "resolution": "720p",
})
endpoint := opts.gateways[logicalIndex%len(opts.gateways)] + "/api/v1/videos/generations"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
if err == nil {
opts.setHeaders(req, logicalIndex, realUpstream)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("X-Async", "true")
var response *http.Response
response, err = client.Do(req)
if err == nil {
if response.StatusCode != http.StatusAccepted {
err = withOperation("video_submit", responseStatusError(response))
} else {
var accepted struct {
TaskID string `json:"taskId"`
}
err = json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&accepted)
_ = response.Body.Close()
taskIDs[index] = strings.TrimSpace(accepted.TaskID)
if err == nil && taskIDs[index] == "" {
err = errors.New("video acceptance response did not include taskId")
}
}
}
}
err = withOperation("video_submit", err)
mu.Lock()
latencies[index] = time.Since(requestStarted)
if !realUpstream && logicalIndex%4 == 0 {
forcedConversions++
}
if err != nil {
failures++
if firstErr == nil {
firstErr = err
}
}
mu.Unlock()
}()
}
wg.Wait()
submissionDuration := time.Since(submitStartedAt)
if firstErr == nil && name == "video-throughput" && submissionDuration > 10*time.Second {
firstErr = fmt.Errorf("1200 video submissions exceeded 10 seconds: %s", submissionDuration)
}
uniqueTaskIDs := map[string]struct{}{}
for _, taskID := range taskIDs {
if taskID != "" {
uniqueTaskIDs[taskID] = struct{}{}
}
}
if firstErr == nil && len(uniqueTaskIDs) != requestCount {
firstErr = fmt.Errorf("video task IDs are not unique: unique=%d requests=%d", len(uniqueTaskIDs), requestCount)
}
pollSlots := make(chan struct{}, 256)
for index, taskID := range taskIDs {
if taskID == "" {
continue
}
index, taskID := index, taskID
logicalIndex := requestOffset + index*requestStride
wg.Add(1)
go func() {
defer wg.Done()
select {
case pollSlots <- struct{}{}:
case <-ctx.Done():
return
}
defer func() { <-pollSlots }()
err := pollVideoTask(ctx, client, opts, taskID, logicalIndex, realUpstream)
mu.Lock()
latencies[index] = time.Since(startedTasks[index])
if err != nil {
failures++
if firstErr == nil {
firstErr = err
}
} else {
completed++
}
mu.Unlock()
}()
}
wg.Wait()
comboSet := map[string]struct{}{}
for _, combo := range combinations {
comboSet[strings.Join(combo, "\x00")] = struct{}{}
}
if firstErr == nil && !realUpstream && len(comboSet) != 128 {
firstErr = fmt.Errorf("video image combinations=%d, want 128", len(comboSet))
}
elapsed := time.Since(startedAt)
result := phaseReport{
Name: name, Requests: requestCount, Completed: completed, Failed: failures,
DurationMS: elapsed.Milliseconds(), SubmissionDurationMS: submissionDuration.Milliseconds(),
UniqueTaskIDs: len(uniqueTaskIDs), UniqueImageCombos: len(comboSet),
ForcedConversionTasks: forcedConversions,
}
if firstErr == nil && completed != requestCount {
firstErr = fmt.Errorf("video completed %d of %d tasks", completed, requestCount)
}
return phaseResult{report: result, latencies: latencies, err: firstErr}
}
func runMixed(
ctx context.Context,
client *http.Client,
opts options,
expectThrottling bool,
) phaseResult {
runCtx, stopRun := context.WithCancel(ctx)
defer stopRun()
startedAt := time.Now()
totalRate := opts.imageRate + opts.videoRate
interval := time.Duration(float64(time.Second) / totalRate)
if interval < time.Millisecond {
interval = time.Millisecond
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
stop := time.NewTimer(opts.mixedDuration)
defer stop.Stop()
slots := make(chan struct{}, 4096)
var wg sync.WaitGroup
var mu sync.Mutex
latencies := make([]time.Duration, 0, int(totalRate*opts.mixedDuration.Seconds()))
requests := 0
completed := 0
failed := 0
throttled := 0
unexpected5xx := 0
var firstErr error
imageAccumulator := float64(0)
generating := true
for generating {
select {
case <-runCtx.Done():
generating = false
if ctx.Err() != nil {
mu.Lock()
if firstErr == nil {
firstErr = ctx.Err()
}
mu.Unlock()
}
case <-stop.C:
generating = false
case <-ticker.C:
select {
case slots <- struct{}{}:
case <-runCtx.Done():
generating = false
continue
}
requestIndex := opts.logicalRequestIndex(requests)
requests++
imageAccumulator += opts.imageRate
isImage := imageAccumulator >= totalRate
if isImage {
imageAccumulator -= totalRate
}
wg.Add(1)
go func() {
defer wg.Done()
defer func() { <-slots }()
var result phaseResult
if isImage {
result = runGemini(runCtx, client, opts, "mixed-image", 1, 256<<10, 256<<10, false, requestIndex, 1)
} else {
result = runVideo(
runCtx,
client,
opts,
"mixed-video",
1,
false,
opts.emulatorFixtureURLs(),
false,
requestIndex,
1,
)
}
latency := time.Since(startedAt)
if len(result.latencies) > 0 {
latency = result.latencies[0]
}
mu.Lock()
defer mu.Unlock()
latencies = append(latencies, latency)
if result.err == nil {
completed++
return
}
var statusErr *httpStatusError
if errors.As(result.err, &statusErr) && statusErr.Status == http.StatusTooManyRequests {
throttled++
if !expectThrottling && firstErr == nil {
firstErr = result.err
failed++
stopRun()
}
return
}
if errors.As(result.err, &statusErr) && statusErr.Status >= 500 {
unexpected5xx++
}
failed++
if firstErr == nil {
firstErr = result.err
stopRun()
}
}()
}
}
wg.Wait()
if firstErr == nil && expectThrottling && throttled == 0 {
firstErr = errors.New("mixed overload did not receive any 429 response")
}
if firstErr == nil && completed+throttled != requests {
firstErr = fmt.Errorf(
"mixed workload accounted for %d of %d requests",
completed+throttled,
requests,
)
}
if firstErr == nil && unexpected5xx > 0 {
firstErr = fmt.Errorf("mixed workload received %d unexpected 5xx responses", unexpected5xx)
}
return phaseResult{
report: phaseReport{
Name: map[bool]string{false: "mixed-soak", true: "mixed-overload"}[expectThrottling],
Requests: requests,
Completed: completed,
Failed: failed,
DurationMS: time.Since(startedAt).Milliseconds(),
Throttled: throttled,
Unexpected5xx: unexpected5xx,
OfferedRatePerSecond: totalRate,
},
latencies: latencies,
err: firstErr,
}
}
func pollVideoTask(ctx context.Context, client *http.Client, opts options, taskID string, index int, realUpstream bool) error {
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
for {
endpoint := opts.gateways[index%len(opts.gateways)] + "/api/v1/ai/result/" + url.PathEscape(taskID)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
if err != nil {
return withOperation("video_poll", err)
}
opts.setHeaders(req, index, realUpstream)
response, err := client.Do(req)
if err == nil {
if response.StatusCode != http.StatusOK {
err = responseStatusError(response)
} else {
var payload map[string]any
err = json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(&payload)
_ = response.Body.Close()
status := strings.ToLower(strings.TrimSpace(fmt.Sprint(payload["status"])))
switch status {
case "succeeded", "success":
mediaURL := findMediaURL(payload)
if mediaURL == "" {
return withOperation("video_poll", fmt.Errorf("video task %s succeeded without a media URL", taskID))
}
return withOperation("video_download", validateVideoAsset(ctx, client, opts, mediaURL, index))
case "failed", "cancelled", "canceled":
return withOperation("video_poll", fmt.Errorf("video task %s finished with status %s", taskID, status))
}
}
}
if err != nil {
return withOperation("video_poll", err)
}
select {
case <-ctx.Done():
return withOperation("video_poll", ctx.Err())
case <-ticker.C:
}
}
}
func findMediaURL(value any) string {
switch typed := value.(type) {
case map[string]any:
preferredKeys := []string{"video_url", "url", "data", "content", "output", "result", "upload"}
visited := make(map[string]struct{}, len(typed))
for _, key := range preferredKeys {
item, ok := typed[key]
if !ok {
continue
}
visited[key] = struct{}{}
normalized := strings.ToLower(strings.TrimSpace(key))
if (normalized == "url" || normalized == "video_url") && acceptanceMediaURL(strings.TrimSpace(fmt.Sprint(item))) {
return strings.TrimSpace(fmt.Sprint(item))
}
if mediaURL := findMediaURL(item); mediaURL != "" {
return mediaURL
}
}
keys := make([]string, 0, len(typed)-len(visited))
for key := range typed {
if _, ok := visited[key]; !ok {
keys = append(keys, key)
}
}
sort.Strings(keys)
for _, key := range keys {
if mediaURL := findMediaURL(typed[key]); mediaURL != "" {
return mediaURL
}
}
case []any:
for _, item := range typed {
if mediaURL := findMediaURL(item); mediaURL != "" {
return mediaURL
}
}
}
return ""
}
func acceptanceMediaURL(value string) bool {
return strings.HasPrefix(value, "/") || strings.HasPrefix(value, "http://") || strings.HasPrefix(value, "https://")
}
func validateVideoAsset(ctx context.Context, client *http.Client, opts options, mediaURL string, index int) error {
requestURL, gatewayRequest, err := acceptanceMediaRequestURL(opts, mediaURL, index)
if err != nil {
return fmt.Errorf("resolve final video: %w", err)
}
request, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil)
if err != nil {
return err
}
if gatewayRequest {
opts.setHeaders(request, index, false)
}
response, err := client.Do(request)
if err != nil {
return fmt.Errorf("download final video: %w", err)
}
defer response.Body.Close()
if response.StatusCode < 200 || response.StatusCode >= 300 {
return fmt.Errorf("download final video: HTTP %d", response.StatusCode)
}
payload, err := io.ReadAll(io.LimitReader(response.Body, 64<<10))
if err != nil {
return fmt.Errorf("read final video: %w", err)
}
contentType := strings.ToLower(strings.TrimSpace(response.Header.Get("Content-Type")))
hasMP4Header := len(payload) >= 12 && string(payload[4:8]) == "ftyp"
if len(payload) < 12 || (!strings.HasPrefix(contentType, "video/") && !hasMP4Header) {
return fmt.Errorf("final video asset is invalid: bytes=%d content_type=%q", len(payload), contentType)
}
return nil
}
func acceptanceMediaRequestURL(opts options, mediaURL string, index int) (string, bool, error) {
mediaURL = strings.TrimSpace(mediaURL)
if mediaURL == "" {
return "", false, errors.New("empty media URL")
}
if strings.HasPrefix(mediaURL, "/") {
if len(opts.gateways) == 0 {
return "", false, errors.New("no gateway URL configured")
}
base, err := url.Parse(opts.gateways[index%len(opts.gateways)])
if err != nil {
return "", false, err
}
reference, err := url.Parse(mediaURL)
if err != nil {
return "", false, err
}
return base.ResolveReference(reference).String(), true, nil
}
parsed, err := url.Parse(mediaURL)
if err != nil {
return "", false, err
}
if opts.gatewayTLSName == "" || !strings.EqualFold(parsed.Hostname(), opts.gatewayTLSName) {
return parsed.String(), false, nil
}
if len(opts.gateways) == 0 {
return "", false, errors.New("no gateway URL configured")
}
base, err := url.Parse(opts.gateways[index%len(opts.gateways)])
if err != nil {
return "", false, err
}
base.Path = parsed.Path
base.RawPath = parsed.RawPath
base.RawQuery = parsed.RawQuery
base.Fragment = parsed.Fragment
return base.String(), true, nil
}
func (o options) setHeaders(request *http.Request, requestIndex int, realUpstream bool) {
if o.gatewayTLSName != "" {
request.Host = o.gatewayTLSName
}
request.Header.Set("Authorization", "Bearer "+o.apiKeys[requestIndex%len(o.apiKeys)])
request.Header.Set(runHeader, o.runID)
request.Header.Set(tokenHeader, o.runToken)
if request.Method == http.MethodPost {
request.Header.Set("Idempotency-Key", fmt.Sprintf("acceptance-%s-%s-%d", o.runID, o.executionID, requestIndex))
}
if realUpstream {
request.Header.Set(upstreamHeader, "real")
}
}
func (o options) emulatorFixtureURLs() []string {
out := make([]string, 0, 16)
for index := 0; index < 4; index++ {
out = append(out, fmt.Sprintf("%s/fixtures/image-%02d.png", o.emulatorURL, index))
}
for index := 4; index < 8; index++ {
out = append(out, fmt.Sprintf("%s/fixtures/image-%02d.jpg", o.emulatorURL, index))
}
for index := 8; index < 12; index++ {
out = append(out, fmt.Sprintf("%s/fixtures/image-%02d.webp", o.emulatorURL, index))
}
for index := 12; index < 16; index++ {
out = append(out, fmt.Sprintf("%s/fixtures/image-%02d-oversized.jpg", o.emulatorURL, index))
}
return out
}
func videoCombinations(images []string, count int) [][]string {
if len(images) < 3 {
return nil
}
if len(images) < 9 {
return [][]string{append([]string(nil), images[:3]...)}
}
out := make([][]string, 0, count)
seen := map[string]struct{}{}
for seed := 0; len(out) < count; seed++ {
imageCount := acceptanceworkload.VideoImageCount(seed)
combo := make([]string, 0, imageCount)
used := map[int]struct{}{}
state := uint64(seed + 1)
for len(combo) < imageCount {
state = state*6364136223846793005 + 1442695040888963407
index := int((state >> 32) % uint64(len(images)))
if _, exists := used[index]; exists {
continue
}
used[index] = struct{}{}
combo = append(combo, images[index])
}
signature := strings.Join(combo, "\x00")
if _, exists := seen[signature]; exists {
continue
}
seen[signature] = struct{}{}
out = append(out, combo)
}
return out
}
func streamGeminiImageHash(reader io.Reader) (int64, string, error) {
buffered := bufio.NewReaderSize(reader, 64<<10)
if err := scanUntil(buffered, []byte(`"inlineData"`), 256<<20); err != nil {
return 0, "", err
}
if err := scanUntil(buffered, []byte(`"data"`), 1<<20); err != nil {
return 0, "", err
}
for {
item, err := buffered.ReadByte()
if err != nil {
return 0, "", err
}
if item == ':' {
break
}
}
for {
item, err := buffered.ReadByte()
if err != nil {
return 0, "", err
}
if item == '"' {
break
}
if item != ' ' && item != '\t' && item != '\r' && item != '\n' {
return 0, "", errors.New("Gemini inlineData.data is not a JSON string")
}
}
jsonString := &base64JSONStringReader{reader: buffered}
digest := sha256.New()
decodedBytes, err := io.Copy(digest, base64.NewDecoder(base64.StdEncoding, jsonString))
if err != nil {
return 0, "", fmt.Errorf("decode Gemini Base64 output: %w", err)
}
if decodedBytes == 0 {
return 0, "", errors.New("Gemini Base64 output is empty")
}
return decodedBytes, hex.EncodeToString(digest.Sum(nil)), nil
}
type base64JSONStringReader struct {
reader *bufio.Reader
done bool
}
func (r *base64JSONStringReader) Read(payload []byte) (int, error) {
if r.done {
return 0, io.EOF
}
written := 0
for written < len(payload) {
item, err := r.reader.ReadByte()
if err != nil {
return written, err
}
if item == '"' {
r.done = true
if written == 0 {
return 0, io.EOF
}
return written, nil
}
if item == '\\' {
return written, errors.New("escaped Base64 JSON strings are not supported")
}
payload[written] = item
written++
}
return written, nil
}
func scanUntil(reader *bufio.Reader, target []byte, limit int64) error {
matched := 0
for read := int64(0); read < limit; read++ {
item, err := reader.ReadByte()
if err != nil {
return err
}
if item == target[matched] {
matched++
if matched == len(target) {
return nil
}
continue
}
if item == target[0] {
matched = 1
} else {
matched = 0
}
}
return errors.New("Gemini inlineData output was not found")
}
func responseStatusError(response *http.Response) error {
defer response.Body.Close()
payload, _ := io.ReadAll(io.LimitReader(response.Body, 512))
return &httpStatusError{Status: response.StatusCode, Body: strings.TrimSpace(string(payload))}
}
func percentileMilliseconds(values []time.Duration, quantile float64) float64 {
filtered := make([]time.Duration, 0, len(values))
for _, value := range values {
if value > 0 {
filtered = append(filtered, value)
}
}
if len(filtered) == 0 {
return 0
}
sort.Slice(filtered, func(i, j int) bool { return filtered[i] < filtered[j] })
index := int(float64(len(filtered)-1) * quantile)
return float64(filtered[index].Microseconds()) / 1000
}
func paddedPNG(size int) []byte {
return paddedPNGVariant(size, 0)
}
func geminiInputVariant(runID string, logicalIndex int) int {
digest := sha256.Sum256([]byte(fmt.Sprintf("%s:%d", strings.TrimSpace(runID), logicalIndex)))
return int(binary.BigEndian.Uint64(digest[:8]) & uint64(^uint(0)>>1))
}
func paddedPNGVariant(size int, variant int) []byte {
base, _ := base64.StdEncoding.DecodeString("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=")
if size <= len(base) {
return base
}
iendOffset := len(base) - 12
paddingLength := size - len(base) - 12
if paddingLength < 0 {
paddingLength = 0
}
chunk := make([]byte, 12+paddingLength)
binary.BigEndian.PutUint32(chunk[:4], uint32(paddingLength))
copy(chunk[4:8], []byte("teST"))
if variant != 0 && paddingLength >= 8 {
binary.BigEndian.PutUint64(chunk[8:16], uint64(variant))
}
binary.BigEndian.PutUint32(chunk[8+paddingLength:], crc32.ChecksumIEEE(chunk[4:8+paddingLength]))
out := make([]byte, 0, size)
out = append(out, base[:iendOffset]...)
out = append(out, chunk...)
out = append(out, base[iendOffset:]...)
return out
}
func redactError(message string, opts options) string {
for _, secret := range append(append([]string(nil), opts.apiKeys...), opts.runToken) {
if secret != "" {
message = strings.ReplaceAll(message, secret, "[REDACTED]")
}
}
return reportURLPattern.ReplaceAllString(message, "[REDACTED_URL]")
}
func acceptanceRootCAs(path string) (*x509.CertPool, error) {
if strings.TrimSpace(path) == "" {
return nil, nil
}
info, err := os.Lstat(path)
if err != nil {
return nil, fmt.Errorf("read Gateway CA file: %w", err)
}
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
return nil, errors.New("Gateway CA file must be a regular non-symlink file")
}
payload, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("read Gateway CA file: %w", err)
}
roots, err := x509.SystemCertPool()
if err != nil || roots == nil {
roots = x509.NewCertPool()
}
if !roots.AppendCertsFromPEM(payload) {
return nil, errors.New("Gateway CA file does not contain a PEM certificate")
}
return roots, nil
}
func splitCSV(value string) []string {
items := strings.Split(value, ",")
out := make([]string, 0, len(items))
for _, item := range items {
if item = strings.TrimSpace(item); item != "" {
out = append(out, item)
}
}
return out
}
func env(name string, fallback string) string {
if value := strings.TrimSpace(os.Getenv(name)); value != "" {
return value
}
return fallback
}
func envDuration(name string, fallback time.Duration) time.Duration {
value := strings.TrimSpace(os.Getenv(name))
if value == "" {
return fallback
}
parsed, err := time.ParseDuration(value)
if err != nil {
return -1
}
return parsed
}
func envFloat(name string, fallback float64) float64 {
value := strings.TrimSpace(os.Getenv(name))
if value == "" {
return fallback
}
parsed, err := strconv.ParseFloat(value, 64)
if err != nil {
return -1
}
return parsed
}