wangbo
215b6b607f
线上 P24 验收证明 Worker 注册容量为 48，但异步 admission 始终只有 8 条活跃租约。将单事务调度窗口覆盖当前 2 x P32 的 64 槽，避免多 Worker 在同一 FIFO 小批次上竞争形成隐性全局上限。\n\nGemini 流式请求补充可回放 GetBody 与精确 Content-Length，使带幂等键的 POST 遇到 HTTP/2 可重试传输错误时可安全重放；同时提前部署验收回调收集器，保证旧 Run side effects 能在接管前收口。\n\n已通过 Go 全量测试、go vet、gofmt、Shell bash -n、ShellCheck、验收与人工发布脚本测试、迁移安全检查，以及独立 PostgreSQL 48 条 admission/租约/River job 原子集成测试。
350 lines
12 KiB
Go
350 lines
12 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"sync/atomic"
|
|
"testing"
|
|
)
|
|
|
|
func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
|
|
payload := paddedPNG(4 << 20)
|
|
response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
|
|
base64.StdEncoding.EncodeToString(payload))
|
|
size, digest, err := streamGeminiImageHash(bytes.NewBufferString(response))
|
|
if err != nil {
|
|
t.Fatalf("stream Gemini output: %v", err)
|
|
}
|
|
expected := sha256.Sum256(payload)
|
|
if size != int64(len(payload)) || digest != hex.EncodeToString(expected[:]) {
|
|
t.Fatalf("size=%d digest=%s", size, digest)
|
|
}
|
|
}
|
|
|
|
func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) {
|
|
first := paddedPNGVariant(256<<10, 1)
|
|
second := paddedPNGVariant(256<<10, 2)
|
|
if len(first) != 256<<10 || len(second) != 256<<10 {
|
|
t.Fatalf("variant sizes=%d/%d", len(first), len(second))
|
|
}
|
|
firstHash := sha256.Sum256(first)
|
|
secondHash := sha256.Sum256(second)
|
|
if firstHash == secondHash {
|
|
t.Fatal("distinct variants have the same SHA-256")
|
|
}
|
|
}
|
|
|
|
func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
|
|
input := paddedPNGVariant(2<<20, 37)
|
|
var body struct {
|
|
Contents []struct {
|
|
Parts []struct {
|
|
InlineData *struct {
|
|
MIMEType string `json:"mimeType"`
|
|
Data string `json:"data"`
|
|
} `json:"inlineData"`
|
|
} `json:"parts"`
|
|
} `json:"contents"`
|
|
GenerationConfig struct {
|
|
ResponseModalities []string `json:"responseModalities"`
|
|
} `json:"generationConfig"`
|
|
}
|
|
stream := streamGeminiRequestBody(input)
|
|
defer stream.Close()
|
|
if err := json.NewDecoder(stream).Decode(&body); err != nil {
|
|
t.Fatalf("decode streamed Gemini request: %v", err)
|
|
}
|
|
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 2 || body.Contents[0].Parts[1].InlineData == nil {
|
|
t.Fatalf("unexpected Gemini body structure: %+v", body)
|
|
}
|
|
inlineData := body.Contents[0].Parts[1].InlineData
|
|
if inlineData.MIMEType != "image/png" {
|
|
t.Fatalf("mime type=%q", inlineData.MIMEType)
|
|
}
|
|
decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
|
|
if err != nil {
|
|
t.Fatalf("decode streamed input: %v", err)
|
|
}
|
|
if !bytes.Equal(decoded, input) {
|
|
t.Fatal("streamed input differs from source")
|
|
}
|
|
if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
|
|
t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
|
|
input := paddedPNGVariant(256<<10, 91)
|
|
request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", input)
|
|
if err != nil {
|
|
t.Fatalf("new Gemini request: %v", err)
|
|
}
|
|
if request.GetBody == nil {
|
|
t.Fatal("streaming Gemini request does not provide GetBody")
|
|
}
|
|
first, err := io.ReadAll(request.Body)
|
|
if err != nil {
|
|
t.Fatalf("read first request body: %v", err)
|
|
}
|
|
_ = request.Body.Close()
|
|
replayedBody, err := request.GetBody()
|
|
if err != nil {
|
|
t.Fatalf("replay Gemini request body: %v", err)
|
|
}
|
|
replayed, err := io.ReadAll(replayedBody)
|
|
if err != nil {
|
|
t.Fatalf("read replayed request body: %v", err)
|
|
}
|
|
_ = replayedBody.Close()
|
|
if !bytes.Equal(first, replayed) {
|
|
t.Fatal("replayed Gemini request body differs from the original")
|
|
}
|
|
if request.ContentLength != int64(len(first)) {
|
|
t.Fatalf("content length=%d, want %d", request.ContentLength, len(first))
|
|
}
|
|
}
|
|
|
|
func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
|
|
images := make([]string, 16)
|
|
for index := range images {
|
|
images[index] = fmt.Sprintf("https://fixtures.example/image-%02d", index)
|
|
}
|
|
combinations := videoCombinations(images, 128)
|
|
if len(combinations) != 128 {
|
|
t.Fatalf("combinations=%d", len(combinations))
|
|
}
|
|
seen := map[string]struct{}{}
|
|
for _, combination := range combinations {
|
|
seen[fmt.Sprint(combination)] = struct{}{}
|
|
if len(combination) != 3 && len(combination) != 6 && len(combination) != 9 {
|
|
t.Fatalf("invalid combination size=%d", len(combination))
|
|
}
|
|
}
|
|
if len(seen) != 128 {
|
|
t.Fatalf("unique combinations=%d", len(seen))
|
|
}
|
|
}
|
|
|
|
func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
|
|
output := paddedPNG(256 << 10)
|
|
encoded := base64.StdEncoding.EncodeToString(output)
|
|
var firstCalls atomic.Int64
|
|
var secondCalls atomic.Int64
|
|
var firstKeyCalls atomic.Int64
|
|
var secondKeyCalls atomic.Int64
|
|
newGateway := func(calls *atomic.Int64) *httptest.Server {
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
calls.Add(1)
|
|
switch r.Header.Get("Authorization") {
|
|
case "Bearer key-1":
|
|
firstKeyCalls.Add(1)
|
|
case "Bearer key-2":
|
|
secondKeyCalls.Add(1)
|
|
default:
|
|
t.Errorf("unexpected authorization header")
|
|
}
|
|
if r.Header.Get(runHeader) != "run-1" || r.Header.Get(tokenHeader) != "token-1" {
|
|
t.Errorf("missing acceptance headers")
|
|
}
|
|
var body map[string]any
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
t.Errorf("decode Gemini body: %v", err)
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"candidates": []any{map[string]any{"content": map[string]any{
|
|
"parts": []any{map[string]any{"inlineData": map[string]any{
|
|
"mimeType": "image/png", "data": encoded,
|
|
}}},
|
|
}}},
|
|
})
|
|
}))
|
|
}
|
|
first := newGateway(&firstCalls)
|
|
defer first.Close()
|
|
second := newGateway(&secondCalls)
|
|
defer second.Close()
|
|
|
|
opts := options{
|
|
gateways: []string{first.URL, second.URL}, apiKeys: []string{"key-1", "key-2"}, runID: "run-1",
|
|
runToken: "token-1", geminiModel: "gemini-image-test",
|
|
}
|
|
result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 256<<10, 256<<10, false, 0, 1)
|
|
if result.err != nil || result.report.Completed != 8 {
|
|
t.Fatalf("result=%+v err=%v", result.report, result.err)
|
|
}
|
|
if firstCalls.Load() != 4 || secondCalls.Load() != 4 {
|
|
t.Fatalf("gateway calls=%d/%d", firstCalls.Load(), secondCalls.Load())
|
|
}
|
|
if firstKeyCalls.Load() != 4 || secondKeyCalls.Load() != 4 {
|
|
t.Fatalf("API key calls=%d/%d", firstKeyCalls.Load(), secondKeyCalls.Load())
|
|
}
|
|
}
|
|
|
|
func TestDistributedShardIndexesCoverWorkloadWithoutOverlap(t *testing.T) {
|
|
first := options{shardIndex: 0, shardCount: 2}
|
|
second := options{shardIndex: 1, shardCount: 2}
|
|
if first.shardRequestCount(1001) != 501 || second.shardRequestCount(1001) != 500 {
|
|
t.Fatalf("shard counts=%d/%d", first.shardRequestCount(1001), second.shardRequestCount(1001))
|
|
}
|
|
seen := map[int]bool{}
|
|
for local := 0; local < first.shardRequestCount(1001); local++ {
|
|
seen[first.logicalRequestIndex(local)] = true
|
|
}
|
|
for local := 0; local < second.shardRequestCount(1001); local++ {
|
|
index := second.logicalRequestIndex(local)
|
|
if seen[index] {
|
|
t.Fatalf("duplicate logical index %d", index)
|
|
}
|
|
seen[index] = true
|
|
}
|
|
if len(seen) != 1001 {
|
|
t.Fatalf("covered indexes=%d", len(seen))
|
|
}
|
|
}
|
|
|
|
func TestGeminiInputVariantIsRunScoped(t *testing.T) {
|
|
if geminiInputVariant("run-a", 7) == geminiInputVariant("run-b", 7) {
|
|
t.Fatal("different Run IDs produced the same input variant")
|
|
}
|
|
if geminiInputVariant("run-a", 7) == geminiInputVariant("run-a", 8) {
|
|
t.Fatal("different logical indexes produced the same input variant")
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceIdempotencyKeyIncludesExecutionAndLogicalIndex(t *testing.T) {
|
|
request := httptest.NewRequest(http.MethodPost, "https://gateway.example/v1", nil)
|
|
opts := options{
|
|
apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1", executionID: "p24-2-gemini-baseline",
|
|
}
|
|
opts.setHeaders(request, 17, false)
|
|
if got := request.Header.Get("Idempotency-Key"); got != "acceptance-run-1-p24-2-gemini-baseline-17" {
|
|
t.Fatalf("idempotency key=%q", got)
|
|
}
|
|
poll := httptest.NewRequest(http.MethodGet, "https://gateway.example/result", nil)
|
|
opts.setHeaders(poll, 17, false)
|
|
if got := poll.Header.Get("Idempotency-Key"); got != "" {
|
|
t.Fatalf("GET request unexpectedly has idempotency key %q", got)
|
|
}
|
|
}
|
|
|
|
func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
|
if request.Header.Get("Authorization") != "" {
|
|
t.Fatal("acceptance credentials leaked to external media host")
|
|
}
|
|
w.Header().Set("Content-Type", "video/mp4")
|
|
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
|
|
}))
|
|
defer server.Close()
|
|
if err := validateVideoAsset(t.Context(), server.Client(), options{}, server.URL+"/result.mp4", 0); err != nil {
|
|
t.Fatalf("validate video asset: %v", err)
|
|
}
|
|
if got := findMediaURL(map[string]any{"content": map[string]any{"video_url": server.URL}}); got != server.URL {
|
|
t.Fatalf("media URL=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestValidateVideoAssetUsesGatewayForMaterializedPath(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
|
if request.Host != "gateway.easyai.local" || request.Header.Get("Authorization") != "Bearer key-1" {
|
|
t.Fatalf("host=%q authorization=%q", request.Host, request.Header.Get("Authorization"))
|
|
}
|
|
if request.URL.Path != "/static/generated/result.mp4" {
|
|
t.Fatalf("path=%q", request.URL.Path)
|
|
}
|
|
w.Header().Set("Content-Type", "video/mp4")
|
|
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
|
|
}))
|
|
defer server.Close()
|
|
opts := options{
|
|
gateways: []string{server.URL}, apiKeys: []string{"key-1"},
|
|
runID: "run-1", runToken: "token-1", gatewayTLSName: "gateway.easyai.local",
|
|
}
|
|
if err := validateVideoAsset(t.Context(), server.Client(), opts, "/static/generated/result.mp4", 0); err != nil {
|
|
t.Fatalf("validate materialized video: %v", err)
|
|
}
|
|
got := findMediaURL(map[string]any{
|
|
"raw": map[string]any{"video_url": "http://internal.invalid/video.mp4"},
|
|
"data": []any{map[string]any{"video_url": "/static/generated/result.mp4"}},
|
|
})
|
|
if got != "/static/generated/result.mp4" {
|
|
t.Fatalf("preferred media URL=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceReportErrorRedactsSecretsAndURLs(t *testing.T) {
|
|
got := redactError(
|
|
`token-1 failed at https://example.invalid/video.mp4?token=signed`,
|
|
options{runToken: "token-1"},
|
|
)
|
|
if got != `[REDACTED] failed at [REDACTED_URL]` {
|
|
t.Fatalf("redacted error=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceReportCannotOverwriteExistingRunArtifact(t *testing.T) {
|
|
path := filepath.Join(t.TempDir(), "load-report.json")
|
|
if err := writeReportExclusive(path, []byte(`{"runId":"first"}`)); err != nil {
|
|
t.Fatalf("write first report: %v", err)
|
|
}
|
|
if err := writeReportExclusive(path, []byte(`{"runId":"second"}`)); err == nil {
|
|
t.Fatal("second report unexpectedly overwrote the first report")
|
|
}
|
|
payload, err := os.ReadFile(path)
|
|
if err != nil {
|
|
t.Fatalf("read report: %v", err)
|
|
}
|
|
if string(payload) != "{\"runId\":\"first\"}\n" {
|
|
t.Fatalf("report changed after rejected overwrite: %s", payload)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceFailurePreservesOperationAndHTTPStatus(t *testing.T) {
|
|
err := withOperation("video_poll", &httpStatusError{Status: http.StatusUnauthorized, Body: "unauthorized"})
|
|
var operationErr *operationError
|
|
if !errors.As(err, &operationErr) || operationErr.Operation != "video_poll" {
|
|
t.Fatalf("operation error=%#v", operationErr)
|
|
}
|
|
var statusErr *httpStatusError
|
|
if !errors.As(err, &statusErr) || statusErr.Status != http.StatusUnauthorized {
|
|
t.Fatalf("HTTP status error=%#v", statusErr)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
|
|
request := httptest.NewRequest(http.MethodGet, "https://127.0.0.1/api/v1/healthz", nil)
|
|
opts := options{
|
|
apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1",
|
|
gatewayTLSName: "ai.example.com",
|
|
}
|
|
opts.setHeaders(request, 0, false)
|
|
if request.Host != "ai.example.com" {
|
|
t.Fatalf("Host=%q", request.Host)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceRootCAsRejectsSymlink(t *testing.T) {
|
|
root := t.TempDir()
|
|
target := filepath.Join(root, "ca.pem")
|
|
if err := os.WriteFile(target, []byte("not a certificate"), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
link := filepath.Join(root, "ca-link.pem")
|
|
if err := os.Symlink(target, link); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := acceptanceRootCAs(link); err == nil {
|
|
t.Fatal("symlink CA file was accepted")
|
|
}
|
|
}
|