perf(acceptance): 使用多钱包身份分片模拟并发

单一验收钱包会把预留和结算串行化，掩盖双节点 Worker 的真实吞吐。验收现在幂等准备 32 个隔离身份和钱包，按请求轮询 API Key，并在 Run 配置中登记允许的 Key/User 身份对；密钥仅经 stdin 传给集群内压测进程。\n\n仍保留真实账务、候选权限、回调、重复扣费和强杀恢复校验。\n\n验证：Go 全量测试、临时 PostgreSQL 集成测试、go vet、OpenAPI 生成、迁移安全检查、bash -n、ShellCheck。
2026-07-31 04:25:09 +08:00
parent 8ce120631f
commit 262090db7b
6 changed files with 263 additions and 42 deletions
@@ -39,7 +39,7 @@ type options struct {
 	gateways       []string
 	gatewayTLSName string
 	emulatorURL    string
-	apiKey         string
+	apiKeys        []string
 	runID          string
 	runToken       string
 	geminiModel    string
@@ -135,7 +135,7 @@ func parseOptions() (options, error) {
 		gateways:       splitCSV(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAYS")),
 		gatewayTLSName: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME")),
 		emulatorURL:    strings.TrimRight(strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_EMULATOR_URL")), "/"),
-		apiKey:         strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_API_KEY")),
+		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")),
@@ -158,7 +158,7 @@ func parseOptions() (options, error) {
 		(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 opts.apiKey == "" || opts.runID == "" || opts.runToken == "" {
+	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 == "" {
@@ -289,7 +289,7 @@ func runGemini(
 			endpoint := opts.gateways[index%len(opts.gateways)] + "/v1beta/models/" + url.PathEscape(opts.geminiModel) + ":generateContent"
 			req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(requestBody))
 			if err == nil {
-				opts.setHeaders(req, realUpstream)
+				opts.setHeaders(req, index, realUpstream)
 				req.Header.Set("Content-Type", "application/json")
 				var response *http.Response
 				response, err = client.Do(req)
@@ -414,7 +414,7 @@ func runVideo(
 			endpoint := opts.gateways[index%len(opts.gateways)] + "/api/v1/videos/generations"
 			req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
 			if err == nil {
-				opts.setHeaders(req, realUpstream)
+				opts.setHeaders(req, index, realUpstream)
 				req.Header.Set("Content-Type", "application/json")
 				req.Header.Set("X-Async", "true")
 				var response *http.Response
@@ -519,7 +519,7 @@ func pollVideoTask(ctx context.Context, client *http.Client, opts options, taskI
 		if err != nil {
 			return err
 		}
-		opts.setHeaders(req, realUpstream)
+		opts.setHeaders(req, index, realUpstream)
 		response, err := client.Do(req)
 		if err == nil {
 			if response.StatusCode != http.StatusOK {
@@ -599,11 +599,11 @@ func validateVideoAsset(ctx context.Context, client *http.Client, mediaURL strin
 	return nil
 }

-func (o options) setHeaders(request *http.Request, realUpstream bool) {
+func (o options) setHeaders(request *http.Request, requestIndex int, realUpstream bool) {
 	if o.gatewayTLSName != "" {
 		request.Host = o.gatewayTLSName
 	}
-	request.Header.Set("Authorization", "Bearer "+o.apiKey)
+	request.Header.Set("Authorization", "Bearer "+o.apiKeys[requestIndex%len(o.apiKeys)])
 	request.Header.Set(runHeader, o.runID)
 	request.Header.Set(tokenHeader, o.runToken)
 	if realUpstream {
@@ -799,7 +799,7 @@ func paddedPNG(size int) []byte {
 }

 func redactError(message string, opts options) string {
-	for _, secret := range []string{opts.apiKey, opts.runToken} {
+	for _, secret := range append(append([]string(nil), opts.apiKeys...), opts.runToken) {
 		if secret != "" {
 			message = strings.ReplaceAll(message, secret, "[REDACTED]")
 		}
@@ -53,9 +53,19 @@ func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
 	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")
 			}
@@ -78,7 +88,7 @@ func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
 	defer second.Close()

 	opts := options{
-		gateways: []string{first.URL, second.URL}, apiKey: "key-1", runID: "run-1",
+		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)
@@ -88,6 +98,9 @@ func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
 	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 TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
@@ -117,10 +130,10 @@ func TestAcceptanceReportErrorRedactsSecretsAndURLs(t *testing.T) {
 func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
 	request := httptest.NewRequest(http.MethodGet, "https://127.0.0.1/api/v1/healthz", nil)
 	opts := options{
-		apiKey: "key-1", runID: "run-1", runToken: "token-1",
+		apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1",
 		gatewayTLSName: "ai.example.com",
 	}
-	opts.setHeaders(request, false)
+	opts.setHeaders(request, 0, false)
 	if request.Host != "ai.example.com" {
 		t.Fatalf("Host=%q", request.Host)
 	}
@@ -583,17 +583,35 @@ func (s *Store) AuthorizeAcceptanceTask(ctx context.Context, runID string, token
 		return "", ErrAcceptanceNotAuthorized
 	}
 	var activeRunID, apiKeyID, userID, tokenHash, status string
+	var participantAuthorized bool
 	err = s.pool.QueryRow(ctx, `
-SELECT id::text, api_key_id, user_id, token_sha256, status
+SELECT id::text, api_key_id, user_id, token_sha256, status,
+       (
+         (api_key_id = $2 AND user_id = $3)
+         OR EXISTS (
+           SELECT 1
+           FROM jsonb_array_elements(
+             CASE
+               WHEN jsonb_typeof(config->'participants') = 'array'
+                 THEN config->'participants'
+               ELSE '[]'::jsonb
+             END
+           ) participant
+           WHERE participant->>'apiKeyId' = $2
+             AND participant->>'userId' = $3
+         )
+       )
 FROM gateway_acceptance_runs
-WHERE id = $1::uuid`, runID).Scan(&activeRunID, &apiKeyID, &userID, &tokenHash, &status)
+WHERE id = $1::uuid`, runID, strings.TrimSpace(user.APIKeyID), strings.TrimSpace(user.ID)).Scan(
+		&activeRunID, &apiKeyID, &userID, &tokenHash, &status, &participantAuthorized,
+	)
 	if err != nil {
 		if errors.Is(err, pgx.ErrNoRows) {
 			return "", ErrAcceptanceNotAuthorized
 		}
 		return "", err
 	}
-	if status != "running" || apiKeyID != strings.TrimSpace(user.APIKeyID) || userID != strings.TrimSpace(user.ID) ||
+	if status != "running" || !participantAuthorized ||
 		subtle.ConstantTimeCompare([]byte(tokenHash), []byte(acceptanceTokenSHA256(token))) != 1 {
 		return "", ErrAcceptanceNotAuthorized
 	}
@@ -45,6 +45,9 @@ WHERE setting_key = $1`, SystemSettingGatewayTrafficMode)
 		EmulatorBaseURL:   "http://acceptance-emulator:8090",
 		CallbackURL:       "http://acceptance-emulator:8090/callbacks",
 		CapacityProfile:   "P24",
+		Config: map[string]any{"participants": []any{
+			map[string]any{"apiKeyId": "acceptance-shard-key", "userId": "acceptance-shard-user"},
+		}},
 	})
 	if err != nil {
 		t.Fatalf("create acceptance run: %v", err)
@@ -70,6 +73,14 @@ WHERE setting_key = $1`, SystemSettingGatewayTrafficMode)
 	if authorizedRunID, err := db.AuthorizeAcceptanceTask(ctx, run.ID, token, user); err != nil || authorizedRunID != run.ID {
 		t.Fatalf("authorize acceptance task run=%q err=%v", authorizedRunID, err)
 	}
+	shardUser := &auth.User{ID: "acceptance-shard-user", APIKeyID: "acceptance-shard-key"}
+	if authorizedRunID, err := db.AuthorizeAcceptanceTask(ctx, run.ID, token, shardUser); err != nil || authorizedRunID != run.ID {
+		t.Fatalf("authorize acceptance shard run=%q err=%v", authorizedRunID, err)
+	}
+	unpairedShard := &auth.User{ID: "acceptance-shard-user", APIKeyID: "acceptance-api-key"}
+	if _, err := db.AuthorizeAcceptanceTask(ctx, run.ID, token, unpairedShard); !errors.Is(err, ErrAcceptanceNotAuthorized) {
+		t.Fatalf("unpaired acceptance shard error=%v, want unauthorized", err)
+	}
 	if baseURL, credential, err := db.AcceptanceCandidateOverride(ctx, run.ID); err != nil ||
 		baseURL != "http://acceptance-emulator:8090" || credential == "" || credential == token {
 		t.Fatalf("candidate override base=%q credential=%q err=%v", baseURL, credential, err)