BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(worker): 实现集群限流与自适应负载

Browse Source 
保留平台模型 RPM、TPM 和并发策略语义，增加 PostgreSQL 集群级租约、饱和候选重选和多平台自动负载，避免突发任务固定等待首个平台。\n\n新增 Worker 实时负载采样、自适应 active/heavy 容量、心跳与管理端指标，并扩展本地 acceptance runner，覆盖三 Worker、同模型三平台 2/4/6 并发和 48 个带图视频突发任务。\n\n验证：go test ./...、go vet ./...、PostgreSQL 跨 Store 集成测试、gofmt、bash -n、ShellCheck 及本地集群 provider-burst 验收通过；48/48 成功，无越限、重复提交、重复计费、重复回调或终态资源泄漏。
This commit is contained in:
wangbo
2026-08-03 00:13:46 +08:00
parent 9a01fd4657
commit c28bf74230
52 changed files with 3700 additions and 272 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/cmd/acceptance-load/main_test.go
+74 -17
View File
@@ -17,6 +17,12 @@ import (
"testing"
)
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
return f(request)
}
func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
payload := paddedPNG(4 << 20)
response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
@@ -44,8 +50,8 @@ func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) {
}
}
func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
input := paddedPNGVariant(2<<20, 37)
func TestStreamGeminiRequestBodyPreservesMultipleInputs(t *testing.T) {
inputs := geminiInputs(3, 2<<20, "multi-image-test", 37)
var body struct {
Contents []struct {
Parts []struct {
@@ -59,24 +65,26 @@ func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
ResponseModalities []string `json:"responseModalities"`
} `json:"generationConfig"`
}
stream := streamGeminiRequestBody(input)
stream := streamGeminiRequestBody(inputs)
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 {
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
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")
for index, input := range inputs {
inlineData := body.Contents[0].Parts[index+1].InlineData
if inlineData == nil || inlineData.MIMEType != "image/png" {
t.Fatalf("image %d inline data=%+v", index, inlineData)
}
decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
if err != nil {
t.Fatalf("decode streamed input %d: %v", index, err)
}
if !bytes.Equal(decoded, input) {
t.Fatalf("streamed input %d differs from source", index)
}
}
if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
@@ -84,8 +92,8 @@ func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
}
func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
input := paddedPNGVariant(256<<10, 91)
request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", input)
inputs := geminiInputs(3, 768<<10, "replay-test", 91)
request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", inputs)
if err != nil {
t.Fatalf("new Gemini request: %v", err)
}
@@ -114,6 +122,37 @@ func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
}
}
func TestGeminiMultiImageRequestUsesThreeFileDataReferences(t *testing.T) {
request, err := newGeminiFileDataRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", []string{
"https://fixtures.example/one.png",
"https://fixtures.example/two.png",
"https://fixtures.example/three.png",
})
if err != nil {
t.Fatalf("new multi-image request: %v", err)
}
var body struct {
Contents []struct {
Parts []struct {
FileData *struct {
FileURI string `json:"fileUri"`
} `json:"fileData"`
} `json:"parts"`
} `json:"contents"`
}
if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
t.Fatalf("decode multi-image request: %v", err)
}
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
t.Fatalf("unexpected multi-image body: %+v", body)
}
for index := 1; index < 4; index++ {
if body.Contents[0].Parts[index].FileData == nil || body.Contents[0].Parts[index].FileData.FileURI == "" {
t.Fatalf("missing fileData at part %d", index)
}
}
}
func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
images := make([]string, 16)
for index := range images {
@@ -133,6 +172,9 @@ func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
if len(seen) != 128 {
t.Fatalf("unique combinations=%d", len(seen))
}
if got := requestedVideoCombinationCount(combinations, 144, 0, 1); got != 131 {
t.Fatalf("requested combinations=%d, want 131", got)
}
}
func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
@@ -178,7 +220,7 @@ func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
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)
result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 1, 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)
}
@@ -254,6 +296,21 @@ func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
}
}
func TestValidateVideoAssetUsesIndependentTransportForExternalMedia(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
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()
poisoned := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return nil, errors.New("gateway-only transport must not be used for external media")
})}
opts := options{gatewayTLSName: "gateway.easyai.local"}
if err := validateVideoAsset(t.Context(), poisoned, opts, server.URL+"/result.mp4", 0); err != nil {
t.Fatalf("validate external video with independent transport: %v", err)
}
}
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" {