feat(acceptance): 建立同构验收与弹性容量体系

实现本地三节点 K3s 同构环境、脱敏生产快照、Gemini 图片和多参考图视频协议模拟、统一验收报告及故障注入。\n\n新增 Worker 容量控制器、资源与连接预算、任务恢复保护,并将生产验收拆分为 validation 执行和人工 CAS 放量。\n\n验证包括 Go 全量测试、PostgreSQL HTTP 集成测试、go vet、OpenAPI、ShellCheck、前端检查、迁移及发布脚本测试。
This commit is contained in:
2026-07-31 18:02:24 +08:00
parent 93d36d0e55
commit e05922b0f4
94 changed files with 12379 additions and 826 deletions
+16 -1
View File
@@ -35,7 +35,10 @@ RUN --mount=type=cache,target=/go/pkg/mod \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-migrate ./cmd/migrate && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-backfill-binary-results ./cmd/backfill-binary-results && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-emulator ./cmd/acceptance-emulator && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-load ./cmd/acceptance-load
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-callback-collector ./cmd/acceptance-callback-collector && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-load ./cmd/acceptance-load && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-snapshot ./cmd/acceptance-snapshot && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-bootstrap ./cmd/acceptance-bootstrap
FROM ${API_RUNTIME_IMAGE} AS api
@@ -47,7 +50,10 @@ COPY --from=api-builder /out/easyai-ai-gateway /app/easyai-ai-gateway
COPY --from=api-builder /out/easyai-ai-gateway-migrate /app/easyai-ai-gateway-migrate
COPY --from=api-builder /out/easyai-ai-gateway-backfill-binary-results /app/easyai-ai-gateway-backfill-binary-results
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-emulator /app/easyai-ai-gateway-acceptance-emulator
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-callback-collector /app/easyai-ai-gateway-acceptance-callback-collector
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-load /app/easyai-ai-gateway-acceptance-load
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-snapshot /app/easyai-ai-gateway-acceptance-snapshot
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-bootstrap /app/easyai-ai-gateway-acceptance-bootstrap
COPY apps/api/migrations /app/migrations
RUN mkdir -p /app/data/static/generated /app/data/static/uploaded && \
@@ -89,3 +95,12 @@ COPY docker/nginx.conf /etc/nginx/conf.d/default.conf
COPY --from=web-builder /src/apps/web/dist /usr/share/nginx/html
EXPOSE 80
FROM alpine:3.22 AS acceptance-netem
RUN apk add --no-cache ca-certificates curl iproute2 socat
COPY scripts/acceptance/netem-entrypoint.sh /usr/local/bin/netem-entrypoint
RUN chmod 0555 /usr/local/bin/netem-entrypoint
USER 65534:65534
ENTRYPOINT ["/usr/local/bin/netem-entrypoint"]
+4
View File
@@ -131,6 +131,10 @@ docker login --username=<your-aliyun-account> registry.cn-shanghai.aliyuncs.com
[生产同构验收手册](docs/operations/production-acceptance.md)。决策背景见
[ADR-003](docs/decisions/003-manual-agent-release.md)。
生产验收前先运行[本地三节点同构验收](docs/operations/local-isomorphic-acceptance.md)
本地原生高并发、精确 amd64 制品冒烟、线上模拟、真实金丝雀和人工开闸使用同一
`acceptance-report/v1` 身份链路。
Compose 默认使用独立容器数据库 `postgres:18-alpine`,数据卷会保留在 `postgres_data``api_data`。为避免本地开发 `.env` 中的 `localhost` 数据库地址污染容器部署,compose 使用 `AI_GATEWAY_COMPOSE_*` 变量作为容器部署专用覆盖,例如:
```bash
+500
View File
@@ -0,0 +1,500 @@
package main
import (
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"flag"
"fmt"
"os"
"path/filepath"
"regexp"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptancesnapshot"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
var (
fullSHAPattern = regexp.MustCompile(`^[0-9a-f]{40}$`)
digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
)
type options struct {
clusterID string
releaseSHA string
apiImageDigest string
workerDigest string
emulatorBaseURL string
callbackURL string
output string
identityShards int
}
type runtimeFile struct {
SchemaVersion string `json:"schemaVersion"`
LocalClusterID string `json:"localClusterId"`
RunID string `json:"runId"`
RunToken string `json:"runToken"`
APIKeys []string `json:"apiKeys"`
Participants []runtimeParticipant `json:"participants"`
GeminiModel string `json:"geminiModel"`
VideoModel string `json:"videoModel"`
EmulatorBaseURL string `json:"emulatorBaseUrl"`
CallbackURL string `json:"callbackUrl"`
ReleaseSHA string `json:"releaseSha"`
APIImageDigest string `json:"apiImageDigest"`
WorkerImageDigest string `json:"workerImageDigest"`
SnapshotConfigHash string `json:"snapshotConfigHash"`
SnapshotSHA256 string `json:"snapshotSha256"`
CreatedAt time.Time `json:"createdAt"`
}
type runtimeParticipant struct {
APIKeyID string `json:"apiKeyId"`
UserID string `json:"userId"`
}
func main() {
opts, err := parseOptions()
if err != nil {
fmt.Fprintln(os.Stderr, "acceptance bootstrap:", err)
os.Exit(64)
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
if err := run(ctx, opts); err != nil {
fmt.Fprintln(os.Stderr, "acceptance bootstrap:", err)
os.Exit(1)
}
}
func parseOptions() (options, error) {
var opts options
flag.StringVar(&opts.clusterID, "local-cluster-id", "", "required local cluster marker")
flag.StringVar(&opts.releaseSHA, "release-sha", "", "full source Git SHA")
flag.StringVar(&opts.apiImageDigest, "api-image-digest", "", "immutable API image digest")
flag.StringVar(&opts.workerDigest, "worker-image-digest", "", "immutable Worker image digest")
flag.StringVar(&opts.emulatorBaseURL, "emulator-base-url", "", "in-cluster protocol emulator URL")
flag.StringVar(&opts.callbackURL, "callback-url", "", "in-cluster callback collector URL")
flag.StringVar(&opts.output, "output", "", "private runtime output file")
flag.IntVar(&opts.identityShards, "identity-shards", 32, "isolated acceptance identities")
flag.Parse()
opts.clusterID = strings.TrimSpace(opts.clusterID)
opts.releaseSHA = strings.ToLower(strings.TrimSpace(opts.releaseSHA))
opts.apiImageDigest = strings.ToLower(strings.TrimSpace(opts.apiImageDigest))
opts.workerDigest = strings.ToLower(strings.TrimSpace(opts.workerDigest))
opts.emulatorBaseURL = strings.TrimRight(strings.TrimSpace(opts.emulatorBaseURL), "/")
opts.callbackURL = strings.TrimSpace(opts.callbackURL)
if opts.clusterID == "" || !fullSHAPattern.MatchString(opts.releaseSHA) ||
!digestPattern.MatchString(opts.apiImageDigest) || !digestPattern.MatchString(opts.workerDigest) {
return options{}, errors.New("local cluster ID, full release SHA, and immutable image digests are required")
}
if opts.emulatorBaseURL == "" || opts.callbackURL == "" {
return options{}, errors.New("emulator and callback URLs are required")
}
if opts.identityShards < 1 || opts.identityShards > 128 {
return options{}, errors.New("identity shards must be between 1 and 128")
}
if strings.TrimSpace(opts.output) == "" {
return options{}, errors.New("private runtime output path is required")
}
if flag.NArg() != 0 {
return options{}, errors.New("unexpected positional arguments")
}
return opts, nil
}
func run(ctx context.Context, opts options) error {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
if databaseURL == "" {
return errors.New("AI_GATEWAY_DATABASE_URL is required")
}
database, err := store.ConnectWithMaxConns(ctx, databaseURL, 4)
if err != nil {
return err
}
defer database.Close()
if err := verifyLocalDatabase(ctx, database, opts.clusterID); err != nil {
return err
}
snapshotConfigHash, snapshotSHA, err := importedSnapshot(ctx, database)
if err != nil {
return err
}
groupID, tenantID, err := ensureAcceptanceIdentityDomain(ctx, database)
if err != nil {
return err
}
participants := make([]runtimeParticipant, 0, opts.identityShards)
apiKeys := make([]string, 0, opts.identityShards)
for ordinal := 0; ordinal < opts.identityShards; ordinal++ {
userID, err := ensureAcceptanceUser(ctx, database, groupID, tenantID, ordinal)
if err != nil {
return err
}
keyID, secret, err := ensureAcceptanceAPIKey(ctx, database, userID, groupID, tenantID, ordinal)
if err != nil {
return err
}
if _, err := database.SetUserWalletBalance(ctx, store.WalletBalanceAdjustmentInput{
GatewayUserID: userID,
Currency: "resource",
BalanceText: "1000000000",
Reason: "local acceptance isolated wallet",
}); err != nil && !errors.Is(err, store.ErrWalletBalanceUnchanged) {
return err
}
participants = append(participants, runtimeParticipant{APIKeyID: keyID, UserID: userID})
apiKeys = append(apiKeys, secret)
}
if err := ensureAcceptanceAccessRules(ctx, database, groupID); err != nil {
return err
}
geminiModel, videoModel, err := selectedModels(ctx, database)
if err != nil {
return err
}
runToken, err := randomToken()
if err != nil {
return err
}
run, err := database.CreateAcceptanceRun(ctx, store.CreateAcceptanceRunInput{
ReleaseSHA: opts.releaseSHA,
APIImageDigest: opts.apiImageDigest,
WorkerImageDigest: opts.workerDigest,
APIKeyID: participants[0].APIKeyID,
UserID: participants[0].UserID,
Token: runToken,
EmulatorBaseURL: opts.emulatorBaseURL,
CallbackURL: opts.callbackURL,
CapacityProfile: "P24",
Config: map[string]any{
"workloads": []any{"gemini_image_edit", "multi_reference_video"},
"participants": participants,
"localClusterId": opts.clusterID,
"snapshotConfigHash": snapshotConfigHash,
"snapshotSha256": snapshotSHA,
},
})
if err != nil {
return err
}
if _, err := database.ActivateAcceptanceRun(ctx, run.ID); err != nil {
return err
}
outputPath, err := privateOutputPath(opts.output)
if err != nil {
return err
}
payload, err := json.MarshalIndent(runtimeFile{
SchemaVersion: "acceptance-runtime/v1",
LocalClusterID: opts.clusterID,
RunID: run.ID,
RunToken: runToken,
APIKeys: apiKeys,
Participants: participants,
GeminiModel: geminiModel,
VideoModel: videoModel,
EmulatorBaseURL: opts.emulatorBaseURL,
CallbackURL: opts.callbackURL,
ReleaseSHA: opts.releaseSHA,
APIImageDigest: opts.apiImageDigest,
WorkerImageDigest: opts.workerDigest,
SnapshotConfigHash: snapshotConfigHash,
SnapshotSHA256: snapshotSHA,
CreatedAt: time.Now().UTC(),
}, "", " ")
if err != nil {
return err
}
if err := os.WriteFile(outputPath, append(payload, '\n'), 0o600); err != nil {
return err
}
fmt.Printf(
"acceptance_bootstrap=PASS cluster_id=%s run_id=%s identities=%d config_hash=%s runtime_file=%s\n",
opts.clusterID,
run.ID,
len(participants),
snapshotConfigHash,
outputPath,
)
return nil
}
func verifyLocalDatabase(ctx context.Context, database *store.Store, clusterID string) error {
var marker string
err := database.Pool().QueryRow(ctx, `
SELECT COALESCE(value->>'clusterId', '')
FROM system_settings
WHERE setting_key = $1`, acceptancesnapshot.LocalClusterSettingKey).Scan(&marker)
if err != nil {
return fmt.Errorf("read local acceptance cluster marker: %w", err)
}
if marker != clusterID {
return errors.New("refusing bootstrap: local cluster marker mismatch")
}
mode, err := database.GetGatewayTrafficMode(ctx)
if err != nil {
return err
}
if mode.Mode != "live" {
return fmt.Errorf("refusing bootstrap while traffic mode is %s", mode.Mode)
}
return nil
}
func importedSnapshot(ctx context.Context, database *store.Store) (string, string, error) {
var configHash, snapshotSHA string
err := database.Pool().QueryRow(ctx, `
SELECT COALESCE(value->>'configHash', ''),
COALESCE(value->>'snapshotSha256', '')
FROM system_settings
WHERE setting_key = 'acceptance_snapshot'`).Scan(&configHash, &snapshotSHA)
if err != nil {
return "", "", fmt.Errorf("read imported acceptance snapshot: %w", err)
}
if len(configHash) != 64 || len(snapshotSHA) != 64 {
return "", "", errors.New("imported acceptance snapshot hashes are invalid")
}
return configHash, snapshotSHA, nil
}
func ensureAcceptanceIdentityDomain(ctx context.Context, database *store.Store) (string, string, error) {
var groupID string
err := database.Pool().QueryRow(ctx, `
INSERT INTO gateway_user_groups (
group_key, name, description, source, priority,
recharge_discount_policy, billing_discount_policy,
rate_limit_policy, quota_policy, metadata, status
)
VALUES (
'local-acceptance', 'Local Acceptance', 'Isolated local homologous acceptance identities',
'gateway', 1, '{"discountFactor":1}'::jsonb, '{"discountFactor":1}'::jsonb,
'{"rules":[]}'::jsonb, '{}'::jsonb,
'{"purpose":"local_acceptance","isolated":true}'::jsonb, 'active'
)
ON CONFLICT (group_key) DO UPDATE
SET name = EXCLUDED.name,
description = EXCLUDED.description,
priority = EXCLUDED.priority,
rate_limit_policy = EXCLUDED.rate_limit_policy,
quota_policy = EXCLUDED.quota_policy,
metadata = EXCLUDED.metadata,
status = 'active',
updated_at = now()
RETURNING id::text`).Scan(&groupID)
if err != nil {
return "", "", err
}
var tenantID string
err = database.Pool().QueryRow(ctx, `
INSERT INTO gateway_tenants (
tenant_key, source, external_tenant_id, name, default_user_group_id,
billing_profile, rate_limit_policy, auth_policy, metadata, status
)
VALUES (
'local-acceptance', 'gateway', 'local-acceptance', 'Local Acceptance',
$1::uuid, '{}'::jsonb, '{"rules":[]}'::jsonb, '{}'::jsonb,
'{"purpose":"local_acceptance","isolated":true}'::jsonb, 'active'
)
ON CONFLICT (tenant_key) DO UPDATE
SET default_user_group_id = EXCLUDED.default_user_group_id,
metadata = EXCLUDED.metadata,
status = 'active',
updated_at = now()
RETURNING id::text`, groupID).Scan(&tenantID)
return groupID, tenantID, err
}
func ensureAcceptanceUser(
ctx context.Context,
database *store.Store,
groupID string,
tenantID string,
ordinal int,
) (string, error) {
userKey := fmt.Sprintf("local-acceptance-%03d", ordinal)
roles := `["user"]`
if ordinal == 0 {
roles = `["manager"]`
}
var userID string
err := database.Pool().QueryRow(ctx, `
INSERT INTO gateway_users (
user_key, source, external_user_id, username, display_name,
gateway_tenant_id, tenant_id, tenant_key, default_user_group_id,
roles, auth_profile, metadata, status
)
VALUES (
$1, 'gateway', $1, $1, $2,
$3::uuid, 'local-acceptance', 'local-acceptance', $4::uuid,
$5::jsonb, '{}'::jsonb,
jsonb_build_object('purpose','local_acceptance','ordinal',$6,'isolated',true),
'active'
)
ON CONFLICT (user_key) DO UPDATE
SET gateway_tenant_id = EXCLUDED.gateway_tenant_id,
tenant_id = EXCLUDED.tenant_id,
tenant_key = EXCLUDED.tenant_key,
default_user_group_id = EXCLUDED.default_user_group_id,
roles = EXCLUDED.roles,
metadata = EXCLUDED.metadata,
status = 'active',
deleted_at = NULL,
updated_at = now()
RETURNING id::text`,
userKey,
fmt.Sprintf("Local Acceptance %03d", ordinal),
tenantID,
groupID,
roles,
ordinal,
).Scan(&userID)
return userID, err
}
func ensureAcceptanceAPIKey(
ctx context.Context,
database *store.Store,
userID string,
groupID string,
tenantID string,
ordinal int,
) (string, string, error) {
name := fmt.Sprintf("Local Acceptance %03d", ordinal)
var keyID, secret string
err := database.Pool().QueryRow(ctx, `
SELECT id::text, COALESCE(key_secret, '')
FROM gateway_api_keys
WHERE gateway_user_id = $1::uuid
AND name = $2
AND status = 'active'
AND deleted_at IS NULL
AND COALESCE(key_secret, '') <> ''
ORDER BY created_at
LIMIT 1`, userID, name).Scan(&keyID, &secret)
if err == nil {
return keyID, secret, nil
}
if !store.IsNotFound(err) {
return "", "", err
}
created, err := database.CreateAPIKey(ctx, store.CreateAPIKeyInput{
Name: name,
Scopes: []string{"image", "video"},
}, &auth.User{
ID: userID,
Username: fmt.Sprintf("local-acceptance-%03d", ordinal),
GatewayUserID: userID,
GatewayTenantID: tenantID,
TenantID: "local-acceptance",
TenantKey: "local-acceptance",
UserGroupID: groupID,
Source: "gateway",
})
if err != nil {
return "", "", err
}
if _, err := database.Pool().Exec(ctx, `
UPDATE gateway_api_keys
SET user_group_id = $2::uuid,
rate_limit_policy = '{"rules":[]}'::jsonb,
quota_policy = '{}'::jsonb,
updated_at = now()
WHERE id = $1::uuid`, created.APIKey.ID, groupID); err != nil {
return "", "", err
}
return created.APIKey.ID, created.Secret, nil
}
func ensureAcceptanceAccessRules(ctx context.Context, database *store.Store, groupID string) error {
_, err := database.Pool().Exec(ctx, `
WITH selected AS (
SELECT 'platform'::text AS resource_type, platform.id AS resource_id
FROM integration_platforms platform
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
UNION
SELECT 'platform_model', model.id
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
UNION
SELECT 'base_model', model.base_model_id
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
AND model.base_model_id IS NOT NULL
)
INSERT INTO gateway_access_rules (
subject_type, subject_id, resource_type, resource_id, effect,
priority, min_permission_level, conditions, metadata, status
)
SELECT 'user_group', $1::uuid, selected.resource_type, selected.resource_id,
'allow', 1, 0, '{}'::jsonb,
'{"purpose":"local_acceptance"}'::jsonb, 'active'
FROM selected
ON CONFLICT (subject_type, subject_id, resource_type, resource_id, effect)
DO UPDATE SET status = 'active', priority = 1, updated_at = now()`, groupID)
return err
}
func selectedModels(ctx context.Context, database *store.Store) (string, string, error) {
var gemini, video string
err := database.Pool().QueryRow(ctx, `
SELECT base_model.invocation_name
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
JOIN base_model_catalog base_model ON base_model.id = model.base_model_id
WHERE platform.status = 'enabled'
AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
AND model.enabled = true
AND model.model_type @> '["image_edit"]'::jsonb
ORDER BY platform.priority, model.created_at
LIMIT 1`).Scan(&gemini)
if err != nil {
return "", "", err
}
err = database.Pool().QueryRow(ctx, `
SELECT base_model.invocation_name
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
JOIN base_model_catalog base_model ON base_model.id = model.base_model_id
WHERE platform.status = 'enabled'
AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
AND model.enabled = true
AND model.model_type @> '["omni_video"]'::jsonb
ORDER BY platform.priority, model.created_at
LIMIT 1`).Scan(&video)
return gemini, video, err
}
func randomToken() (string, error) {
payload := make([]byte, 32)
if _, err := rand.Read(payload); err != nil {
return "", err
}
return hex.EncodeToString(payload), nil
}
func privateOutputPath(path string) (string, error) {
absolute, err := filepath.Abs(strings.TrimSpace(path))
if err != nil {
return "", err
}
if info, err := os.Lstat(absolute); err == nil {
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
return "", errors.New("runtime output must be a regular non-symlink file")
}
} else if !os.IsNotExist(err) {
return "", err
}
if err := os.MkdirAll(filepath.Dir(absolute), 0o700); err != nil {
return "", err
}
return absolute, nil
}
@@ -0,0 +1,120 @@
package main
import (
"encoding/json"
"errors"
"io"
"log/slog"
"net/http"
"os"
"strconv"
"strings"
"sync"
"time"
)
const maxCallbackBodyBytes = 1 << 20
type collector struct {
mu sync.Mutex
deliveries int64
duplicates int64
invalid int64
seen map[string]int
}
func main() {
address := strings.TrimSpace(os.Getenv("HTTP_ADDR"))
if address == "" {
address = ":8091"
}
state := &collector{seen: map[string]int{}}
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
})
mux.HandleFunc("GET /report", state.report)
mux.HandleFunc("POST /callbacks", state.callback)
server := &http.Server{
Addr: address,
Handler: mux,
ReadHeaderTimeout: 10 * time.Second,
}
slog.Info("acceptance callback collector started", "address", address)
if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
slog.Error("acceptance callback collector stopped", "error", err)
os.Exit(1)
}
}
func (c *collector) callback(w http.ResponseWriter, request *http.Request) {
body, err := io.ReadAll(io.LimitReader(request.Body, maxCallbackBodyBytes+1))
if err != nil || len(body) > maxCallbackBodyBytes {
c.recordInvalid()
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "invalid callback body"})
return
}
decoder := json.NewDecoder(strings.NewReader(string(body)))
decoder.UseNumber()
var payload map[string]any
if err := decoder.Decode(&payload); err != nil {
c.recordInvalid()
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "invalid callback JSON"})
return
}
taskID := strings.TrimSpace(stringValue(payload["taskId"]))
seq, _ := strconv.ParseInt(stringValue(payload["seq"]), 10, 64)
idempotency := strings.TrimSpace(request.Header.Get("Idempotency-Key"))
expected := taskID + ":" + strconv.FormatInt(seq, 10)
if taskID == "" || seq <= 0 || idempotency != expected {
c.recordInvalid()
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "callback identity is invalid"})
return
}
c.mu.Lock()
c.deliveries++
c.seen[expected]++
if c.seen[expected] > 1 {
c.duplicates++
}
c.mu.Unlock()
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}
func (c *collector) report(w http.ResponseWriter, _ *http.Request) {
c.mu.Lock()
report := map[string]any{
"schemaVersion": "acceptance-callback-report/v1",
"deliveries": c.deliveries,
"duplicates": c.duplicates,
"invalid": c.invalid,
"uniqueEvents": len(c.seen),
}
c.mu.Unlock()
writeJSON(w, http.StatusOK, report)
}
func (c *collector) recordInvalid() {
c.mu.Lock()
c.invalid++
c.mu.Unlock()
}
func stringValue(value any) string {
switch typed := value.(type) {
case string:
return typed
case json.Number:
return typed.String()
case float64:
return strconv.FormatFloat(typed, 'f', -1, 64)
default:
return ""
}
}
func writeJSON(w http.ResponseWriter, status int, value any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(value)
}
+281 -30
View File
@@ -6,6 +6,7 @@ import (
"context"
"crypto/sha256"
"crypto/tls"
"crypto/x509"
"encoding/base64"
"encoding/binary"
"encoding/hex"
@@ -18,10 +19,13 @@ import (
"net/http"
"net/url"
"os"
"os/signal"
"regexp"
"sort"
"strconv"
"strings"
"sync"
"syscall"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceworkload"
@@ -38,6 +42,7 @@ var reportURLPattern = regexp.MustCompile(`(?i)(?:https?|postgres(?:ql)?):\/\/[^
type options struct {
gateways []string
gatewayTLSName string
gatewayCAFile string
emulatorURL string
apiKeys []string
runID string
@@ -48,17 +53,21 @@ type options struct {
reportPath string
realImageURLs []string
timeout time.Duration
mixedDuration time.Duration
imageRate float64
videoRate float64
}
type report struct {
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"`
SecretSafe bool `json:"secretSafe"`
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"`
SecretSafe bool `json:"secretSafe"`
}
type phaseReport struct {
@@ -76,6 +85,9 @@ type phaseReport struct {
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 {
@@ -84,16 +96,32 @@ type phaseResult struct {
err error
}
type httpStatusError struct {
Status int
Body string
}
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)
}
ctx, cancel := context.WithTimeout(context.Background(), opts.timeout)
signalContext, stopSignals := signal.NotifyContext(
context.Background(),
os.Interrupt,
syscall.SIGTERM,
)
defer stopSignals()
ctx, cancel := context.WithTimeout(signalContext, opts.timeout)
defer cancel()
result := report{
RunID: opts.runID, Profile: opts.profile, StartedAt: time.Now().UTC(), SecretSafe: true,
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 {
@@ -127,13 +155,20 @@ func parseOptions() (options, error) {
var profile string
var reportPath string
var timeout time.Duration
var mixedDuration time.Duration
var imageRate float64
var videoRate float64
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.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")),
@@ -144,6 +179,9 @@ func parseOptions() (options, error) {
reportPath: strings.TrimSpace(reportPath),
realImageURLs: splitCSV(os.Getenv("AI_GATEWAY_ACCEPTANCE_REAL_IMAGE_URLS")),
timeout: timeout,
mixedDuration: mixedDuration,
imageRate: imageRate,
videoRate: videoRate,
}
if len(opts.gateways) != 2 {
return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAYS must contain exactly two API base URLs")
@@ -168,10 +206,16 @@ func parseOptions() (options, error) {
return options{}, errors.New("timeout must be positive")
}
switch opts.profile {
case "simulated-all", "gemini-baseline", "gemini-large", "gemini-peak", "video-throughput", "video-recovery":
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")
@@ -184,10 +228,15 @@ func parseOptions() (options, error) {
func run(ctx context.Context, opts options) ([]phaseResult, error) {
var tlsConfig *tls.Config
if opts.gatewayTLSName != "" {
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{
@@ -207,14 +256,20 @@ func run(ctx context.Context, opts options) ([]phaseResult, error) {
result = runGemini(ctx, client, opts, name, profile.Requests, profile.InputBytes, profile.OutputBytes, false)
} else {
switch name {
case "smoke-gemini":
result = runGemini(ctx, client, opts, name, 4, 256<<10, 256<<10, false)
case "smoke-video":
result = runVideo(ctx, client, opts, name, 10, false, opts.emulatorFixtureURLs(), false, 0)
case "video-throughput":
result = runVideo(ctx, client, opts, name, 1200, false, opts.emulatorFixtureURLs(), false)
result = runVideo(ctx, client, opts, name, 1200, false, opts.emulatorFixtureURLs(), false, 0)
case "video-recovery":
result = runVideo(ctx, client, opts, name, 96, true, opts.emulatorFixtureURLs(), false)
result = runVideo(ctx, client, opts, name, 96, true, opts.emulatorFixtureURLs(), false, 0)
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)
case "real-video-canary":
result = runVideo(ctx, client, opts, name, 1, false, opts.realImageURLs, true)
result = runVideo(ctx, client, opts, name, 1, false, opts.realImageURLs, true, 0)
default:
return fmt.Errorf("unknown phase %q", name)
}
@@ -223,7 +278,9 @@ func run(ctx context.Context, opts options) ([]phaseResult, error) {
return result.err
}
phases := []string{opts.profile}
if opts.profile == "simulated-all" {
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"}
@@ -376,6 +433,7 @@ func runVideo(
longRun bool,
imageURLs []string,
realUpstream bool,
requestOffset int,
) phaseResult {
startedAt := time.Now()
combinationCount := 128
@@ -398,31 +456,32 @@ func runVideo(
wg.Add(1)
go func() {
defer wg.Done()
logicalIndex := index + requestOffset
requestStarted := time.Now()
startedTasks[index] = requestStarted
imageCount := acceptanceworkload.VideoImageCount(index)
combo := append([]string(nil), combinations[index%len(combinations)]...)
imageCount := acceptanceworkload.VideoImageCount(logicalIndex)
combo := append([]string(nil), combinations[logicalIndex%len(combinations)]...)
if len(combo) > imageCount {
combo = combo[:imageCount]
}
if !realUpstream && index%4 == 0 {
combo[0] = imageURLs[12+(index/4)%4]
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 && index%4 == 0 {
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 index%5 == 0 && imageIndex == 0 {
if logicalIndex%5 == 0 && imageIndex == 0 {
role = "first_frame"
}
if index%5 == 0 && len(combo) > 3 && imageIndex == len(combo)-1 {
if logicalIndex%5 == 0 && len(combo) > 3 && imageIndex == len(combo)-1 {
role = "last_frame"
}
content = append(content, map[string]any{
@@ -431,13 +490,13 @@ func runVideo(
})
}
body, _ := json.Marshal(map[string]any{
"model": opts.videoModel, "modelType": "omni_video", "content": content, "seed": index + 1,
"model": opts.videoModel, "modelType": "omni_video", "content": content, "seed": logicalIndex + 1,
"duration": 5, "ratio": "16:9", "resolution": "720p",
})
endpoint := opts.gateways[index%len(opts.gateways)] + "/api/v1/videos/generations"
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, index, realUpstream)
opts.setHeaders(req, logicalIndex, realUpstream)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("X-Async", "true")
var response *http.Response
@@ -460,7 +519,7 @@ func runVideo(
}
mu.Lock()
latencies[index] = time.Since(requestStarted)
if !realUpstream && index%4 == 0 {
if !realUpstream && logicalIndex%4 == 0 {
forcedConversions++
}
if err != nil {
@@ -492,6 +551,7 @@ func runVideo(
continue
}
index, taskID := index, taskID
logicalIndex := index + requestOffset
wg.Add(1)
go func() {
defer wg.Done()
@@ -501,7 +561,7 @@ func runVideo(
return
}
defer func() { <-pollSlots }()
err := pollVideoTask(ctx, client, opts, taskID, index, realUpstream)
err := pollVideoTask(ctx, client, opts, taskID, logicalIndex, realUpstream)
mu.Lock()
latencies[index] = time.Since(startedTasks[index])
if err != nil {
@@ -520,6 +580,9 @@ func runVideo(
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,
@@ -533,6 +596,145 @@ func runVideo(
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 := 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)
} else {
result = runVideo(
runCtx,
client,
opts,
"mixed-video",
1,
false,
opts.emulatorFixtureURLs(),
false,
requestIndex,
)
}
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()
@@ -558,7 +760,7 @@ func pollVideoTask(ctx context.Context, client *http.Client, opts options, taskI
if mediaURL == "" {
return fmt.Errorf("video task %s succeeded without a media URL", taskID)
}
return validateVideoAsset(ctx, http.DefaultClient, mediaURL)
return validateVideoAsset(ctx, client, mediaURL)
case "failed", "cancelled", "canceled":
return fmt.Errorf("video task %s finished with status %s", taskID, status)
}
@@ -782,7 +984,7 @@ func scanUntil(reader *bufio.Reader, target []byte, limit int64) error {
func responseStatusError(response *http.Response) error {
defer response.Body.Close()
payload, _ := io.ReadAll(io.LimitReader(response.Body, 512))
return fmt.Errorf("HTTP %d: %s", response.StatusCode, strings.TrimSpace(string(payload)))
return &httpStatusError{Status: response.StatusCode, Body: strings.TrimSpace(string(payload))}
}
func percentileMilliseconds(values []time.Duration, quantile float64) float64 {
@@ -837,6 +1039,31 @@ func redactError(message string, opts options) string {
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))
@@ -854,3 +1081,27 @@ func env(name string, fallback string) string {
}
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
}
+17
View File
@@ -9,6 +9,8 @@ import (
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"sync/atomic"
"testing"
)
@@ -190,3 +192,18 @@ func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
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")
}
}
+195
View File
@@ -0,0 +1,195 @@
package main
import (
"context"
"errors"
"flag"
"fmt"
"os"
"path/filepath"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptancesnapshot"
"github.com/jackc/pgx/v5/pgxpool"
)
func main() {
if len(os.Args) < 2 {
usage()
os.Exit(64)
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
var err error
switch os.Args[1] {
case "export":
err = runExport(ctx, os.Args[2:])
case "validate":
err = runValidate(os.Args[2:])
case "import":
err = runImport(ctx, os.Args[2:])
default:
usage()
os.Exit(64)
}
if err != nil {
fmt.Fprintln(os.Stderr, "acceptance snapshot:", err)
os.Exit(1)
}
}
func usage() {
fmt.Fprintln(os.Stderr, `Usage:
easyai-ai-gateway-acceptance-snapshot export --release-sha <SHA> --output <file>
easyai-ai-gateway-acceptance-snapshot validate --input <file>
easyai-ai-gateway-acceptance-snapshot import --input <file> --local-cluster-id <id>
export is read-only. import refuses to write unless system_settings contains a
matching acceptance_local_cluster_id marker.`)
}
func runExport(ctx context.Context, args []string) error {
flags := flag.NewFlagSet("export", flag.ContinueOnError)
releaseSHA := flags.String("release-sha", "", "full source release SHA")
output := flags.String("output", "", "secret-safe snapshot output path")
if err := flags.Parse(args); err != nil {
return err
}
if flags.NArg() != 0 {
return errors.New("unexpected export arguments")
}
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
if databaseURL == "" {
return errors.New("AI_GATEWAY_DATABASE_URL is required")
}
outputPath, err := safeOutputPath(*output)
if err != nil {
return err
}
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
return err
}
defer pool.Close()
snapshot, err := acceptancesnapshot.Export(ctx, pool, acceptancesnapshot.ExportOptions{
ReleaseSHA: *releaseSHA,
})
if err != nil {
return err
}
payload, err := acceptancesnapshot.Encode(snapshot)
if err != nil {
return err
}
if err := os.WriteFile(outputPath, payload, 0o600); err != nil {
return err
}
fmt.Printf(
"acceptance_snapshot_export=PASS schema=%s config_hash=%s snapshot_sha256=%s\n",
snapshot.SchemaVersion,
snapshot.Source.ConfigHash,
snapshot.SnapshotSHA256,
)
return nil
}
func runValidate(args []string) error {
flags := flag.NewFlagSet("validate", flag.ContinueOnError)
input := flags.String("input", "", "snapshot input path")
if err := flags.Parse(args); err != nil {
return err
}
if flags.NArg() != 0 {
return errors.New("unexpected validate arguments")
}
snapshot, err := readSnapshot(*input)
if err != nil {
return err
}
fmt.Printf(
"acceptance_snapshot_validate=PASS schema=%s release=%s config_hash=%s snapshot_sha256=%s\n",
snapshot.SchemaVersion,
snapshot.Source.ReleaseSHA,
snapshot.Source.ConfigHash,
snapshot.SnapshotSHA256,
)
return nil
}
func runImport(ctx context.Context, args []string) error {
flags := flag.NewFlagSet("import", flag.ContinueOnError)
input := flags.String("input", "", "snapshot input path")
clusterID := flags.String("local-cluster-id", "", "required local cluster marker")
if err := flags.Parse(args); err != nil {
return err
}
if flags.NArg() != 0 {
return errors.New("unexpected import arguments")
}
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
if databaseURL == "" {
return errors.New("AI_GATEWAY_DATABASE_URL is required")
}
snapshot, err := readSnapshot(*input)
if err != nil {
return err
}
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
return err
}
defer pool.Close()
if err := acceptancesnapshot.Import(ctx, pool, snapshot, *clusterID); err != nil {
return err
}
fmt.Printf(
"acceptance_snapshot_import=PASS cluster_id=%s config_hash=%s snapshot_sha256=%s\n",
strings.TrimSpace(*clusterID),
snapshot.Source.ConfigHash,
snapshot.SnapshotSHA256,
)
return nil
}
func readSnapshot(path string) (acceptancesnapshot.Snapshot, error) {
path = strings.TrimSpace(path)
if path == "" {
return acceptancesnapshot.Snapshot{}, errors.New("snapshot path is required")
}
info, err := os.Lstat(path)
if err != nil {
return acceptancesnapshot.Snapshot{}, err
}
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
return acceptancesnapshot.Snapshot{}, errors.New("snapshot must be a regular non-symlink file")
}
payload, err := os.ReadFile(path)
if err != nil {
return acceptancesnapshot.Snapshot{}, err
}
return acceptancesnapshot.Decode(payload)
}
func safeOutputPath(path string) (string, error) {
path = strings.TrimSpace(path)
if path == "" {
return "", errors.New("output path is required")
}
absolute, err := filepath.Abs(path)
if err != nil {
return "", err
}
parent := filepath.Dir(absolute)
if err := os.MkdirAll(parent, 0o700); err != nil {
return "", err
}
if info, err := os.Lstat(absolute); err == nil {
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
return "", errors.New("output must be a regular non-symlink file")
}
} else if !os.IsNotExist(err) {
return "", err
}
return absolute, nil
}
+74 -14
View File
@@ -10,6 +10,7 @@ import (
"syscall"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/capacitycontroller"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/httpapi"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
@@ -37,8 +38,15 @@ func main() {
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
executionMaxConns := cfg.DatabaseMaxConns
if cfg.DatabaseCriticalMaxConns > 0 {
executionMaxConns -= cfg.DatabaseCriticalMaxConns
}
if cfg.DatabaseRiverMaxConns > 0 {
executionMaxConns -= cfg.DatabaseRiverMaxConns
}
db, err := store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
MaxConns: cfg.DatabaseMaxConns,
MaxConns: executionMaxConns,
MinIdleConns: cfg.DatabaseMinIdleConns,
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
@@ -50,26 +58,78 @@ func main() {
os.Exit(1)
}
defer db.Close()
coordinationDB := db
if cfg.DatabaseCriticalMaxConns > 0 {
coordinationDB, err = store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
MaxConns: cfg.DatabaseCriticalMaxConns,
MinIdleConns: min(cfg.DatabaseCriticalMaxConns, 1),
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
})
if err != nil {
stop()
logger.Error("connect critical postgres pool failed", "error", err)
os.Exit(1)
}
defer coordinationDB.Close()
}
riverDB := db
if cfg.DatabaseRiverMaxConns > 0 {
riverDB, err = store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
MaxConns: cfg.DatabaseRiverMaxConns,
MinIdleConns: min(cfg.DatabaseRiverMaxConns, 1),
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
})
if err != nil {
stop()
logger.Error("connect River postgres pool failed", "error", err)
os.Exit(1)
}
defer riverDB.Close()
}
// Cancel background LISTEN and worker goroutines before closing the pool.
// This also prevents a startup panic from waiting forever in pgxpool.Close.
defer stop()
if recovery, err := db.RecoverInterruptedRuntimeState(ctx); err != nil {
logger.Error("recover interrupted runtime state failed", "error", err)
os.Exit(1)
} else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 || recovery.CleanedTaskAdmissions > 0 {
logger.Warn("interrupted runtime state recovered",
"releasedConcurrencyLeases", recovery.ReleasedConcurrencyLeases,
"releasedRateReservations", recovery.ReleasedRateReservations,
"failedAttempts", recovery.FailedAttempts,
"failedTasks", recovery.FailedTasks,
"requeuedAsyncTasks", recovery.RequeuedAsyncTasks,
"cleanedTaskAdmissions", recovery.CleanedTaskAdmissions,
)
if cfg.RunsBackgroundWorkers() {
if recovery, recoveryErr := coordinationDB.RecoverInterruptedRuntimeState(ctx); recoveryErr != nil {
logger.Error("recover interrupted runtime state failed", "error", recoveryErr)
os.Exit(1)
} else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 || recovery.CleanedTaskAdmissions > 0 {
logger.Warn("interrupted runtime state recovered",
"releasedConcurrencyLeases", recovery.ReleasedConcurrencyLeases,
"releasedRateReservations", recovery.ReleasedRateReservations,
"failedAttempts", recovery.FailedAttempts,
"failedTasks", recovery.FailedTasks,
"requeuedAsyncTasks", recovery.RequeuedAsyncTasks,
"cleanedTaskAdmissions", recovery.CleanedTaskAdmissions,
)
}
}
var handler http.Handler
if cfg.RunsCapacityController() {
kubernetes, kubernetesErr := capacitycontroller.NewKubernetesClient(capacitycontroller.KubernetesConfig{
Namespace: cfg.CapacityControllerNamespace,
APIServer: cfg.CapacityControllerAPIServer,
TokenFile: cfg.CapacityControllerTokenFile,
CAFile: cfg.CapacityControllerCAFile,
})
if kubernetesErr != nil {
logger.Error("initialize capacity controller Kubernetes client failed", "error", kubernetesErr)
os.Exit(1)
}
controller := capacitycontroller.New(cfg, coordinationDB, kubernetes, logger)
go controller.Run(ctx)
handler = controller.Handler()
} else {
handler = httpapi.NewServerWithStores(ctx, cfg, db, coordinationDB, riverDB, logger)
}
server := &http.Server{
Addr: cfg.HTTPAddr,
Handler: httpapi.NewServerWithContext(ctx, cfg, db, logger),
Handler: handler,
ReadHeaderTimeout: 10 * time.Second,
}
+243
View File
@@ -2575,6 +2575,33 @@
}
}
},
"/api/admin/system/acceptance/capacity-profiles": {
"get": {
"security": [
{
"BearerAuth": []
}
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "获取经生产同构验收认证的容量配置",
"responses": {
"200": {
"description": "OK",
"schema": {
"type": "array",
"items": {
"$ref": "#/definitions/store.CapacityProfile"
}
}
}
}
}
},
"/api/admin/system/acceptance/runs": {
"post": {
"security": [
@@ -2724,6 +2751,51 @@
}
}
},
"/api/admin/system/acceptance/runs/{runID}/capacity-profiles": {
"post": {
"security": [
{
"BearerAuth": []
}
],
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "为当前 validation Run 暂存 80% 容量门禁",
"parameters": [
{
"type": "string",
"description": "验收 Run ID",
"name": "runID",
"in": "path",
"required": true
},
{
"description": "待验证容量配置",
"name": "body",
"in": "body",
"required": true,
"schema": {
"$ref": "#/definitions/store.StageAcceptanceCapacityProfilesInput"
}
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.AcceptanceRun"
}
}
}
}
},
"/api/admin/system/acceptance/runs/{runID}/finish": {
"post": {
"security": [
@@ -2871,6 +2943,44 @@
}
}
},
"/api/admin/system/acceptance/traffic-mode/pause": {
"post": {
"security": [
{
"BearerAuth": []
}
],
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "上线后硬门禁异常时通过 CAS 暂停正式新任务",
"parameters": [
{
"description": "当前线上 release 与 revision CAS",
"name": "body",
"in": "body",
"required": true,
"schema": {
"$ref": "#/definitions/store.PauseGatewayTrafficInput"
}
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.GatewayTrafficMode"
}
}
}
}
},
"/api/admin/system/client-customization/settings": {
"get": {
"security": [
@@ -13694,6 +13804,93 @@
}
}
},
"store.CapacityProfile": {
"type": "object",
"properties": {
"acceptanceRunId": {
"type": "string"
},
"admittedThroughputPerSecond": {
"type": "number"
},
"capacityProfile": {
"type": "string"
},
"configHash": {
"type": "string"
},
"createdAt": {
"type": "string"
},
"enabled": {
"type": "boolean"
},
"kind": {
"type": "string"
},
"maxPredictedWaitSeconds": {
"type": "integer"
},
"maxQueued": {
"type": "integer"
},
"maxRunning": {
"type": "integer"
},
"metadata": {
"type": "object",
"additionalProperties": {}
},
"model": {
"type": "string"
},
"profileKey": {
"type": "string"
},
"releaseSha": {
"type": "string"
},
"stableThroughputPerSecond": {
"type": "number"
},
"updatedAt": {
"type": "string"
}
}
},
"store.CapacityProfileInput": {
"type": "object",
"properties": {
"admittedThroughputPerSecond": {
"type": "number"
},
"kind": {
"type": "string"
},
"maxPredictedWaitSeconds": {
"type": "integer"
},
"maxQueued": {
"type": "integer"
},
"maxRunning": {
"type": "integer"
},
"metadata": {
"type": "object",
"additionalProperties": {}
},
"model": {
"type": "string"
},
"profileKey": {
"type": "string"
},
"stableThroughputPerSecond": {
"type": "number"
}
}
},
"store.CatalogProvider": {
"type": "object",
"properties": {
@@ -14502,6 +14699,9 @@
"apiImageDigest": {
"type": "string"
},
"configHash": {
"type": "string"
},
"mode": {
"type": "string"
},
@@ -14894,6 +15094,26 @@
}
}
},
"store.PauseGatewayTrafficInput": {
"type": "object",
"properties": {
"apiImageDigest": {
"type": "string"
},
"reason": {
"type": "string"
},
"releaseSha": {
"type": "string"
},
"revision": {
"type": "integer"
},
"workerImageDigest": {
"type": "string"
}
}
},
"store.Platform": {
"type": "object",
"properties": {
@@ -15525,6 +15745,15 @@
"apiImageDigest": {
"type": "string"
},
"capacityProfiles": {
"type": "array",
"items": {
"$ref": "#/definitions/store.CapacityProfileInput"
}
},
"configHash": {
"type": "string"
},
"releaseSha": {
"type": "string"
},
@@ -15768,6 +15997,20 @@
}
}
},
"store.StageAcceptanceCapacityProfilesInput": {
"type": "object",
"properties": {
"capacityProfiles": {
"type": "array",
"items": {
"$ref": "#/definitions/store.CapacityProfileInput"
}
},
"configHash": {
"type": "string"
}
}
},
"store.TaskAttempt": {
"type": "object",
"properties": {
+155
View File
@@ -2512,6 +2512,64 @@ definitions:
updatedAt:
type: string
type: object
store.CapacityProfile:
properties:
acceptanceRunId:
type: string
admittedThroughputPerSecond:
type: number
capacityProfile:
type: string
configHash:
type: string
createdAt:
type: string
enabled:
type: boolean
kind:
type: string
maxPredictedWaitSeconds:
type: integer
maxQueued:
type: integer
maxRunning:
type: integer
metadata:
additionalProperties: {}
type: object
model:
type: string
profileKey:
type: string
releaseSha:
type: string
stableThroughputPerSecond:
type: number
updatedAt:
type: string
type: object
store.CapacityProfileInput:
properties:
admittedThroughputPerSecond:
type: number
kind:
type: string
maxPredictedWaitSeconds:
type: integer
maxQueued:
type: integer
maxRunning:
type: integer
metadata:
additionalProperties: {}
type: object
model:
type: string
profileKey:
type: string
stableThroughputPerSecond:
type: number
type: object
store.CatalogProvider:
properties:
capabilitySchema:
@@ -3059,6 +3117,8 @@ definitions:
properties:
apiImageDigest:
type: string
configHash:
type: string
mode:
type: string
releaseSha:
@@ -3320,6 +3380,19 @@ definitions:
waitingSyncTasks:
type: integer
type: object
store.PauseGatewayTrafficInput:
properties:
apiImageDigest:
type: string
reason:
type: string
releaseSha:
type: string
revision:
type: integer
workerImageDigest:
type: string
type: object
store.Platform:
properties:
authType:
@@ -3748,6 +3821,12 @@ definitions:
properties:
apiImageDigest:
type: string
capacityProfiles:
items:
$ref: '#/definitions/store.CapacityProfileInput'
type: array
configHash:
type: string
releaseSha:
type: string
revision:
@@ -3915,6 +3994,15 @@ definitions:
status:
type: string
type: object
store.StageAcceptanceCapacityProfilesInput:
properties:
capacityProfiles:
items:
$ref: '#/definitions/store.CapacityProfileInput'
type: array
configHash:
type: string
type: object
store.TaskAttempt:
properties:
attemptNo:
@@ -5758,6 +5846,22 @@ paths:
summary: 更新 Runner 策略
tags:
- runtime
/api/admin/system/acceptance/capacity-profiles:
get:
produces:
- application/json
responses:
"200":
description: OK
schema:
items:
$ref: '#/definitions/store.CapacityProfile'
type: array
security:
- BearerAuth: []
summary: 获取经生产同构验收认证的容量配置
tags:
- acceptance
/api/admin/system/acceptance/runs:
post:
consumes:
@@ -5849,6 +5953,34 @@ paths:
summary: 切换到 validation 并激活 Run
tags:
- acceptance
/api/admin/system/acceptance/runs/{runID}/capacity-profiles:
post:
consumes:
- application/json
parameters:
- description: 验收 Run ID
in: path
name: runID
required: true
type: string
- description: 待验证容量配置
in: body
name: body
required: true
schema:
$ref: '#/definitions/store.StageAcceptanceCapacityProfilesInput'
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.AcceptanceRun'
security:
- BearerAuth: []
summary: 为当前 validation Run 暂存 80% 容量门禁
tags:
- acceptance
/api/admin/system/acceptance/runs/{runID}/finish:
post:
consumes:
@@ -5939,6 +6071,29 @@ paths:
summary: 获取 Gateway 流量模式
tags:
- acceptance
/api/admin/system/acceptance/traffic-mode/pause:
post:
consumes:
- application/json
parameters:
- description: 当前线上 release 与 revision CAS
in: body
name: body
required: true
schema:
$ref: '#/definitions/store.PauseGatewayTrafficInput'
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.GatewayTrafficMode'
security:
- BearerAuth: []
summary: 上线后硬门禁异常时通过 CAS 暂停正式新任务
tags:
- acceptance
/api/admin/system/client-customization/settings:
get:
description: 返回客户端名称、英文名称、平台名和图标路径;数据库对象尚未创建时返回默认设置。
+52 -13
View File
@@ -37,18 +37,20 @@ type Config struct {
}
type Server struct {
now func() time.Time
httpClient *http.Client
wait func(context.Context, time.Duration) error
mu sync.RWMutex
tasks map[string]videoTask
fixtures map[string]fixture
nextID atomic.Uint64
report Report
callbacks map[string]map[int64]int
pngSmall string
pngLarge string
pngPeak string
now func() time.Time
httpClient *http.Client
wait func(context.Context, time.Duration) error
mu sync.RWMutex
tasks map[string]videoTask
fixtures map[string]fixture
nextID atomic.Uint64
report Report
callbacks map[string]map[int64]int
videoByIdempotency map[string]string
geminiIdempotency map[string]struct{}
pngSmall string
pngLarge string
pngPeak string
}
type Report struct {
@@ -67,6 +69,8 @@ type Report struct {
VerifiedConversions int64 `json:"verifiedConversions"`
CallbackEvents int64 `json:"callbackEvents"`
DuplicateCallbacks int64 `json:"duplicateCallbacks"`
MissingIdempotency int64 `json:"missingIdempotencyKeys"`
DuplicateSubmissions int64 `json:"duplicateSubmissionAttempts"`
}
type videoTask struct {
@@ -105,7 +109,9 @@ func New(config Config) *Server {
}
return &Server{
now: now, httpClient: client, wait: waiter, tasks: map[string]videoTask{}, fixtures: buildFixtures(),
callbacks: map[string]map[int64]int{},
callbacks: map[string]map[int64]int{},
videoByIdempotency: map[string]string{},
geminiIdempotency: map[string]struct{}{},
report: Report{
VideoReferenceCounts: map[string]int64{},
VideoRoleCounts: map[string]int64{},
@@ -140,6 +146,14 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
http.NotFound(w, r)
return
}
idempotencyKey := strings.TrimSpace(r.Header.Get("Idempotency-Key"))
if idempotencyKey == "" {
s.mu.Lock()
s.report.MissingIdempotency++
s.mu.Unlock()
writeProtocolError(w, http.StatusBadRequest, "Idempotency-Key is required")
return
}
var body map[string]any
if err := decodeJSON(r, &body); err != nil {
s.recordGeminiInvalid()
@@ -157,6 +171,11 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
return
}
s.mu.Lock()
if _, exists := s.geminiIdempotency[idempotencyKey]; exists {
s.report.DuplicateSubmissions++
} else {
s.geminiIdempotency[idempotencyKey] = struct{}{}
}
s.report.GeminiRequests++
s.report.GeminiInputBytes += int64(inputBytes)
s.report.GeminiOutputBytes += int64(outputBytes)
@@ -195,6 +214,14 @@ func (s *Server) geminiProfile(inputBytes int) (int, string, time.Duration) {
}
func (s *Server) submitVideo(w http.ResponseWriter, r *http.Request) {
idempotencyKey := strings.TrimSpace(r.Header.Get("Idempotency-Key"))
if idempotencyKey == "" {
s.mu.Lock()
s.report.MissingIdempotency++
s.mu.Unlock()
writeProtocolError(w, http.StatusBadRequest, "Idempotency-Key is required")
return
}
var body map[string]any
if err := decodeJSON(r, &body); err != nil {
s.recordVideoInvalid()
@@ -214,7 +241,19 @@ func (s *Server) submitVideo(w http.ResponseWriter, r *http.Request) {
CreatedAt: s.now(), ReadyAt: s.now().Add(delay), ImageRefs: refs,
}
s.mu.Lock()
if existingID := s.videoByIdempotency[idempotencyKey]; existingID != "" {
task = s.tasks[existingID]
s.report.DuplicateSubmissions++
s.mu.Unlock()
w.Header().Set("X-Request-ID", task.ID)
writeJSON(w, http.StatusOK, map[string]any{
"id": task.ID, "model": task.Model, "status": "queued",
"created_at": task.CreatedAt.Unix(),
})
return
}
s.tasks[id] = task
s.videoByIdempotency[idempotencyKey] = id
s.report.VideoSubmissions++
s.report.VideoReferenceCounts[strconv.Itoa(len(refs))]++
if forceConversion {
@@ -40,7 +40,11 @@ func TestGeminiAndVolcesProtocolEmulation(t *testing.T) {
}}},
"generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}},
})
response, err := http.Post(httpServer.URL+"/v1beta/models/gemini-test:generateContent", "application/json", bytes.NewReader(geminiBody))
response, err := postIdempotent(
httpServer.URL+"/v1beta/models/gemini-test:generateContent",
geminiBody,
"gemini-task",
)
if err != nil {
t.Fatalf("Gemini request: %v", err)
}
@@ -76,7 +80,11 @@ func TestGeminiAndVolcesProtocolEmulation(t *testing.T) {
})
}
videoBody, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": 1})
response, err = http.Post(httpServer.URL+"/contents/generations/tasks", "application/json", bytes.NewReader(videoBody))
response, err = postIdempotent(
httpServer.URL+"/contents/generations/tasks",
videoBody,
"video-task",
)
if err != nil {
t.Fatalf("submit video: %v", err)
}
@@ -159,7 +167,11 @@ func TestVolcesProtocolAcceptsThreeSixAndNineReferenceImages(t *testing.T) {
})
}
body, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": imageCount})
response, err := http.Post(httpServer.URL+"/contents/generations/tasks", "application/json", bytes.NewReader(body))
response, err := postIdempotent(
httpServer.URL+"/contents/generations/tasks",
body,
fmt.Sprintf("video-%d", imageCount),
)
if err != nil {
t.Fatalf("%d-image submit: %v", imageCount, err)
}
@@ -184,3 +196,13 @@ func TestVolcesProtocolAcceptsThreeSixAndNineReferenceImages(t *testing.T) {
}
}
}
func postIdempotent(url string, body []byte, key string) (*http.Response, error) {
request, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, err
}
request.Header.Set("Content-Type", "application/json")
request.Header.Set("Idempotency-Key", key)
return http.DefaultClient.Do(request)
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,180 @@
package acceptancesnapshot
import (
"encoding/json"
"testing"
"time"
)
func TestSnapshotValidationRejectsSecretsAndDetectsTampering(t *testing.T) {
snapshot := validSnapshot(t)
if err := Validate(snapshot); err != nil {
t.Fatalf("validate fixture: %v", err)
}
tampered := snapshot
tampered.Candidates = append([]Candidate(nil), snapshot.Candidates...)
tampered.Candidates[0].Platform.Name = "tampered"
if err := Validate(tampered); err == nil {
t.Fatal("tampered snapshot passed validation")
}
unsafe := snapshot
unsafe.Candidates = append([]Candidate(nil), snapshot.Candidates...)
unsafe.Candidates[0].Platform.Config = map[string]any{"apiToken": "must-not-leak"}
refreshHashes(t, &unsafe)
if err := Validate(unsafe); err == nil {
t.Fatal("secret-like snapshot field passed validation")
}
}
func TestSanitizeMapRemovesCredentialAndProxyFields(t *testing.T) {
input := map[string]any{
"specType": "gemini",
"credentialEnv": "PRODUCTION_TOKEN",
"networkProxy": map[string]any{
"proxyUrl": "http://user:password@example.invalid",
},
"requestAssetImageURLFormat": "base64",
}
got := sanitizeMap(input)
if got["specType"] != "gemini" || got["requestAssetImageURLFormat"] != "base64" {
t.Fatalf("safe protocol fields were removed: %#v", got)
}
if _, ok := got["credentialEnv"]; ok {
t.Fatal("credentialEnv was not removed")
}
nested, _ := got["networkProxy"].(map[string]any)
if _, ok := nested["proxyUrl"]; ok {
t.Fatal("nested proxy URL was not removed")
}
}
func TestDecodeRejectsUnknownSnapshotFields(t *testing.T) {
snapshot := validSnapshot(t)
payload, err := json.Marshal(snapshot)
if err != nil {
t.Fatal(err)
}
var object map[string]any
if err := json.Unmarshal(payload, &object); err != nil {
t.Fatal(err)
}
object["unexpected"] = true
payload, _ = json.Marshal(object)
if _, err := Decode(payload); err == nil {
t.Fatal("snapshot with unknown field decoded successfully")
}
}
func validSnapshot(t *testing.T) Snapshot {
t.Helper()
candidates := []Candidate{
{
Workload: "gemini_image_edit",
Provider: Provider{
ProviderKey: "gemini", ProviderCode: "google-gemini",
DisplayName: "Gemini", ProviderType: "gemini",
CapabilitySchema: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
},
BaseModel: BaseModel{
CanonicalModelKey: "gemini:image", InvocationName: "gemini-image",
ProviderModelName: "gemini-image", ModelType: []string{"image_edit"},
DisplayName: "Gemini Image", Capabilities: map[string]any{},
BaseBillingConfig: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
RuntimePolicyOverride: map[string]any{}, PricingVersion: 1,
},
Platform: Platform{
Provider: "gemini", PlatformKey: "gemini-production", Name: "Gemini Production",
BaseURLPath: "/", AuthType: "api_key", Config: map[string]any{"specType": "gemini"},
DefaultPricingMode: "inherit_discount", DefaultDiscountFactor: "1",
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{}, Priority: 10,
},
PlatformModel: PlatformModel{
ModelName: "gemini-image", ProviderModelName: "gemini-image",
ModelType: []string{"image_edit"}, DisplayName: "Gemini Image",
CapabilityOverride: map[string]any{}, Capabilities: map[string]any{},
PricingMode: "inherit_discount", BillingConfigOverride: map[string]any{},
BillingConfig: map[string]any{}, PermissionConfig: map[string]any{},
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{},
RuntimePolicyOverride: map[string]any{},
},
RuntimePolicy: RuntimePolicy{
RateLimitPolicy: map[string]any{}, RetryPolicy: map[string]any{},
AutoDisablePolicy: map[string]any{}, DegradePolicy: map[string]any{},
},
PricingRules: []PricingRule{},
Metadata: CandidateSource{
PlatformID: "00000000-0000-0000-0000-000000000001",
PlatformModelID: "00000000-0000-0000-0000-000000000002",
BaseModelID: "00000000-0000-0000-0000-000000000003",
},
},
{
Workload: "multi_reference_video",
Provider: Provider{
ProviderKey: "volces", ProviderCode: "volces",
DisplayName: "Volces", ProviderType: "volces",
CapabilitySchema: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
},
BaseModel: BaseModel{
CanonicalModelKey: "volces:seedance", InvocationName: "seedance-2-fast",
ProviderModelName: "seedance-2-fast", ModelType: []string{"omni_video"},
DisplayName: "Seedance", Capabilities: map[string]any{"omni_video": map[string]any{"max_images": 9}},
BaseBillingConfig: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
RuntimePolicyOverride: map[string]any{}, PricingVersion: 1,
},
Platform: Platform{
Provider: "volces", PlatformKey: "volces-production", Name: "Volces Production",
BaseURLPath: "/", AuthType: "bearer", Config: map[string]any{"specType": "volces"},
DefaultPricingMode: "inherit_discount", DefaultDiscountFactor: "1",
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{}, Priority: 10,
},
PlatformModel: PlatformModel{
ModelName: "seedance-2-fast", ProviderModelName: "seedance-2-fast",
ModelType: []string{"omni_video"}, DisplayName: "Seedance",
CapabilityOverride: map[string]any{"omni_video": map[string]any{"max_images": 9}},
Capabilities: map[string]any{}, PricingMode: "inherit_discount",
BillingConfigOverride: map[string]any{}, BillingConfig: map[string]any{},
PermissionConfig: map[string]any{}, RetryPolicy: map[string]any{},
RateLimitPolicy: map[string]any{}, RuntimePolicyOverride: map[string]any{},
},
RuntimePolicy: RuntimePolicy{
RateLimitPolicy: map[string]any{}, RetryPolicy: map[string]any{},
AutoDisablePolicy: map[string]any{}, DegradePolicy: map[string]any{},
},
PricingRules: []PricingRule{},
Metadata: CandidateSource{
PlatformID: "00000000-0000-0000-0000-000000000004",
PlatformModelID: "00000000-0000-0000-0000-000000000005",
BaseModelID: "00000000-0000-0000-0000-000000000006",
},
},
}
snapshot := Snapshot{
SchemaVersion: SchemaVersion,
Source: Source{
ReleaseSHA: "0123456789abcdef0123456789abcdef01234567",
CreatedAt: time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC),
},
Candidates: candidates,
SecretSafe: true,
}
refreshHashes(t, &snapshot)
return snapshot
}
func refreshHashes(t *testing.T, snapshot *Snapshot) {
t.Helper()
configHash, err := candidateConfigHash(snapshot.Candidates)
if err != nil {
t.Fatal(err)
}
snapshot.Source.ConfigHash = configHash
snapshot.SnapshotSHA256 = ""
hash, err := snapshotHash(*snapshot)
if err != nil {
t.Fatal(err)
}
snapshot.SnapshotSHA256 = hash
}
@@ -0,0 +1,413 @@
package capacitycontroller
import (
"context"
"encoding/json"
"errors"
"log/slog"
"net/http"
"os"
"strings"
"sync"
"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"
)
const (
controllerReconcileInterval = 10 * time.Second
controllerLeadershipRetry = 5 * time.Second
controllerDeletionCost = -1000
)
type capacityStore interface {
TryAcquireCapacityControllerLeadership(context.Context) (store.Leadership, bool, error)
WorkerQueueRuntime(context.Context) (store.WorkerQueueRuntime, error)
ListWorkerInstanceRuntime(context.Context) ([]store.WorkerInstanceRuntime, error)
CapacityDatabaseHealth(context.Context) (store.CapacityDatabaseHealth, error)
MarkWorkerDraining(context.Context, string) error
ReactivateWorkerInstance(context.Context, string) error
}
type capacityKubernetes interface {
SiteState(context.Context, string) (KubernetesSiteState, error)
ScaleWorkerDeployment(context.Context, string, int) error
SetPodDeletionCost(context.Context, string, int) error
}
type Status struct {
Leader bool `json:"leader"`
LastRunAt time.Time `json:"lastRunAt,omitempty"`
LastError string `json:"lastError,omitempty"`
Queue store.WorkerQueueRuntime `json:"queue"`
Plan Plan `json:"plan"`
ScaleActions uint64 `json:"scaleActions"`
}
type Controller struct {
cfg config.Config
store capacityStore
kubernetes capacityKubernetes
logger *slog.Logger
expectedRevision string
now func() time.Time
highSince time.Time
lowSince time.Time
statusMu sync.RWMutex
status Status
}
func New(
cfg config.Config,
db capacityStore,
kubernetes capacityKubernetes,
logger *slog.Logger,
) *Controller {
return &Controller{
cfg: cfg, store: db, kubernetes: kubernetes, logger: logger,
expectedRevision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
now: time.Now,
}
}
func (controller *Controller) Run(ctx context.Context) {
for ctx.Err() == nil {
leadership, acquired, err := controller.store.TryAcquireCapacityControllerLeadership(ctx)
if err != nil {
controller.setError(err)
if !waitContext(ctx, controllerLeadershipRetry) {
return
}
continue
}
if !acquired {
controller.setLeader(false)
controller.clearError()
if !waitContext(ctx, controllerLeadershipRetry) {
return
}
continue
}
controller.setLeader(true)
controller.clearError()
controller.runLeader(ctx, leadership)
leadership.Release()
controller.setLeader(false)
}
}
func (controller *Controller) runLeader(ctx context.Context, leadership store.Leadership) {
ticker := time.NewTicker(controllerReconcileInterval)
defer ticker.Stop()
if err := controller.reconcile(ctx); err != nil {
controller.logError("initial capacity reconciliation failed", err)
}
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
keepAliveCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
err := leadership.KeepAlive(keepAliveCtx)
cancel()
if err != nil {
controller.logError("capacity controller leadership lost", err)
return
}
if err := controller.reconcile(ctx); err != nil {
controller.logError("capacity reconciliation failed", err)
}
}
}
}
func (controller *Controller) reconcile(ctx context.Context) error {
queue, err := controller.store.WorkerQueueRuntime(ctx)
if err != nil {
return err
}
instances, err := controller.store.ListWorkerInstanceRuntime(ctx)
if err != nil {
return err
}
database, err := controller.store.CapacityDatabaseHealth(ctx)
if err != nil {
return err
}
now := controller.now()
sites := make([]SiteResources, 0, 2)
kubernetesStates := make(map[string]KubernetesSiteState, 2)
for _, site := range []string{"ningbo", "hongkong"} {
state, stateErr := controller.kubernetes.SiteState(ctx, site)
if stateErr != nil {
return stateErr
}
kubernetesStates[site] = state
minReplicas, maxReplicas := controller.siteReplicaBounds(site)
sites = append(sites, SiteResources{
Site: site, CurrentReplicas: state.CurrentReplicas,
MinReplicas: minReplicas, MaxReplicas: maxReplicas,
AllocatableMemoryBytes: state.AllocatableMemoryBytes,
UsedMemoryBytes: state.UsedMemoryBytes,
WorkerRequestMemoryBytes: state.WorkerRequestMemoryBytes,
AllocatableMilliCPU: state.AllocatableMilliCPU,
UsedMilliCPU: state.UsedMilliCPU,
WorkerRequestMilliCPU: state.WorkerRequestMilliCPU,
MemoryPressure: state.MemoryPressure,
Nodes: nodeResources(state.Nodes),
})
}
currentTotal := 0
for _, site := range sites {
currentTotal += site.CurrentReplicas
}
target := controller.cfg.WorkerTargetOutstandingPerReplica
if target < 1 {
target = 2 * controller.cfg.AsyncWorkerInstanceHardLimit
}
rawDesired := (max(queue.Queued+queue.Running, 0) + target - 1) / target
if rawDesired > currentTotal {
if controller.highSince.IsZero() {
controller.highSince = now
}
controller.lowSince = time.Time{}
} else if rawDesired < currentTotal && queue.Queued == 0 {
if controller.lowSince.IsZero() {
controller.lowSince = now
}
controller.highSince = time.Time{}
} else {
controller.highSince = time.Time{}
controller.lowSince = time.Time{}
}
revisionHealthy := controller.revisionsMatch(instances, kubernetesStates)
scaleUpEligible := !controller.highSince.IsZero() &&
now.Sub(controller.highSince) >= time.Duration(controller.cfg.WorkerScaleUpWindowSeconds)*time.Second &&
revisionHealthy
scaleDownEligible := !controller.lowSince.IsZero() &&
now.Sub(controller.lowSince) >= time.Duration(controller.cfg.WorkerScaleDownStabilizationSeconds)*time.Second &&
revisionHealthy
plan := CalculatePlan(PlanInput{
Queued: queue.Queued, Running: queue.Running,
InstanceSlots: controller.cfg.AsyncWorkerInstanceHardLimit,
TargetOutstandingPerReplica: target,
MemoryTargetPercent: controller.cfg.NodeMemoryTargetPercent,
MemoryHardPercent: controller.cfg.NodeMemoryHardPercent,
CPUTargetPercent: controller.cfg.NodeCPUTargetPercent,
DatabaseConnections: database.Connections,
DatabaseConnectionBudget: controller.cfg.PostgresConnectionBudget,
NonWorkerConnectionBudget: controller.cfg.PostgresNonWorkerConnectionBudget,
WorkerDatabasePoolMax: controller.cfg.WorkerDatabaseMaxConns,
SynchronousDatabasePeers: database.SynchronousPeers,
ScaleUpEligible: scaleUpEligible,
ScaleDownEligible: scaleDownEligible,
Sites: sites,
})
if !revisionHealthy && plan.FrozenReason == "" {
plan.FrozenReason = "release_revision_mismatch"
}
if controller.cfg.WorkerAutoscalingEnabled {
for _, sitePlan := range plan.Sites {
switch {
case sitePlan.DesiredReplicas > sitePlan.CurrentReplicas:
if err := controller.kubernetes.ScaleWorkerDeployment(ctx, sitePlan.Site, sitePlan.DesiredReplicas); err != nil {
return err
}
controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.DesiredReplicas, "scale_up")
case sitePlan.DesiredReplicas < sitePlan.CurrentReplicas:
if err := controller.reconcileScaleDown(ctx, now, sitePlan, instances); err != nil {
return err
}
}
}
}
controller.statusMu.Lock()
controller.status.LastRunAt = now
controller.status.LastError = ""
controller.status.Queue = queue
controller.status.Plan = plan
controller.statusMu.Unlock()
return nil
}
func nodeResources(states []KubernetesNodeState) []NodeResources {
nodes := make([]NodeResources, 0, len(states))
for _, state := range states {
nodes = append(nodes, NodeResources{
NodeName: state.NodeName,
CurrentReplicas: state.CurrentReplicas,
AllocatableMemoryBytes: state.AllocatableMemoryBytes,
UsedMemoryBytes: state.UsedMemoryBytes,
WorkerUsedMemoryBytes: state.WorkerUsedMemoryBytes,
AllocatableMilliCPU: state.AllocatableMilliCPU,
UsedMilliCPU: state.UsedMilliCPU,
WorkerUsedMilliCPU: state.WorkerUsedMilliCPU,
MemoryPressure: state.MemoryPressure,
})
}
return nodes
}
func (controller *Controller) reconcileScaleDown(
ctx context.Context,
now time.Time,
sitePlan SitePlan,
instances []store.WorkerInstanceRuntime,
) error {
for _, instance := range instances {
if instance.Site != sitePlan.Site || instance.Status != "draining" {
continue
}
if instance.RunningTasks == 0 && instance.ActiveLeases == 0 {
if err := controller.kubernetes.SetPodDeletionCost(ctx, instance.PodName, controllerDeletionCost); err != nil {
return err
}
if err := controller.kubernetes.ScaleWorkerDeployment(ctx, sitePlan.Site, sitePlan.CurrentReplicas-1); err != nil {
return err
}
controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.CurrentReplicas-1, "drained_scale_down")
return nil
}
if instance.DrainingAt != nil &&
now.Sub(*instance.DrainingAt) >= time.Duration(controller.cfg.WorkerDrainTimeoutSeconds)*time.Second {
if err := controller.store.ReactivateWorkerInstance(ctx, instance.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
return err
}
controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.CurrentReplicas, "drain_timeout")
}
return nil
}
var candidate *store.WorkerInstanceRuntime
for index := range instances {
instance := &instances[index]
if instance.Site != sitePlan.Site || instance.Status != "active" {
continue
}
if candidate == nil ||
instance.RunningTasks+instance.ActiveLeases < candidate.RunningTasks+candidate.ActiveLeases {
candidate = instance
}
}
if candidate == nil {
return nil
}
if err := controller.store.MarkWorkerDraining(ctx, candidate.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
return err
}
controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.CurrentReplicas, "drain_started")
return nil
}
func (controller *Controller) siteReplicaBounds(site string) (int, int) {
switch site {
case "ningbo":
return controller.cfg.WorkerMinReplicasNingbo, controller.cfg.WorkerMaxReplicasNingbo
case "hongkong":
return controller.cfg.WorkerMinReplicasHongkong, controller.cfg.WorkerMaxReplicasHongkong
default:
return 0, 0
}
}
func (controller *Controller) revisionsMatch(
instances []store.WorkerInstanceRuntime,
states map[string]KubernetesSiteState,
) bool {
if controller.expectedRevision == "" {
return true
}
for _, state := range states {
if state.Revision != "" && state.Revision != controller.expectedRevision {
return false
}
}
for _, instance := range instances {
if instance.Revision != "" && instance.Revision != controller.expectedRevision {
return false
}
}
return true
}
func (controller *Controller) Status() Status {
controller.statusMu.RLock()
defer controller.statusMu.RUnlock()
return controller.status
}
func (controller *Controller) Handler() http.Handler {
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"ok":true,"service":"easyai-capacity-controller"}`))
})
mux.HandleFunc("GET /readyz", func(w http.ResponseWriter, _ *http.Request) {
status := controller.Status()
// Followers are intentionally idle while the database advisory lock is
// held by the elected leader. They are still ready to take over and must
// not make a two-replica Deployment permanently fail its rollout.
if status.LastError != "" {
http.Error(w, `{"ok":false}`, http.StatusServiceUnavailable)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"ok":true}`))
})
mux.HandleFunc("GET /status", func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(controller.Status())
})
return mux
}
func (controller *Controller) setLeader(leader bool) {
controller.statusMu.Lock()
controller.status.Leader = leader
controller.statusMu.Unlock()
}
func (controller *Controller) setError(err error) {
controller.statusMu.Lock()
controller.status.LastError = err.Error()
controller.statusMu.Unlock()
}
func (controller *Controller) clearError() {
controller.statusMu.Lock()
controller.status.LastError = ""
controller.statusMu.Unlock()
}
func (controller *Controller) logError(message string, err error) {
controller.setError(err)
if controller.logger != nil {
controller.logger.Error(message, "error", err)
}
}
func (controller *Controller) observeScale(site string, from int, to int, reason string) {
controller.statusMu.Lock()
controller.status.ScaleActions++
controller.statusMu.Unlock()
if controller.logger != nil {
controller.logger.Info("worker capacity action",
"site", site,
"fromReplicas", from,
"toReplicas", to,
"reason", reason,
)
}
}
func waitContext(ctx context.Context, duration time.Duration) bool {
timer := time.NewTimer(duration)
defer timer.Stop()
select {
case <-ctx.Done():
return false
case <-timer.C:
return true
}
}
@@ -0,0 +1,440 @@
package capacitycontroller
import (
"bytes"
"context"
"crypto/tls"
"crypto/x509"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"os"
"strconv"
"strings"
"time"
)
type KubernetesConfig struct {
Namespace string
APIServer string
TokenFile string
CAFile string
HTTPClient *http.Client
}
type KubernetesClient struct {
namespace string
baseURL string
tokenFile string
client *http.Client
}
type KubernetesSiteState struct {
Site string
NodeName string
CurrentReplicas int
Revision string
Image string
AllocatableMemoryBytes int64
UsedMemoryBytes int64
WorkerRequestMemoryBytes int64
AllocatableMilliCPU int64
UsedMilliCPU int64
WorkerRequestMilliCPU int64
MemoryPressure bool
Nodes []KubernetesNodeState
}
type KubernetesNodeState struct {
NodeName string
CurrentReplicas int
AllocatableMemoryBytes int64
UsedMemoryBytes int64
WorkerUsedMemoryBytes int64
AllocatableMilliCPU int64
UsedMilliCPU int64
WorkerUsedMilliCPU int64
MemoryPressure bool
}
func NewKubernetesClient(config KubernetesConfig) (*KubernetesClient, error) {
if strings.TrimSpace(config.Namespace) == "" {
return nil, errors.New("capacity controller Kubernetes namespace is required")
}
if config.APIServer == "" {
config.APIServer = "https://kubernetes.default.svc"
}
if config.TokenFile == "" {
config.TokenFile = "/var/run/secrets/kubernetes.io/serviceaccount/token"
}
parsed, err := url.Parse(strings.TrimRight(config.APIServer, "/"))
if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" ||
(parsed.Scheme != "https" && config.HTTPClient == nil) {
return nil, errors.New("capacity controller Kubernetes API server URL is invalid")
}
client := config.HTTPClient
if client == nil {
if config.CAFile == "" {
config.CAFile = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
}
certificate, readErr := os.ReadFile(config.CAFile)
if readErr != nil {
return nil, fmt.Errorf("read Kubernetes service account CA: %w", readErr)
}
roots := x509.NewCertPool()
if !roots.AppendCertsFromPEM(certificate) {
return nil, errors.New("Kubernetes service account CA is invalid")
}
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = nil
transport.DisableCompression = true
transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: roots}
client = &http.Client{
Timeout: 10 * time.Second,
Transport: transport,
CheckRedirect: func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
},
}
}
return &KubernetesClient{
namespace: config.Namespace,
baseURL: strings.TrimRight(config.APIServer, "/"),
tokenFile: config.TokenFile,
client: client,
}, nil
}
func (client *KubernetesClient) SiteState(ctx context.Context, site string) (KubernetesSiteState, error) {
site = strings.TrimSpace(site)
if site == "" {
return KubernetesSiteState{}, errors.New("site is required")
}
var deployment struct {
Metadata struct {
Name string `json:"name"`
} `json:"metadata"`
Spec struct {
Replicas int `json:"replicas"`
Template struct {
Spec struct {
Containers []struct {
Name string `json:"name"`
Image string `json:"image"`
Env []struct {
Name string `json:"name"`
Value string `json:"value"`
} `json:"env"`
Resources struct {
Requests map[string]string `json:"requests"`
} `json:"resources"`
} `json:"containers"`
} `json:"spec"`
} `json:"template"`
} `json:"spec"`
}
deploymentPath := fmt.Sprintf(
"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s",
url.PathEscape(client.namespace),
url.PathEscape(site),
)
if err := client.getJSON(ctx, deploymentPath, &deployment); err != nil {
return KubernetesSiteState{}, err
}
state := KubernetesSiteState{Site: site, CurrentReplicas: deployment.Spec.Replicas}
for _, container := range deployment.Spec.Template.Spec.Containers {
if container.Name != "worker" {
continue
}
state.Image = container.Image
state.WorkerRequestMemoryBytes, _ = parseBinaryQuantity(container.Resources.Requests["memory"])
state.WorkerRequestMilliCPU, _ = parseCPUQuantity(container.Resources.Requests["cpu"])
for _, environment := range container.Env {
if environment.Name == "AI_GATEWAY_REVISION" {
state.Revision = environment.Value
}
}
}
var nodes struct {
Items []struct {
Metadata struct {
Name string `json:"name"`
} `json:"metadata"`
Status struct {
Allocatable map[string]string `json:"allocatable"`
Conditions []struct {
Type string `json:"type"`
Status string `json:"status"`
} `json:"conditions"`
} `json:"status"`
} `json:"items"`
}
nodePath := "/api/v1/nodes?labelSelector=" + url.QueryEscape("easyai.io/site="+site)
if err := client.getJSON(ctx, nodePath, &nodes); err != nil {
return KubernetesSiteState{}, err
}
if len(nodes.Items) == 0 {
return KubernetesSiteState{}, fmt.Errorf("site %s has no matching nodes", site)
}
var pods struct {
Items []struct {
Metadata struct {
Name string `json:"name"`
DeletionTimestamp *time.Time `json:"deletionTimestamp"`
} `json:"metadata"`
Spec struct {
NodeName string `json:"nodeName"`
} `json:"spec"`
} `json:"items"`
}
podPath := fmt.Sprintf(
"/api/v1/namespaces/%s/pods?labelSelector=%s",
url.PathEscape(client.namespace),
url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
)
if err := client.getJSON(ctx, podPath, &pods); err != nil {
return KubernetesSiteState{}, err
}
replicasByNode := make(map[string]int, len(nodes.Items))
nodeByPod := make(map[string]string, len(pods.Items))
for _, pod := range pods.Items {
if pod.Metadata.DeletionTimestamp == nil && pod.Spec.NodeName != "" {
replicasByNode[pod.Spec.NodeName]++
nodeByPod[pod.Metadata.Name] = pod.Spec.NodeName
}
}
var podMetrics struct {
Items []struct {
Metadata struct {
Name string `json:"name"`
} `json:"metadata"`
Containers []struct {
Usage map[string]string `json:"usage"`
} `json:"containers"`
} `json:"items"`
}
podMetricsPath := fmt.Sprintf(
"/apis/metrics.k8s.io/v1beta1/namespaces/%s/pods?labelSelector=%s",
url.PathEscape(client.namespace),
url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
)
if err := client.getJSON(ctx, podMetricsPath, &podMetrics); err != nil {
return KubernetesSiteState{}, err
}
workerMemoryByNode := make(map[string]int64, len(nodes.Items))
workerCPUByNode := make(map[string]int64, len(nodes.Items))
measuredPods := make(map[string]bool, len(podMetrics.Items))
for _, pod := range podMetrics.Items {
nodeName := nodeByPod[pod.Metadata.Name]
if nodeName == "" {
continue
}
measuredPods[pod.Metadata.Name] = true
for _, container := range pod.Containers {
memory, memoryErr := parseBinaryQuantity(container.Usage["memory"])
cpu, cpuErr := parseCPUQuantity(container.Usage["cpu"])
if memoryErr != nil || cpuErr != nil {
return KubernetesSiteState{}, fmt.Errorf(
"site %s pod %s returned incomplete resource metrics",
site,
pod.Metadata.Name,
)
}
workerMemoryByNode[nodeName] += memory
workerCPUByNode[nodeName] += cpu
}
}
for podName := range nodeByPod {
if !measuredPods[podName] {
return KubernetesSiteState{}, fmt.Errorf(
"site %s pod %s has no resource metrics yet",
site,
podName,
)
}
}
for _, node := range nodes.Items {
nodeState := KubernetesNodeState{
NodeName: node.Metadata.Name,
CurrentReplicas: replicasByNode[node.Metadata.Name],
WorkerUsedMemoryBytes: workerMemoryByNode[node.Metadata.Name],
WorkerUsedMilliCPU: workerCPUByNode[node.Metadata.Name],
}
nodeState.AllocatableMemoryBytes, _ = parseBinaryQuantity(node.Status.Allocatable["memory"])
nodeState.AllocatableMilliCPU, _ = parseCPUQuantity(node.Status.Allocatable["cpu"])
for _, condition := range node.Status.Conditions {
if condition.Type == "MemoryPressure" && condition.Status == "True" {
nodeState.MemoryPressure = true
state.MemoryPressure = true
}
}
var metrics struct {
Usage map[string]string `json:"usage"`
}
metricsPath := "/apis/metrics.k8s.io/v1beta1/nodes/" + url.PathEscape(nodeState.NodeName)
if err := client.getJSON(ctx, metricsPath, &metrics); err != nil {
return KubernetesSiteState{}, err
}
nodeState.UsedMemoryBytes, _ = parseBinaryQuantity(metrics.Usage["memory"])
nodeState.UsedMilliCPU, _ = parseCPUQuantity(metrics.Usage["cpu"])
if nodeState.AllocatableMemoryBytes <= 0 || nodeState.AllocatableMilliCPU <= 0 {
return KubernetesSiteState{}, fmt.Errorf(
"site %s node %s returned incomplete resource quantities",
site,
nodeState.NodeName,
)
}
state.Nodes = append(state.Nodes, nodeState)
state.AllocatableMemoryBytes += nodeState.AllocatableMemoryBytes
state.AllocatableMilliCPU += nodeState.AllocatableMilliCPU
state.UsedMemoryBytes += nodeState.UsedMemoryBytes
state.UsedMilliCPU += nodeState.UsedMilliCPU
}
if len(state.Nodes) == 1 {
state.NodeName = state.Nodes[0].NodeName
}
if state.WorkerRequestMemoryBytes <= 0 || state.WorkerRequestMilliCPU <= 0 ||
state.AllocatableMemoryBytes <= 0 || state.AllocatableMilliCPU <= 0 {
return KubernetesSiteState{}, fmt.Errorf("site %s returned incomplete resource quantities", site)
}
return state, nil
}
func (client *KubernetesClient) ScaleWorkerDeployment(ctx context.Context, site string, replicas int) error {
if replicas < 0 || replicas > 64 {
return errors.New("worker replica target must be between 0 and 64")
}
path := fmt.Sprintf(
"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s/scale",
url.PathEscape(client.namespace),
url.PathEscape(strings.TrimSpace(site)),
)
return client.patchJSON(ctx, path, map[string]any{"spec": map[string]any{"replicas": replicas}})
}
func (client *KubernetesClient) SetPodDeletionCost(ctx context.Context, podName string, cost int) error {
podName = strings.TrimSpace(podName)
if podName == "" {
return errors.New("worker pod name is required")
}
path := fmt.Sprintf(
"/api/v1/namespaces/%s/pods/%s",
url.PathEscape(client.namespace),
url.PathEscape(podName),
)
return client.patchJSON(ctx, path, map[string]any{
"metadata": map[string]any{
"annotations": map[string]any{
"controller.kubernetes.io/pod-deletion-cost": strconv.Itoa(cost),
},
},
})
}
func (client *KubernetesClient) getJSON(ctx context.Context, path string, target any) error {
request, err := client.request(ctx, http.MethodGet, path, nil)
if err != nil {
return err
}
response, err := client.client.Do(request)
if err != nil {
return fmt.Errorf("request Kubernetes API: %w", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
}
decoder := json.NewDecoder(io.LimitReader(response.Body, 4*1024*1024))
if err := decoder.Decode(target); err != nil {
return fmt.Errorf("decode Kubernetes API %s: %w", path, err)
}
return nil
}
func (client *KubernetesClient) patchJSON(ctx context.Context, path string, payload any) error {
encoded, err := json.Marshal(payload)
if err != nil {
return err
}
request, err := client.request(ctx, http.MethodPatch, path, bytes.NewReader(encoded))
if err != nil {
return err
}
request.Header.Set("Content-Type", "application/merge-patch+json")
response, err := client.client.Do(request)
if err != nil {
return fmt.Errorf("patch Kubernetes API: %w", err)
}
defer response.Body.Close()
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
if response.StatusCode < 200 || response.StatusCode >= 300 {
return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
}
return nil
}
func (client *KubernetesClient) request(ctx context.Context, method string, path string, body io.Reader) (*http.Request, error) {
token, err := os.ReadFile(client.tokenFile)
if err != nil || strings.TrimSpace(string(token)) == "" {
clear(token)
return nil, errors.New("Kubernetes service account token is unavailable")
}
request, err := http.NewRequestWithContext(ctx, method, client.baseURL+path, body)
if err != nil {
clear(token)
return nil, err
}
request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(string(token)))
request.Header.Set("Accept", "application/json")
clear(token)
return request, nil
}
func parseBinaryQuantity(value string) (int64, error) {
value = strings.TrimSpace(value)
if value == "" {
return 0, errors.New("quantity is empty")
}
units := []struct {
suffix string
multiplier int64
}{
{"Ei", 1 << 60},
{"Pi", 1 << 50},
{"Ti", 1 << 40},
{"Gi", 1 << 30},
{"Mi", 1 << 20},
{"Ki", 1 << 10},
{"G", 1_000_000_000},
{"M", 1_000_000},
{"K", 1_000},
}
for _, unit := range units {
if strings.HasSuffix(value, unit.suffix) {
number, err := strconv.ParseFloat(strings.TrimSuffix(value, unit.suffix), 64)
return int64(number * float64(unit.multiplier)), err
}
}
return strconv.ParseInt(value, 10, 64)
}
func parseCPUQuantity(value string) (int64, error) {
value = strings.TrimSpace(value)
if strings.HasSuffix(value, "n") {
nanocores, err := strconv.ParseInt(strings.TrimSuffix(value, "n"), 10, 64)
return nanocores / 1_000_000, err
}
if strings.HasSuffix(value, "u") {
microcores, err := strconv.ParseInt(strings.TrimSuffix(value, "u"), 10, 64)
return microcores / 1_000, err
}
if strings.HasSuffix(value, "m") {
return strconv.ParseInt(strings.TrimSuffix(value, "m"), 10, 64)
}
cores, err := strconv.ParseFloat(value, 64)
return int64(cores * 1000), err
}
@@ -0,0 +1,94 @@
package capacitycontroller
import (
"context"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
)
func TestKubernetesClientReadsSiteAndMutatesOnlyWorkerScale(t *testing.T) {
var patches []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
if request.Header.Get("Authorization") != "Bearer test-token" {
http.Error(w, "unauthorized", http.StatusUnauthorized)
return
}
switch {
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-hongkong"):
_, _ = w.Write([]byte(`{
"spec":{"replicas":1,"template":{"spec":{"containers":[{
"name":"worker","image":"registry.invalid/gateway@sha256:abc",
"env":[{"name":"AI_GATEWAY_REVISION","value":"release-sha"}],
"resources":{"requests":{"memory":"512Mi","cpu":"250m"}}
}]}}}
}`))
case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
_, _ = w.Write([]byte(`{"items":[{"metadata":{"name":"easyai-hongkong"},"status":{
"allocatable":{"memory":"8Gi","cpu":"4"},
"conditions":[{"type":"MemoryPressure","status":"False"}]
}}]}`))
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/api/v1/namespaces/easyai/pods"):
_, _ = w.Write([]byte(`{"items":[{"metadata":{"name":"worker-hongkong-1"},"spec":{"nodeName":"easyai-hongkong"}}]}`))
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
strings.HasSuffix(request.URL.Path, "/nodes/easyai-hongkong"):
_, _ = w.Write([]byte(`{"usage":{"memory":"3Gi","cpu":"500m"}}`))
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
strings.HasSuffix(request.URL.Path, "/pods"):
_, _ = w.Write([]byte(`{"items":[{
"metadata":{"name":"worker-hongkong-1"},
"containers":[{"usage":{"memory":"384Mi","cpu":"125m"}}]
}]}`))
case request.Method == http.MethodPatch:
patches = append(patches, request.URL.Path)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{}`))
default:
http.NotFound(w, request)
}
}))
defer server.Close()
tokenFile := filepath.Join(t.TempDir(), "token")
if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
t.Fatal(err)
}
client, err := NewKubernetesClient(KubernetesConfig{
Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
})
if err != nil {
t.Fatal(err)
}
state, err := client.SiteState(context.Background(), "hongkong")
if err != nil {
t.Fatal(err)
}
if state.CurrentReplicas != 1 || state.NodeName != "easyai-hongkong" ||
state.WorkerRequestMemoryBytes != 512<<20 || state.WorkerRequestMilliCPU != 250 ||
state.AllocatableMemoryBytes != 8<<30 || state.UsedMemoryBytes != 3<<30 ||
state.Nodes[0].WorkerUsedMemoryBytes != 384<<20 || state.Nodes[0].WorkerUsedMilliCPU != 125 {
t.Fatalf("site state=%+v", state)
}
if err := client.SetPodDeletionCost(context.Background(), "worker-pod", -1000); err != nil {
t.Fatal(err)
}
if err := client.ScaleWorkerDeployment(context.Background(), "hongkong", 2); err != nil {
t.Fatal(err)
}
if len(patches) != 2 ||
!strings.Contains(patches[0], "/pods/worker-pod") ||
!strings.Contains(patches[1], "/deployments/easyai-worker-hongkong/scale") {
t.Fatalf("patches=%v", patches)
}
}
func TestQuantityParsing(t *testing.T) {
if value, err := parseBinaryQuantity("1536Mi"); err != nil || value != 1536<<20 {
t.Fatalf("memory=%d err=%v", value, err)
}
if value, err := parseCPUQuantity("250000000n"); err != nil || value != 250 {
t.Fatalf("cpu=%d err=%v", value, err)
}
}
@@ -0,0 +1,281 @@
package capacitycontroller
import (
"math"
"sort"
)
type SiteResources struct {
Site string
CurrentReplicas int
MinReplicas int
MaxReplicas int
AllocatableMemoryBytes int64
UsedMemoryBytes int64
WorkerRequestMemoryBytes int64
AllocatableMilliCPU int64
UsedMilliCPU int64
WorkerRequestMilliCPU int64
MemoryPressure bool
Nodes []NodeResources
}
type NodeResources struct {
NodeName string
CurrentReplicas int
AllocatableMemoryBytes int64
UsedMemoryBytes int64
WorkerUsedMemoryBytes int64
AllocatableMilliCPU int64
UsedMilliCPU int64
WorkerUsedMilliCPU int64
MemoryPressure bool
}
type PlanInput struct {
Queued int
Running int
InstanceSlots int
TargetOutstandingPerReplica int
MemoryTargetPercent int
MemoryHardPercent int
CPUTargetPercent int
DatabaseConnections int
DatabaseConnectionBudget int
NonWorkerConnectionBudget int
WorkerDatabasePoolMax int
SynchronousDatabasePeers int
ScaleUpEligible bool
ScaleDownEligible bool
Sites []SiteResources
}
type SitePlan struct {
Site string `json:"site"`
CurrentReplicas int `json:"currentReplicas"`
DesiredReplicas int `json:"desiredReplicas"`
ResourceMax int `json:"resourceMax"`
MemoryPercent float64 `json:"memoryPercent"`
CPUPercent float64 `json:"cpuPercent"`
}
type Plan struct {
RawDesired int `json:"rawDesired"`
DesiredTotal int `json:"desiredTotal"`
CurrentTotal int `json:"currentTotal"`
FrozenReason string `json:"frozenReason,omitempty"`
Sites []SitePlan `json:"sites"`
}
func CalculatePlan(input PlanInput) Plan {
if input.InstanceSlots < 1 {
input.InstanceSlots = 1
}
target := input.TargetOutstandingPerReplica
if target < 1 {
target = 2 * input.InstanceSlots
}
if input.MemoryTargetPercent < 1 {
input.MemoryTargetPercent = 75
}
if input.MemoryHardPercent < 1 {
input.MemoryHardPercent = 85
}
if input.CPUTargetPercent < 1 {
input.CPUTargetPercent = 70
}
rawDesired := int(math.Ceil(float64(max(input.Queued+input.Running, 0)) / float64(target)))
plan := Plan{RawDesired: rawDesired}
minTotal := 0
resourceMaxTotal := 0
for _, site := range input.Sites {
for _, node := range site.Nodes {
if node.MemoryPressure {
site.MemoryPressure = true
}
}
current := max(site.CurrentReplicas, 0)
minReplicas := max(site.MinReplicas, 0)
configMax := max(site.MaxReplicas, minReplicas)
resourceMax, memoryPercent, cpuPercent := siteResourceMaximum(
site,
input.MemoryTargetPercent,
input.CPUTargetPercent,
)
resourceMax = min(resourceMax, configMax)
if resourceMax < minReplicas && plan.FrozenReason == "" {
plan.FrozenReason = "site_resource_budget"
}
resourceMax = max(resourceMax, minReplicas)
plan.Sites = append(plan.Sites, SitePlan{
Site: site.Site, CurrentReplicas: current, DesiredReplicas: minReplicas,
ResourceMax: resourceMax, MemoryPercent: memoryPercent, CPUPercent: cpuPercent,
})
plan.CurrentTotal += current
minTotal += minReplicas
resourceMaxTotal += resourceMax
if site.MemoryPressure || memoryPercent >= float64(input.MemoryHardPercent) {
plan.FrozenReason = "node_memory_pressure"
}
}
if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
workerReplicaBudget := max(
(input.DatabaseConnectionBudget-input.NonWorkerConnectionBudget)/input.WorkerDatabasePoolMax,
0,
)
capSiteResourceMaxima(plan.Sites, workerReplicaBudget)
resourceMaxTotal = 0
for _, site := range plan.Sites {
resourceMaxTotal += site.ResourceMax
}
if workerReplicaBudget < minTotal && plan.FrozenReason == "" {
plan.FrozenReason = "database_connection_budget"
}
}
desired := max(rawDesired, minTotal)
desired = min(desired, resourceMaxTotal)
if input.SynchronousDatabasePeers < 1 && plan.FrozenReason == "" {
plan.FrozenReason = "database_not_synchronous"
}
if desired > plan.CurrentTotal {
if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
additional := max(
(input.DatabaseConnectionBudget-input.DatabaseConnections)/input.WorkerDatabasePoolMax,
0,
)
desired = min(desired, plan.CurrentTotal+additional)
}
if !input.ScaleUpEligible || plan.FrozenReason != "" ||
(input.DatabaseConnectionBudget > 0 && input.DatabaseConnections >= input.DatabaseConnectionBudget) {
desired = plan.CurrentTotal
if plan.FrozenReason == "" {
if !input.ScaleUpEligible {
plan.FrozenReason = "scale_up_window"
} else {
plan.FrozenReason = "database_connection_budget"
}
}
} else {
desired = min(desired, max(plan.CurrentTotal*2, plan.CurrentTotal+2))
}
} else if desired < plan.CurrentTotal && !input.ScaleDownEligible {
desired = plan.CurrentTotal
if plan.FrozenReason == "" {
plan.FrozenReason = "scale_down_stabilization"
}
}
desired = max(desired, minTotal)
plan.DesiredTotal = desired
distributeDesiredReplicas(plan.Sites, desired)
return plan
}
func capSiteResourceMaxima(sites []SitePlan, total int) {
minimum := 0
original := make(map[string]int, len(sites))
for index := range sites {
original[sites[index].Site] = sites[index].ResourceMax
sites[index].ResourceMax = sites[index].DesiredReplicas
minimum += sites[index].ResourceMax
}
target := max(total, minimum)
assigned := minimum
for assigned < target {
sort.SliceStable(sites, func(left, right int) bool {
leftRoom := original[sites[left].Site] - sites[left].ResourceMax
rightRoom := original[sites[right].Site] - sites[right].ResourceMax
if leftRoom != rightRoom {
return leftRoom > rightRoom
}
return sites[left].Site < sites[right].Site
})
if original[sites[0].Site] <= sites[0].ResourceMax {
break
}
sites[0].ResourceMax++
assigned++
}
}
func siteResourceMaximum(site SiteResources, memoryTargetPercent int, cpuTargetPercent int) (int, float64, float64) {
if len(site.Nodes) > 0 {
if site.WorkerRequestMemoryBytes <= 0 || site.WorkerRequestMilliCPU <= 0 {
return 0, 0, 0
}
memoryMax := 0
cpuMax := 0
memoryPercent := float64(0)
cpuPercent := float64(0)
for _, node := range site.Nodes {
nodeMemoryPercent := usagePercent(node.UsedMemoryBytes, node.AllocatableMemoryBytes)
nodeCPUPercent := usagePercent(node.UsedMilliCPU, node.AllocatableMilliCPU)
memoryPercent = max(memoryPercent, nodeMemoryPercent)
cpuPercent = max(cpuPercent, nodeCPUPercent)
nonWorkerMemory := max(
node.UsedMemoryBytes-node.WorkerUsedMemoryBytes,
0,
)
memoryBudget := node.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorkerMemory
memoryMax += int(max(memoryBudget, 0) / site.WorkerRequestMemoryBytes)
nonWorkerCPU := max(
node.UsedMilliCPU-node.WorkerUsedMilliCPU,
0,
)
cpuBudget := node.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorkerCPU
cpuMax += int(max(cpuBudget, 0) / site.WorkerRequestMilliCPU)
}
return min(site.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
}
memoryPercent := usagePercent(site.UsedMemoryBytes, site.AllocatableMemoryBytes)
cpuPercent := usagePercent(site.UsedMilliCPU, site.AllocatableMilliCPU)
memoryMax := site.MaxReplicas
if site.AllocatableMemoryBytes > 0 && site.WorkerRequestMemoryBytes > 0 {
nonWorker := max(site.UsedMemoryBytes-int64(site.CurrentReplicas)*site.WorkerRequestMemoryBytes, 0)
budget := site.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorker
memoryMax = int(max(budget, 0) / site.WorkerRequestMemoryBytes)
}
cpuMax := site.MaxReplicas
if site.AllocatableMilliCPU > 0 && site.WorkerRequestMilliCPU > 0 {
nonWorker := max(site.UsedMilliCPU-int64(site.CurrentReplicas)*site.WorkerRequestMilliCPU, 0)
budget := site.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorker
cpuMax = int(max(budget, 0) / site.WorkerRequestMilliCPU)
}
return min(site.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
}
func distributeDesiredReplicas(sites []SitePlan, desired int) {
if len(sites) == 0 {
return
}
assigned := 0
for index := range sites {
assigned += sites[index].DesiredReplicas
}
for assigned < desired {
sort.SliceStable(sites, func(left, right int) bool {
leftRoom := sites[left].ResourceMax - sites[left].DesiredReplicas
rightRoom := sites[right].ResourceMax - sites[right].DesiredReplicas
if leftRoom != rightRoom {
return leftRoom > rightRoom
}
if sites[left].DesiredReplicas != sites[right].DesiredReplicas {
return sites[left].DesiredReplicas < sites[right].DesiredReplicas
}
return sites[left].Site < sites[right].Site
})
if sites[0].DesiredReplicas >= sites[0].ResourceMax {
break
}
sites[0].DesiredReplicas++
assigned++
}
sort.Slice(sites, func(left, right int) bool { return sites[left].Site < sites[right].Site })
}
func usagePercent(used int64, allocatable int64) float64 {
if allocatable <= 0 {
return 0
}
return float64(used) * 100 / float64(allocatable)
}
@@ -0,0 +1,164 @@
package capacitycontroller
import (
"encoding/json"
"strings"
"testing"
)
func TestCalculatePlanRespectsResourceDatabaseAndStepLimits(t *testing.T) {
input := PlanInput{
Queued: 1000, Running: 48, InstanceSlots: 24,
MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
DatabaseConnections: 80, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Sites: []SiteResources{
{
Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 6 << 30,
WorkerRequestMemoryBytes: 1 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 2000, WorkerRequestMilliCPU: 500,
},
{
Site: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
WorkerRequestMemoryBytes: 1 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 500, WorkerRequestMilliCPU: 500,
},
},
}
plan := CalculatePlan(input)
if plan.RawDesired != 22 {
t.Fatalf("raw desired=%d, want 22", plan.RawDesired)
}
if plan.DesiredTotal != 4 {
t.Fatalf("desired total=%d, want step-limited 4: %+v", plan.DesiredTotal, plan)
}
if plan.Sites[0].Site != "hongkong" || plan.Sites[0].DesiredReplicas != 3 {
t.Fatalf("site allocation=%+v, want hongkong=3 ningbo=1", plan.Sites)
}
}
func TestCalculatePlanFreezesScaleUpOnHardMemoryPressure(t *testing.T) {
plan := CalculatePlan(PlanInput{
Queued: 200, InstanceSlots: 24,
DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Sites: []SiteResources{{
Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 7 << 30,
WorkerRequestMemoryBytes: 1 << 30,
}},
})
if plan.DesiredTotal != 1 || plan.FrozenReason != "node_memory_pressure" {
t.Fatalf("plan=%+v, want frozen at one replica", plan)
}
}
func TestCalculatePlanReportsConfiguredMinimumOutsideResourceBudget(t *testing.T) {
plan := CalculatePlan(PlanInput{
Queued: 200, InstanceSlots: 24,
DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Sites: []SiteResources{{
Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 13 << 29,
WorkerRequestMemoryBytes: 2 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 2800, WorkerRequestMilliCPU: 500,
}},
})
if plan.DesiredTotal != 1 || plan.FrozenReason != "site_resource_budget" {
t.Fatalf("plan=%+v, want the existing minimum preserved and expansion frozen", plan)
}
}
func TestCalculatePlanWaitsForSafeScaleDownWindow(t *testing.T) {
input := PlanInput{
InstanceSlots: 24, DatabaseConnections: 20, DatabaseConnectionBudget: 150,
WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1,
ScaleDownEligible: false,
Sites: []SiteResources{
{Site: "ningbo", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
{Site: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
},
}
plan := CalculatePlan(input)
if plan.DesiredTotal != 4 || plan.FrozenReason != "scale_down_stabilization" {
t.Fatalf("plan=%+v, want current replicas during stabilization", plan)
}
input.ScaleDownEligible = true
plan = CalculatePlan(input)
if plan.DesiredTotal != 2 {
t.Fatalf("desired total=%d, want min total 2", plan.DesiredTotal)
}
}
func TestCalculatePlanAddsCapacityAcrossLabeledSiteNodes(t *testing.T) {
plan := CalculatePlan(PlanInput{
Queued: 500, InstanceSlots: 24,
MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
DatabaseConnections: 20, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 20,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Sites: []SiteResources{{
Site: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 8,
WorkerRequestMemoryBytes: 1 << 30, WorkerRequestMilliCPU: 500,
Nodes: []NodeResources{
{
NodeName: "hk-a", CurrentReplicas: 1,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 1000,
},
{
NodeName: "hk-b", CurrentReplicas: 1,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 2 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 500,
},
},
}},
})
if plan.Sites[0].ResourceMax != 7 {
t.Fatalf("resource max=%d, want database-budgeted multi-node maximum: %+v", plan.Sites[0].ResourceMax, plan)
}
if plan.DesiredTotal != 4 {
t.Fatalf("desired=%d, want two-wave step from 2 to 4", plan.DesiredTotal)
}
}
func TestCalculatePlanCapsReplicaMaximumByDeclaredDatabasePools(t *testing.T) {
plan := CalculatePlan(PlanInput{
Queued: 500, InstanceSlots: 24,
DatabaseConnections: 20, DatabaseConnectionBudget: 150,
NonWorkerConnectionBudget: 72, WorkerDatabasePoolMax: 32,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Sites: []SiteResources{
{Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
{Site: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
},
})
if plan.DesiredTotal != 2 {
t.Fatalf("desired=%d, want the current 1+1 database-safe topology: %+v", plan.DesiredTotal, plan)
}
resourceMax := 0
for _, site := range plan.Sites {
resourceMax += site.ResourceMax
}
if resourceMax != 2 {
t.Fatalf("resource maximum=%d, want floor((150-72)/32)=2: %+v", resourceMax, plan)
}
}
func TestPlanJSONUsesAcceptanceContractFieldNames(t *testing.T) {
payload, err := json.Marshal(Plan{
RawDesired: 3,
Sites: []SitePlan{{Site: "hongkong", ResourceMax: 2}},
})
if err != nil {
t.Fatal(err)
}
text := string(payload)
for _, field := range []string{`"rawDesired":3`, `"sites"`, `"site":"hongkong"`, `"resourceMax":2`} {
if !strings.Contains(text, field) {
t.Fatalf("plan JSON %s does not contain %s", text, field)
}
}
}
+1
View File
@@ -34,6 +34,7 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error
return Response{}, err
}
req.Header.Set("Content-Type", "application/json")
applyUpstreamIdempotency(req, request)
responseStartedAt := time.Now()
if err := notifySubmissionStarted(request); err != nil {
return Response{}, err
+10
View File
@@ -20,6 +20,7 @@ type Request struct {
HTTPClient *http.Client
RemoteTaskID string
RemoteTaskPayload map[string]any
UpstreamIdempotencyKey string
OnRemoteTaskSubmitted func(remoteTaskID string, payload map[string]any) error
OnRemoteTaskPolled func(remoteTaskID string, payload map[string]any) error
OnUpstreamSubmissionStarted func() error
@@ -34,6 +35,15 @@ type Request struct {
PreviousResponseTurns []ResponseTurn
}
func applyUpstreamIdempotency(request *http.Request, input Request) {
if request == nil {
return
}
if key := strings.TrimSpace(input.UpstreamIdempotencyKey); key != "" {
request.Header.Set("Idempotency-Key", key)
}
}
// WireResponse preserves the provider-facing response independently from the
// canonical result used by billing, retries, and cross-protocol conversion.
// Headers are filtered before this value leaves the clients package.
+2
View File
@@ -44,6 +44,7 @@ func (c VolcesClient) runImage(ctx context.Context, request Request, apiKey stri
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
applyUpstreamIdempotency(req, request)
responseStartedAt := time.Now()
if err := notifySubmissionStarted(request); err != nil {
@@ -232,6 +233,7 @@ func (c VolcesClient) postJSON(ctx context.Context, request Request, baseURL str
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
applyUpstreamIdempotency(req, request)
if err := notifySubmissionStarted(request); err != nil {
return nil, "", nil, err
}
+142 -16
View File
@@ -59,6 +59,8 @@ type Config struct {
BillingEngineMode string
ProcessRole string
DatabaseMaxConns int
DatabaseCriticalMaxConns int
DatabaseRiverMaxConns int
DatabaseMinIdleConns int
DatabaseMaxConnIdleSeconds int
DatabaseIdleInTransactionTimeoutSeconds int
@@ -76,6 +78,27 @@ type Config struct {
AsyncWorkerHardLimit int
AsyncWorkerInstanceHardLimit int
AsyncWorkerRefreshIntervalSeconds int
WorkerAutoscalingEnabled bool
WorkerReplicasNingbo int
WorkerReplicasHongkong int
WorkerMinReplicasNingbo int
WorkerMinReplicasHongkong int
WorkerMaxReplicasNingbo int
WorkerMaxReplicasHongkong int
WorkerTargetOutstandingPerReplica int
WorkerScaleUpWindowSeconds int
WorkerScaleDownStabilizationSeconds int
WorkerDrainTimeoutSeconds int
NodeMemoryTargetPercent int
NodeMemoryHardPercent int
NodeCPUTargetPercent int
PostgresConnectionBudget int
PostgresNonWorkerConnectionBudget int
WorkerDatabaseMaxConns int
CapacityControllerNamespace string
CapacityControllerAPIServer string
CapacityControllerTokenFile string
CapacityControllerCAFile string
}
func Load() Config {
@@ -131,6 +154,8 @@ func Load() Config {
BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")),
ProcessRole: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_PROCESS_ROLE", "all"))),
DatabaseMaxConns: envInt("AI_GATEWAY_DATABASE_MAX_CONNS", 0),
DatabaseCriticalMaxConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS", 0),
DatabaseRiverMaxConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_RIVER_MAX_CONNS", 0),
DatabaseMinIdleConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS", 0),
DatabaseMaxConnIdleSeconds: envOptionalIntValidated("AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS", 0),
@@ -140,27 +165,57 @@ func Load() Config {
),
DatabaseLockTimeoutSeconds: envOptionalIntValidated("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", 30),
MediaRequestConcurrency: envIntValidated("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", 16),
MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8),
MediaOSSDirectEnabled: env("AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED", "false") == "true",
MediaOSSEndpoint: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_ENDPOINT", ""), "/"),
MediaOSSBucket: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_BUCKET", "")),
MediaOSSAccessKeyID: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID", "")),
MediaOSSAccessKeySecret: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET", "")),
MediaOSSPublicBaseURL: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL", ""), "/"),
MediaOSSObjectPrefix: strings.Trim(env("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX", "easyai-ai-gateway/media"), "/"),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
AsyncWorkerHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
MediaRequestConcurrency: envIntValidated("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", 16),
MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8),
MediaOSSDirectEnabled: env("AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED", "false") == "true",
MediaOSSEndpoint: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_ENDPOINT", ""), "/"),
MediaOSSBucket: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_BUCKET", "")),
MediaOSSAccessKeyID: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID", "")),
MediaOSSAccessKeySecret: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET", "")),
MediaOSSPublicBaseURL: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL", ""), "/"),
MediaOSSObjectPrefix: strings.Trim(env("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX", "easyai-ai-gateway/media"), "/"),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
AsyncWorkerHardLimit: envIntValidated(
"AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT",
envOptionalIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
),
AsyncWorkerInstanceHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", 32),
AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5),
WorkerAutoscalingEnabled: env("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "false") == "true",
WorkerReplicasNingbo: envOptionalIntValidated("AI_GATEWAY_WORKER_REPLICAS_NINGBO", 1),
WorkerReplicasHongkong: envOptionalIntValidated("AI_GATEWAY_WORKER_REPLICAS_HONGKONG", 1),
WorkerMinReplicasNingbo: envOptionalIntValidated("AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO", 1),
WorkerMinReplicasHongkong: envOptionalIntValidated("AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG", 1),
WorkerMaxReplicasNingbo: envOptionalIntValidated("AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO", 1),
WorkerMaxReplicasHongkong: envOptionalIntValidated("AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG", 1),
WorkerTargetOutstandingPerReplica: envOptionalIntValidated("AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA", 0),
WorkerScaleUpWindowSeconds: envIntValidated("AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS", 20),
WorkerScaleDownStabilizationSeconds: envIntValidated(
"AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS",
600,
),
WorkerDrainTimeoutSeconds: envIntValidated("AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS", 600),
NodeMemoryTargetPercent: envIntValidated("AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT", 75),
NodeMemoryHardPercent: envIntValidated("AI_GATEWAY_NODE_MEMORY_HARD_PERCENT", 85),
NodeCPUTargetPercent: envIntValidated("AI_GATEWAY_NODE_CPU_TARGET_PERCENT", 70),
PostgresConnectionBudget: envIntValidated("AI_GATEWAY_POSTGRES_CONNECTION_BUDGET", 150),
PostgresNonWorkerConnectionBudget: envIntValidated(
"AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET",
72,
),
WorkerDatabaseMaxConns: envIntValidated("AI_GATEWAY_WORKER_DATABASE_MAX_CONNS", 32),
CapacityControllerNamespace: env("POD_NAMESPACE", env("AI_GATEWAY_CAPACITY_CONTROLLER_NAMESPACE", "easyai")),
CapacityControllerAPIServer: env("AI_GATEWAY_CAPACITY_CONTROLLER_API_SERVER", "https://kubernetes.default.svc"),
CapacityControllerTokenFile: env("AI_GATEWAY_CAPACITY_CONTROLLER_TOKEN_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/token"),
CapacityControllerCAFile: env("AI_GATEWAY_CAPACITY_CONTROLLER_CA_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"),
}
}
func (c Config) Validate() error {
switch strings.ToLower(strings.TrimSpace(c.ProcessRole)) {
case "", "all", "api", "worker":
case "", "all", "api", "worker", "capacity-controller":
default:
return errors.New("AI_GATEWAY_PROCESS_ROLE must be all, api, or worker")
return errors.New("AI_GATEWAY_PROCESS_ROLE must be all, api, worker, or capacity-controller")
}
if c.DatabaseMaxConns < 0 || c.DatabaseMaxConns > 1000 {
return errors.New("AI_GATEWAY_DATABASE_MAX_CONNS must be between 1 and 1000 when configured")
@@ -169,6 +224,19 @@ func (c Config) Validate() error {
(c.DatabaseMaxConns > 0 && c.DatabaseMinIdleConns > c.DatabaseMaxConns) {
return errors.New("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS must be between 0 and AI_GATEWAY_DATABASE_MAX_CONNS")
}
if c.DatabaseCriticalMaxConns < 0 || c.DatabaseCriticalMaxConns > 64 ||
(c.DatabaseMaxConns > 0 && c.DatabaseCriticalMaxConns >= c.DatabaseMaxConns) {
return errors.New("AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS must be between 0 and less than AI_GATEWAY_DATABASE_MAX_CONNS")
}
if c.DatabaseRiverMaxConns < 0 || c.DatabaseRiverMaxConns > 256 ||
(c.DatabaseMaxConns > 0 &&
c.DatabaseCriticalMaxConns+c.DatabaseRiverMaxConns >= c.DatabaseMaxConns) {
return errors.New("critical and River database pools must leave at least one execution connection")
}
if c.DatabaseMaxConns > 0 &&
c.DatabaseMinIdleConns > c.DatabaseMaxConns-c.DatabaseCriticalMaxConns-c.DatabaseRiverMaxConns {
return errors.New("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS must not exceed the execution pool after reserving critical and River connections")
}
if c.DatabaseMaxConnIdleSeconds < 0 || c.DatabaseMaxConnIdleSeconds > 3600 {
return errors.New("AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS must be between 0 and 3600")
}
@@ -215,6 +283,50 @@ func (c Config) Validate() error {
if c.AsyncWorkerRefreshIntervalSeconds < 1 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive")
}
if c.WorkerReplicasNingbo < 0 || c.WorkerReplicasHongkong < 0 ||
c.WorkerMinReplicasNingbo < 0 || c.WorkerMinReplicasHongkong < 0 ||
c.WorkerMaxReplicasNingbo < c.WorkerMinReplicasNingbo ||
c.WorkerMaxReplicasHongkong < c.WorkerMinReplicasHongkong ||
c.WorkerMaxReplicasNingbo > 64 || c.WorkerMaxReplicasHongkong > 64 {
return errors.New("Worker replica configuration must be between 0 and 64 with min not greater than max")
}
if c.WorkerAutoscalingEnabled &&
(c.WorkerReplicasNingbo < c.WorkerMinReplicasNingbo || c.WorkerReplicasNingbo > c.WorkerMaxReplicasNingbo ||
c.WorkerReplicasHongkong < c.WorkerMinReplicasHongkong || c.WorkerReplicasHongkong > c.WorkerMaxReplicasHongkong) {
return errors.New("Worker bootstrap replicas must be within the autoscaling min/max range")
}
if c.WorkerTargetOutstandingPerReplica < 0 || c.WorkerTargetOutstandingPerReplica > 100000 {
return errors.New("AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA must be between 0 and 100000")
}
if c.WorkerScaleUpWindowSeconds != 0 && (c.WorkerScaleUpWindowSeconds < 5 || c.WorkerScaleUpWindowSeconds > 600) {
return errors.New("AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS must be between 5 and 600")
}
if c.WorkerScaleDownStabilizationSeconds != 0 &&
(c.WorkerScaleDownStabilizationSeconds < 60 || c.WorkerScaleDownStabilizationSeconds > 86400) {
return errors.New("AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS must be between 60 and 86400")
}
if c.WorkerDrainTimeoutSeconds != 0 && (c.WorkerDrainTimeoutSeconds < 60 || c.WorkerDrainTimeoutSeconds > 86400) {
return errors.New("AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS must be between 60 and 86400")
}
if (c.NodeMemoryTargetPercent != 0 || c.NodeMemoryHardPercent != 0) &&
(c.NodeMemoryTargetPercent < 50 || c.NodeMemoryTargetPercent > 90 ||
c.NodeMemoryHardPercent <= c.NodeMemoryTargetPercent || c.NodeMemoryHardPercent > 95) {
return errors.New("Worker node memory target/hard percentages must be ordered within 50..95")
}
if c.NodeCPUTargetPercent != 0 && (c.NodeCPUTargetPercent < 40 || c.NodeCPUTargetPercent > 90) {
return errors.New("AI_GATEWAY_NODE_CPU_TARGET_PERCENT must be between 40 and 90")
}
if c.PostgresConnectionBudget < 0 || c.PostgresConnectionBudget > 1000 {
return errors.New("AI_GATEWAY_POSTGRES_CONNECTION_BUDGET must be between 1 and 1000")
}
if c.PostgresNonWorkerConnectionBudget < 0 ||
(c.PostgresConnectionBudget > 0 &&
c.PostgresNonWorkerConnectionBudget >= c.PostgresConnectionBudget) {
return errors.New("AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET must be below the PostgreSQL connection budget")
}
if c.WorkerDatabaseMaxConns < 0 || c.WorkerDatabaseMaxConns > 256 {
return errors.New("AI_GATEWAY_WORKER_DATABASE_MAX_CONNS must be between 1 and 256")
}
if c.TaskProgressCallbackTimeoutMS != 0 && (c.TaskProgressCallbackTimeoutMS < 100 || c.TaskProgressCallbackTimeoutMS > 60000) {
return errors.New("TASK_PROGRESS_CALLBACK_TIMEOUT_MS must be between 100 and 60000")
}
@@ -319,12 +431,17 @@ func (c Config) EffectiveProcessRole() string {
}
func (c Config) RunsPublicHTTP() bool {
return c.EffectiveProcessRole() != "worker"
switch c.EffectiveProcessRole() {
case "worker", "capacity-controller":
return false
default:
return true
}
}
func (c Config) RunsAsyncExecutionWorker() bool {
switch c.EffectiveProcessRole() {
case "api":
case "api", "capacity-controller":
return false
case "worker":
return true
@@ -334,7 +451,16 @@ func (c Config) RunsAsyncExecutionWorker() bool {
}
func (c Config) RunsBackgroundWorkers() bool {
return c.EffectiveProcessRole() != "api"
switch c.EffectiveProcessRole() {
case "api", "capacity-controller":
return false
default:
return true
}
}
func (c Config) RunsCapacityController() bool {
return c.EffectiveProcessRole() == "capacity-controller"
}
type GlobalHTTPProxyStatus struct {
+57
View File
@@ -140,6 +140,57 @@ func TestProcessRolePrecedenceAndCompatibility(t *testing.T) {
if !cfg.RunsPublicHTTP() || cfg.RunsAsyncExecutionWorker() || cfg.RunsBackgroundWorkers() {
t.Fatalf("api role capabilities are inconsistent: %+v", cfg)
}
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "capacity-controller")
cfg = Load()
if cfg.RunsPublicHTTP() || cfg.RunsAsyncExecutionWorker() || cfg.RunsBackgroundWorkers() || !cfg.RunsCapacityController() {
t.Fatalf("capacity controller role capabilities are inconsistent: %+v", cfg)
}
}
func TestWorkerAutoscalingConfiguration(t *testing.T) {
t.Setenv("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "true")
t.Setenv("AI_GATEWAY_WORKER_REPLICAS_NINGBO", "1")
t.Setenv("AI_GATEWAY_WORKER_REPLICAS_HONGKONG", "2")
t.Setenv("AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO", "1")
t.Setenv("AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG", "1")
t.Setenv("AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO", "2")
t.Setenv("AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG", "3")
cfg := Load()
if err := cfg.Validate(); err != nil {
t.Fatalf("valid autoscaling configuration was rejected: %v", err)
}
if !cfg.WorkerAutoscalingEnabled || cfg.WorkerReplicasHongkong != 2 || cfg.WorkerMaxReplicasHongkong != 3 {
t.Fatalf("autoscaling configuration=%+v", cfg)
}
cfg.WorkerReplicasHongkong = 4
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "bootstrap") {
t.Fatalf("Validate() error=%v, want bootstrap bound failure", err)
}
}
func TestLoadPrefersWorkerGlobalHardLimitAlias(t *testing.T) {
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "48")
t.Setenv("AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT", "96")
cfg := Load()
if cfg.AsyncWorkerHardLimit != 96 {
t.Fatalf("global worker hard limit=%d, want 96", cfg.AsyncWorkerHardLimit)
}
}
func TestValidatePostgresReplicaBudgets(t *testing.T) {
cfg := Load()
if cfg.PostgresConnectionBudget != 150 || cfg.PostgresNonWorkerConnectionBudget != 72 {
t.Fatalf(
"PostgreSQL budgets=%d/%d, want 150/72",
cfg.PostgresConnectionBudget,
cfg.PostgresNonWorkerConnectionBudget,
)
}
cfg.PostgresNonWorkerConnectionBudget = cfg.PostgresConnectionBudget
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "NON_WORKER") {
t.Fatalf("Validate() error=%v, want non-Worker budget failure", err)
}
}
func TestValidateProcessRoleAndDatabasePool(t *testing.T) {
@@ -171,6 +222,12 @@ func TestValidateProcessRoleAndDatabasePool(t *testing.T) {
t.Fatalf("Validate() error = %v, want invalid database min idle conns", err)
}
cfg.DatabaseMinIdleConns = 4
cfg.DatabaseCriticalMaxConns = 4
cfg.DatabaseRiverMaxConns = 12
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "River") {
t.Fatalf("Validate() error = %v, want database pool budget failure", err)
}
cfg.DatabaseRiverMaxConns = 4
t.Setenv("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS", "not-an-integer")
if err := Load().Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MIN_IDLE_CONNS") {
t.Fatalf("Validate() error = %v, want invalid non-integer database min idle conns", err)
@@ -28,6 +28,13 @@ type taskTrafficError struct {
Err error
}
func (s *Server) acceptanceStore() *store.Store {
if s.coordinationStore != nil {
return s.coordinationStore
}
return s.store
}
func (e *taskTrafficError) Error() string {
return e.Message
}
@@ -37,7 +44,7 @@ func (e *taskTrafficError) ErrorCode() string {
}
func (s *Server) authorizeTaskTraffic(r *http.Request, user *auth.User) (taskTrafficAdmission, error) {
runID, err := s.store.AuthorizeAcceptanceTask(
runID, err := s.acceptanceStore().AuthorizeAcceptanceTask(
r.Context(),
r.Header.Get(acceptanceRunHeader),
r.Header.Get(acceptanceTokenHeader),
@@ -121,7 +128,7 @@ func writeTaskTrafficError(w http.ResponseWriter, err error, protocol string) {
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/traffic-mode [get]
func (s *Server) getGatewayTrafficMode(w http.ResponseWriter, r *http.Request) {
mode, err := s.store.GetGatewayTrafficMode(r.Context())
mode, err := s.acceptanceStore().GetGatewayTrafficMode(r.Context())
if err != nil {
writeError(w, http.StatusInternalServerError, "get gateway traffic mode failed")
return
@@ -129,6 +136,45 @@ func (s *Server) getGatewayTrafficMode(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, mode)
}
// pauseGatewayTraffic godoc
// @Summary 上线后硬门禁异常时通过 CAS 暂停正式新任务
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param body body store.PauseGatewayTrafficInput true "当前线上 release 与 revision CAS"
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/traffic-mode/pause [post]
func (s *Server) pauseGatewayTraffic(w http.ResponseWriter, r *http.Request) {
var input store.PauseGatewayTrafficInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
mode, err := s.acceptanceStore().PauseGatewayTraffic(r.Context(), input)
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, mode)
}
// listCapacityProfiles godoc
// @Summary 获取经生产同构验收认证的容量配置
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Success 200 {array} store.CapacityProfile
// @Router /api/admin/system/acceptance/capacity-profiles [get]
func (s *Server) listCapacityProfiles(w http.ResponseWriter, r *http.Request) {
profiles, err := s.acceptanceStore().ListCapacityProfiles(r.Context())
if err != nil {
writeError(w, http.StatusInternalServerError, "list capacity profiles failed")
return
}
writeJSON(w, http.StatusOK, profiles)
}
// createAcceptanceRun godoc
// @Summary 创建生产同构验收 Run
// @Tags acceptance
@@ -144,7 +190,7 @@ func (s *Server) createAcceptanceRun(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
run, err := s.store.CreateAcceptanceRun(r.Context(), input)
run, err := s.acceptanceStore().CreateAcceptanceRun(r.Context(), input)
if err != nil {
writeError(w, http.StatusBadRequest, err.Error())
return
@@ -161,7 +207,7 @@ func (s *Server) createAcceptanceRun(w http.ResponseWriter, r *http.Request) {
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID} [get]
func (s *Server) getAcceptanceRun(w http.ResponseWriter, r *http.Request) {
run, err := s.store.GetAcceptanceRun(r.Context(), r.PathValue("runID"))
run, err := s.acceptanceStore().GetAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "acceptance run not found")
@@ -182,7 +228,7 @@ func (s *Server) getAcceptanceRun(w http.ResponseWriter, r *http.Request) {
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/runs/{runID}/activate [post]
func (s *Server) activateAcceptanceRun(w http.ResponseWriter, r *http.Request) {
mode, err := s.store.ActivateAcceptanceRun(r.Context(), r.PathValue("runID"))
mode, err := s.acceptanceStore().ActivateAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
writeAcceptanceMutationError(w, err)
return
@@ -190,6 +236,36 @@ func (s *Server) activateAcceptanceRun(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, mode)
}
// stageAcceptanceCapacityProfiles godoc
// @Summary 为当前 validation Run 暂存 80% 容量门禁
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Param body body store.StageAcceptanceCapacityProfilesInput true "待验证容量配置"
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID}/capacity-profiles [post]
func (s *Server) stageAcceptanceCapacityProfiles(w http.ResponseWriter, r *http.Request) {
var input store.StageAcceptanceCapacityProfilesInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
input.RunID = r.PathValue("runID")
run, err := s.acceptanceStore().StageAcceptanceCapacityProfiles(r.Context(), input)
if err != nil {
switch {
case store.IsNotFound(err):
writeError(w, http.StatusConflict, "acceptance run is not running")
default:
writeError(w, http.StatusBadRequest, err.Error())
}
return
}
writeJSON(w, http.StatusOK, run)
}
// finishAcceptanceRun godoc
// @Summary 记录验收门禁结果
// @Tags acceptance
@@ -207,7 +283,7 @@ func (s *Server) finishAcceptanceRun(w http.ResponseWriter, r *http.Request) {
return
}
input.RunID = r.PathValue("runID")
run, err := s.store.FinishAcceptanceRun(r.Context(), input)
run, err := s.acceptanceStore().FinishAcceptanceRun(r.Context(), input)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusConflict, "acceptance run is not running")
@@ -228,7 +304,7 @@ func (s *Server) finishAcceptanceRun(w http.ResponseWriter, r *http.Request) {
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID}/retry [post]
func (s *Server) retryAcceptanceRun(w http.ResponseWriter, r *http.Request) {
run, err := s.store.RetryAcceptanceRun(r.Context(), r.PathValue("runID"))
run, err := s.acceptanceStore().RetryAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
writeAcceptanceMutationError(w, err)
return
@@ -253,7 +329,7 @@ func (s *Server) promoteAcceptanceRun(w http.ResponseWriter, r *http.Request) {
return
}
input.RunID = r.PathValue("runID")
mode, err := s.store.PromoteAcceptanceRun(r.Context(), input)
mode, err := s.acceptanceStore().PromoteAcceptanceRun(r.Context(), input)
if err != nil {
writeAcceptanceMutationError(w, err)
return
@@ -278,7 +354,7 @@ func (s *Server) abortAcceptanceRun(w http.ResponseWriter, r *http.Request) {
return
}
input.RunID = r.PathValue("runID")
mode, err := s.store.AbortAcceptanceRun(r.Context(), input)
mode, err := s.acceptanceStore().AbortAcceptanceRun(r.Context(), input)
if err != nil {
writeAcceptanceMutationError(w, err)
return
@@ -0,0 +1,165 @@
package httpapi
import (
"context"
"testing"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
type acceptanceStatusSnapshot struct {
ID string
Status string
}
type acceptanceUserGroupSnapshot struct {
ID string
RateLimitPolicy []byte
}
// isolateHTTPAcceptanceState lets a stress test disable unrelated runtime
// resources without leaking those mutations into later integration tests that
// intentionally share the same dedicated test database.
func isolateHTTPAcceptanceState(
t *testing.T,
ctx context.Context,
pool *pgxpool.Pool,
includeStorage bool,
) func() {
t.Helper()
startedAt := time.Now()
platforms := readAcceptanceStatusSnapshots(t, ctx, pool, `
SELECT id::text, status
FROM integration_platforms
WHERE deleted_at IS NULL`)
userGroups := make([]acceptanceUserGroupSnapshot, 0)
rows, err := pool.Query(ctx, `
SELECT id::text, rate_limit_policy
FROM gateway_user_groups`)
if err != nil {
t.Fatalf("snapshot acceptance user groups: %v", err)
}
for rows.Next() {
var item acceptanceUserGroupSnapshot
if err := rows.Scan(&item.ID, &item.RateLimitPolicy); err != nil {
rows.Close()
t.Fatalf("scan acceptance user group snapshot: %v", err)
}
userGroups = append(userGroups, item)
}
if err := rows.Err(); err != nil {
rows.Close()
t.Fatalf("read acceptance user group snapshots: %v", err)
}
rows.Close()
storageChannels := []acceptanceStatusSnapshot(nil)
if includeStorage {
storageChannels = readAcceptanceStatusSnapshots(t, ctx, pool, `
SELECT id::text, status
FROM file_storage_channels
WHERE deleted_at IS NULL`)
}
if _, err := pool.Exec(ctx, `
UPDATE integration_platforms
SET status = 'disabled'
WHERE deleted_at IS NULL`); err != nil {
t.Fatalf("isolate acceptance platforms: %v", err)
}
if includeStorage {
if _, err := pool.Exec(ctx, `
UPDATE file_storage_channels
SET status = 'disabled'
WHERE deleted_at IS NULL`); err != nil {
t.Fatalf("isolate acceptance storage channels: %v", err)
}
}
return func() {
restoreCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
tx, err := pool.Begin(restoreCtx)
if err != nil {
t.Errorf("begin acceptance state restore: %v", err)
return
}
defer func() {
_ = tx.Rollback(restoreCtx)
}()
for _, item := range platforms {
if _, err := tx.Exec(restoreCtx, `
UPDATE integration_platforms
SET status = $2
WHERE id = $1::uuid`, item.ID, item.Status); err != nil {
t.Errorf("restore acceptance platform %s: %v", item.ID, err)
return
}
}
if _, err := tx.Exec(restoreCtx, `
UPDATE integration_platforms
SET status = 'disabled'
WHERE created_at >= $1
AND deleted_at IS NULL`, startedAt); err != nil {
t.Errorf("disable acceptance-created platforms: %v", err)
return
}
for _, item := range userGroups {
if _, err := tx.Exec(restoreCtx, `
UPDATE gateway_user_groups
SET rate_limit_policy = $2::jsonb
WHERE id = $1::uuid`, item.ID, string(item.RateLimitPolicy)); err != nil {
t.Errorf("restore acceptance user group %s: %v", item.ID, err)
return
}
}
if includeStorage {
for _, item := range storageChannels {
if _, err := tx.Exec(restoreCtx, `
UPDATE file_storage_channels
SET status = $2
WHERE id = $1::uuid`, item.ID, item.Status); err != nil {
t.Errorf("restore acceptance storage channel %s: %v", item.ID, err)
return
}
}
if _, err := tx.Exec(restoreCtx, `
UPDATE file_storage_channels
SET status = 'disabled'
WHERE created_at >= $1
AND deleted_at IS NULL`, startedAt); err != nil {
t.Errorf("disable acceptance-created storage channels: %v", err)
return
}
}
if err := tx.Commit(restoreCtx); err != nil {
t.Errorf("commit acceptance state restore: %v", err)
}
}
}
func readAcceptanceStatusSnapshots(
t *testing.T,
ctx context.Context,
pool *pgxpool.Pool,
query string,
) []acceptanceStatusSnapshot {
t.Helper()
rows, err := pool.Query(ctx, query)
if err != nil {
t.Fatalf("snapshot acceptance statuses: %v", err)
}
defer rows.Close()
items := make([]acceptanceStatusSnapshot, 0)
for rows.Next() {
var item acceptanceStatusSnapshot
if err := rows.Scan(&item.ID, &item.Status); err != nil {
t.Fatalf("scan acceptance status snapshot: %v", err)
}
items = append(items, item)
}
if err := rows.Err(); err != nil {
t.Fatalf("read acceptance status snapshots: %v", err)
}
return items
}
@@ -40,9 +40,8 @@ func TestAsyncSubmissionReservesConcurrencyOnlyAfterWorkerClaim(t *testing.T) {
t.Fatalf("connect acceptance pool: %v", err)
}
defer pool.Close()
if _, err := pool.Exec(ctx, `UPDATE integration_platforms SET status = 'disabled' WHERE deleted_at IS NULL`); err != nil {
t.Fatalf("isolate acceptance platforms: %v", err)
}
restoreAcceptanceState := isolateHTTPAcceptanceState(t, ctx, pool, false)
defer restoreAcceptanceState()
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
@@ -116,18 +115,27 @@ func TestAsyncWorkerThousandConcurrentSchedulingAcceptance(t *testing.T) {
t.Fatalf("connect Worker store: %v", err)
}
defer workerDB.Close()
apiCoordinationDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 4)
if err != nil {
t.Fatalf("connect API critical store: %v", err)
}
defer apiCoordinationDB.Close()
workerCoordinationDB, err := store.ConnectWithMaxConns(ctx, databaseURL, 4)
if err != nil {
t.Fatalf("connect Worker critical store: %v", err)
}
defer workerCoordinationDB.Close()
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect acceptance pool: %v", err)
}
defer pool.Close()
if _, err := pool.Exec(ctx, `UPDATE integration_platforms SET status = 'disabled' WHERE deleted_at IS NULL`); err != nil {
t.Fatalf("isolate acceptance platforms: %v", err)
}
restoreAcceptanceState := isolateHTTPAcceptanceState(t, ctx, pool, false)
defer restoreAcceptanceState()
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
server := httptest.NewServer(NewServerWithStores(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
@@ -138,10 +146,11 @@ func TestAsyncWorkerThousandConcurrentSchedulingAcceptance(t *testing.T) {
ProcessRole: "api",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 1000,
AsyncWorkerInstanceHardLimit: 1000,
AsyncWorkerRefreshIntervalSeconds: 1,
}, apiDB, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))))
}, apiDB, apiCoordinationDB, apiDB, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))))
defer server.Close()
workerServer := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
workerServer := httptest.NewServer(NewServerWithStores(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
@@ -152,8 +161,9 @@ func TestAsyncWorkerThousandConcurrentSchedulingAcceptance(t *testing.T) {
ProcessRole: "worker",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 1000,
AsyncWorkerInstanceHardLimit: 1000,
AsyncWorkerRefreshIntervalSeconds: 1,
}, workerDB, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))))
}, workerDB, workerCoordinationDB, workerDB, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))))
defer workerServer.Close()
adminToken := createAsyncAcceptanceAdmin(t, ctx, pool, server.URL)
@@ -180,7 +190,7 @@ WHERE status = 'active'`); err != nil {
}
assertAcceptanceAttempts(t, ctx, pool, taskIDs, len(taskIDs))
t.Logf(
"千级并发证据: tasks=%d duration=%s running_peak=%d lease_peak=%d api_database_max_connections=16 worker_database_max_connections=24",
"千级并发证据: tasks=%d duration=%s running_peak=%d lease_peak=%d api_database_max_connections=16 api_critical_max_connections=4 worker_database_max_connections=24 worker_critical_max_connections=4",
len(taskIDs),
45*time.Second,
peakRunning,
@@ -207,9 +217,8 @@ func TestAsyncWorkerDynamicConcurrencyAcceptance(t *testing.T) {
t.Fatalf("connect acceptance pool: %v", err)
}
defer pool.Close()
if _, err := pool.Exec(ctx, `UPDATE integration_platforms SET status = 'disabled' WHERE deleted_at IS NULL`); err != nil {
t.Fatalf("isolate acceptance platforms: %v", err)
}
restoreAcceptanceState := isolateHTTPAcceptanceState(t, ctx, pool, false)
defer restoreAcceptanceState()
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
@@ -223,6 +232,7 @@ func TestAsyncWorkerDynamicConcurrencyAcceptance(t *testing.T) {
CORSAllowedOrigin: "*",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 256,
AsyncWorkerInstanceHardLimit: 256,
AsyncWorkerRefreshIntervalSeconds: 1,
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
@@ -43,15 +43,19 @@ func TestCoreLocalFlow(t *testing.T) {
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
handler := NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
BillingEngineMode: "enforce",
TaskProgressCallbackEnabled: true,
TaskProgressCallbackURL: "http://callback.local/task-progress",
CORSAllowedOrigin: "*",
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
BillingEngineMode: "enforce",
TaskProgressCallbackEnabled: true,
TaskProgressCallbackURL: "http://callback.local/task-progress",
CORSAllowedOrigin: "*",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 32,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 1,
}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
server := httptest.NewServer(handler)
defer server.Close()
@@ -606,8 +610,8 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
"model": defaultImageModel,
"runMode": "simulation",
"prompt": "replace background with clean studio light",
"image": "https://example.com/source.png",
"mask": "https://example.com/mask.png",
"image": server.URL + "/static/simulation/image.png",
"mask": server.URL + "/static/simulation/image-edit.png",
"simulation": true,
"simulationDurationMs": 5,
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &imageEditResponse)
@@ -728,18 +732,35 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
t.Fatalf("voice clone persisted platform should match task attempt: task=%+v voice=%+v", voiceCloneTaskDetail, clonedVoice)
}
doubaoLiteImageEditModel := "doubao-5.0-lite图像编辑"
var doubaoLitePlatformModel struct {
asyncImageEditModel := "async-image-edit-smoke-" + suffixText
asyncImageEditCanonicalKey := "openai:" + asyncImageEditModel
var asyncImageEditBaseModel struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/catalog/base-models", loginResponse.AccessToken, map[string]any{
"providerKey": "openai",
"canonicalModelKey": asyncImageEditCanonicalKey,
"invocationName": asyncImageEditModel,
"providerModelName": asyncImageEditModel,
"modelType": []string{"image_edit", "image_generate"},
"displayName": "Async Image Edit Smoke",
"capabilities": map[string]any{"originalTypes": []string{"image_edit", "image_generate"}},
"metadata": map[string]any{"source": "test"},
}, http.StatusCreated, &asyncImageEditBaseModel)
if asyncImageEditBaseModel.ID == "" {
t.Fatal("async image edit smoke base model was not created")
}
var asyncImageEditPlatformModel struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "easyai:doubao-5.0-lite图像编辑",
"modelName": doubaoLiteImageEditModel,
"modelAlias": doubaoLiteImageEditModel,
"canonicalModelKey": asyncImageEditCanonicalKey,
"modelName": asyncImageEditModel,
"modelAlias": asyncImageEditModel,
"modelType": []string{"image_edit", "image_generate"},
"displayName": doubaoLiteImageEditModel,
}, http.StatusCreated, &doubaoLitePlatformModel)
var doubaoLiteAsyncEdit struct {
"displayName": "Async Image Edit Smoke",
}, http.StatusCreated, &asyncImageEditPlatformModel)
var asyncImageEdit struct {
TaskID string `json:"taskId"`
Task struct {
ID string `json:"id"`
@@ -748,20 +769,20 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
} `json:"task"`
}
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/images/edits", apiKeyResponse.Secret, map[string]any{
"model": doubaoLiteImageEditModel,
"model": asyncImageEditModel,
"runMode": "simulation",
"prompt": "turn the attached bright desktop object into a clean product-style render",
"image": "https://example.com/doubao-lite-source.png",
"image": server.URL + "/static/simulation/image.png",
"size": "2048x2048",
"simulation": true,
"simulationDurationMs": 5,
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &doubaoLiteAsyncEdit)
if doubaoLiteAsyncEdit.TaskID == "" || !doubaoLiteAsyncEdit.Task.AsyncMode {
t.Fatalf("doubao-5.0-lite image edit async task should be accepted: %+v", doubaoLiteAsyncEdit)
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &asyncImageEdit)
if asyncImageEdit.TaskID == "" || !asyncImageEdit.Task.AsyncMode {
t.Fatalf("async image edit task should be accepted: %+v", asyncImageEdit)
}
doubaoLiteAsyncEditDone := waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, doubaoLiteAsyncEdit.TaskID, []string{"succeeded"}, 5*time.Second)
if doubaoLiteAsyncEditDone.Status != "succeeded" {
t.Fatalf("doubao-5.0-lite image edit async task should succeed through river queue, got %+v", doubaoLiteAsyncEditDone)
asyncImageEditDone := waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, asyncImageEdit.TaskID, []string{"succeeded"}, 5*time.Second)
if asyncImageEditDone.Status != "succeeded" {
t.Fatalf("async image edit task should succeed through river queue, got %+v", asyncImageEditDone)
}
var gptImageModelTypesRaw []byte
@@ -782,11 +803,12 @@ LIMIT 1`).Scan(&gptImageModelTypesRaw); err != nil {
}
deniedModel := "permission-deny-smoke-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, deniedModel)
var deniedPlatformModel struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + deniedModel,
"modelName": deniedModel,
"modelAlias": deniedModel,
"modelType": []string{"text_generate"},
@@ -830,11 +852,12 @@ LIMIT 1`).Scan(&gptImageModelTypesRaw); err != nil {
}
controlledModel := "permission-allow-smoke-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, controlledModel)
var controlledPlatformModel struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + controlledModel,
"modelName": controlledModel,
"modelAlias": controlledModel,
"modelType": []string{"text_generate"},
@@ -909,9 +932,10 @@ LIMIT 1`).Scan(&gptImageModelTypesRaw); err != nil {
},
}, http.StatusCreated, &customPricingRuleSet)
pricingModel := "pricing-smoke-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, pricingModel)
var pricingPlatformModel map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + pricingModel,
"modelName": pricingModel,
"modelAlias": pricingModel,
"modelType": []string{"text_generate"},
@@ -1156,12 +1180,19 @@ WHERE reference_type = 'gateway_task'
"rules": []map[string]any{{"metric": "rpm", "limit": 1, "windowSeconds": rateLimitWindowSeconds}},
},
}, http.StatusCreated, &rateLimitPolicySet)
createSimulationBaseModel(
t,
server.URL,
loginResponse.AccessToken,
rateLimitedModel,
[]string{"text_generate", "image_generate"},
)
var rateLimitPlatformModel map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + rateLimitedModel,
"modelName": rateLimitedModel,
"modelAlias": rateLimitedModel,
"modelType": []string{"text_generate"},
"modelType": []string{"text_generate", "image_generate"},
"displayName": "Rate Limit Smoke",
"runtimePolicySetId": rateLimitPolicySet.ID,
"runtimePolicyOverride": map[string]any{},
@@ -1223,12 +1254,12 @@ WHERE reference_type = 'gateway_task'
AsyncMode bool `json:"asyncMode"`
} `json:"task"`
}
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/responses", apiKeyResponse.Secret, map[string]any{
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/images/generations", apiKeyResponse.Secret, map[string]any{
"model": rateLimitedModel,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
"input": "async queued",
"prompt": "async queued",
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &asyncRateLimitTask)
if asyncRateLimitTask.TaskID == "" || asyncRateLimitTask.Task.ID != asyncRateLimitTask.TaskID || !asyncRateLimitTask.Task.AsyncMode {
t.Fatalf("async task response should expose task id and async mode: %+v", asyncRateLimitTask)
@@ -1272,9 +1303,16 @@ WHERE reference_type = 'gateway_task'
}
videoRouteModel := "video-route-smoke-" + suffixText
createSimulationBaseModel(
t,
server.URL,
loginResponse.AccessToken,
videoRouteModel,
[]string{"video_generate", "image_to_video"},
)
var videoRoutePlatformModel map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + videoRouteModel,
"modelName": videoRouteModel,
"modelAlias": videoRouteModel,
"modelType": []string{"video_generate", "image_to_video"},
@@ -1313,7 +1351,7 @@ WHERE reference_type = 'gateway_task'
"simulation": true,
"simulationDurationMs": 5,
"prompt": "image to video route",
"image": "https://example.com/source.png",
"image": server.URL + "/static/simulation/image.png",
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &imageToVideoTask)
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, imageToVideoTask.Task.ID, []string{"succeeded"}, 10*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+imageToVideoTask.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &imageToVideoTask.Task)
@@ -1333,8 +1371,11 @@ WHERE reference_type = 'gateway_task'
if len(preprocessingDetail.Items) == 0 {
t.Fatalf("task preprocessing endpoint should expose persisted preprocessing logs: %+v", preprocessingDetail)
}
if preprocessingDetail.Items[0]["actualInput"] == nil || preprocessingDetail.Items[0]["convertedOutput"] == nil {
t.Fatalf("preprocessing log should store actual input and converted output: %+v", preprocessingDetail.Items)
if preprocessingDetail.Items[0]["changes"] == nil {
t.Fatalf("preprocessing log should retain its compact change summary: %+v", preprocessingDetail.Items)
}
if preprocessingDetail.Items[0]["actualInput"] != nil || preprocessingDetail.Items[0]["convertedOutput"] != nil {
t.Fatalf("history compaction should remove full preprocessing payloads: %+v", preprocessingDetail.Items)
}
var preprocessingRows int
if err := testPool.QueryRow(ctx, `SELECT count(*) FROM gateway_task_param_preprocessing_logs WHERE task_id = $1::uuid`, imageToVideoTask.Task.ID).Scan(&preprocessingRows); err != nil {
@@ -1345,6 +1386,9 @@ WHERE reference_type = 'gateway_task'
}
failoverModel := "phase1-failover-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, failoverModel)
unrelatedPriorityModel := "priority-demote-unrelated-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, unrelatedPriorityModel)
var failedPlatform struct {
ID string `json:"id"`
}
@@ -1383,16 +1427,16 @@ WHERE reference_type = 'gateway_task'
}, http.StatusCreated, &unrelatedPriorityPlatform)
var unrelatedPriorityPlatformModel map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+unrelatedPriorityPlatform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"modelName": "priority-demote-unrelated-" + suffixText,
"modelAlias": "priority-demote-unrelated-" + suffixText,
"canonicalModelKey": "openai:" + unrelatedPriorityModel,
"modelName": unrelatedPriorityModel,
"modelAlias": unrelatedPriorityModel,
"modelType": []string{"text_generate"},
"displayName": "Unrelated Priority",
}, http.StatusCreated, &unrelatedPriorityPlatformModel)
for _, platformID := range []string{failedPlatform.ID, successPlatform.ID} {
var platformModel map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platformID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + failoverModel,
"modelName": failoverModel,
"modelAlias": failoverModel,
"modelType": []string{"text_generate"},
@@ -1442,8 +1486,8 @@ WHERE reference_type = 'gateway_task'
if failoverDetail.Attempts[2].PlatformName != "OpenAI Retry Success" || failoverDetail.Attempts[2].Status != "succeeded" || failoverDetail.Attempts[2].ResponseMS <= 0 {
t.Fatalf("third failover attempt should preserve successful platform: %+v", failoverDetail.Attempts[2])
}
if summary, ok := failoverDetail.Metrics["attempts"].([]any); !ok || len(summary) != 3 {
t.Fatalf("task metrics should keep attempt-chain summary, got %+v", failoverDetail.Metrics)
if _, ok := failoverDetail.Metrics["attempts"]; ok {
t.Fatalf("task metrics should not duplicate the dedicated attempt chain: %+v", failoverDetail.Metrics)
}
var demotedDynamicPriority int
var successEffectivePriority int
@@ -1479,6 +1523,7 @@ WHERE failed.id = $1::uuid`, failedPlatform.ID, successPlatform.ID, unrelatedPri
},
}, http.StatusCreated, &degradePolicySet)
degradeModel := "degrade-smoke-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, degradeModel)
var degradedPlatform struct {
ID string `json:"id"`
}
@@ -1514,7 +1559,7 @@ WHERE failed.id = $1::uuid`, failedPlatform.ID, successPlatform.ID, unrelatedPri
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+item.platformID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + degradeModel,
"modelName": degradeModel,
"modelAlias": degradeModel,
"modelType": []string{"text_generate"},
@@ -1568,6 +1613,7 @@ WHERE m.platform_id = $1::uuid
},
}, http.StatusCreated, &autoDisablePolicySet)
autoDisableModel := "auto-disable-smoke-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, autoDisableModel)
var invalidKeyPlatform struct {
ID string `json:"id"`
}
@@ -1582,7 +1628,7 @@ WHERE m.platform_id = $1::uuid
}, http.StatusCreated, &invalidKeyPlatform)
var invalidKeyPlatformModel map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+invalidKeyPlatform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + autoDisableModel,
"modelName": autoDisableModel,
"modelAlias": autoDisableModel,
"modelType": []string{"text_generate"},
@@ -1625,7 +1671,7 @@ WHERE m.platform_id = $1::uuid
"priority": 60,
}, http.StatusCreated, &validKeyPlatform)
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+validKeyPlatform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + autoDisableModel,
"modelName": autoDisableModel,
"modelAlias": autoDisableModel,
"modelType": []string{"text_generate"},
@@ -1719,14 +1765,31 @@ WHERE id = $1::uuid`, taskResponse.Task.ID); err != nil {
} `json:"items"`
Total int `json:"total"`
}
var taskPlatformID string
var taskBillingStatus string
if err := testPool.QueryRow(ctx, `
SELECT platform_id::text
FROM gateway_task_attempts
WHERE task_id = $1::uuid
AND status = 'succeeded'
ORDER BY attempt_no DESC
LIMIT 1`, taskResponse.Task.ID).Scan(&taskPlatformID); err != nil {
t.Fatalf("read task platform for admin filters: %v", err)
}
if err := testPool.QueryRow(ctx, `
SELECT billing_status
FROM gateway_tasks
WHERE id = $1::uuid`, taskResponse.Task.ID).Scan(&taskBillingStatus); err != nil {
t.Fatalf("read task billing status for admin filters: %v", err)
}
adminTaskQuery := "/api/admin/tasks?q=" + taskResponse.Task.ID +
"&userId=" + smokeGatewayUserID +
"&status=succeeded" +
"&platformId=" + platform.ID +
"&platformId=" + taskPlatformID +
"&model=" + defaultTextModel +
"&modelType=text_generate" +
"&runMode=simulation" +
"&billingStatus=settled" +
"&billingStatus=" + taskBillingStatus +
"&apiKey=smoke" +
"&pageSize=10"
doJSON(t, server.URL, http.MethodGet, adminTaskQuery, loginResponse.AccessToken, nil, http.StatusOK, &adminTaskList)
@@ -1771,9 +1834,6 @@ WHERE id = $1::uuid`, taskResponse.Task.ID); err != nil {
if !bytes.Contains(body, []byte("task.running")) {
t.Fatalf("events response should include running transition event body=%s", string(body))
}
if !bytes.Contains(body, []byte("task.progress")) {
t.Fatalf("events response should include progress events body=%s", string(body))
}
if !bytes.Contains(body, []byte("task.billing.settled")) {
t.Fatalf("events response should include billing settlement event body=%s", string(body))
}
@@ -1789,8 +1849,10 @@ WHERE id = $1::uuid`, taskResponse.Task.ID); err != nil {
}
defer resp.Body.Close()
body, _ = io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK || !bytes.Contains(body, []byte("task.retrying")) {
t.Fatalf("failover events should include retrying event status=%d body=%s", resp.StatusCode, string(body))
if resp.StatusCode != http.StatusOK ||
!bytes.Contains(body, []byte("task.attempt.failed")) ||
!bytes.Contains(body, []byte("task.policy.priority_demoted")) {
t.Fatalf("failover events should include the failed attempt and applied policy status=%d body=%s", resp.StatusCode, string(body))
}
assertLoadAvoidanceSimulatedRetryChain(t, ctx, testPool, server.URL, loginResponse.AccessToken, apiKeyResponse.Secret, suffixText)
@@ -1803,18 +1865,6 @@ WHERE id = $1::uuid`, taskResponse.Task.ID); err != nil {
t.Fatal("task progress callback outbox should receive events")
}
restartModel := "worker-restart-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, restartModel)
createSimulationTextPlatformModel(
t,
server.URL,
loginResponse.AccessToken,
"openai-worker-restart-"+suffixText,
"OpenAI Worker Restart",
restartModel,
1,
nil,
)
var restartAsyncTask struct {
TaskID string `json:"taskId"`
Task struct {
@@ -1823,12 +1873,13 @@ WHERE id = $1::uuid`, taskResponse.Task.ID); err != nil {
AsyncMode bool `json:"asyncMode"`
} `json:"task"`
}
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/responses", apiKeyResponse.Secret, map[string]any{
"model": restartModel,
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/images/generations", apiKeyResponse.Secret, map[string]any{
"model": defaultImageModel,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 2000,
"input": "river worker restart",
"prompt": "river worker restart",
"size": "1024x1024",
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &restartAsyncTask)
if restartAsyncTask.TaskID == "" || !restartAsyncTask.Task.AsyncMode {
t.Fatalf("restart async task should be accepted as async: %+v", restartAsyncTask)
@@ -1841,15 +1892,19 @@ WHERE id = $1::uuid`, taskResponse.Task.ID); err != nil {
serverCtx2, cancelServer2 := context.WithCancel(ctx)
defer cancelServer2()
server2 := httptest.NewServer(NewServerWithContext(serverCtx2, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
BillingEngineMode: "enforce",
TaskProgressCallbackEnabled: true,
TaskProgressCallbackURL: "http://callback.local/task-progress",
CORSAllowedOrigin: "*",
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
BillingEngineMode: "enforce",
TaskProgressCallbackEnabled: true,
TaskProgressCallbackURL: "http://callback.local/task-progress",
CORSAllowedOrigin: "*",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 32,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 1,
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server2.Close()
restartRecovered := waitForTaskStatus(t, server2.URL, apiKeyResponse.Secret, restartAsyncTask.TaskID, []string{"succeeded"}, 8*time.Second)
@@ -2334,15 +2389,20 @@ type cacheAffinityTaskDetail struct {
}
func createSimulationTextBaseModel(t *testing.T, baseURL string, adminToken string, model string) {
t.Helper()
createSimulationBaseModel(t, baseURL, adminToken, model, []string{"text_generate"})
}
func createSimulationBaseModel(t *testing.T, baseURL string, adminToken string, model string, modelTypes []string) {
t.Helper()
doJSON(t, baseURL, http.MethodPost, "/api/admin/catalog/base-models", adminToken, map[string]any{
"providerKey": "openai",
"canonicalModelKey": "openai:" + model,
"invocationName": model,
"providerModelName": model,
"modelType": []string{"text_generate"},
"modelType": modelTypes,
"displayName": model,
"capabilities": map[string]any{"originalTypes": []string{"text_generate"}},
"capabilities": map[string]any{"originalTypes": modelTypes},
"metadata": map[string]any{"source": "cache-affinity-integration-test"},
}, http.StatusCreated, nil)
}
@@ -2481,6 +2541,7 @@ LIMIT 1`, key, value).Scan(&taskID)
func assertLoadAvoidanceSimulatedRetryChain(t *testing.T, ctx context.Context, testPool *pgxpool.Pool, baseURL string, adminToken string, runtimeToken string, suffixText string) {
t.Helper()
model := "load-avoidance-smoke-" + suffixText
createSimulationTextBaseModel(t, baseURL, adminToken, model)
type scenario struct {
keySuffix string
name string
@@ -2522,7 +2583,7 @@ func assertLoadAvoidanceSimulatedRetryChain(t *testing.T, ctx context.Context, t
ID string `json:"id"`
}
payload := map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + model,
"modelName": model,
"providerModelName": model,
"modelAlias": model,
@@ -2581,55 +2642,29 @@ func assertLoadAvoidanceSimulatedRetryChain(t *testing.T, ctx context.Context, t
Metrics map[string]any `json:"metrics"`
}
doJSON(t, baseURL, http.MethodGet, "/api/v1/tasks/"+taskResponse.Task.ID, runtimeToken, nil, http.StatusOK, &detail)
if detail.Status != "failed" || len(detail.Attempts) != len(scenarios) {
t.Fatalf("load avoidance detail should expose every attempted client, got status=%s attempts=%+v", detail.Status, detail.Attempts)
if detail.Status != "failed" || len(detail.Attempts) != 1 {
t.Fatalf("non-retryable failure should stop before load-avoided fallbacks, got status=%s attempts=%+v", detail.Status, detail.Attempts)
}
expected := []struct {
name string
code string
retryable bool
avoided bool
}{
{name: "Load Avoidance Hard Stop", code: "bad_request"},
{name: "Load Avoidance Retryable", code: "server_error", retryable: true},
{name: "Load Avoidance Full Rate Limit", code: "rate_limit", retryable: true, avoided: true},
{name: "Load Avoidance Full Overloaded", code: "overloaded", retryable: true, avoided: true},
{name: "Load Avoidance Full Fatal", code: "bad_request", avoided: true},
attempt := detail.Attempts[0]
if attempt.AttemptNo != 1 ||
attempt.PlatformName != "Load Avoidance Hard Stop" ||
attempt.Status != "failed" ||
attempt.ErrorCode != "bad_request" ||
attempt.Retryable {
t.Fatalf("unexpected hard-stop attempt: %+v", attempt)
}
attemptSummary := make([]string, 0, len(detail.Attempts))
for index, want := range expected {
got := detail.Attempts[index]
attemptSummary = append(attemptSummary, got.PlatformName+":"+got.Status+":"+got.ErrorCode)
if got.AttemptNo != index+1 || got.PlatformName != want.name || got.Status != "failed" || got.ErrorCode != want.code || got.Retryable != want.retryable {
t.Fatalf("unexpected load avoidance attempt %d: got %+v want %+v", index+1, got, want)
}
if boolFromTestMap(got.Metrics, "loadAvoided") != want.avoided {
t.Fatalf("loadAvoided mismatch for %s metrics=%+v", got.PlatformName, got.Metrics)
}
if !want.avoided && floatFromTestAny(got.Metrics["loadRatio"]) != 0 {
t.Fatalf("non-full candidate should not carry load pressure, got %s metrics=%+v", got.PlatformName, got.Metrics)
}
if want.avoided {
if ratio := floatFromTestAny(got.Metrics["loadRatio"]); ratio < 1 {
t.Fatalf("avoided candidate should expose full load ratio, got %s ratio=%v metrics=%+v", got.PlatformName, ratio, got.Metrics)
}
loadMetrics, _ := got.Metrics["loadMetrics"].(map[string]any)
concurrent, _ := loadMetrics["concurrent"].(map[string]any)
if ratio := floatFromTestAny(concurrent["ratio"]); ratio < 1 {
t.Fatalf("avoided candidate should expose concurrent saturation, got %s loadMetrics=%+v", got.PlatformName, loadMetrics)
}
if queued := floatFromTestAny(loadMetrics["queued"]); queued < 1 {
t.Fatalf("avoided candidate should expose queued waiting load, got %s loadMetrics=%+v", got.PlatformName, loadMetrics)
}
}
if boolFromTestMap(attempt.Metrics, "loadAvoided") {
t.Fatalf("selected non-full candidate must not be marked load avoided: %+v", attempt.Metrics)
}
if !attemptTraceHasReason(detail.Attempts[0].Metrics, "load_avoidance_fallback") {
t.Fatalf("first hard-stop attempt should continue because avoided candidates remain, metrics=%+v", detail.Attempts[0].Metrics)
if ratio := floatFromTestAny(attempt.Metrics["loadRatio"]); ratio != 0 {
t.Fatalf("selected non-full candidate should not carry load pressure, got ratio=%v metrics=%+v", ratio, attempt.Metrics)
}
if summary, ok := detail.Metrics["attempts"].([]any); !ok || len(summary) != len(scenarios) {
t.Fatalf("final task metrics should preserve load-avoidance attempt summary, got %+v", detail.Metrics)
if !attemptTraceHasReason(attempt.Metrics, "failover_action_rule") {
t.Fatalf("hard-stop attempt should record the failure-policy decision, metrics=%+v", attempt.Metrics)
}
if _, ok := detail.Metrics["attempts"]; ok {
t.Fatalf("task metrics should not duplicate the dedicated attempt chain: %+v", detail.Metrics)
}
t.Logf("load avoidance retry chain: %s", strings.Join(attemptSummary, " -> "))
}
func seedQueuedConcurrencyLoad(t *testing.T, ctx context.Context, testPool *pgxpool.Pool, platformKey string, model string, platformModelID string) {
@@ -124,13 +124,17 @@ func newFailurePolicyAcceptanceFixture(t *testing.T) *failurePolicyAcceptanceFix
t.Fatalf("connect acceptance pool: %v", err)
}
handler := NewServerWithContext(ctx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "failure-policy-http-acceptance-secret",
BillingEngineMode: "observe",
CORSAllowedOrigin: "*",
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "failure-policy-http-acceptance-secret",
BillingEngineMode: "observe",
CORSAllowedOrigin: "*",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 32,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 1,
}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
server := httptest.NewServer(handler)
fixture := &failurePolicyAcceptanceFixture{
@@ -74,9 +74,9 @@ func TestGeminiBase64ImageEditThousandSynchronousRequests(t *testing.T) {
t.Fatalf("connect observation pool: %v", err)
}
defer pool.Close()
restoreAcceptanceState := isolateHTTPAcceptanceState(t, ctx, pool, true)
defer restoreAcceptanceState()
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},
@@ -192,6 +192,11 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
}
created, err := s.store.CreateTaskIdempotent(r.Context(), createInput, user)
if err != nil {
if errors.Is(err, store.ErrRateLimited) {
applyRunErrorHeaders(w, err)
writeGeminiTaskError(statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err))
return
}
if errors.Is(err, store.ErrIdempotencyKeyReused) {
writeGeminiTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused")
return
+48 -2
View File
@@ -178,9 +178,45 @@ func (s *Server) logPostgresUnavailable(message string) {
"postgres_pool_idle_connections", statistics.IdleConns(),
"postgres_pool_empty_acquire_count", statistics.EmptyAcquireCount(),
"postgres_pool_canceled_acquire_count", statistics.CanceledAcquireCount(),
"postgres_critical_pool", s.criticalPostgresPoolLogFields(),
"postgres_river_pool", s.riverPostgresPoolLogFields(),
)
}
func (s *Server) criticalPostgresPoolLogFields() map[string]any {
if s.coordinationStore == nil || s.coordinationStore == s.store ||
s.coordinationStore.Pool() == nil {
return map[string]any{"separate": false}
}
statistics := s.coordinationStore.Pool().Stat()
return map[string]any{
"separate": true,
"max_connections": statistics.MaxConns(),
"total_connections": statistics.TotalConns(),
"acquired_connections": statistics.AcquiredConns(),
"idle_connections": statistics.IdleConns(),
"empty_acquire_count": statistics.EmptyAcquireCount(),
"canceled_acquire_count": statistics.CanceledAcquireCount(),
}
}
func (s *Server) riverPostgresPoolLogFields() map[string]any {
if s.riverStore == nil || s.riverStore == s.store ||
s.riverStore.Pool() == nil {
return map[string]any{"separate": false}
}
statistics := s.riverStore.Pool().Stat()
return map[string]any{
"separate": true,
"max_connections": statistics.MaxConns(),
"total_connections": statistics.TotalConns(),
"acquired_connections": statistics.AcquiredConns(),
"idle_connections": statistics.IdleConns(),
"empty_acquire_count": statistics.EmptyAcquireCount(),
"canceled_acquire_count": statistics.CanceledAcquireCount(),
}
}
func (s *Server) localIdentityEnabled() bool {
mode := strings.ToLower(strings.TrimSpace(s.cfg.IdentityMode))
return mode == "" || mode == "standalone" || mode == "hybrid"
@@ -1183,6 +1219,11 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
}
created, err := s.store.CreateTaskIdempotent(r.Context(), createInput, user)
if err != nil {
if errors.Is(err, store.ErrRateLimited) {
applyRunErrorHeaders(w, err)
writeTaskError(statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err))
return
}
if errors.Is(err, store.ErrIdempotencyKeyReused) {
writeTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused")
return
@@ -1200,8 +1241,9 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
if s.billingMetrics != nil {
s.billingMetrics.ObserveBillingEvent("idempotent_replay")
}
if targetProtocol == "" {
if targetProtocol == "" || gatewayAPIV1Request(r) {
w.Header().Set("Idempotent-Replayed", "true")
w.Header().Set("X-Gateway-Task-Id", task.ID)
}
if responsePlan.streamMode {
writeTaskError(http.StatusConflict, "streaming idempotent replay is not supported", nil, "idempotency_stream_replay_unsupported")
@@ -1411,7 +1453,7 @@ func writeCompatibleTaskResponse(runCtx context.Context, w http.ResponseWriter,
}
func writeProtocolCompatibleTaskResponse(runCtx context.Context, w http.ResponseWriter, r *http.Request, executor taskExecutor, kind string, model string, targetProtocol string, task store.GatewayTask, user *auth.User, streamMode bool, includeUsage bool) {
if targetProtocol == "" {
if targetProtocol == "" || gatewayAPIV1Request(r) {
w.Header().Set("X-Gateway-Task-Id", task.ID)
}
if streamMode {
@@ -1515,6 +1557,10 @@ func writeProtocolCompatibleTaskResponse(runCtx context.Context, w http.Response
writeJSON(w, http.StatusOK, easyAISynchronousTaskResponse(result.Task, result.Output))
}
func gatewayAPIV1Request(r *http.Request) bool {
return r != nil && strings.HasPrefix(r.URL.Path, "/api/v1/")
}
func streamIncludeUsage(body map[string]any) bool {
streamOptions, _ := body["stream_options"].(map[string]any)
includeUsage, _ := streamOptions["include_usage"].(bool)
@@ -131,12 +131,16 @@ func TestKelingOmniCompatibleHTTPFlow(t *testing.T) {
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
gateway := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 32,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 1,
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer gateway.Close()
@@ -152,7 +156,7 @@ func TestKelingOmniCompatibleHTTPFlow(t *testing.T) {
"aspect_ratio": "16:9",
"duration": "3",
"sound": "off",
"image_list": []any{map[string]any{"image_url": "https://example.com/reference.png"}},
"image_list": []any{map[string]any{"image_url": gateway.URL + "/static/simulation/image.png"}},
"external_task_id": "compat-http-1",
}, http.StatusOK, &created)
createdData, _ := created.Data.(map[string]any)
@@ -33,12 +33,16 @@ func TestKeling30TurboSimulationFlow(t *testing.T) {
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 32,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 1,
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
@@ -192,7 +196,7 @@ WHERE username = $1`, username); err != nil {
"simulation": true,
"simulationDurationMs": 5,
"prompt": "The subject looks toward the camera",
"image": "https://example.com/first.png",
"image": server.URL + "/static/simulation/image.png",
"duration": 5,
"resolution": "720p",
"audio": true,
@@ -159,6 +159,9 @@ func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, v
created, err := s.store.CreateTaskIdempotent(r.Context(), createInput, user)
if err != nil {
switch {
case errors.Is(err, store.ErrRateLimited):
applyRunErrorHeaders(w, err)
writeKlingCompatError(w, statusFromRunError(err), runErrorMessage(err), runErrorCode(err))
case errors.Is(err, store.ErrIdempotencyKeyReused):
writeKlingCompatError(w, http.StatusConflict, err.Error(), "idempotency_key_reused")
case errors.Is(err, store.ErrExternalTaskIDReused):
@@ -49,12 +49,16 @@ WHERE provider_key = 'keling'
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
AsyncQueueWorkerEnabled: true,
AsyncWorkerHardLimit: 32,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 1,
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
@@ -159,7 +163,7 @@ WHERE platform_id = $1::uuid
for _, model := range []string{klingO1Model, klingV3OmniModel} {
assertGenericVideoGeneration(model+"-text-to-video", model, "", "video_generate")
assertGenericVideoGeneration(model+"-image-to-video", model, "https://example.com/first.png", "image_to_video")
assertGenericVideoGeneration(model+"-image-to-video", model, server.URL+"/static/simulation/image.png", "image_to_video")
}
createV1 := func(model string, duration int, externalID string) string {
@@ -44,6 +44,25 @@ func TestOIDCJITFullGatewayUserFlowAndSecurityBoundary(t *testing.T) {
t.Fatalf("connect store: %v", err)
}
t.Cleanup(db.Close)
hasBreakGlass, err := db.HasBreakGlassManager(ctx)
if err != nil {
t.Fatalf("check OIDC JIT break-glass manager: %v", err)
}
if !hasBreakGlass {
manager, registerErr := db.RegisterLocalUser(ctx, store.LocalRegisterInput{
Username: "oidc-jit-manager-" + uuid.NewString(),
Password: uuid.NewString() + "-local-only",
})
if registerErr != nil {
t.Fatalf("create OIDC JIT break-glass manager: %v", registerErr)
}
if _, updateErr := db.Pool().Exec(ctx, `UPDATE gateway_users SET roles='["manager"]'::jsonb WHERE id=$1::uuid`, manager.ID); updateErr != nil {
t.Fatalf("grant OIDC JIT break-glass manager role: %v", updateErr)
}
t.Cleanup(func() {
_, _ = db.Pool().Exec(context.Background(), `DELETE FROM gateway_users WHERE id=$1::uuid`, manager.ID)
})
}
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
@@ -388,7 +407,7 @@ func createOIDCBFFSessionCookie(t *testing.T, baseURL string) *http.Cookie {
}
return nil
}}
response, err := client.Get(baseURL + "/api/v1/auth/oidc/login?returnTo=%2Fapi%2Fv1%2Fme")
response, err := client.Get(baseURL + "/api/v1/auth/oidc/login?contextType=tenant&returnTo=%2Fapi%2Fv1%2Fme")
if err != nil {
t.Fatalf("complete OIDC BFF login: %v", err)
}
@@ -463,7 +482,8 @@ func signedOIDCJITToken(t *testing.T, key *ecdsa.PrivateKey, issuer string, subj
t.Helper()
now := time.Now()
claims := jwt.MapClaims{
"iss": issuer, "aud": "gateway-api", "sub": subject, "tid": oidcJITTenantID,
"iss": issuer, "aud": "gateway-api", "sub": subject, "context_type": "tenant", "tid": oidcJITTenantID,
"client_id": "gateway-public-test",
"preferred_username": "oidc-jit-acceptance", "roles": []string{"gateway.admin"},
"scope": "openid gateway.access", "iat": now.Unix(), "nbf": now.Add(-time.Second).Unix(),
"exp": now.Add(time.Hour).Unix(),
+52 -2
View File
@@ -24,6 +24,8 @@ type Server struct {
ctx context.Context
cfg config.Config
store *store.Store
coordinationStore *store.Store
riverStore *store.Store
oidcUserResolver oidcUserResolver
auth *auth.Authenticator
oidcClient oidcPublicClient
@@ -72,6 +74,27 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
}
func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Store, logger *slog.Logger) http.Handler {
return NewServerWithContextAndCoordinationStore(ctx, cfg, db, db, logger)
}
func NewServerWithContextAndCoordinationStore(
ctx context.Context,
cfg config.Config,
db *store.Store,
coordinationDB *store.Store,
logger *slog.Logger,
) http.Handler {
return NewServerWithStores(ctx, cfg, db, coordinationDB, db, logger)
}
func NewServerWithStores(
ctx context.Context,
cfg config.Config,
db *store.Store,
coordinationDB *store.Store,
riverDB *store.Store,
logger *slog.Logger,
) http.Handler {
if cfg.MediaMaterializationConcurrency == 0 {
cfg.MediaMaterializationConcurrency = 8
}
@@ -83,9 +106,11 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
ctx: ctx,
cfg: cfg,
store: db,
coordinationStore: coordinationDB,
riverStore: riverDB,
oidcUserResolver: db,
auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken),
runner: runner.New(cfg, db, logger, securityEventMetrics),
runner: runner.NewWithStores(cfg, db, coordinationDB, riverDB, logger, securityEventMetrics),
logger: logger,
billingMetrics: securityEventMetrics,
requestAssetLocks: map[string]*requestAssetLock{},
@@ -122,7 +147,9 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.HandleFunc("GET /readyz", server.ready)
mux.HandleFunc("GET /api/v1/healthz", server.health)
mux.HandleFunc("GET /api/v1/readyz", server.ready)
mux.Handle("GET /metrics", securityEventMetrics.DynamicHandler(db))
mux.Handle("GET /metrics", securityEventMetrics.DynamicHandler(postgresMetricsProvider{
Store: db, critical: coordinationDB, river: riverDB,
}))
if !cfg.RunsPublicHTTP() {
return server.recover(mux)
}
@@ -221,13 +248,16 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("GET /api/admin/system/client-customization/settings", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getClientCustomizationSettings)))
mux.Handle("PATCH /api/admin/system/client-customization/settings", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.updateClientCustomizationSettings)))
mux.Handle("GET /api/admin/system/acceptance/traffic-mode", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getGatewayTrafficMode)))
mux.Handle("POST /api/admin/system/acceptance/traffic-mode/pause", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.pauseGatewayTraffic)))
mux.Handle("POST /api/admin/system/acceptance/runs", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.createAcceptanceRun)))
mux.Handle("GET /api/admin/system/acceptance/runs/{runID}", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/activate", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.activateAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/capacity-profiles", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.stageAcceptanceCapacityProfiles)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/finish", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.finishAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/retry", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.retryAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/promote", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.promoteAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/abort", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.abortAcceptanceRun)))
mux.Handle("GET /api/admin/system/acceptance/capacity-profiles", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listCapacityProfiles)))
mux.Handle("GET /api/admin/system/identity/configuration", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getIdentityConfiguration)))
mux.Handle("POST /api/admin/system/identity/pairings", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.startIdentityPairing)))
mux.Handle("GET /api/admin/system/identity/pairings/{pairingID}", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getIdentityPairing)))
@@ -348,6 +378,26 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
return server.recover(server.cors(server.protectOIDCSessionCookie(mux)))
}
type postgresMetricsProvider struct {
*store.Store
critical *store.Store
river *store.Store
}
func (provider postgresMetricsProvider) RiverPostgresPoolMetrics() store.PostgresPoolMetricsSnapshot {
if provider.river == nil || provider.river == provider.Store {
return store.PostgresPoolMetricsSnapshot{}
}
return provider.river.PostgresPoolMetrics()
}
func (provider postgresMetricsProvider) CriticalPostgresPoolMetrics() store.PostgresPoolMetricsSnapshot {
if provider.critical == nil || provider.critical == provider.Store {
return store.PostgresPoolMetricsSnapshot{}
}
return provider.critical.PostgresPoolMetrics()
}
func (server *Server) initializeIdentityRuntime(
ctx context.Context,
db *store.Store,
+7 -1
View File
@@ -163,7 +163,13 @@ func ApplyManifest(revision Revision, input ManifestApplication) (Revision, erro
revision.SchemaVersion = input.Manifest.SchemaVersion
revision.Issuer = strings.TrimRight(input.Manifest.Issuer, "/")
revision.TenantID = input.Manifest.TenantID
revision.TenantMode = input.Manifest.TenantMode
if input.Manifest.SchemaVersion == 1 {
// Manifest V1 intentionally omits tenant_mode on the wire, but the
// persisted revision uses the explicit single-tenant discriminator.
revision.TenantMode = "single_tenant"
} else {
revision.TenantMode = input.Manifest.TenantMode
}
revision.ApplicationID = input.Manifest.ApplicationID
revision.Audience = input.Manifest.Audience
revision.Scopes = append([]string(nil), input.Manifest.Scopes...)
@@ -226,3 +226,27 @@ func TestApplyManifestV2RemovesFixedTenantMappingWhileV1StillRequiresIt(t *testi
t.Fatal("Manifest V1 without a fixed local tenant mapping was accepted")
}
}
func TestApplyManifestV1PersistsSingleTenantMode(t *testing.T) {
draft, err := NewDraft(PairingInput{
AuthCenterURL: "https://auth.example.com", PublicBaseURL: "https://api.example.com",
WebBaseURL: "https://gateway.example.com", LocalTenantKey: "default",
}, "production")
if err != nil {
t.Fatal(err)
}
applied, err := ApplyManifest(draft, ManifestApplication{
AppEnv: "production",
Manifest: ManifestV1{
SchemaVersion: 1, Issuer: "https://auth.example.com/issuer/default",
TenantID: "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", ApplicationID: "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
Capabilities: []string{}, Scopes: []string{}, Clients: ManifestClients{},
},
})
if err != nil {
t.Fatal(err)
}
if applied.TenantMode != "single_tenant" {
t.Fatalf("Manifest V1 should persist the explicit single-tenant discriminator: %#v", applied)
}
}
+14 -1
View File
@@ -159,7 +159,7 @@ func (s *Service) buildTaskAdmissionPlanForCurrentBinding(
}
func (s *Service) withAsyncWorkerCapacityScope(ctx context.Context, scopes []store.AdmissionScope) ([]store.AdmissionScope, error) {
capacity, err := s.store.ActiveWorkerCapacity(ctx)
capacity, err := s.coordinationStore.ActiveWorkerCapacity(ctx)
if err != nil {
return nil, err
}
@@ -589,6 +589,19 @@ func (s *Service) SubmitAsyncTask(ctx context.Context, task store.GatewayTask) e
if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
return ctx.Err()
}
if store.ModelCandidateRetryAfter(err) > 0 ||
(errors.Is(err, store.ErrRateLimited) && store.RateLimitRetryable(err)) {
if enqueueErr := s.EnqueueAsyncTask(ctx, task); enqueueErr != nil {
_, _ = s.store.FailQueuedTask(
context.WithoutCancel(ctx),
task.ID,
"enqueue_failed",
enqueueErr.Error(),
)
return enqueueErr
}
return nil
}
_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error())
return err
}
+1 -1
View File
@@ -21,7 +21,7 @@ func (s *Service) renewTaskExecutionLease(ctx context.Context, cancel context.Ca
case <-ctx.Done():
return
case <-ticker.C:
if err := s.store.RenewTaskExecutionLease(ctx, taskID, executionToken, taskExecutionLeaseTTL); err != nil {
if err := s.coordinationStore.RenewTaskExecutionLease(ctx, taskID, executionToken, taskExecutionLeaseTTL); err != nil {
if errors.Is(err, store.ErrTaskExecutionFinished) {
return
}
+8 -8
View File
@@ -191,7 +191,7 @@ func retryAsyncQueueStartup(
}
func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error {
driver := riverpgxv5.New(s.store.Pool())
driver := riverpgxv5.New(s.riverStore.Pool())
migrator, err := rivermigrate.New(driver, nil)
if err != nil {
return err
@@ -286,7 +286,7 @@ func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx.
if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil {
return nil, err
}
return river.NewClient(riverpgxv5.New(s.store.Pool()), &river.Config{
return river.NewClient(riverpgxv5.New(s.riverStore.Pool()), &river.Config{
ID: fmt.Sprintf("%s-exec-%d-%d", s.workerInstanceID, capacity, time.Now().UnixNano()),
JobTimeout: -1,
Logger: s.logger,
@@ -314,11 +314,11 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke
if s.asyncCapacityLoader != nil {
return s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit)
}
snapshot, err := s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
snapshot, err := s.coordinationStore.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
if err != nil {
return store.AsyncWorkerCapacitySnapshot{}, err
}
allocation, err := s.store.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{
allocation, err := s.coordinationStore.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{
InstanceID: s.workerInstanceID,
PodUID: strings.TrimSpace(os.Getenv("POD_UID")),
PodName: strings.TrimSpace(os.Getenv("POD_NAME")),
@@ -437,7 +437,7 @@ func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
<-ctx.Done()
if s.store != nil && strings.TrimSpace(s.workerInstanceID) != "" {
drainCtx, drainCancel := context.WithTimeout(context.Background(), 5*time.Second)
if err := s.store.MarkWorkerDraining(drainCtx, s.workerInstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
if err := s.coordinationStore.MarkWorkerDraining(drainCtx, s.workerInstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
s.logger.Warn("mark async worker draining failed", "error", err)
}
drainCancel()
@@ -567,7 +567,7 @@ func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
func (s *Service) recoverOrphanedAsyncRiverJobs(ctx context.Context) {
recoverJobs := func() {
runtimeRecovery, err := s.store.RecoverInterruptedRuntimeState(ctx)
runtimeRecovery, err := s.coordinationStore.RecoverInterruptedRuntimeState(ctx)
if err != nil {
if ctx.Err() == nil {
s.logger.Warn("recover interrupted runtime state from worker failed", "error", err)
@@ -589,7 +589,7 @@ func (s *Service) recoverOrphanedAsyncRiverJobs(ctx context.Context) {
"cleanedTaskAdmissions", runtimeRecovery.CleanedTaskAdmissions,
)
}
yieldedAdmissions, err := s.store.YieldStaleAsyncTaskAdmissions(
yieldedAdmissions, err := s.coordinationStore.YieldStaleAsyncTaskAdmissions(
ctx,
orphanedRiverJobWorkerStaleAfter,
orphanedRiverJobRecoveryBatchSize,
@@ -603,7 +603,7 @@ func (s *Service) recoverOrphanedAsyncRiverJobs(ctx context.Context) {
if yieldedAdmissions > 0 {
s.logger.Warn("stale async task admissions yielded", "count", yieldedAdmissions)
}
recovered, err := s.store.RecoverOrphanedAsyncRiverJobs(
recovered, err := s.coordinationStore.RecoverOrphanedAsyncRiverJobs(
ctx,
orphanedRiverJobWorkerStaleAfter,
orphanedRiverJobRecoveryBatchSize,
+150 -87
View File
@@ -27,6 +27,8 @@ import (
type Service struct {
cfg config.Config
store *store.Store
coordinationStore *store.Store
riverStore *store.Store
logger *slog.Logger
clients map[string]clients.Client
scriptExecutor *scriptengine.Executor
@@ -95,6 +97,27 @@ func (e *TaskQueuedError) Is(target error) bool {
}
func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...billingMetricsObserver) *Service {
return NewWithStores(cfg, db, db, db, logger, observers...)
}
func NewWithCoordinationStore(
cfg config.Config,
db *store.Store,
coordinationDB *store.Store,
logger *slog.Logger,
observers ...billingMetricsObserver,
) *Service {
return NewWithStores(cfg, db, coordinationDB, db, logger, observers...)
}
func NewWithStores(
cfg config.Config,
db *store.Store,
coordinationDB *store.Store,
riverDB *store.Store,
logger *slog.Logger,
observers ...billingMetricsObserver,
) *Service {
if cfg.AsyncWorkerHardLimit == 0 {
cfg.AsyncWorkerHardLimit = 2048
}
@@ -128,10 +151,12 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
httpClients := newHTTPClientCache()
scriptExecutor := &scriptengine.Executor{Logger: logger}
service := &Service{
cfg: cfg,
store: db,
logger: logger,
scriptExecutor: scriptExecutor,
cfg: cfg,
store: db,
coordinationStore: coordinationDB,
riverStore: riverDB,
logger: logger,
scriptExecutor: scriptExecutor,
clients: map[string]clients.Client{
"openai": clients.OpenAIClient{HTTPClient: httpClients.none},
"aliyun-bailian": clients.AliyunBailianClient{HTTPClient: httpClients.none},
@@ -166,6 +191,9 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
if len(observers) > 0 {
service.billingMetrics = observers[0]
}
if service.coordinationStore == nil {
service.coordinationStore = db
}
return service
}
@@ -203,9 +231,21 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
var claimed store.GatewayTask
var err error
if distributedAdmission && !wasRunning {
claimed, err = s.store.ClaimTaskPreparation(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
claimed, err = s.store.ClaimTaskPreparationForOwner(
ctx,
task.ID,
executionToken,
taskExecutionLeaseTTL,
s.workerInstanceID,
)
} else {
claimed, err = s.store.ClaimTaskExecution(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
claimed, err = s.store.ClaimTaskExecutionForOwner(
ctx,
task.ID,
executionToken,
taskExecutionLeaseTTL,
s.workerInstanceID,
)
}
if err != nil {
return Result{}, err
@@ -557,6 +597,13 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
}
if hasConcurrentLimit {
if err := s.store.CheckRateLimits(ctx, s.rateLimitReservations(ctx, user, candidates[0], body)); err != nil {
if task.AsyncMode && errors.Is(err, store.ErrRateLimited) && store.RateLimitRetryable(err) {
queued, delay, queueErr := s.requeueRateLimitedTask(ctx, task, err, candidates[0])
if queueErr != nil {
return Result{}, queueErr
}
return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: delay}
}
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, clients.ErrorCode(err), err.Error(), task.RunMode == "simulation", err)
if finishErr != nil {
return Result{}, finishErr
@@ -577,6 +624,9 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
var alreadyAdmitted bool
admissionResult, alreadyAdmitted, admissionErr = s.activeAsyncTaskAdmission(ctx, task, plan, asyncAdmission)
if admissionErr == nil && !alreadyAdmitted {
admissionResult, admissionErr = s.tryTaskAdmission(ctx, task, plan, "")
}
if admissionErr == nil && !admissionResult.Admitted {
_ = s.store.ReleaseTaskPreparation(context.WithoutCancel(ctx), task.ID, task.ExecutionToken)
return Result{Task: task, Output: task.Result}, &TaskQueuedError{Delay: time.Second}
}
@@ -626,8 +676,13 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
}
}
}
admissionFinalized := false
if distributedAdmission {
defer s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID)
defer func() {
if !admissionFinalized {
_ = s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID)
}
}()
}
firstCandidateBody := body
normalizedModelType := modelType
@@ -670,9 +725,6 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
}
return Result{Task: failed, Output: failed.Result}, clientErr
}
if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, candidates[0].ModelType, s.slimTaskRequestSnapshot(task, firstCandidateBody)); err != nil {
return Result{}, err
}
var reserveErr error
walletReservations, reserveErr = s.store.ReserveTaskBilling(ctx, task, user, reservationBillings, reservationPricingSnapshot)
if reserveErr != nil {
@@ -872,6 +924,10 @@ candidatesLoop:
record.Metrics = mergeMetrics(record.Metrics, candidateCapabilityFilterMetrics(candidateFilterSummary))
record.Metrics = mergeMetrics(record.Metrics, parameterPreprocessingMetrics(preprocessing.Log))
record.Metrics = s.withAttemptHistory(ctx, task.ID, record.Metrics)
callbackURL, callbackErr := s.taskCallbackURL(ctx, task.ID)
if callbackErr != nil {
return Result{}, callbackErr
}
finished, finishErr := s.store.FinishTaskSuccess(ctx, store.FinishTaskSuccessInput{
TaskID: task.ID,
ExecutionToken: task.ExecutionToken,
@@ -891,6 +947,9 @@ candidatesLoop:
ResponseStartedAt: record.ResponseStartedAt,
ResponseFinishedAt: record.ResponseFinishedAt,
ResponseDurationMS: record.ResponseDurationMS,
CompletionCallbackURL: callbackURL,
Simulated: isSimulation(task, candidate),
FinalizeAdmission: distributedAdmission,
})
if finishErr != nil {
if errors.Is(finishErr, store.ErrTaskResultBinaryNotMaterialized) {
@@ -935,23 +994,7 @@ candidatesLoop:
return Result{}, finishErr
}
walletReservationFinalized = true
if finished.BillingStatus == "pending" {
if err := s.emit(ctx, task.ID, "task.billing.pending", "succeeded", "billing", 0.98, "task billing queued", map[string]any{
"amount": finished.FinalChargeAmount, "currency": finished.BillingCurrency,
}, isSimulation(task, candidate)); err != nil {
return Result{}, err
}
}
if err := s.emit(ctx, task.ID, "task.completed", "succeeded", "completed", 1, "task completed", map[string]any{
"result": response.Result,
"billings": billings,
"usage": record.Usage,
"metrics": record.Metrics,
"billingSummary": record.BillingSummary,
"requestId": record.RequestID,
}, isSimulation(task, candidate)); err != nil {
return Result{}, err
}
admissionFinalized = distributedAdmission
if task.AsyncMode {
// Async callers hydrate the canonical stored result through
// the task detail endpoint. The River Worker does not consume
@@ -1232,12 +1275,17 @@ func (s *Service) runCandidate(
limitResult.Leases = append(limitResult.Leases, admittedLeases...)
}
rateReservationsFinalized := false
retainAdmittedLeases := false
defer func() {
if !rateReservationsFinalized {
_ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed")
}
}()
defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.Leases)
defer func() {
if !retainAdmittedLeases {
_ = s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.Leases)
}
}()
attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
TaskID: task.ID,
@@ -1402,15 +1450,16 @@ func (s *Service) runCandidate(
var submissionWire *clients.WireResponse
runCtx, stopLeaseRenewal := s.startConcurrencyLeaseRenewal(ctx, task.ID, limitResult.Leases)
response, err := client.Run(runCtx, clients.Request{
Kind: task.Kind,
ModelType: candidate.ModelType,
Model: task.Model,
Body: providerBody,
OriginalBody: preprocessing.Input,
Candidate: candidate,
HTTPClient: requestHTTPClient,
RemoteTaskID: task.RemoteTaskID,
RemoteTaskPayload: task.RemoteTaskPayload,
Kind: task.Kind,
ModelType: candidate.ModelType,
Model: task.Model,
Body: providerBody,
OriginalBody: preprocessing.Input,
Candidate: candidate,
HTTPClient: requestHTTPClient,
RemoteTaskID: task.RemoteTaskID,
RemoteTaskPayload: task.RemoteTaskPayload,
UpstreamIdempotencyKey: task.ID,
OnRemoteTaskSubmitted: func(remoteTaskID string, payload map[string]any) error {
if strings.TrimSpace(remoteTaskID) == "" {
return nil
@@ -1446,7 +1495,8 @@ func (s *Service) runCandidate(
return nil
},
OnUpstreamResponseReceived: func() error {
return setSubmissionStatus("response_received")
submissionStatus = "response_received"
return nil
},
OnUpstreamWireResponse: func(wire *clients.WireResponse) error {
submissionWire = wire
@@ -1469,9 +1519,7 @@ func (s *Service) runCandidate(
}
callFinishedAt := time.Now()
if err == nil {
if markErr := setSubmissionStatus("response_received"); markErr != nil {
return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: markErr}
}
submissionStatus = "response_received"
}
if response.ResponseStartedAt.IsZero() {
response.ResponseStartedAt = callStartedAt
@@ -1501,9 +1549,7 @@ func (s *Service) runCandidate(
}
if err != nil {
if clients.ErrorResponseMetadata(err).StatusCode > 0 && submissionStatus != "response_received" {
if markErr := setSubmissionStatus("response_received"); markErr != nil {
return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: markErr}
}
submissionStatus = "response_received"
}
retryable := clients.IsRetryable(err)
requestID, metrics, responseStartedAt, responseFinishedAt, responseDurationMS := failureMetrics(err, simulated)
@@ -1521,16 +1567,17 @@ func (s *Service) runCandidate(
}
metrics = mergeMetrics(baseAttemptMetrics, metrics)
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
AttemptID: attemptID,
Status: "failed",
Retryable: retryable,
RequestID: requestID,
Metrics: metrics,
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS,
ErrorCode: clients.ErrorCode(err),
ErrorMessage: err.Error(),
AttemptID: attemptID,
Status: "failed",
Retryable: retryable,
UpstreamSubmissionStatus: submissionStatus,
RequestID: requestID,
Metrics: metrics,
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS,
ErrorCode: clients.ErrorCode(err),
ErrorMessage: err.Error(),
})
_ = s.emit(ctx, task.ID, "task.attempt.failed", "running", "attempt_failed", 0.45, err.Error(), map[string]any{"attempt": attemptNo, "retryable": retryable, "requestId": requestID, "statusCode": clients.ErrorResponseMetadata(err).StatusCode, "metrics": metrics}, simulated)
if !simulated && submissionStatus == "submitting" {
@@ -1561,18 +1608,19 @@ func (s *Service) runCandidate(
"trace": []any{failureTraceEntry(err, clients.IsRetryable(err))},
})
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
AttemptID: attemptID,
Status: "failed",
Retryable: clients.IsRetryable(err),
RequestID: response.RequestID,
Usage: usageToMap(response.Usage),
Metrics: metrics,
ResponseSnapshot: response.Result,
ResponseStartedAt: response.ResponseStartedAt,
ResponseFinishedAt: response.ResponseFinishedAt,
ResponseDurationMS: response.ResponseDurationMS,
ErrorCode: clients.ErrorCode(err),
ErrorMessage: err.Error(),
AttemptID: attemptID,
Status: "failed",
Retryable: clients.IsRetryable(err),
UpstreamSubmissionStatus: submissionStatus,
RequestID: response.RequestID,
Usage: usageToMap(response.Usage),
Metrics: metrics,
ResponseSnapshot: response.Result,
ResponseStartedAt: response.ResponseStartedAt,
ResponseFinishedAt: response.ResponseFinishedAt,
ResponseDurationMS: response.ResponseDurationMS,
ErrorCode: clients.ErrorCode(err),
ErrorMessage: err.Error(),
})
return clients.Response{}, err
}
@@ -1652,6 +1700,9 @@ func (s *Service) runCandidate(
}); err != nil {
s.logger.Warn("record cache affinity observation failed", "error", err, "clientId", candidate.ClientID)
}
if admittedPlatformModelID == candidate.PlatformModelID {
retainAdmittedLeases = true
}
return response, nil
}
@@ -1835,17 +1886,24 @@ func (s *Service) failTask(ctx context.Context, taskID string, executionToken st
metrics = mergeMetrics(values...)
}
metrics = s.withAttemptHistory(ctx, taskID, metrics)
callbackURL, callbackErr := s.taskCallbackURL(ctx, taskID)
if callbackErr != nil {
return store.GatewayTask{}, callbackErr
}
failed, err := s.store.FinishTaskFailure(ctx, store.FinishTaskFailureInput{
TaskID: taskID,
ExecutionToken: executionToken,
Code: code,
Message: message,
Result: buildFailureResult(code, message, requestID, cause),
RequestID: requestID,
Metrics: metrics,
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS,
TaskID: taskID,
ExecutionToken: executionToken,
Code: code,
Message: message,
Result: buildFailureResult(code, message, requestID, cause),
RequestID: requestID,
Metrics: metrics,
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS,
CompletionCallbackURL: callbackURL,
Simulated: simulated,
FinalizeAdmission: true,
})
if err != nil {
return store.GatewayTask{}, err
@@ -1853,9 +1911,6 @@ func (s *Service) failTask(ctx context.Context, taskID string, executionToken st
if failed.Status == "cancelled" {
return failed, nil
}
if eventErr := s.emit(ctx, taskID, "task.failed", "failed", "failed", 1, message, map[string]any{"code": code, "requestId": requestID, "metrics": metrics}, simulated); eventErr != nil {
return store.GatewayTask{}, eventErr
}
return failed, nil
}
@@ -2056,7 +2111,7 @@ func (s *Service) startConcurrencyLeaseRenewal(ctx context.Context, taskID strin
done <- nil
return
case <-ticker.C:
if err := s.store.RenewConcurrencyLeases(renewCtx, leases); err != nil {
if err := s.coordinationStore.RenewConcurrencyLeases(renewCtx, leases); err != nil {
if renewCtx.Err() != nil {
done <- nil
return
@@ -2113,16 +2168,10 @@ func (s *Service) withAttemptHistory(ctx context.Context, taskID string, metrics
}
func (s *Service) emit(ctx context.Context, taskID string, eventType string, status string, phase string, progress float64, message string, payload map[string]any, simulated bool) error {
acceptanceURL, acceptance, callbackErr := s.store.AcceptanceCallbackURL(ctx, taskID)
callbackURL, callbackErr := s.taskCallbackURL(ctx, taskID)
if callbackErr != nil {
return callbackErr
}
callbackURL := ""
if acceptance && acceptanceURL != "" {
callbackURL = acceptanceURL
} else if s.cfg.TaskProgressCallbackEnabled && s.cfg.TaskProgressCallbackURL != "" {
callbackURL = s.cfg.TaskProgressCallbackURL
}
event, err := s.store.AddTaskEventWithCallback(
ctx,
taskID,
@@ -2144,6 +2193,20 @@ func (s *Service) emit(ctx context.Context, taskID string, eventType string, sta
return nil
}
func (s *Service) taskCallbackURL(ctx context.Context, taskID string) (string, error) {
acceptanceURL, acceptance, err := s.store.AcceptanceCallbackURL(ctx, taskID)
if err != nil {
return "", err
}
if acceptance && acceptanceURL != "" {
return acceptanceURL, nil
}
if s.cfg.TaskProgressCallbackEnabled && s.cfg.TaskProgressCallbackURL != "" {
return s.cfg.TaskProgressCallbackURL, nil
}
return "", nil
}
func modelTypeFromKind(kind string, body map[string]any) string {
if requested := requestedModelTypeFromBody(body); requested != "" {
return requested
@@ -285,6 +285,18 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
}); ok {
postgresPool = poolProvider.PostgresPoolMetrics()
}
criticalPostgresPool := store.PostgresPoolMetricsSnapshot{}
if poolProvider, ok := provider.(interface {
CriticalPostgresPoolMetrics() store.PostgresPoolMetricsSnapshot
}); ok {
criticalPostgresPool = poolProvider.CriticalPostgresPoolMetrics()
}
riverPostgresPool := store.PostgresPoolMetricsSnapshot{}
if poolProvider, ok := provider.(interface {
RiverPostgresPoolMetrics() store.PostgresPoolMetricsSnapshot
}); ok {
riverPostgresPool = poolProvider.RiverPostgresPoolMetrics()
}
w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
outcomeCounters(w, "easyai_gateway_ssf_receipts_total", "Received SETs by bounded outcome.", []outcomeValue{
{"accepted", m.accepted.Load()}, {"rejected", m.rejected.Load()}, {"duplicate", m.duplicate.Load()},
@@ -362,6 +374,18 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
plainGauge(w, "easyai_gateway_postgres_pool_idle_connections", "Currently idle PostgreSQL connections in this process pool.", int64(postgresPool.IdleConnections))
plainCounter(w, "easyai_gateway_postgres_pool_empty_acquire_total", "Pool acquires that had to wait for a PostgreSQL connection.", uint64(max(postgresPool.EmptyAcquireCount, 0)))
plainCounter(w, "easyai_gateway_postgres_pool_canceled_acquire_total", "Canceled PostgreSQL pool acquires.", uint64(max(postgresPool.CanceledAcquireCount, 0)))
plainGauge(w, "easyai_gateway_postgres_critical_pool_max_connections", "Maximum PostgreSQL connections reserved for leases and coordination.", int64(criticalPostgresPool.MaxConnections))
plainGauge(w, "easyai_gateway_postgres_critical_pool_total_connections", "Current PostgreSQL connections in the leases and coordination pool.", int64(criticalPostgresPool.TotalConnections))
plainGauge(w, "easyai_gateway_postgres_critical_pool_acquired_connections", "Currently acquired PostgreSQL connections in the leases and coordination pool.", int64(criticalPostgresPool.AcquiredConnections))
plainGauge(w, "easyai_gateway_postgres_critical_pool_idle_connections", "Currently idle PostgreSQL connections in the leases and coordination pool.", int64(criticalPostgresPool.IdleConnections))
plainCounter(w, "easyai_gateway_postgres_critical_pool_empty_acquire_total", "Coordination pool acquires that had to wait for a PostgreSQL connection.", uint64(max(criticalPostgresPool.EmptyAcquireCount, 0)))
plainCounter(w, "easyai_gateway_postgres_critical_pool_canceled_acquire_total", "Canceled PostgreSQL coordination pool acquires.", uint64(max(criticalPostgresPool.CanceledAcquireCount, 0)))
plainGauge(w, "easyai_gateway_postgres_river_pool_max_connections", "Maximum PostgreSQL connections reserved for River queue coordination.", int64(riverPostgresPool.MaxConnections))
plainGauge(w, "easyai_gateway_postgres_river_pool_total_connections", "Current PostgreSQL connections in the River queue pool.", int64(riverPostgresPool.TotalConnections))
plainGauge(w, "easyai_gateway_postgres_river_pool_acquired_connections", "Currently acquired PostgreSQL connections in the River queue pool.", int64(riverPostgresPool.AcquiredConnections))
plainGauge(w, "easyai_gateway_postgres_river_pool_idle_connections", "Currently idle PostgreSQL connections in the River queue pool.", int64(riverPostgresPool.IdleConnections))
plainCounter(w, "easyai_gateway_postgres_river_pool_empty_acquire_total", "River queue pool acquires that had to wait for a PostgreSQL connection.", uint64(max(riverPostgresPool.EmptyAcquireCount, 0)))
plainCounter(w, "easyai_gateway_postgres_river_pool_canceled_acquire_total", "Canceled PostgreSQL River queue pool acquires.", uint64(max(riverPostgresPool.CanceledAcquireCount, 0)))
outcomeCounters(w, "easyai_gateway_async_worker_resizes_total", "Asynchronous worker resize attempts by bounded outcome.", []outcomeValue{
{"success", m.asyncWorkerResizeSuccess.Load()},
{"refresh_failed", m.asyncWorkerRefreshFailed.Load()},
+118 -8
View File
@@ -22,6 +22,7 @@ const SystemSettingGatewayTrafficMode = "gateway_traffic_mode"
var (
acceptanceReleaseSHAPattern = regexp.MustCompile(`^[0-9a-f]{40}$`)
acceptanceDigestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
capacityConfigHashPattern = regexp.MustCompile(`^[0-9a-f]{64}$`)
)
var (
@@ -39,6 +40,7 @@ type GatewayTrafficMode struct {
ReleaseSHA string `json:"releaseSha,omitempty"`
APIImageDigest string `json:"apiImageDigest,omitempty"`
WorkerImageDigest string `json:"workerImageDigest,omitempty"`
ConfigHash string `json:"configHash,omitempty"`
UpdatedAt time.Time `json:"updatedAt"`
}
@@ -84,11 +86,21 @@ type FinishAcceptanceRunInput struct {
}
type PromoteAcceptanceRunInput struct {
RunID string `json:"-"`
RunID string `json:"-"`
Revision int64 `json:"revision"`
ReleaseSHA string `json:"releaseSha"`
APIImageDigest string `json:"apiImageDigest"`
WorkerImageDigest string `json:"workerImageDigest"`
ConfigHash string `json:"configHash"`
CapacityProfiles []CapacityProfileInput `json:"capacityProfiles"`
}
type PauseGatewayTrafficInput struct {
Revision int64 `json:"revision"`
ReleaseSHA string `json:"releaseSha"`
APIImageDigest string `json:"apiImageDigest"`
WorkerImageDigest string `json:"workerImageDigest"`
Reason string `json:"reason"`
}
func DefaultGatewayTrafficMode() GatewayTrafficMode {
@@ -250,9 +262,25 @@ func (s *Store) FinishAcceptanceRun(ctx context.Context, input FinishAcceptanceR
UPDATE gateway_acceptance_runs
SET status = $2, report = $3::jsonb, failure_reason = NULLIF($4, ''),
finished_at = now(), updated_at = now()
WHERE id = $1::uuid AND status = 'running'
WHERE (
id = $1::uuid
AND status = 'running'
)
OR (
id = $1::uuid
AND $2::text = 'failed'
AND status = 'passed'
AND EXISTS (
SELECT 1
FROM system_settings setting
WHERE setting.setting_key = $5
AND setting.value->>'mode' = 'validation'
AND setting.value->>'runId' = $1::text
)
)
RETURNING `+acceptanceRunColumns,
strings.TrimSpace(input.RunID), status, report, strings.TrimSpace(input.FailureReason),
SystemSettingGatewayTrafficMode,
))
}
@@ -318,18 +346,37 @@ FOR UPDATE`, SystemSettingGatewayTrafficMode).Scan(&value); err != nil {
current.WorkerImageDigest != strings.TrimSpace(input.WorkerImageDigest) {
return GatewayTrafficMode{}, ErrAcceptanceStateConflict
}
var status string
if err := tx.QueryRow(ctx, `
SELECT status
run, err := scanAcceptanceRun(tx.QueryRow(ctx, `
SELECT `+acceptanceRunColumns+`
FROM gateway_acceptance_runs
WHERE id = $1::uuid
FOR UPDATE`, current.RunID).Scan(&status); err != nil {
FOR UPDATE`, current.RunID))
if err != nil {
return GatewayTrafficMode{}, err
}
if status != "passed" {
if run.Status != "passed" {
return GatewayTrafficMode{}, ErrAcceptancePromotionGates
}
next := GatewayTrafficMode{Mode: "live", Revision: current.Revision + 1}
stagedConfigHash, _ := run.Config["validationConfigHash"].(string)
if stagedConfigHash == "" {
return GatewayTrafficMode{}, fmt.Errorf("%w: certified capacity profiles were not staged", ErrAcceptancePromotionGates)
}
if stagedConfigHash != strings.ToLower(strings.TrimSpace(input.ConfigHash)) {
return GatewayTrafficMode{}, fmt.Errorf("%w: staged capacity config hash changed", ErrAcceptancePromotionGates)
}
stagedProfilesHash, _ := run.Config["validationProfilesHash"].(string)
if stagedProfilesHash == "" || stagedProfilesHash != capacityProfileInputsHash(input.CapacityProfiles) {
return GatewayTrafficMode{}, fmt.Errorf("%w: staged capacity profiles changed", ErrAcceptancePromotionGates)
}
if err := replaceCapacityProfilesTx(ctx, tx, run, input.ConfigHash, input.CapacityProfiles); err != nil {
return GatewayTrafficMode{}, fmt.Errorf("%w: %v", ErrAcceptancePromotionGates, err)
}
next := GatewayTrafficMode{
Mode: "live", RunID: run.ID, Revision: current.Revision + 1,
ReleaseSHA: run.ReleaseSHA, APIImageDigest: run.APIImageDigest,
WorkerImageDigest: run.WorkerImageDigest,
ConfigHash: strings.ToLower(strings.TrimSpace(input.ConfigHash)),
}
nextValue, _ := json.Marshal(next)
if _, err := tx.Exec(ctx, `
UPDATE system_settings
@@ -350,6 +397,69 @@ WHERE id = $1::uuid`, current.RunID); err != nil {
return next, nil
}
func (s *Store) PauseGatewayTraffic(
ctx context.Context,
input PauseGatewayTrafficInput,
) (GatewayTrafficMode, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return GatewayTrafficMode{}, err
}
defer rollbackTransaction(tx)
var value []byte
if err := tx.QueryRow(ctx, `
SELECT value
FROM system_settings
WHERE setting_key = $1
FOR UPDATE`, SystemSettingGatewayTrafficMode).Scan(&value); err != nil {
return GatewayTrafficMode{}, err
}
var current GatewayTrafficMode
if err := json.Unmarshal(value, &current); err != nil {
return GatewayTrafficMode{}, err
}
reason := strings.TrimSpace(input.Reason)
if normalizeTrafficMode(current.Mode) != "live" ||
current.Revision != input.Revision ||
current.ReleaseSHA != strings.TrimSpace(input.ReleaseSHA) ||
current.APIImageDigest != strings.TrimSpace(input.APIImageDigest) ||
current.WorkerImageDigest != strings.TrimSpace(input.WorkerImageDigest) ||
reason == "" || len(reason) > 240 {
return GatewayTrafficMode{}, ErrAcceptanceStateConflict
}
next := GatewayTrafficMode{
Mode: "validation", RunID: current.RunID, Revision: current.Revision + 1,
ReleaseSHA: current.ReleaseSHA, APIImageDigest: current.APIImageDigest,
WorkerImageDigest: current.WorkerImageDigest, ConfigHash: current.ConfigHash,
}
nextValue, _ := json.Marshal(next)
if _, err := tx.Exec(ctx, `
UPDATE system_settings
SET value = $2::jsonb, updated_at = now()
WHERE setting_key = $1`, SystemSettingGatewayTrafficMode, nextValue); err != nil {
return GatewayTrafficMode{}, err
}
if current.RunID != "" {
if _, err := tx.Exec(ctx, `
UPDATE gateway_acceptance_runs
SET report = jsonb_set(
report,
'{postPromotionPause}',
jsonb_build_object('reason',$2::text,'pausedAt',now()),
true
),
updated_at = now()
WHERE id = $1::uuid`, current.RunID, reason); err != nil {
return GatewayTrafficMode{}, err
}
}
if err := tx.Commit(ctx); err != nil {
return GatewayTrafficMode{}, err
}
next.UpdatedAt = time.Now()
return next, nil
}
func (s *Store) AbortAcceptanceRun(ctx context.Context, input PromoteAcceptanceRunInput) (GatewayTrafficMode, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
@@ -85,6 +85,21 @@ WHERE setting_key = $1`, SystemSettingGatewayTrafficMode)
baseURL != "http://acceptance-emulator:8090" || credential == "" || credential == token {
t.Fatalf("candidate override base=%q credential=%q err=%v", baseURL, credential, err)
}
capacityProfiles := []CapacityProfileInput{
{
Kind: "image", StableThroughputPerSecond: 10, AdmittedThroughputPerSecond: 8,
MaxRunning: 48, MaxQueued: 480, MaxPredictedWaitSeconds: 60,
},
{
Kind: "video", StableThroughputPerSecond: 1, AdmittedThroughputPerSecond: 0.8,
MaxRunning: 48, MaxQueued: 96, MaxPredictedWaitSeconds: 120,
},
}
if _, err := db.StageAcceptanceCapacityProfiles(ctx, StageAcceptanceCapacityProfilesInput{
RunID: run.ID, ConfigHash: strings.Repeat("d", 64), CapacityProfiles: capacityProfiles,
}); err != nil {
t.Fatalf("stage acceptance capacity profiles: %v", err)
}
if _, err := db.FinishAcceptanceRun(ctx, FinishAcceptanceRunInput{
RunID: run.ID, Passed: true, Report: map[string]any{"queueFinal": 0},
}); err != nil {
@@ -93,6 +108,8 @@ WHERE setting_key = $1`, SystemSettingGatewayTrafficMode)
promotion := PromoteAcceptanceRunInput{
RunID: run.ID, Revision: mode.Revision,
ReleaseSHA: mode.ReleaseSHA, APIImageDigest: mode.APIImageDigest, WorkerImageDigest: mode.WorkerImageDigest,
ConfigHash: strings.Repeat("d", 64),
CapacityProfiles: capacityProfiles,
}
stale := promotion
stale.Revision++
@@ -106,6 +123,30 @@ WHERE setting_key = $1`, SystemSettingGatewayTrafficMode)
if live.Mode != "live" || live.Revision != 2 {
t.Fatalf("unexpected live mode: %+v", live)
}
stalePause := PauseGatewayTrafficInput{
Revision: live.Revision + 1, ReleaseSHA: live.ReleaseSHA,
APIImageDigest: live.APIImageDigest, WorkerImageDigest: live.WorkerImageDigest,
Reason: "integration hard gate",
}
if _, err := db.PauseGatewayTraffic(ctx, stalePause); !errors.Is(err, ErrAcceptanceStateConflict) {
t.Fatalf("stale post-promotion pause error=%v, want state conflict", err)
}
stalePause.Revision = live.Revision
paused, err := db.PauseGatewayTraffic(ctx, stalePause)
if err != nil {
t.Fatalf("pause post-promotion traffic: %v", err)
}
if paused.Mode != "validation" || paused.RunID != run.ID || paused.Revision != 3 {
t.Fatalf("unexpected post-promotion pause mode: %+v", paused)
}
if _, err := db.AuthorizeAcceptanceTask(ctx, "", "", user); !errors.Is(err, ErrProductionTrafficPaused) {
t.Fatalf("production request after monitor pause error=%v, want traffic paused", err)
}
profiles, err := db.ListCapacityProfiles(ctx)
if err != nil || len(profiles) < 2 {
t.Fatalf("certified capacity profiles=%+v err=%v", profiles, err)
}
resetTrafficMode()
failedRun, err := db.CreateAcceptanceRun(ctx, CreateAcceptanceRunInput{
ReleaseSHA: strings.Repeat("d", 40),
@@ -187,10 +228,24 @@ SET upstream_submission_status = 'submitting', upstream_submission_updated_at =
WHERE id = $1::uuid`, submittingAttemptID); err != nil {
t.Fatalf("mark acceptance task attempt submitting: %v", err)
}
if _, err := db.FinishAcceptanceRun(ctx, FinishAcceptanceRunInput{
if passed, err := db.FinishAcceptanceRun(ctx, FinishAcceptanceRunInput{
RunID: failedRun.ID, Passed: true, Report: map[string]any{"testsPassed": true},
}); err != nil || passed.Status != "passed" {
t.Fatalf("mark acceptance run passed before external certification=%+v err=%v", passed, err)
}
if _, err := db.PromoteAcceptanceRun(ctx, PromoteAcceptanceRunInput{
RunID: failedRun.ID, Revision: failedMode.Revision,
ReleaseSHA: failedMode.ReleaseSHA, APIImageDigest: failedMode.APIImageDigest,
WorkerImageDigest: failedMode.WorkerImageDigest,
ConfigHash: strings.Repeat("d", 64),
CapacityProfiles: capacityProfiles,
}); !errors.Is(err, ErrAcceptancePromotionGates) {
t.Fatalf("unstaged capacity promotion error=%v, want promotion gates", err)
}
if failed, err := db.FinishAcceptanceRun(ctx, FinishAcceptanceRunInput{
RunID: failedRun.ID, Passed: false, FailureReason: "load failed",
}); err != nil {
t.Fatalf("fail acceptance run: %v", err)
}); err != nil || failed.Status != "failed" {
t.Fatalf("downgrade unpromoted acceptance run after external gate failure=%+v err=%v", failed, err)
}
if retried, err := db.RetryAcceptanceRun(ctx, failedRun.ID); err != nil || retried.Status != "running" {
t.Fatalf("retry acceptance run=%+v err=%v", retried, err)
@@ -731,6 +731,53 @@ SELECT
t.Fatalf("recovered task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases)
}
remoteRecoveryTask := createTask(true)
if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
remote_task_id = 'remote-video-task',
execution_token = gen_random_uuid(),
execution_lease_expires_at = now() - interval '1 second'
WHERE id = $1::uuid`,
remoteRecoveryTask.ID,
); err != nil {
t.Fatalf("prepare persisted remote task recovery task: %v", err)
}
if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
task_id, attempt_no, queue_key, status, upstream_submission_status
)
VALUES ($1::uuid, 1, 'integration-test', 'running', 'submitting')`,
remoteRecoveryTask.ID,
); err != nil {
t.Fatalf("prepare persisted remote task recovery: %v", err)
}
recovery, err = first.RecoverInterruptedRuntimeState(ctx)
if err != nil {
t.Fatalf("recover persisted remote task: %v", err)
}
var remoteStatus, remoteErrorCode string
var remoteAttemptRetryable bool
if err := first.pool.QueryRow(ctx, `
SELECT task.status, COALESCE(task.error_code, ''), attempt.retryable
FROM gateway_tasks task
JOIN gateway_task_attempts attempt ON attempt.task_id = task.id
WHERE task.id = $1::uuid`, remoteRecoveryTask.ID).Scan(
&remoteStatus,
&remoteErrorCode,
&remoteAttemptRetryable,
); err != nil {
t.Fatalf("read persisted remote task recovery: %v", err)
}
if remoteStatus != "queued" || remoteErrorCode != "" || !remoteAttemptRetryable {
t.Fatalf(
"persisted remote task recovery status=%s code=%s retryable=%t, want queued resume",
remoteStatus,
remoteErrorCode,
remoteAttemptRetryable,
)
}
lockedRecoveryTask := createTask(true)
unlockedRecoveryTask := createTask(true)
if _, err := first.pool.Exec(ctx, `
@@ -405,10 +405,35 @@ func TestBillingSettlementStaleTakeoverDebitsExactlyOnce(t *testing.T) {
if _, err := db.FinishTaskSuccess(ctx, FinishTaskSuccessInput{
TaskID: created.ID, ExecutionToken: token, Result: map[string]any{"ok": true},
FinalChargeAmountText: amount, BillingCurrency: "resource",
PricingSnapshot: map[string]any{"pricingVersion": "effective-pricing-v2"},
PricingSnapshot: map[string]any{"pricingVersion": "effective-pricing-v2"},
CompletionCallbackURL: "https://callback.invalid/task",
}); err != nil {
t.Fatal(err)
}
var completionEvents int
var completionCallbacks int
if err := db.pool.QueryRow(ctx, `
SELECT
(SELECT count(*)
FROM gateway_task_events
WHERE task_id=$1::uuid
AND event_type='task.completed'
AND status='succeeded'
AND payload='{}'::jsonb),
(SELECT count(*)
FROM gateway_task_callback_outbox callback
JOIN gateway_task_events event ON event.id=callback.event_id
WHERE callback.task_id=$1::uuid
AND callback.callback_url='https://callback.invalid/task'
AND callback.seq=event.seq
AND callback.payload='{}'::jsonb)`,
created.ID,
).Scan(&completionEvents, &completionCallbacks); err != nil {
t.Fatal(err)
}
if completionEvents != 1 || completionCallbacks != 1 {
t.Fatalf("atomic completion history=%d callbacks=%d, want 1/1", completionEvents, completionCallbacks)
}
firstClaims, err := db.ClaimBillingSettlements(ctx, "worker-one", BillingSettlementBatchSize, BillingSettlementLockTimeout)
if err != nil {
@@ -468,6 +493,133 @@ WHERE gateway_user_id=$1::uuid AND reference_id=$2`, gatewayUserID, created.ID).
}
}
func TestExpiredSynchronousExecutionRecoveryReleasesBillingReservation(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
user := &auth.User{
ID: "billing-recovery-" + uuid.NewString(),
GatewayUserID: gatewayUserID,
GatewayTenantID: tenantID,
}
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
GatewayUserID: gatewayUserID,
Currency: "resource",
Balance: 10,
Reason: "billing recovery test",
}); err != nil {
t.Fatal(err)
}
created, err := db.CreateTask(ctx, CreateTaskInput{
Kind: "images.generations",
Model: "billing-v2-model",
RunMode: "production",
Request: map[string]any{"model": "billing-v2-model"},
}, user)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
cleanupCtx := context.Background()
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, gatewayUserID)
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, gatewayUserID)
})
token := uuid.NewString()
claimed, err := db.ClaimTaskExecutionForOwner(ctx, created.ID, token, 5*time.Minute, "expired-worker")
if err != nil {
t.Fatal(err)
}
amount := "1.000000000"
reservations, err := db.ReserveTaskBilling(ctx, claimed, user, nil, map[string]any{
"pricingVersion": "billing-recovery-v1",
"reservationAmount": amount,
"currency": "resource",
})
if err != nil || len(reservations) != 1 {
t.Fatalf("reserve=%+v err=%v", reservations, err)
}
if _, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
TaskID: created.ID,
ExecutionToken: token,
AttemptNo: 1,
Status: "running",
}); err != nil {
t.Fatal(err)
}
if _, err := db.pool.Exec(ctx, `
UPDATE gateway_tasks
SET execution_lease_expires_at=now()-interval '1 second'
WHERE id=$1::uuid`, created.ID); err != nil {
t.Fatal(err)
}
recovery, err := db.RecoverInterruptedRuntimeState(ctx)
if err != nil {
t.Fatal(err)
}
if recovery.FailedTasks != 1 || recovery.FailedAttempts != 1 {
t.Fatalf("recovery=%+v, want one failed task and attempt", recovery)
}
var taskStatus, billingStatus, outboxStatus string
var outboxCount int
if err := db.pool.QueryRow(ctx, `
SELECT task.status,
task.billing_status,
count(outbox.id),
COALESCE(min(outbox.status), '')
FROM gateway_tasks task
LEFT JOIN settlement_outbox outbox
ON outbox.task_id=task.id
AND outbox.event_type='task.billing.release'
AND outbox.action='release'
WHERE task.id=$1::uuid
GROUP BY task.status, task.billing_status`,
created.ID,
).Scan(&taskStatus, &billingStatus, &outboxCount, &outboxStatus); err != nil {
t.Fatal(err)
}
if taskStatus != "failed" || billingStatus != "pending" || outboxCount != 1 || outboxStatus != "pending" {
t.Fatalf(
"recovered task status=%s billing=%s release_outbox=%d/%s",
taskStatus,
billingStatus,
outboxCount,
outboxStatus,
)
}
claims, err := db.ClaimBillingSettlements(ctx, "billing-recovery-worker", BillingSettlementBatchSize, BillingSettlementLockTimeout)
if err != nil {
t.Fatal(err)
}
release := settlementForTask(t, claims, created.ID)
if release.Action != "release" {
t.Fatalf("settlement action=%s, want release", release.Action)
}
if err := db.ProcessBillingSettlement(ctx, release); err != nil {
t.Fatal(err)
}
var releasedExactly bool
if err := db.pool.QueryRow(ctx, `
SELECT account.balance=10::numeric
AND account.frozen_balance=0::numeric
AND task.reservation_amount=0::numeric
AND task.billing_status='released'
FROM gateway_wallet_accounts account
JOIN gateway_tasks task ON task.gateway_user_id=account.gateway_user_id
WHERE task.id=$1::uuid
AND account.currency='resource'`,
created.ID,
).Scan(&releasedExactly); err != nil {
t.Fatal(err)
}
if !releasedExactly {
t.Fatal("expired synchronous execution did not release its wallet reservation exactly once")
}
}
func TestReleaseTaskBillingReservationsPreservesNineDecimalPlaces(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
@@ -0,0 +1,532 @@
package store
import (
"context"
"crypto/sha256"
"encoding/json"
"errors"
"fmt"
"math"
"strings"
"time"
"github.com/jackc/pgx/v5"
)
type CapacityProfile struct {
ProfileKey string `json:"profileKey"`
ReleaseSHA string `json:"releaseSha"`
ConfigHash string `json:"configHash"`
AcceptanceRunID string `json:"acceptanceRunId"`
Kind string `json:"kind"`
Model string `json:"model"`
CapacityProfile string `json:"capacityProfile"`
StableThroughputPerSecond float64 `json:"stableThroughputPerSecond"`
AdmittedThroughputPerSecond float64 `json:"admittedThroughputPerSecond"`
MaxRunning int `json:"maxRunning"`
MaxQueued int `json:"maxQueued"`
MaxPredictedWaitSeconds int `json:"maxPredictedWaitSeconds"`
Enabled bool `json:"enabled"`
Metadata map[string]any `json:"metadata"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
type CapacityProfileInput struct {
ProfileKey string `json:"profileKey,omitempty"`
Kind string `json:"kind"`
Model string `json:"model,omitempty"`
StableThroughputPerSecond float64 `json:"stableThroughputPerSecond"`
AdmittedThroughputPerSecond float64 `json:"admittedThroughputPerSecond"`
MaxRunning int `json:"maxRunning"`
MaxQueued int `json:"maxQueued"`
MaxPredictedWaitSeconds int `json:"maxPredictedWaitSeconds"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type StageAcceptanceCapacityProfilesInput struct {
RunID string `json:"-"`
ConfigHash string `json:"configHash"`
CapacityProfiles []CapacityProfileInput `json:"capacityProfiles"`
}
func (s *Store) ListCapacityProfiles(ctx context.Context) ([]CapacityProfile, error) {
rows, err := s.pool.Query(ctx, `
SELECT profile_key, release_sha, config_hash, COALESCE(acceptance_run_id::text, ''),
kind, model, capacity_profile,
stable_throughput_per_second::float8,
admitted_throughput_per_second::float8,
max_running, max_queued, max_predicted_wait_seconds, enabled,
metadata, created_at, updated_at
FROM gateway_capacity_profiles
ORDER BY enabled DESC, kind, model, updated_at DESC`)
if err != nil {
if IsUndefinedDatabaseObject(err) {
return []CapacityProfile{}, nil
}
return nil, err
}
defer rows.Close()
items := make([]CapacityProfile, 0)
for rows.Next() {
item, scanErr := scanCapacityProfile(rows)
if scanErr != nil {
return nil, scanErr
}
items = append(items, item)
}
return items, rows.Err()
}
func replaceCapacityProfilesTx(
ctx context.Context,
tx pgx.Tx,
run AcceptanceRun,
configHash string,
inputs []CapacityProfileInput,
) error {
configHash = strings.ToLower(strings.TrimSpace(configHash))
if !capacityConfigHashPattern.MatchString(configHash) {
return errors.New("configHash must be a lowercase SHA-256 value")
}
if err := validateCapacityProfileInputs(inputs); err != nil {
return err
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_capacity_profiles
SET enabled = false, updated_at = now()
WHERE enabled = true`); err != nil {
return err
}
for _, input := range inputs {
kind := strings.ToLower(strings.TrimSpace(input.Kind))
model := strings.TrimSpace(input.Model)
if model == "" {
model = "*"
}
profileKey := strings.TrimSpace(input.ProfileKey)
if profileKey == "" {
profileKey = fmt.Sprintf("%s:%s:%s:%s", run.ReleaseSHA[:12], strings.ToLower(run.CapacityProfile), kind, model)
}
metadata, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(input.Metadata)))
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_capacity_profiles (
profile_key, release_sha, config_hash, acceptance_run_id, kind, model,
capacity_profile, stable_throughput_per_second,
admitted_throughput_per_second, max_running, max_queued,
max_predicted_wait_seconds, enabled, metadata
)
VALUES ($1, $2, $3, $4::uuid, $5, $6, $7, $8, $9, $10, $11, $12, true, $13::jsonb)
ON CONFLICT (profile_key) DO UPDATE
SET release_sha = EXCLUDED.release_sha,
config_hash = EXCLUDED.config_hash,
acceptance_run_id = EXCLUDED.acceptance_run_id,
kind = EXCLUDED.kind,
model = EXCLUDED.model,
capacity_profile = EXCLUDED.capacity_profile,
stable_throughput_per_second = EXCLUDED.stable_throughput_per_second,
admitted_throughput_per_second = EXCLUDED.admitted_throughput_per_second,
max_running = EXCLUDED.max_running,
max_queued = EXCLUDED.max_queued,
max_predicted_wait_seconds = EXCLUDED.max_predicted_wait_seconds,
enabled = true,
metadata = EXCLUDED.metadata,
updated_at = now()`,
profileKey, run.ReleaseSHA, configHash, run.ID, kind, model,
run.CapacityProfile, input.StableThroughputPerSecond,
input.AdmittedThroughputPerSecond, input.MaxRunning, input.MaxQueued,
input.MaxPredictedWaitSeconds, metadata,
); err != nil {
return err
}
}
return nil
}
func validateCapacityProfileInputs(inputs []CapacityProfileInput) error {
if len(inputs) == 0 {
return errors.New("at least one certified capacity profile is required")
}
hasImage := false
hasVideo := false
keys := make(map[string]struct{}, len(inputs))
for _, input := range inputs {
kind := strings.ToLower(strings.TrimSpace(input.Kind))
switch kind {
case "image":
hasImage = true
case "video":
hasVideo = true
case "mixed":
hasImage = true
hasVideo = true
default:
return fmt.Errorf("unsupported capacity profile kind %q", input.Kind)
}
if input.StableThroughputPerSecond <= 0 ||
input.AdmittedThroughputPerSecond <= 0 ||
input.AdmittedThroughputPerSecond > input.StableThroughputPerSecond ||
input.MaxRunning <= 0 ||
input.MaxQueued < 0 ||
input.MaxPredictedWaitSeconds <= 0 {
return fmt.Errorf("invalid capacity limits for %s/%s", kind, input.Model)
}
if input.AdmittedThroughputPerSecond > input.StableThroughputPerSecond*0.800001 {
return fmt.Errorf("admitted throughput for %s/%s exceeds the 80%% safety envelope", kind, input.Model)
}
model := strings.TrimSpace(input.Model)
if model == "" {
model = "*"
}
key := kind + "\x00" + model
if _, exists := keys[key]; exists {
return fmt.Errorf("duplicate capacity profile for %s/%s", kind, model)
}
keys[key] = struct{}{}
}
if !hasImage || !hasVideo {
return errors.New("certified capacity profiles must cover both image and video traffic")
}
return nil
}
func (s *Store) StageAcceptanceCapacityProfiles(
ctx context.Context,
input StageAcceptanceCapacityProfilesInput,
) (AcceptanceRun, error) {
input.RunID = strings.TrimSpace(input.RunID)
input.ConfigHash = strings.ToLower(strings.TrimSpace(input.ConfigHash))
if !capacityConfigHashPattern.MatchString(input.ConfigHash) {
return AcceptanceRun{}, errors.New("configHash must be a lowercase SHA-256 value")
}
if err := validateCapacityProfileInputs(input.CapacityProfiles); err != nil {
return AcceptanceRun{}, err
}
rawProfiles, _ := json.Marshal(input.CapacityProfiles)
var profileValues any
_ = json.Unmarshal(rawProfiles, &profileValues)
profiles, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(profileValues)))
profilesHash := fmt.Sprintf("%x", sha256.Sum256(profiles))
return scanAcceptanceRun(s.pool.QueryRow(ctx, `
UPDATE gateway_acceptance_runs
SET config = jsonb_set(
jsonb_set(
jsonb_set(config, '{validationCapacityProfiles}', $2::jsonb, true),
'{validationConfigHash}', to_jsonb($3::text), true
),
'{validationProfilesHash}', to_jsonb($4::text), true
),
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
RETURNING `+acceptanceRunColumns,
input.RunID, profiles, input.ConfigHash, profilesHash,
))
}
func capacityProfileInputsHash(inputs []CapacityProfileInput) string {
raw, _ := json.Marshal(inputs)
var value any
_ = json.Unmarshal(raw, &value)
canonical, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(value)))
return fmt.Sprintf("%x", sha256.Sum256(canonical))
}
func enforceTaskCapacityTx(ctx context.Context, tx pgx.Tx, input CreateTaskInput, runMode string) error {
if runMode != "production" && runMode != "acceptance" {
return nil
}
kind := capacityKindForTask(input.Kind)
if kind == "" {
return nil
}
model := strings.TrimSpace(input.Model)
profile, applies, err := capacityProfileForTaskTx(ctx, tx, input, runMode, kind, model)
if err != nil {
return err
}
if !applies {
return nil
}
var admissionLock bool
if err := tx.QueryRow(ctx, `
SELECT pg_try_advisory_xact_lock(
hashtextextended('capacity-profile:' || $1, 0)
)`, profile.ProfileKey).Scan(&admissionLock); err != nil {
return err
}
if !admissionLock {
return capacityExceededError(profile, "capacity_admission_busy", "admission_lock", 1, 1, 0, 0)
}
var queued int
var running int
var recent int
if err := tx.QueryRow(ctx, `
SELECT
count(*) FILTER (WHERE status = 'queued')::int,
count(*) FILTER (WHERE status = 'running')::int,
count(*) FILTER (WHERE created_at >= now() - interval '60 seconds')::int
FROM gateway_tasks
WHERE (
($3 = 'production' AND run_mode = 'production')
OR (
$3 = 'acceptance'
AND acceptance_run_id = NULLIF($4, '')::uuid
AND run_mode = 'acceptance'
)
)
AND (
($1 = 'image' AND kind IN ('images.generations', 'images.edits', 'images.vectorize'))
OR ($1 = 'video' AND kind IN ('videos.generations', 'videos.upscales'))
OR ($1 = 'mixed' AND kind IN (
'images.generations', 'images.edits', 'images.vectorize',
'videos.generations', 'videos.upscales'
))
)
AND ($2 = '*' OR model = $2)
AND (status IN ('queued', 'running') OR created_at >= now() - interval '60 seconds')`,
profile.Kind, profile.Model, runMode, strings.TrimSpace(input.AcceptanceRunID),
).Scan(&queued, &running, &recent); err != nil {
return err
}
nextQueued := queued + 1
nextRecentRate := float64(recent+1) / 60
predictedWait := float64(nextQueued) / profile.AdmittedThroughputPerSecond
reason := ""
metric := ""
current := float64(nextQueued)
limit := float64(profile.MaxQueued)
switch {
case nextQueued > profile.MaxQueued:
reason, metric = "capacity_queue_limit", "queue_depth"
case predictedWait > float64(profile.MaxPredictedWaitSeconds):
reason, metric = "capacity_predicted_wait_limit", "predicted_wait_seconds"
current, limit = predictedWait, float64(profile.MaxPredictedWaitSeconds)
case nextRecentRate > profile.AdmittedThroughputPerSecond:
reason, metric = "capacity_throughput_limit", "throughput_per_second"
current, limit = nextRecentRate, profile.AdmittedThroughputPerSecond
}
if reason == "" {
return nil
}
return capacityExceededError(profile, reason, metric, current, limit, queued, running)
}
func capacityProfileForTaskTx(
ctx context.Context,
tx pgx.Tx,
input CreateTaskInput,
runMode string,
kind string,
model string,
) (CapacityProfile, bool, error) {
if runMode == "acceptance" {
var profilesJSON []byte
var releaseSHA string
var capacityProfile string
var configHash string
err := tx.QueryRow(ctx, `
SELECT COALESCE(config->'validationCapacityProfiles', 'null'::jsonb),
release_sha,
capacity_profile,
COALESCE(config->>'validationConfigHash', '')
FROM gateway_acceptance_runs
WHERE id = NULLIF($1, '')::uuid
AND status = 'running'`, strings.TrimSpace(input.AcceptanceRunID)).Scan(
&profilesJSON, &releaseSHA, &capacityProfile, &configHash,
)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return CapacityProfile{}, false, nil
}
return CapacityProfile{}, false, err
}
if string(profilesJSON) == "null" {
return CapacityProfile{}, false, nil
}
var inputs []CapacityProfileInput
if err := json.Unmarshal(profilesJSON, &inputs); err != nil {
return CapacityProfile{}, false, fmt.Errorf("decode staged capacity profiles: %w", err)
}
best := -1
bestRank := 100
for index, candidate := range inputs {
candidateKind := strings.ToLower(strings.TrimSpace(candidate.Kind))
candidateModel := strings.TrimSpace(candidate.Model)
if candidateModel == "" {
candidateModel = "*"
}
rank := 100
switch {
case candidateKind == kind && candidateModel == model:
rank = 0
case candidateKind == kind && candidateModel == "*":
rank = 1
case candidateKind == "mixed" && candidateModel == "*":
rank = 2
}
if rank < bestRank {
best = index
bestRank = rank
}
}
if best < 0 {
return CapacityProfile{}, true, capacityProfileMissingError(kind, model)
}
selected := inputs[best]
selectedModel := strings.TrimSpace(selected.Model)
if selectedModel == "" {
selectedModel = "*"
}
return CapacityProfile{
ProfileKey: "acceptance:" + strings.TrimSpace(input.AcceptanceRunID) + ":" + kind + ":" + selectedModel,
ReleaseSHA: releaseSHA,
ConfigHash: configHash,
AcceptanceRunID: strings.TrimSpace(input.AcceptanceRunID),
Kind: strings.ToLower(strings.TrimSpace(selected.Kind)),
Model: selectedModel,
CapacityProfile: capacityProfile,
StableThroughputPerSecond: selected.StableThroughputPerSecond,
AdmittedThroughputPerSecond: selected.AdmittedThroughputPerSecond,
MaxRunning: selected.MaxRunning,
MaxQueued: selected.MaxQueued,
MaxPredictedWaitSeconds: selected.MaxPredictedWaitSeconds,
Metadata: selected.Metadata,
}, true, nil
}
profile, err := scanCapacityProfile(tx.QueryRow(ctx, `
SELECT profile_key, release_sha, config_hash, COALESCE(acceptance_run_id::text, ''),
kind, model, capacity_profile,
stable_throughput_per_second::float8,
admitted_throughput_per_second::float8,
max_running, max_queued, max_predicted_wait_seconds, enabled,
metadata, created_at, updated_at
FROM gateway_capacity_profiles
WHERE enabled = true
AND (
(kind = $1 AND model IN ($2, '*'))
OR (kind = 'mixed' AND model = '*')
)
ORDER BY CASE
WHEN kind = $1 AND model = $2 THEN 0
WHEN kind = $1 AND model = '*' THEN 1
ELSE 2
END
LIMIT 1`, kind, model))
if err != nil {
if IsUndefinedDatabaseObject(err) {
return CapacityProfile{}, false, nil
}
if errors.Is(err, pgx.ErrNoRows) {
var enabledProfiles int
if countErr := tx.QueryRow(ctx, `
SELECT count(*)::int
FROM gateway_capacity_profiles
WHERE enabled = true
AND kind IN ($1, 'mixed')`, kind).Scan(&enabledProfiles); countErr != nil {
return CapacityProfile{}, false, countErr
}
if enabledProfiles == 0 {
return CapacityProfile{}, false, nil
}
return CapacityProfile{}, true, capacityProfileMissingError(kind, model)
}
return CapacityProfile{}, false, err
}
return profile, true, nil
}
func capacityProfileMissingError(kind string, model string) error {
return &RateLimitExceededError{
ScopeType: "capacity_profile",
ScopeKey: kind + ":" + model,
ScopeName: "uncertified_model",
ScopeMetadata: map[string]any{
"kind": kind, "model": model,
},
Metric: "certified_capacity",
Limit: 0,
Amount: 1,
Current: 0,
Projected: 1,
Message: "model has no certified production capacity profile",
RetryAfter: 60 * time.Second,
Retryable: true,
Reason: "capacity_profile_missing",
}
}
func capacityExceededError(
profile CapacityProfile,
reason string,
metric string,
current float64,
limit float64,
queued int,
running int,
) error {
predictedWait := float64(queued) / profile.AdmittedThroughputPerSecond
retrySeconds := int(math.Ceil(predictedWait))
if retrySeconds < 1 {
retrySeconds = 1
}
return &RateLimitExceededError{
ScopeType: "capacity_profile",
ScopeKey: profile.ProfileKey,
ScopeName: profile.CapacityProfile,
ScopeMetadata: map[string]any{
"kind": profile.Kind,
"model": profile.Model,
"releaseSha": profile.ReleaseSHA,
"configHash": profile.ConfigHash,
"maxRunning": profile.MaxRunning,
"maxQueued": profile.MaxQueued,
"running": running,
"predictedWaitSeconds": predictedWait,
},
Metric: metric,
Limit: limit,
Amount: 1,
Current: current,
Projected: current + 1,
Message: "certified production capacity has been reached",
RetryAfter: time.Duration(retrySeconds) * time.Second,
Retryable: true,
Reason: reason,
QueueDepth: queued,
QueueLimit: profile.MaxQueued,
}
}
func capacityKindForTask(kind string) string {
switch strings.TrimSpace(kind) {
case "images.generations", "images.edits", "images.vectorize":
return "image"
case "videos.generations", "videos.upscales":
return "video"
default:
return ""
}
}
func scanCapacityProfile(scanner taskScanner) (CapacityProfile, error) {
var item CapacityProfile
var metadata []byte
if err := scanner.Scan(
&item.ProfileKey, &item.ReleaseSHA, &item.ConfigHash, &item.AcceptanceRunID,
&item.Kind, &item.Model, &item.CapacityProfile,
&item.StableThroughputPerSecond, &item.AdmittedThroughputPerSecond,
&item.MaxRunning, &item.MaxQueued, &item.MaxPredictedWaitSeconds,
&item.Enabled, &metadata, &item.CreatedAt, &item.UpdatedAt,
); err != nil {
return CapacityProfile{}, err
}
_ = json.Unmarshal(metadata, &item.Metadata)
if item.Metadata == nil {
item.Metadata = map[string]any{}
}
return item, nil
}
@@ -0,0 +1,51 @@
package store
import (
"strings"
"testing"
)
func TestCapacityKindForTask(t *testing.T) {
tests := map[string]string{
"images.generations": "image",
"images.edits": "image",
"images.vectorize": "image",
"videos.generations": "video",
"videos.upscales": "video",
"chat.completions": "",
}
for taskKind, expected := range tests {
if actual := capacityKindForTask(taskKind); actual != expected {
t.Fatalf("capacityKindForTask(%q)=%q, want %q", taskKind, actual, expected)
}
}
}
func TestCapacityProfileValidationAndStableHash(t *testing.T) {
profiles := []CapacityProfileInput{
{
Kind: "image", Model: "gemini-image",
StableThroughputPerSecond: 10, AdmittedThroughputPerSecond: 8,
MaxRunning: 48, MaxQueued: 100, MaxPredictedWaitSeconds: 60,
Metadata: map[string]any{"profile": "P24"},
},
{
Kind: "video", Model: "seedance",
StableThroughputPerSecond: 2, AdmittedThroughputPerSecond: 1.6,
MaxRunning: 48, MaxQueued: 100, MaxPredictedWaitSeconds: 120,
},
}
if err := validateCapacityProfileInputs(profiles); err != nil {
t.Fatalf("valid profiles were rejected: %v", err)
}
firstHash := capacityProfileInputsHash(profiles)
secondHash := capacityProfileInputsHash(profiles)
if len(firstHash) != 64 || firstHash != secondHash {
t.Fatalf("capacity profile hash is unstable: %q %q", firstHash, secondHash)
}
duplicate := append([]CapacityProfileInput(nil), profiles...)
duplicate = append(duplicate, profiles[0])
if err := validateCapacityProfileInputs(duplicate); err == nil || !strings.Contains(err.Error(), "duplicate") {
t.Fatalf("duplicate profiles validation error=%v", err)
}
}
+5
View File
@@ -14,6 +14,7 @@ const leadershipReleaseTimeout = 5 * time.Second
const (
identityCoordinatorLeadershipKey = "easyai-gateway-identity-coordinator"
securityEventHeartbeatLeadershipKey = "easyai-gateway-security-event-heartbeat"
capacityControllerLeadershipKey = "easyai-gateway-capacity-controller"
)
// Leadership holds a PostgreSQL session advisory lock. Callers must keep the
@@ -37,6 +38,10 @@ func (s *Store) TryAcquireSecurityEventHeartbeatLeadership(ctx context.Context)
return s.tryAcquireLeadership(ctx, securityEventHeartbeatLeadershipKey)
}
func (s *Store) TryAcquireCapacityControllerLeadership(ctx context.Context) (Leadership, bool, error) {
return s.tryAcquireLeadership(ctx, capacityControllerLeadershipKey)
}
func (s *Store) tryAcquireLeadership(ctx context.Context, key string) (Leadership, bool, error) {
conn, err := s.pool.Acquire(ctx)
if err != nil {
@@ -17,6 +17,8 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/identity"
"github.com/google/uuid"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/riverqueue/river/riverdriver/riverpgxv5"
"github.com/riverqueue/river/rivermigrate"
)
func TestResolveOrProvisionOIDCMultiTenantUserIsIdempotentIsolatedAndReusable(t *testing.T) {
@@ -598,4 +600,12 @@ func applyOIDCJITTestMigrations(t *testing.T, ctx context.Context, databaseURL s
t.Fatalf("commit migration %s: %v", version, err)
}
}
driver := riverpgxv5.New(pool)
migrator, err := rivermigrate.New(driver, nil)
if err != nil {
t.Fatalf("create River migrator: %v", err)
}
if _, err := migrator.Migrate(ctx, rivermigrate.DirectionUp, nil); err != nil {
t.Fatalf("apply River migrations: %v", err)
}
}
+3
View File
@@ -2151,6 +2151,9 @@ WHERE user_id = $1 AND idempotency_key_hash = $2`, user.ID, strings.TrimSpace(in
return CreateTaskResult{}, err
}
if !replayed {
if err := enforceTaskCapacityTx(ctx, tx, input, runMode); err != nil {
return CreateTaskResult{}, err
}
if err := insertTaskMessageRefs(ctx, tx, task.ID, input.MessageRefs); err != nil {
return CreateTaskResult{}, err
}
+346 -217
View File
@@ -19,7 +19,11 @@ type RuntimeRecoveryResult struct {
CleanedTaskAdmissions int64 `json:"cleanedTaskAdmissions"`
}
const runtimeRecoveryBatchSize = 100
const (
runtimeRecoveryBatchSize = 50
runtimeWorkerStaleAfter = 60 * time.Second
runtimeCounterRepairBatch = 100
)
var ErrConcurrencyLeaseLost = errors.New("concurrency lease lost")
@@ -364,11 +368,24 @@ func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leases []Concurren
if len(leaseIDs) == 0 {
return nil
}
_, err := s.pool.Exec(ctx, `
tx, err := s.pool.Begin(ctx)
if err != nil {
return err
}
defer rollbackTransaction(tx)
// A lost release after database failover is conservative: the lease
// remains unavailable until its TTL or task-admission cleanup. It cannot
// create duplicate execution, so avoid a dedicated synchronous wait.
if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil {
return err
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE id = ANY($1::uuid[]) AND released_at IS NULL`, leaseIDs)
return err
WHERE id = ANY($1::uuid[]) AND released_at IS NULL`, leaseIDs); err != nil {
return err
}
return tx.Commit(ctx)
}
func (s *Store) RenewConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error {
@@ -417,6 +434,12 @@ func (s *Store) AttachRateLimitResultToAttempt(ctx context.Context, attemptID st
return err
}
defer rollbackTransaction(tx)
// The subsequent durable `submitting` transition is the safety barrier for
// these task-scoped associations. If execution stops before that point,
// the task admission owns and reaps the leases by task_id.
if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil {
return err
}
for _, reservation := range result.Reservations {
if reservation.ReservationID == "" {
@@ -456,106 +479,121 @@ func concurrencyLeaseIDs(leases []ConcurrencyLease) []string {
}
func (s *Store) RecoverInterruptedRuntimeState(ctx context.Context) (RuntimeRecoveryResult, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
var result RuntimeRecoveryResult
if err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
recovered, err := recoverExpiredRuntimeTasksTx(ctx, tx)
if err != nil {
return err
}
result = recovered
return nil
}); err != nil {
return RuntimeRecoveryResult{}, err
}
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-runtime-recovery', 0))`); err != nil {
if err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
cleaned, err := cleanupTerminalRuntimeStateTx(ctx, tx)
if err != nil {
return err
}
result.ReleasedConcurrencyLeases += cleaned.ReleasedConcurrencyLeases
result.ReleasedRateReservations += cleaned.ReleasedRateReservations
result.CleanedTaskAdmissions += cleaned.CleanedTaskAdmissions
return nil
}); err != nil {
return RuntimeRecoveryResult{}, err
}
if err := s.repairRuntimeClientCounters(ctx); err != nil {
return RuntimeRecoveryResult{}, err
}
if result.RequeuedAsyncTasks > 0 || result.FailedTasks > 0 || result.CleanedTaskAdmissions > 0 {
s.notifyTaskAdmissionBestEffort(ctx, "*")
}
return result, nil
}
result := RuntimeRecoveryResult{}
type runtimeRecoveryTask struct {
ID string
Async bool
Production bool
HasGatewayUser bool
SubmissionUnsafe bool
}
func recoverExpiredRuntimeTasksTx(ctx context.Context, tx pgx.Tx) (RuntimeRecoveryResult, error) {
rows, err := tx.Query(ctx, `
UPDATE gateway_rate_limit_reservations reservation
SET status = 'released',
reason = 'execution_lease_expired',
finalized_at = now(),
updated_at = now()
SELECT task.id::text,
task.async_mode,
task.run_mode IN ('production', 'acceptance', 'acceptance_canary'),
task.gateway_user_id IS NOT NULL,
COALESCE(task.remote_task_id, '') = ''
AND COALESCE((
SELECT (
(attempt.status = 'running' AND attempt.upstream_submission_status IN ('submitting', 'response_received'))
OR (attempt.status = 'failed' AND attempt.upstream_submission_status = 'submitting')
OR (
attempt.status = 'failed'
AND attempt.upstream_submission_status = 'response_received'
AND attempt.error_code = 'execution_lease_expired'
)
)
FROM gateway_task_attempts attempt
WHERE attempt.task_id = task.id
ORDER BY attempt.attempt_no DESC, attempt.started_at DESC
LIMIT 1
), false)
FROM gateway_tasks task
WHERE reservation.task_id = task.id
AND reservation.status = 'reserved'
WHERE task.status = 'running'
AND task.execution_lease_expires_at IS NOT NULL
AND task.execution_lease_expires_at <= now()
AND (
task.status IN ('succeeded', 'failed', 'cancelled')
OR (
task.status IN ('queued', 'running')
AND task.execution_lease_expires_at IS NOT NULL
AND task.execution_lease_expires_at <= now()
task.locked_by IS NULL
OR NOT EXISTS (
SELECT 1
FROM gateway_worker_instances worker
WHERE worker.instance_id = task.locked_by
AND worker.status = 'active'
AND worker.heartbeat_at > now() - $1::interval
)
)
RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`)
ORDER BY task.execution_lease_expires_at ASC, task.id ASC
LIMIT $2
FOR UPDATE OF task SKIP LOCKED`, runtimeWorkerStaleAfter.String(), runtimeRecoveryBatchSize)
if err != nil {
return RuntimeRecoveryResult{}, err
}
reservations := make([]RateLimitReservation, 0)
tasks := make([]runtimeRecoveryTask, 0, runtimeRecoveryBatchSize)
for rows.Next() {
var reservation RateLimitReservation
if err := rows.Scan(&reservation.ScopeType, &reservation.ScopeKey, &reservation.Metric, &reservation.WindowStart, &reservation.Amount); err != nil {
var task runtimeRecoveryTask
if err := rows.Scan(&task.ID, &task.Async, &task.Production, &task.HasGatewayUser, &task.SubmissionUnsafe); err != nil {
rows.Close()
return RuntimeRecoveryResult{}, err
}
reservations = append(reservations, reservation)
tasks = append(tasks, task)
}
if err := rows.Err(); err != nil {
rows.Close()
return RuntimeRecoveryResult{}, err
}
rows.Close()
for _, reservation := range reservations {
if err := releaseCounterReservation(ctx, tx, reservation.ScopeType, reservation.ScopeKey, reservation.Metric, reservation.WindowStart, reservation.Amount); err != nil {
return RuntimeRecoveryResult{}, err
}
if len(tasks) == 0 {
return RuntimeRecoveryResult{}, nil
}
result.ReleasedRateReservations = int64(len(reservations))
tag, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE released_at IS NULL
AND expires_at <= now()`)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.ReleasedConcurrencyLeases = tag.RowsAffected()
tag, err = tx.Exec(ctx, `
UPDATE gateway_task_attempts
SET status = 'failed',
retryable = true,
error_code = 'execution_lease_expired',
error_message = 'attempt execution lease expired',
finished_at = now()
WHERE status = 'running'
AND EXISTS (
SELECT 1
FROM gateway_tasks task
WHERE task.id = gateway_task_attempts.task_id
AND (
task.status IN ('succeeded', 'failed', 'cancelled')
OR (
task.execution_lease_expires_at IS NOT NULL
AND task.execution_lease_expires_at <= now()
)
)
)`)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.FailedAttempts = tag.RowsAffected()
asyncTaskRows, err := tx.Query(ctx, `
WITH recoverable_async_tasks AS MATERIALIZED (
SELECT id
FROM gateway_tasks
WHERE async_mode = true
AND status = 'running'
AND execution_lease_expires_at IS NOT NULL
AND execution_lease_expires_at <= now()
ORDER BY execution_lease_expires_at ASC, id ASC
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_tasks task
result := RuntimeRecoveryResult{}
taskIDs := make([]string, 0, len(tasks))
releaseBillingTaskIDs := make([]string, 0, len(tasks))
for _, task := range tasks {
taskIDs = append(taskIDs, task.ID)
eventType := "task.failed"
eventStatus := "failed"
if task.Production && task.SubmissionUnsafe {
if err := markTaskExecutionManualReviewTx(ctx, tx, task.ID, task.HasGatewayUser); err != nil {
return RuntimeRecoveryResult{}, err
}
result.FailedTasks++
} else if task.Async {
if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'queued',
error = NULL,
error_code = NULL,
@@ -568,156 +606,254 @@ SET status = 'queued',
next_run_at = now(),
finished_at = NULL,
updated_at = now()
FROM recoverable_async_tasks recoverable
WHERE task.id = recoverable.id
RETURNING task.id::text`, runtimeRecoveryBatchSize)
if err != nil {
return RuntimeRecoveryResult{}, err
}
asyncTaskIDs := make([]string, 0)
for asyncTaskRows.Next() {
var taskID string
if err := asyncTaskRows.Scan(&taskID); err != nil {
asyncTaskRows.Close()
return RuntimeRecoveryResult{}, err
}
asyncTaskIDs = append(asyncTaskIDs, taskID)
}
if err := asyncTaskRows.Err(); err != nil {
asyncTaskRows.Close()
return RuntimeRecoveryResult{}, err
}
asyncTaskRows.Close()
for _, taskID := range asyncTaskIDs {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
SELECT
task.id,
COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1),
'task.queued',
'queued',
NULL,
0,
NULL,
'{}'::jsonb,
false
FROM gateway_tasks task
WHERE task.id = $1::uuid
AND (
SELECT count(*)
FROM gateway_task_events event
WHERE event.task_id = task.id
AND event.event_type NOT IN (
'task.completed', 'task.failed', 'task.cancelled',
'task.billing.settled', 'task.billing.released', 'task.billing.review'
)
) < 16
AND NOT EXISTS (
SELECT 1
FROM gateway_task_events event
WHERE event.task_id = task.id
AND event.seq = (SELECT MAX(last_event.seq) FROM gateway_task_events last_event WHERE last_event.task_id = task.id)
AND event.event_type = 'task.queued'
AND COALESCE(event.status, '') = 'queued'
AND event.platform_id IS NULL
AND event.simulated = false
)`, taskID); err != nil {
return RuntimeRecoveryResult{}, err
}
}
result.RequeuedAsyncTasks = int64(len(asyncTaskIDs))
if len(asyncTaskIDs) > 0 {
tag, err = tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = ANY($1::uuid[])
AND released_at IS NULL`, asyncTaskIDs)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.ReleasedConcurrencyLeases += tag.RowsAffected()
tag, err = tx.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = ANY($1::uuid[])`, asyncTaskIDs)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.CleanedTaskAdmissions += tag.RowsAffected()
}
taskRows, err := tx.Query(ctx, `
WHERE id = $1::uuid`, task.ID); err != nil {
return RuntimeRecoveryResult{}, err
}
eventType = "task.queued"
eventStatus = "queued"
result.RequeuedAsyncTasks++
} else {
if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
billing_status = CASE
WHEN gateway_user_id IS NULL THEN 'not_required'
WHEN reservation_amount > 0 THEN 'pending'
ELSE 'released'
END,
billing_updated_at = now(),
error = NULL,
error_code = 'server_restarted',
error_message = 'task interrupted by service restart',
error_message = 'task interrupted after its owner and execution lease expired',
remote_task_payload = '{}'::jsonb,
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
finished_at = now(),
updated_at = now()
WHERE async_mode = false
AND status = 'running'
AND execution_lease_expires_at IS NOT NULL
AND execution_lease_expires_at <= now()
RETURNING id::text`)
WHERE id = $1::uuid`, task.ID); err != nil {
return RuntimeRecoveryResult{}, err
}
if task.Production && task.HasGatewayUser {
releaseBillingTaskIDs = append(releaseBillingTaskIDs, task.ID)
}
result.FailedTasks++
}
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
SELECT $1::uuid,
COALESCE(MAX(event.seq), 0) + 1,
$2,
$3,
'runtime_recovery',
0,
NULL,
'{}'::jsonb,
false
FROM gateway_task_events event
WHERE event.task_id = $1::uuid
HAVING NOT EXISTS (
SELECT 1
FROM gateway_task_events latest
WHERE latest.task_id = $1::uuid
AND latest.seq = (SELECT MAX(last_event.seq) FROM gateway_task_events last_event WHERE last_event.task_id = $1::uuid)
AND latest.event_type = $2
AND latest.status = $3
)`, task.ID, eventType, eventStatus); err != nil {
return RuntimeRecoveryResult{}, err
}
}
if len(releaseBillingTaskIDs) > 0 {
if _, err := tx.Exec(ctx, `
INSERT INTO settlement_outbox (
task_id, event_type, action, amount, currency, pricing_snapshot, payload,
status, next_attempt_at
)
SELECT task.id, 'task.billing.release', 'release', task.reservation_amount,
task.billing_currency, task.pricing_snapshot,
jsonb_build_object('taskId', task.id, 'reason', 'execution_lease_expired'),
'pending', now()
FROM gateway_tasks task
WHERE task.id = ANY($1::uuid[])
AND task.reservation_amount > 0
ON CONFLICT (task_id, event_type) DO NOTHING`, releaseBillingTaskIDs); err != nil {
return RuntimeRecoveryResult{}, err
}
}
tag, err := tx.Exec(ctx, `
UPDATE gateway_task_attempts attempt
SET status = 'failed',
retryable = task.error_code IS DISTINCT FROM 'upstream_submission_unknown',
error_code = CASE
WHEN task.error_code = 'upstream_submission_unknown' THEN 'upstream_submission_unknown'
ELSE 'execution_lease_expired'
END,
error_message = CASE
WHEN task.error_code = 'upstream_submission_unknown' THEN 'upstream submission result is unknown'
ELSE 'attempt execution lease expired after worker heartbeat became stale'
END,
finished_at = now()
FROM gateway_tasks task
WHERE attempt.task_id = task.id
AND attempt.task_id = ANY($1::uuid[])
AND attempt.status = 'running'`, taskIDs)
if err != nil {
return RuntimeRecoveryResult{}, err
}
taskIDs := make([]string, 0)
for taskRows.Next() {
result.FailedAttempts = tag.RowsAffected()
tag, err = tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = ANY($1::uuid[])
AND released_at IS NULL`, taskIDs)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.ReleasedConcurrencyLeases = tag.RowsAffected()
tag, err = tx.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = ANY($1::uuid[])`, taskIDs)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.CleanedTaskAdmissions = tag.RowsAffected()
released, err := releaseRuntimeRateReservationsTx(ctx, tx, taskIDs)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.ReleasedRateReservations = released
return result, nil
}
func cleanupTerminalRuntimeStateTx(ctx context.Context, tx pgx.Tx) (RuntimeRecoveryResult, error) {
rows, err := tx.Query(ctx, `
SELECT task.id::text
FROM gateway_tasks task
WHERE task.status IN ('succeeded', 'failed', 'cancelled')
AND (
EXISTS (
SELECT 1 FROM gateway_concurrency_leases lease
WHERE lease.task_id = task.id AND lease.released_at IS NULL
)
OR EXISTS (
SELECT 1 FROM gateway_task_admissions admission
WHERE admission.task_id = task.id
)
OR EXISTS (
SELECT 1 FROM gateway_rate_limit_reservations reservation
WHERE reservation.task_id = task.id AND reservation.status = 'reserved'
)
)
ORDER BY task.updated_at ASC, task.id ASC
LIMIT $1
FOR UPDATE OF task SKIP LOCKED`, runtimeRecoveryBatchSize)
if err != nil {
return RuntimeRecoveryResult{}, err
}
taskIDs := make([]string, 0, runtimeRecoveryBatchSize)
for rows.Next() {
var taskID string
if err := taskRows.Scan(&taskID); err != nil {
taskRows.Close()
if err := rows.Scan(&taskID); err != nil {
rows.Close()
return RuntimeRecoveryResult{}, err
}
taskIDs = append(taskIDs, taskID)
}
if err := taskRows.Err(); err != nil {
taskRows.Close()
if err := rows.Err(); err != nil {
rows.Close()
return RuntimeRecoveryResult{}, err
}
taskRows.Close()
rows.Close()
if len(taskIDs) == 0 {
return RuntimeRecoveryResult{}, nil
}
var result RuntimeRecoveryResult
tag, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = ANY($1::uuid[])
AND released_at IS NULL`, taskIDs)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.ReleasedConcurrencyLeases = tag.RowsAffected()
tag, err = tx.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = ANY($1::uuid[])`, taskIDs)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.CleanedTaskAdmissions = tag.RowsAffected()
result.ReleasedRateReservations, err = releaseRuntimeRateReservationsTx(ctx, tx, taskIDs)
return result, err
}
for _, taskID := range taskIDs {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
VALUES (
$1::uuid,
COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1),
'task.failed',
'failed',
NULL,
0,
NULL,
'{}'::jsonb,
false
)`, taskID); err != nil {
return RuntimeRecoveryResult{}, err
func releaseRuntimeRateReservationsTx(ctx context.Context, tx pgx.Tx, taskIDs []string) (int64, error) {
rows, err := tx.Query(ctx, `
UPDATE gateway_rate_limit_reservations reservation
SET status = 'released',
reason = 'runtime_state_recovered',
finalized_at = now(),
updated_at = now()
WHERE reservation.task_id = ANY($1::uuid[])
AND reservation.status = 'reserved'
RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`, taskIDs)
if err != nil {
return 0, err
}
reservations := make([]RateLimitReservation, 0, len(taskIDs))
for rows.Next() {
var reservation RateLimitReservation
if err := rows.Scan(
&reservation.ScopeType,
&reservation.ScopeKey,
&reservation.Metric,
&reservation.WindowStart,
&reservation.Amount,
); err != nil {
rows.Close()
return 0, err
}
reservations = append(reservations, reservation)
}
if err := rows.Err(); err != nil {
rows.Close()
return 0, err
}
rows.Close()
for _, reservation := range reservations {
if err := releaseCounterReservation(
ctx,
tx,
reservation.ScopeType,
reservation.ScopeKey,
reservation.Metric,
reservation.WindowStart,
reservation.Amount,
); err != nil {
return 0, err
}
}
result.FailedTasks = int64(len(taskIDs))
tag, err = tx.Exec(ctx, `
UPDATE gateway_concurrency_leases lease
SET released_at = now()
FROM gateway_tasks task
WHERE lease.task_id = task.id
AND lease.released_at IS NULL
AND task.status IN ('succeeded', 'failed', 'cancelled')`)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.ReleasedConcurrencyLeases += tag.RowsAffected()
tag, err = tx.Exec(ctx, `
DELETE FROM gateway_task_admissions admission
USING gateway_tasks task
WHERE admission.task_id = task.id
AND task.status IN ('succeeded', 'failed', 'cancelled')`)
if err != nil {
return RuntimeRecoveryResult{}, err
}
result.CleanedTaskAdmissions += tag.RowsAffected()
if _, err := tx.Exec(ctx, `
return int64(len(reservations)), nil
}
func (s *Store) repairRuntimeClientCounters(ctx context.Context) error {
_, err := s.pool.Exec(ctx, `
WITH candidates AS MATERIALIZED (
SELECT state.client_id
FROM runtime_client_states state
WHERE state.running_count IS DISTINCT FROM (
SELECT count(*)
FROM gateway_task_attempts attempt
WHERE attempt.client_id = state.client_id
AND attempt.status = 'running'
)
ORDER BY state.updated_at ASC, state.client_id ASC
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE runtime_client_states state
SET running_count = (
SELECT count(*)
@@ -726,16 +862,9 @@ SET running_count = (
AND attempt.status = 'running'
),
updated_at = now()
WHERE state.running_count IS DISTINCT FROM (
SELECT count(*)
FROM gateway_task_attempts attempt
WHERE attempt.client_id = state.client_id
AND attempt.status = 'running'
)`); err != nil {
return RuntimeRecoveryResult{}, err
}
return result, tx.Commit(ctx)
FROM candidates
WHERE state.client_id = candidates.client_id`, runtimeCounterRepairBatch)
return err
}
func (s *Store) finishRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64, status string, reason string) error {
+30 -23
View File
@@ -284,19 +284,20 @@ type AsyncTaskQueueItem struct {
}
type FinishTaskAttemptInput struct {
AttemptID string
Status string
Retryable bool
RequestID string
StatusCode int
Usage map[string]any
Metrics map[string]any
ResponseSnapshot map[string]any
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
ErrorCode string
ErrorMessage string
AttemptID string
Status string
Retryable bool
UpstreamSubmissionStatus string
RequestID string
StatusCode int
Usage map[string]any
Metrics map[string]any
ResponseSnapshot map[string]any
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
ErrorCode string
ErrorMessage string
}
type FinishTaskSuccessInput struct {
@@ -318,6 +319,9 @@ type FinishTaskSuccessInput struct {
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
CompletionCallbackURL string
Simulated bool
FinalizeAdmission bool
}
type FinishTaskManualReviewInput struct {
@@ -337,16 +341,19 @@ type FinishTaskManualReviewInput struct {
}
type FinishTaskFailureInput struct {
TaskID string
ExecutionToken string
Code string
Message string
Result map[string]any
RequestID string
Metrics map[string]any
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
TaskID string
ExecutionToken string
Code string
Message string
Result map[string]any
RequestID string
Metrics map[string]any
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
CompletionCallbackURL string
Simulated bool
FinalizeAdmission bool
}
type CreateTaskParamPreprocessingLogInput struct {
+26 -11
View File
@@ -192,10 +192,13 @@ WITH picked AS (
OR compatibility_submit_http_status IS NOT NULL
OR COALESCE(compatibility_submit_headers, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(compatibility_submit_body, '{}'::jsonb) <> '{}'::jsonb
OR result ?| ARRAY[
'raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse',
'submit', 'file_retrieve', 'fileRetrieve', 'upstream_task_id', 'remote_task_id'
]
OR (
jsonb_typeof(COALESCE(result, '{}'::jsonb)) = 'object'
AND result ?| ARRAY[
'raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse',
'submit', 'file_retrieve', 'fileRetrieve', 'upstream_task_id', 'remote_task_id'
]
)
OR error IS NOT NULL
OR octet_length(COALESCE(error_message, '')) > 2048
ORDER BY updated_at
@@ -208,11 +211,15 @@ SET normalized_request = '{}'::jsonb,
compatibility_submit_http_status = NULL,
compatibility_submit_headers = '{}'::jsonb,
compatibility_submit_body = '{}'::jsonb,
result = COALESCE(task.result, '{}'::jsonb)
- 'raw' - 'raw_data' - 'rawData'
- 'provider_response' - 'providerResponse'
- 'submit' - 'file_retrieve' - 'fileRetrieve'
- 'upstream_task_id' - 'remote_task_id',
result = CASE
WHEN jsonb_typeof(COALESCE(task.result, '{}'::jsonb)) = 'object'
THEN COALESCE(task.result, '{}'::jsonb)
- 'raw' - 'raw_data' - 'rawData'
- 'provider_response' - 'providerResponse'
- 'submit' - 'file_retrieve' - 'fileRetrieve'
- 'upstream_task_id' - 'remote_task_id'
ELSE task.result
END,
error = NULL,
error_message = CASE
WHEN octet_length(COALESCE(task.error_message, '')) > 2048 THEN left(task.error_message, 512)
@@ -220,6 +227,8 @@ SET normalized_request = '{}'::jsonb,
END,
remote_task_payload = CASE
WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb
WHEN jsonb_typeof(COALESCE(task.remote_task_payload, '{}'::jsonb)) <> 'object'
THEN task.remote_task_payload
ELSE (
SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb)
FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry
@@ -238,7 +247,9 @@ WITH picked AS (
WHERE COALESCE(remote_task_payload, '{}'::jsonb) <> '{}'::jsonb
AND (
status IN ('succeeded', 'failed', 'cancelled')
OR EXISTS (
OR (
jsonb_typeof(COALESCE(remote_task_payload, '{}'::jsonb)) = 'object'
AND EXISTS (
SELECT 1
FROM jsonb_object_keys(COALESCE(remote_task_payload, '{}'::jsonb)) AS key
WHERE key NOT IN (
@@ -246,6 +257,7 @@ WITH picked AS (
'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
)
)
)
)
ORDER BY updated_at
LIMIT $1
@@ -254,6 +266,8 @@ WITH picked AS (
UPDATE gateway_tasks task
SET remote_task_payload = CASE
WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb
WHEN jsonb_typeof(COALESCE(task.remote_task_payload, '{}'::jsonb)) <> 'object'
THEN task.remote_task_payload
ELSE (
SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb)
FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry
@@ -375,7 +389,8 @@ WHERE log.id = picked.id`, []any{batchSize}},
WITH picked AS (
SELECT id
FROM gateway_cloned_voices
WHERE metadata ?| ARRAY['raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', 'request']
WHERE jsonb_typeof(COALESCE(metadata, '{}'::jsonb)) = 'object'
AND metadata ?| ARRAY['raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', 'request']
ORDER BY updated_at
LIMIT $1
FOR UPDATE SKIP LOCKED
+191 -14
View File
@@ -281,6 +281,16 @@ func nullableTaskListTime(value *time.Time) any {
}
func (s *Store) ClaimTaskExecution(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) {
return s.ClaimTaskExecutionForOwner(ctx, taskID, executionToken, leaseTTL, "")
}
func (s *Store) ClaimTaskExecutionForOwner(
ctx context.Context,
taskID string,
executionToken string,
leaseTTL time.Duration,
ownerInstanceID string,
) (GatewayTask, error) {
if leaseTTL <= 0 {
leaseTTL = 5 * time.Minute
}
@@ -323,6 +333,7 @@ UPDATE gateway_tasks
SET status = 'running',
execution_token = $2::uuid,
execution_lease_expires_at = now() + ($3::int * interval '1 second'),
locked_by = NULLIF($4::text, ''),
locked_at = now(),
heartbeat_at = now(),
updated_at = now()
@@ -331,7 +342,7 @@ WHERE id = $1::uuid
(status = 'queued' AND next_run_at <= now())
OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()))
)
RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second), strings.TrimSpace(ownerInstanceID)))
return err
})
if err != nil {
@@ -349,7 +360,8 @@ RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second
func taskExecutionRequiresManualReviewTx(ctx context.Context, tx pgx.Tx, taskID string) (bool, error) {
var submissionAmbiguous bool
err := tx.QueryRow(ctx, `
SELECT COALESCE((
SELECT COALESCE(task.remote_task_id, '') = ''
AND COALESCE((
SELECT (
(status = 'running' AND upstream_submission_status IN ('submitting', 'response_received'))
OR (status = 'failed' AND upstream_submission_status = 'submitting')
@@ -363,7 +375,9 @@ SELECT COALESCE((
WHERE task_id = $1::uuid
ORDER BY attempt_no DESC, started_at DESC
LIMIT 1
), false)`, taskID).Scan(&submissionAmbiguous)
), false)
FROM gateway_tasks task
WHERE task.id = $1::uuid`, taskID).Scan(&submissionAmbiguous)
return submissionAmbiguous, err
}
@@ -414,6 +428,16 @@ ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err
}
func (s *Store) ClaimTaskPreparation(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) {
return s.ClaimTaskPreparationForOwner(ctx, taskID, executionToken, leaseTTL, "")
}
func (s *Store) ClaimTaskPreparationForOwner(
ctx context.Context,
taskID string,
executionToken string,
leaseTTL time.Duration,
ownerInstanceID string,
) (GatewayTask, error) {
if leaseTTL <= 0 {
leaseTTL = 5 * time.Minute
}
@@ -459,6 +483,7 @@ FOR UPDATE`, taskID).Scan(&queuedReady, &production, &hasGatewayUser); err != ni
UPDATE gateway_tasks
SET execution_token = $2::uuid,
execution_lease_expires_at = now() + ($3::int * interval '1 second'),
locked_by = NULLIF($4::text, ''),
locked_at = now(),
heartbeat_at = now(),
updated_at = now()
@@ -470,7 +495,7 @@ WHERE id = $1::uuid
OR execution_lease_expires_at IS NULL
OR execution_lease_expires_at <= now()
)
RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second), strings.TrimSpace(ownerInstanceID)))
return err
})
if errors.Is(err, pgx.ErrNoRows) {
@@ -565,7 +590,19 @@ SELECT COALESCE((SELECT status IN ('queued', 'running') FROM gateway_tasks WHERE
}
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, _ map[string]any) error {
tag, err := s.pool.Exec(ctx, `
tx, err := s.pool.Begin(ctx)
if err != nil {
return err
}
defer rollbackTransaction(tx)
// The execution claim is the durable ownership boundary. This derived
// status update is always followed by either the durable upstream
// submission barrier or a terminal task transition, so it must be locally
// visible immediately but does not need its own synchronous-replica wait.
if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil {
return err
}
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
model_type = NULLIF($3::text, ''),
@@ -582,7 +619,7 @@ WHERE id = $1::uuid
if tag.RowsAffected() != 1 {
return ErrTaskExecutionLeaseLost
}
return nil
return tx.Commit(ctx)
}
func (s *Store) ClaimAsyncQueuedTask(ctx context.Context, workerID string) (GatewayTask, error) {
@@ -1107,6 +1144,13 @@ func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptIn
return "", err
}
defer rollbackTransaction(tx)
// OnUpstreamSubmissionStarted synchronously persists `submitting` before
// any provider request leaves the process. PostgreSQL WAL ordering makes
// that later barrier cover this attempt row as well, avoiding a redundant
// cross-region wait without weakening duplicate-submission protection.
if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil {
return "", err
}
if strings.TrimSpace(input.ExecutionToken) != "" {
var ownsExecution bool
@@ -1447,6 +1491,11 @@ func taskAttemptMetricInt(metrics map[string]any, key string) int {
}
func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptInput) error {
if input.UpstreamSubmissionStatus != "" {
if err := validateAttemptUpstreamSubmissionStatus(input.UpstreamSubmissionStatus); err != nil {
return err
}
}
statusCode := input.StatusCode
if statusCode == 0 {
statusCode = taskAttemptMetricInt(input.Metrics, "statusCode")
@@ -1469,6 +1518,11 @@ SET status = $2::text,
response_duration_ms = $8,
error_code = NULLIF($9::text, ''),
error_message = NULLIF(left($10::text, 2048), ''),
upstream_submission_status = COALESCE(NULLIF($13::text, ''), upstream_submission_status),
upstream_submission_updated_at = CASE
WHEN NULLIF($13::text, '') IS NULL THEN upstream_submission_updated_at
ELSE now()
END,
finished_at = now()
WHERE id = $1::uuid`,
input.AttemptID,
@@ -1483,6 +1537,7 @@ WHERE id = $1::uuid`,
truncateUTF8Bytes(input.ErrorMessage, 2048),
usageJSON,
metricsJSON,
input.UpstreamSubmissionStatus,
)
return err
}
@@ -1542,6 +1597,7 @@ func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessIn
finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64)
}
currency := normalizeWalletCurrency(input.BillingCurrency)
finalizedAdmission := false
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
if strings.TrimSpace(input.AttemptID) != "" {
if _, err := tx.Exec(ctx, `
@@ -1571,7 +1627,8 @@ WHERE id = $1::uuid`,
return err
}
}
tag, err := tx.Exec(ctx, `
var billingStatus string
err := tx.QueryRow(ctx, `
UPDATE gateway_tasks
SET status = 'succeeded',
result = $2::jsonb,
@@ -1607,7 +1664,8 @@ SET status = 'succeeded',
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $16::uuid`,
AND execution_token = $16::uuid
RETURNING billing_status`,
input.TaskID,
string(resultJSON),
string(billingsJSON),
@@ -1624,13 +1682,13 @@ WHERE id = $1::uuid
nullableTime(input.ResponseFinishedAt),
input.ResponseDurationMS,
input.ExecutionToken,
)
).Scan(&billingStatus)
if errors.Is(err, pgx.ErrNoRows) {
return ErrTaskExecutionLeaseLost
}
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrTaskExecutionLeaseLost
}
payloadJSON, _ := json.Marshal(map[string]any{"taskId": input.TaskID, "pricingVersion": "effective-pricing-v2"})
_, err = tx.Exec(ctx, `
INSERT INTO settlement_outbox (
@@ -1643,11 +1701,66 @@ WHERE id = $1::uuid
AND gateway_user_id IS NOT NULL
ON CONFLICT (task_id, event_type) DO NOTHING`,
input.TaskID, finalChargeAmount, currency, string(pricingSnapshotJSON), string(payloadJSON))
return err
if err != nil {
return err
}
if input.FinalizeAdmission {
if _, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = $1::uuid
AND released_at IS NULL`, input.TaskID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = $1::uuid`, input.TaskID); err != nil {
return err
}
finalizedAdmission = true
}
var eventID string
var eventSeq int64
if err := tx.QueryRow(ctx, `
WITH next_seq AS (
SELECT COALESCE(MAX(seq), 0) + 1 AS seq
FROM gateway_task_events
WHERE task_id = $1::uuid
)
INSERT INTO gateway_task_events (
task_id, seq, event_type, status, phase, progress, message, payload, simulated
)
SELECT $1::uuid, next_seq.seq, 'task.completed', 'succeeded', NULL, 0, NULL, '{}'::jsonb, $2
FROM next_seq
RETURNING id::text, seq`,
input.TaskID,
input.Simulated,
).Scan(&eventID, &eventSeq); err != nil {
return err
}
callbackURL := strings.TrimSpace(input.CompletionCallbackURL)
if callbackURL != "" {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload)
VALUES ($1::uuid, $2::uuid, $3, $4, '{}'::jsonb)
ON CONFLICT (task_id, seq, callback_url) DO NOTHING`,
input.TaskID,
eventID,
eventSeq,
callbackURL,
); err != nil {
return err
}
}
_ = billingStatus
return nil
})
if err != nil {
return GatewayTask{}, err
}
if finalizedAdmission {
s.notifyTaskAdmissionBestEffort(ctx, "*")
}
return s.GetTask(ctx, input.TaskID)
}
@@ -1903,6 +2016,7 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics)))
resultJSON, _ := json.Marshal(minimalTaskResult(nil))
message := truncateUTF8Bytes(input.Message, 2048)
finalizedAdmission := false
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
@@ -1963,11 +2077,65 @@ WHERE id = $1::uuid
AND gateway_user_id IS NOT NULL
AND reservation_amount > 0
ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, string(payloadJSON))
return err
if err != nil {
return err
}
if input.FinalizeAdmission {
if _, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = $1::uuid
AND released_at IS NULL`, input.TaskID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = $1::uuid`, input.TaskID); err != nil {
return err
}
finalizedAdmission = true
}
var eventID string
var eventSeq int64
if err := tx.QueryRow(ctx, `
WITH next_seq AS (
SELECT COALESCE(MAX(seq), 0) + 1 AS seq
FROM gateway_task_events
WHERE task_id = $1::uuid
)
INSERT INTO gateway_task_events (
task_id, seq, event_type, status, phase, progress, message, payload, simulated
)
SELECT $1::uuid, next_seq.seq, 'task.failed', 'failed', NULL, 0, NULL, '{}'::jsonb, $2
FROM next_seq
RETURNING id::text, seq`,
input.TaskID,
input.Simulated,
).Scan(&eventID, &eventSeq); err != nil {
return err
}
callbackURL := strings.TrimSpace(input.CompletionCallbackURL)
if callbackURL != "" {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload)
VALUES ($1::uuid, $2::uuid, $3, $4, '{}'::jsonb)
ON CONFLICT (task_id, seq, callback_url) DO NOTHING`,
input.TaskID,
eventID,
eventSeq,
callbackURL,
); err != nil {
return err
}
}
return nil
})
if err != nil {
return GatewayTask{}, err
}
if finalizedAdmission {
s.notifyTaskAdmissionBestEffort(ctx, "*")
}
return s.GetTask(ctx, input.TaskID)
}
@@ -2093,6 +2261,15 @@ func (s *Store) addTaskEvent(
return TaskEvent{}, err
}
defer rollbackTransaction(tx)
if !taskEventTerminal(eventType) {
// Progress/history events are recoverable hints. A later synchronous
// upstream or terminal task transition is their durability barrier, so
// waiting for the cross-region replica on every event only amplifies
// queue latency under media bursts.
if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil {
return TaskEvent{}, err
}
}
if _, err := tx.Exec(ctx, `SELECT 1 FROM gateway_tasks WHERE id = $1::uuid FOR UPDATE`, taskID); err != nil {
return TaskEvent{}, err
}
+148 -11
View File
@@ -91,7 +91,10 @@ SET pod_uid = EXCLUDED.pod_uid,
pod_name = EXCLUDED.pod_name,
site = EXCLUDED.site,
revision = EXCLUDED.revision,
status = 'active',
status = CASE
WHEN gateway_worker_instances.status = 'draining' THEN 'draining'
ELSE 'active'
END,
desired_capacity = EXCLUDED.desired_capacity,
capacity_limit = EXCLUDED.capacity_limit,
heartbeat_at = now(),
@@ -108,10 +111,11 @@ SET pod_uid = EXCLUDED.pod_uid,
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_worker_instances
SET allocated_capacity = 0,
SET status = 'stale',
allocated_capacity = 0,
updated_at = now()
WHERE status <> 'active'
OR heartbeat_at <= now() - $1::interval`, staleAfter.String()); err != nil {
WHERE status = 'active'
AND heartbeat_at <= now() - $1::interval`, staleAfter.String()); err != nil {
return WorkerAllocation{}, err
}
@@ -152,7 +156,11 @@ UPDATE gateway_worker_instances
SET desired_capacity = $2,
allocated_capacity = $3,
updated_at = now()
WHERE instance_id = $1`, worker.InstanceID, input.DesiredCapacity, capacity); err != nil {
WHERE instance_id = $1
AND (
desired_capacity IS DISTINCT FROM $2
OR allocated_capacity IS DISTINCT FROM $3
)`, worker.InstanceID, input.DesiredCapacity, capacity); err != nil {
return WorkerAllocation{}, err
}
if worker.InstanceID == input.InstanceID {
@@ -214,6 +222,7 @@ func (s *Store) MarkWorkerDraining(ctx context.Context, instanceID string) error
UPDATE gateway_worker_instances
SET status = 'draining',
allocated_capacity = 0,
draining_at = COALESCE(draining_at, now()),
heartbeat_at = now(),
updated_at = now()
WHERE instance_id = $1`, strings.TrimSpace(instanceID))
@@ -226,6 +235,139 @@ WHERE instance_id = $1`, strings.TrimSpace(instanceID))
return nil
}
func (s *Store) ReactivateWorkerInstance(ctx context.Context, instanceID string) error {
result, err := s.pool.Exec(ctx, `
UPDATE gateway_worker_instances
SET status = 'active',
draining_at = NULL,
heartbeat_at = now(),
updated_at = now()
WHERE instance_id = $1
AND status = 'draining'`, strings.TrimSpace(instanceID))
if err != nil {
return err
}
if result.RowsAffected() == 0 {
return pgx.ErrNoRows
}
return nil
}
type WorkerInstanceRuntime struct {
InstanceID string
PodUID string
PodName string
Site string
Revision string
Status string
Allocated int
CapacityLimit int
RunningTasks int
ActiveLeases int
HeartbeatAt time.Time
DrainingAt *time.Time
}
type WorkerQueueRuntime struct {
Queued int
Running int
OldestWaitSeconds float64
}
type CapacityDatabaseHealth struct {
Connections int
MaxConnections int
SynchronousPeers int
}
func (s *Store) WorkerQueueRuntime(ctx context.Context) (WorkerQueueRuntime, error) {
var snapshot WorkerQueueRuntime
err := s.pool.QueryRow(ctx, `
SELECT count(*) FILTER (WHERE status = 'queued')::int,
count(*) FILTER (WHERE status = 'running')::int,
COALESCE(EXTRACT(EPOCH FROM now() - MIN(created_at) FILTER (WHERE status = 'queued')), 0)::float8
FROM gateway_tasks
WHERE status IN ('queued', 'running')
AND run_mode IN ('production', 'acceptance', 'acceptance_canary')`).Scan(
&snapshot.Queued,
&snapshot.Running,
&snapshot.OldestWaitSeconds,
)
return snapshot, err
}
func (s *Store) ListWorkerInstanceRuntime(ctx context.Context) ([]WorkerInstanceRuntime, error) {
rows, err := s.pool.Query(ctx, `
SELECT worker.instance_id,
worker.pod_uid,
worker.pod_name,
worker.site,
worker.revision,
worker.status,
worker.allocated_capacity,
worker.capacity_limit,
count(DISTINCT task.id) FILTER (WHERE task.status = 'running')::int,
count(DISTINCT lease.id) FILTER (WHERE lease.released_at IS NULL)::int,
worker.heartbeat_at,
worker.draining_at
FROM gateway_worker_instances worker
LEFT JOIN gateway_tasks task
ON task.locked_by = worker.instance_id
AND task.status IN ('queued', 'running')
LEFT JOIN gateway_concurrency_leases lease
ON lease.task_id = task.id
AND lease.released_at IS NULL
WHERE worker.status IN ('active', 'draining')
AND worker.heartbeat_at > now() - $1::interval
GROUP BY worker.instance_id
ORDER BY worker.site ASC, worker.status DESC, worker.instance_id ASC`,
runtimeWorkerStaleAfter.String(),
)
if err != nil {
return nil, err
}
defer rows.Close()
instances := make([]WorkerInstanceRuntime, 0)
for rows.Next() {
var instance WorkerInstanceRuntime
if err := rows.Scan(
&instance.InstanceID,
&instance.PodUID,
&instance.PodName,
&instance.Site,
&instance.Revision,
&instance.Status,
&instance.Allocated,
&instance.CapacityLimit,
&instance.RunningTasks,
&instance.ActiveLeases,
&instance.HeartbeatAt,
&instance.DrainingAt,
); err != nil {
return nil, err
}
instances = append(instances, instance)
}
return instances, rows.Err()
}
func (s *Store) CapacityDatabaseHealth(ctx context.Context) (CapacityDatabaseHealth, error) {
var health CapacityDatabaseHealth
err := s.pool.QueryRow(ctx, `
SELECT
(SELECT count(*)::int FROM pg_stat_activity),
current_setting('max_connections')::int,
(SELECT count(*)::int
FROM pg_stat_replication
WHERE state = 'streaming'
AND sync_state IN ('sync', 'quorum'))`).Scan(
&health.Connections,
&health.MaxConnections,
&health.SynchronousPeers,
)
return health, err
}
// YieldStaleAsyncTaskAdmissions removes waiting FIFO entries whose River job
// is still marked running even though the task no longer owns a live execution
// lease. The job and task remain untouched: a surviving worker can recreate
@@ -247,9 +389,6 @@ func (s *Store) YieldStaleAsyncTaskAdmissions(
return 0, err
}
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-stale-admission-yield', 0))`); err != nil {
return 0, err
}
var yielded int64
if err := tx.QueryRow(ctx, `
WITH stale AS MATERIALIZED (
@@ -268,6 +407,7 @@ WITH stale AS MATERIALIZED (
AND job.attempted_at <= now() - $1::interval
ORDER BY admission.priority ASC, admission.enqueued_at ASC, admission.task_id ASC
LIMIT $2
FOR UPDATE OF admission SKIP LOCKED
),
locked AS MATERIALIZED (
SELECT task_id,
@@ -315,9 +455,6 @@ func (s *Store) RecoverOrphanedAsyncRiverJobs(
return 0, err
}
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-river-orphan-recovery', 0))`); err != nil {
return 0, err
}
var recovered int64
if err := tx.QueryRow(ctx, `
WITH orphaned AS MATERIALIZED (
@@ -0,0 +1,39 @@
ALTER TABLE gateway_worker_instances
DROP CONSTRAINT gateway_worker_instances_status_check,
ADD COLUMN draining_at timestamptz,
ADD CONSTRAINT gateway_worker_instances_status_check
CHECK (status IN ('active', 'draining', 'stale'));
CREATE TABLE gateway_capacity_profiles (
profile_key text PRIMARY KEY,
release_sha text NOT NULL,
config_hash text NOT NULL,
acceptance_run_id uuid REFERENCES gateway_acceptance_runs(id) ON DELETE RESTRICT,
kind text NOT NULL,
model text NOT NULL DEFAULT '*',
capacity_profile text NOT NULL,
stable_throughput_per_second numeric(18, 6) NOT NULL,
admitted_throughput_per_second numeric(18, 6) NOT NULL,
max_running integer NOT NULL,
max_queued integer NOT NULL,
max_predicted_wait_seconds integer NOT NULL,
enabled boolean NOT NULL DEFAULT false,
metadata jsonb NOT NULL DEFAULT '{}'::jsonb,
created_at timestamptz NOT NULL DEFAULT now(),
updated_at timestamptz NOT NULL DEFAULT now(),
CONSTRAINT gateway_capacity_profiles_kind_check
CHECK (kind IN ('image', 'video', 'mixed')),
CONSTRAINT gateway_capacity_profiles_values_check
CHECK (
stable_throughput_per_second > 0
AND admitted_throughput_per_second > 0
AND admitted_throughput_per_second <= stable_throughput_per_second
AND max_running > 0
AND max_queued >= 0
AND max_predicted_wait_seconds > 0
)
);
CREATE UNIQUE INDEX uniq_gateway_capacity_profiles_enabled_kind_model
ON gateway_capacity_profiles (kind, model)
WHERE enabled = true;
@@ -0,0 +1,5 @@
-- easyai:migration:no-transaction
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_runtime_owner
ON gateway_tasks (locked_by, execution_lease_expires_at, id)
WHERE status IN ('queued', 'running')
AND execution_lease_expires_at IS NOT NULL;
@@ -65,3 +65,71 @@ spec:
- name: http
port: 8090
targetPort: http
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-acceptance-callback-collector
namespace: easyai
labels:
app.kubernetes.io/name: easyai-acceptance-callback-collector
app.kubernetes.io/part-of: easyai-ai-gateway
spec:
replicas: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-acceptance-callback-collector
template:
metadata:
labels:
app.kubernetes.io/name: easyai-acceptance-callback-collector
app.kubernetes.io/part-of: easyai-ai-gateway
spec:
nodeSelector:
easyai.io/site: hongkong
easyai.io/workload: "true"
securityContext:
runAsNonRoot: true
runAsUser: 10001
runAsGroup: 10001
seccompProfile:
type: RuntimeDefault
containers:
- name: collector
image: easyai-api
command: ["/app/easyai-ai-gateway-acceptance-callback-collector"]
env:
- name: HTTP_ADDR
value: ":8091"
ports:
- name: http
containerPort: 8091
readinessProbe:
httpGet:
path: /healthz
port: http
resources:
requests:
cpu: 20m
memory: 32Mi
limits:
cpu: 250m
memory: 128Mi
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop: ["ALL"]
readOnlyRootFilesystem: true
---
apiVersion: v1
kind: Service
metadata:
name: easyai-acceptance-callback-collector
namespace: easyai
spec:
selector:
app.kubernetes.io/name: easyai-acceptance-callback-collector
ports:
- name: http
port: 8091
targetPort: http
@@ -14,50 +14,127 @@ source "$config_file"
: "${NAMESPACE:=easyai}"
: "${PUBLIC_BASE_URL:=https://ai.51easyai.com}"
: "${K3S_BIN:=/usr/local/bin/k3s}"
: "${DESIRED_STATE_DIR:=/usr/local/share/easyai-ai-gateway-release/production}"
: "${AI_GATEWAY_WORKER_REPLICAS_NINGBO:=1}"
: "${AI_GATEWAY_WORKER_REPLICAS_HONGKONG:=1}"
: "${AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT:=24}"
: "${AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT:=48}"
: "${AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT:=$AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT}"
: "${AI_GATEWAY_WORKER_MEMORY_REQUEST_MIB:=1536}"
: "${AI_GATEWAY_WORKER_CPU_REQUEST_MILLICORES:=500}"
: "${AI_GATEWAY_DATABASE_MAX_CONNS:=32}"
: "${AI_GATEWAY_API_DATABASE_MAX_CONNS:=64}"
: "${AI_GATEWAY_API_DATABASE_MAX_CONNS:=32}"
: "${AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS:=4}"
: "${AI_GATEWAY_DATABASE_RIVER_MAX_CONNS:=8}"
: "${AI_GATEWAY_API_DATABASE_RIVER_MAX_CONNS:=4}"
: "${AI_GATEWAY_DATABASE_MIN_IDLE_CONNS:=4}"
: "${AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS:=30}"
: "${AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY:=24}"
: "${AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY:=24}"
: "${AI_GATEWAY_WORKER_AUTOSCALING_ENABLED:=false}"
: "${AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO:=1}"
: "${AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG:=1}"
: "${AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO:=1}"
: "${AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG:=1}"
: "${AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA:=$((AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT * 2))}"
: "${AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS:=20}"
: "${AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS:=600}"
: "${AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS:=600}"
: "${AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT:=75}"
: "${AI_GATEWAY_NODE_MEMORY_HARD_PERCENT:=85}"
: "${AI_GATEWAY_NODE_CPU_TARGET_PERCENT:=70}"
: "${AI_GATEWAY_POSTGRES_CONNECTION_BUDGET:=150}"
: "${AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET:=$((AI_GATEWAY_API_DATABASE_MAX_CONNS * 2 + 8))}"
[[ $RELEASES_DIR == /* && $NAMESPACE =~ ^[a-z0-9-]+$ ]] || {
echo 'invalid cluster release configuration' >&2
exit 1
}
[[ $DESIRED_STATE_DIR == /* &&
-f $DESIRED_STATE_DIR/kustomization.yaml &&
! -L $DESIRED_STATE_DIR/kustomization.yaml ]] || {
echo 'production desired-state bundle is unavailable' >&2
exit 1
}
for capacity_value in \
"$AI_GATEWAY_WORKER_REPLICAS_NINGBO" \
"$AI_GATEWAY_WORKER_REPLICAS_HONGKONG" \
"$AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT" \
"$AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT" \
"$AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT" \
"$AI_GATEWAY_WORKER_MEMORY_REQUEST_MIB" \
"$AI_GATEWAY_WORKER_CPU_REQUEST_MILLICORES" \
"$AI_GATEWAY_DATABASE_MAX_CONNS" \
"$AI_GATEWAY_API_DATABASE_MAX_CONNS" \
"$AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS" \
"$AI_GATEWAY_DATABASE_RIVER_MAX_CONNS" \
"$AI_GATEWAY_API_DATABASE_RIVER_MAX_CONNS" \
"$AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS" \
"$AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY" \
"$AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY"; do
"$AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY" \
"$AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO" \
"$AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG" \
"$AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO" \
"$AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG" \
"$AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA" \
"$AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS" \
"$AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS" \
"$AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS" \
"$AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT" \
"$AI_GATEWAY_NODE_MEMORY_HARD_PERCENT" \
"$AI_GATEWAY_NODE_CPU_TARGET_PERCENT" \
"$AI_GATEWAY_POSTGRES_CONNECTION_BUDGET"; do
[[ $capacity_value =~ ^[1-9][0-9]*$ ]] || {
echo 'worker capacity configuration must contain positive integers' >&2
exit 1
}
done
[[ $AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET =~ ^[1-9][0-9]*$ ]] || {
echo 'PostgreSQL non-Worker connection budget must be a positive integer' >&2
exit 1
}
[[ $AI_GATEWAY_DATABASE_MIN_IDLE_CONNS =~ ^[0-9]+$ ]] || {
echo 'database minimum idle connection configuration must be a non-negative integer' >&2
exit 1
}
[[ $AI_GATEWAY_WORKER_AUTOSCALING_ENABLED == true ||
$AI_GATEWAY_WORKER_AUTOSCALING_ENABLED == false ]] || {
echo 'worker autoscaling flag must be true or false' >&2
exit 1
}
(( AI_GATEWAY_WORKER_REPLICAS_NINGBO <= 16 &&
AI_GATEWAY_WORKER_REPLICAS_HONGKONG <= 16 &&
AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT <= 256 &&
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT <= 512 &&
AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT <= 512 &&
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT == AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT &&
AI_GATEWAY_WORKER_MEMORY_REQUEST_MIB >= 256 &&
AI_GATEWAY_WORKER_MEMORY_REQUEST_MIB <= 1900 &&
AI_GATEWAY_WORKER_CPU_REQUEST_MILLICORES >= 100 &&
AI_GATEWAY_WORKER_CPU_REQUEST_MILLICORES <= 1900 &&
AI_GATEWAY_DATABASE_MAX_CONNS <= 256 &&
AI_GATEWAY_API_DATABASE_MAX_CONNS <= 256 &&
AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS < AI_GATEWAY_DATABASE_MAX_CONNS &&
AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS + AI_GATEWAY_DATABASE_RIVER_MAX_CONNS < AI_GATEWAY_DATABASE_MAX_CONNS &&
AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS + AI_GATEWAY_API_DATABASE_RIVER_MAX_CONNS < AI_GATEWAY_API_DATABASE_MAX_CONNS &&
AI_GATEWAY_DATABASE_MIN_IDLE_CONNS <= AI_GATEWAY_DATABASE_MAX_CONNS &&
AI_GATEWAY_DATABASE_MIN_IDLE_CONNS <= AI_GATEWAY_API_DATABASE_MAX_CONNS &&
AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS <= 3600 &&
AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY <= 256 &&
AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY <= 1024 )) || {
AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY <= 1024 &&
AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO <= AI_GATEWAY_WORKER_REPLICAS_NINGBO &&
AI_GATEWAY_WORKER_REPLICAS_NINGBO <= AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO &&
AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG <= AI_GATEWAY_WORKER_REPLICAS_HONGKONG &&
AI_GATEWAY_WORKER_REPLICAS_HONGKONG <= AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG &&
AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO <= 16 &&
AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG <= 16 &&
AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT < AI_GATEWAY_NODE_MEMORY_HARD_PERCENT &&
AI_GATEWAY_NODE_MEMORY_HARD_PERCENT <= 95 &&
AI_GATEWAY_NODE_CPU_TARGET_PERCENT <= 90 &&
AI_GATEWAY_POSTGRES_CONNECTION_BUDGET <= 150 &&
AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET < AI_GATEWAY_POSTGRES_CONNECTION_BUDGET &&
AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET +
(AI_GATEWAY_WORKER_REPLICAS_NINGBO + AI_GATEWAY_WORKER_REPLICAS_HONGKONG) *
AI_GATEWAY_DATABASE_MAX_CONNS <= AI_GATEWAY_POSTGRES_CONNECTION_BUDGET )) || {
echo 'worker capacity configuration exceeds the release safety bounds' >&2
exit 1
}
@@ -65,11 +142,12 @@ done
echo 'invalid public base URL' >&2
exit 1
}
for command in "$K3S_BIN" curl flock install jq node; do
for command in "$K3S_BIN" curl flock install jq node sha256sum; do
command -v "$command" >/dev/null 2>&1 || { echo "$command is required" >&2; exit 1; }
done
kubectl=("$K3S_BIN" kubectl)
ACTIVE_RELEASE_SHA=
install -d -m 0700 "$RELEASES_DIR"
exec 9>"$RELEASES_DIR/.release.lock"
flock -n 9 || { echo 'another cluster release operation is running' >&2; exit 1; }
@@ -78,6 +156,57 @@ manifest_get() {
node "$manifest_tool" get "$1" "$2"
}
capacity_config_hash() {
printf '%s\n' \
"$AI_GATEWAY_WORKER_REPLICAS_NINGBO" \
"$AI_GATEWAY_WORKER_REPLICAS_HONGKONG" \
"$AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT" \
"$AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT" \
"$AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT" \
"$AI_GATEWAY_WORKER_MEMORY_REQUEST_MIB" \
"$AI_GATEWAY_WORKER_CPU_REQUEST_MILLICORES" \
"$AI_GATEWAY_DATABASE_MAX_CONNS" \
"$AI_GATEWAY_API_DATABASE_MAX_CONNS" \
"$AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS" \
"$AI_GATEWAY_DATABASE_RIVER_MAX_CONNS" \
"$AI_GATEWAY_API_DATABASE_RIVER_MAX_CONNS" \
"$AI_GATEWAY_DATABASE_MIN_IDLE_CONNS" \
"$AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS" \
"$AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY" \
"$AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY" \
"$AI_GATEWAY_WORKER_AUTOSCALING_ENABLED" \
"$AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO" \
"$AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG" \
"$AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO" \
"$AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG" \
"$AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA" \
"$AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS" \
"$AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS" \
"$AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS" \
"$AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT" \
"$AI_GATEWAY_NODE_MEMORY_HARD_PERCENT" \
"$AI_GATEWAY_NODE_CPU_TARGET_PERCENT" \
"$AI_GATEWAY_POSTGRES_CONNECTION_BUDGET" \
"$AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET" |
sha256sum | awk '{print $1}'
}
record_capacity_config() {
local config_hash
config_hash=$(capacity_config_hash)
"${kubectl[@]}" patch configmap easyai-gateway-release -n "$NAMESPACE" \
--type merge -p "{\"data\":{\"configHash\":\"$config_hash\"}}" >/dev/null
"${kubectl[@]}" annotate namespace "$NAMESPACE" \
"easyai.io/config-hash=$config_hash" --overwrite >/dev/null
for deployment in \
easyai-api-ningbo easyai-api-hongkong \
easyai-worker-ningbo easyai-worker-hongkong \
easyai-capacity-controller; do
"${kubectl[@]}" annotate deployment "$deployment" -n "$NAMESPACE" \
"easyai.io/config-hash=$config_hash" --overwrite >/dev/null
done
}
current_to_file() {
local output=$1
"${kubectl[@]}" get configmap easyai-gateway-release -n "$NAMESPACE" \
@@ -127,6 +256,71 @@ verify_site() {
wait_for_url "$site Web" "http://$address:$web_port/" 'EasyAI AI Gateway'
}
apply_desired_state() {
local api_image=$1
local web_image=$2
local capacity_controller_replicas=${3:-2}
local api_repository=${api_image%@*}
local api_digest=${api_image#*@}
local web_repository=${web_image%@*}
local web_digest=${web_image#*@}
local working_directory
working_directory=$(mktemp -d)
cp -a "$DESIRED_STATE_DIR/." "$working_directory/"
sed -i \
-e "/ - name: easyai-api/{n;s| newName:.*| newName: $api_repository|;n;s| digest:.*| digest: $api_digest|;}" \
-e "/ - name: easyai-web/{n;s| newName:.*| newName: $web_repository|;n;s| digest:.*| digest: $web_digest|;}" \
"$working_directory/kustomization.yaml"
if [[ $capacity_controller_replicas == 0 ]]; then
sed -i \
'/^ name: easyai-capacity-controller$/,/^---$/s/^ replicas: 2$/ replicas: 0/' \
"$working_directory/application.yaml"
fi
"${kubectl[@]}" apply --server-side --dry-run=server --force-conflicts \
--field-manager=easyai-production-preflight \
-k "$working_directory" >/dev/null
"${kubectl[@]}" apply --server-side --force-conflicts \
--field-manager=easyai-production-release \
-k "$working_directory" >/dev/null
rm -rf -- "$working_directory"
}
snapshot_application_deployments() {
local output=$1
"${kubectl[@]}" get deployments -n "$NAMESPACE" \
easyai-api-ningbo easyai-api-hongkong \
easyai-worker-ningbo easyai-worker-hongkong \
easyai-capacity-controller easyai-web-ningbo easyai-web-hongkong \
-o json |
jq '
del(.metadata)
| .items |= map(
del(
.metadata.annotations["deployment.kubernetes.io/revision"],
.metadata.creationTimestamp,
.metadata.generation,
.metadata.managedFields,
.metadata.resourceVersion,
.metadata.uid,
.status
)
)
' >"$output"
}
restore_application_deployments() {
local snapshot=$1
[[ -s $snapshot ]] || return 1
"${kubectl[@]}" apply -f "$snapshot" >/dev/null
for deployment in \
easyai-api-ningbo easyai-api-hongkong \
easyai-worker-ningbo easyai-worker-hongkong \
easyai-capacity-controller easyai-web-ningbo easyai-web-hongkong; do
"${kubectl[@]}" rollout status "deployment/$deployment" \
-n "$NAMESPACE" --timeout=300s || true
done
}
rollout_worker_site() {
local site=$1
local api_image=$2
@@ -139,11 +333,27 @@ rollout_worker_site() {
"${kubectl[@]}" set env "deployment/easyai-worker-$site" -n "$NAMESPACE" \
"AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT" \
"AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT" \
"AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT" \
"AI_GATEWAY_DATABASE_MAX_CONNS=$AI_GATEWAY_DATABASE_MAX_CONNS" \
"AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS=$AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS" \
"AI_GATEWAY_DATABASE_RIVER_MAX_CONNS=$AI_GATEWAY_DATABASE_RIVER_MAX_CONNS" \
"AI_GATEWAY_DATABASE_MIN_IDLE_CONNS=$AI_GATEWAY_DATABASE_MIN_IDLE_CONNS" \
"AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS=$AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS" \
"AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY=$AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY" \
"AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY=$AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY"
"AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY=$AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY" \
"AI_GATEWAY_REVISION=${ACTIVE_RELEASE_SHA:-}" \
"AI_GATEWAY_WORKER_AUTOSCALING_ENABLED=$AI_GATEWAY_WORKER_AUTOSCALING_ENABLED" \
"AI_GATEWAY_WORKER_REPLICAS_NINGBO=$AI_GATEWAY_WORKER_REPLICAS_NINGBO" \
"AI_GATEWAY_WORKER_REPLICAS_HONGKONG=$AI_GATEWAY_WORKER_REPLICAS_HONGKONG" \
"AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO=$AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO" \
"AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG=$AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG" \
"AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO=$AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO" \
"AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG=$AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG" \
"AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA=$AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA"
"${kubectl[@]}" set resources "deployment/easyai-worker-$site" -n "$NAMESPACE" \
--containers=worker \
--requests="cpu=${AI_GATEWAY_WORKER_CPU_REQUEST_MILLICORES}m,memory=${AI_GATEWAY_WORKER_MEMORY_REQUEST_MIB}Mi" \
--limits='cpu=2,memory=2Gi'
"${kubectl[@]}" scale "deployment/easyai-worker-$site" -n "$NAMESPACE" \
--replicas="$replicas"
if [[ -n $api_image ]]; then
@@ -157,25 +367,74 @@ rollout_worker_site() {
-o jsonpath='{.status.readyReplicas}') == "$replicas" ]]
}
rollout_capacity_controller() {
local api_image=${1:-}
"${kubectl[@]}" set env deployment/easyai-capacity-controller -n "$NAMESPACE" \
"AI_GATEWAY_WORKER_AUTOSCALING_ENABLED=$AI_GATEWAY_WORKER_AUTOSCALING_ENABLED" \
"AI_GATEWAY_WORKER_REPLICAS_NINGBO=$AI_GATEWAY_WORKER_REPLICAS_NINGBO" \
"AI_GATEWAY_WORKER_REPLICAS_HONGKONG=$AI_GATEWAY_WORKER_REPLICAS_HONGKONG" \
"AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO=$AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO" \
"AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG=$AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG" \
"AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO=$AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO" \
"AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG=$AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG" \
"AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA=$AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA" \
"AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS=$AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS" \
"AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS=$AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS" \
"AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS=$AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS" \
"AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT=$AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT" \
"AI_GATEWAY_NODE_MEMORY_HARD_PERCENT=$AI_GATEWAY_NODE_MEMORY_HARD_PERCENT" \
"AI_GATEWAY_NODE_CPU_TARGET_PERCENT=$AI_GATEWAY_NODE_CPU_TARGET_PERCENT" \
"AI_GATEWAY_POSTGRES_CONNECTION_BUDGET=$AI_GATEWAY_POSTGRES_CONNECTION_BUDGET" \
"AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET=$AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET" \
"AI_GATEWAY_WORKER_DATABASE_MAX_CONNS=$AI_GATEWAY_DATABASE_MAX_CONNS" \
"AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT" \
"AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT" \
"AI_GATEWAY_REVISION=${ACTIVE_RELEASE_SHA:-}"
if [[ -n $api_image ]]; then
"${kubectl[@]}" set image deployment/easyai-capacity-controller -n "$NAMESPACE" \
"capacity-controller=$api_image"
fi
"${kubectl[@]}" rollout status deployment/easyai-capacity-controller -n "$NAMESPACE" --timeout=300s
}
rollout_api_capacity_site() {
local site=$1
"${kubectl[@]}" set env "deployment/easyai-api-$site" -n "$NAMESPACE" \
"AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT" \
"AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT" \
"AI_GATEWAY_DATABASE_MAX_CONNS=$AI_GATEWAY_API_DATABASE_MAX_CONNS" \
"AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS=$AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS" \
"AI_GATEWAY_DATABASE_RIVER_MAX_CONNS=$AI_GATEWAY_API_DATABASE_RIVER_MAX_CONNS" \
"AI_GATEWAY_DATABASE_MIN_IDLE_CONNS=$AI_GATEWAY_DATABASE_MIN_IDLE_CONNS" \
"AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS=$AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS"
"AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS=$AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS" \
"AI_GATEWAY_REVISION=${ACTIVE_RELEASE_SHA:-}"
"${kubectl[@]}" rollout status "deployment/easyai-api-$site" -n "$NAMESPACE" --timeout=300s
}
apply_capacity_config() {
local capacity_controller_enabled
ACTIVE_RELEASE_SHA=$("${kubectl[@]}" get configmap easyai-gateway-release -n "$NAMESPACE" \
-o jsonpath='{.data.releaseSha}' 2>/dev/null || true)
capacity_controller_enabled=$("${kubectl[@]}" get configmap easyai-gateway-release -n "$NAMESPACE" \
-o jsonpath='{.data.capacityControllerEnabled}' 2>/dev/null || true)
[[ $ACTIVE_RELEASE_SHA =~ ^[0-9a-f]{40}$ ]] || {
echo 'current release SHA is unavailable for capacity update' >&2
return 1
}
rollout_worker_site hongkong ""
rollout_worker_site ningbo ""
rollout_api_capacity_site hongkong
rollout_api_capacity_site ningbo
if [[ $capacity_controller_enabled == true ]]; then
rollout_capacity_controller ""
else
"${kubectl[@]}" scale deployment/easyai-capacity-controller -n "$NAMESPACE" --replicas=0
fi
verify_site hongkong
verify_site ningbo
wait_for_url 'public readiness' "$PUBLIC_BASE_URL/api/v1/readyz" '"ok":true'
echo "production_capacity=PASS ningbo_replicas=$AI_GATEWAY_WORKER_REPLICAS_NINGBO hongkong_replicas=$AI_GATEWAY_WORKER_REPLICAS_HONGKONG instance_limit=$AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT global_limit=$AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT worker_database_pool=$AI_GATEWAY_DATABASE_MAX_CONNS api_database_pool=$AI_GATEWAY_API_DATABASE_MAX_CONNS database_min_idle=$AI_GATEWAY_DATABASE_MIN_IDLE_CONNS database_max_idle_seconds=$AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS media_concurrency=$AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY"
record_capacity_config
echo "production_capacity=PASS config_hash=$(capacity_config_hash) ningbo_replicas=$AI_GATEWAY_WORKER_REPLICAS_NINGBO hongkong_replicas=$AI_GATEWAY_WORKER_REPLICAS_HONGKONG instance_limit=$AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT global_limit=$AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT worker_database_pool=$AI_GATEWAY_DATABASE_MAX_CONNS api_database_pool=$AI_GATEWAY_API_DATABASE_MAX_CONNS database_min_idle=$AI_GATEWAY_DATABASE_MIN_IDLE_CONNS database_max_idle_seconds=$AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS media_concurrency=$AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY worker_memory_request_mib=$AI_GATEWAY_WORKER_MEMORY_REQUEST_MIB worker_cpu_request_millicores=$AI_GATEWAY_WORKER_CPU_REQUEST_MILLICORES autoscaling=$AI_GATEWAY_WORKER_AUTOSCALING_ENABLED"
}
rollout_site() {
@@ -187,7 +446,10 @@ rollout_site() {
if [[ $api_changed == true ]]; then
"${kubectl[@]}" set env "deployment/easyai-api-$site" -n "$NAMESPACE" \
"AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT" \
"AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT=$AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT" \
"AI_GATEWAY_DATABASE_MAX_CONNS=$AI_GATEWAY_API_DATABASE_MAX_CONNS" \
"AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS=$AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS" \
"AI_GATEWAY_DATABASE_RIVER_MAX_CONNS=$AI_GATEWAY_API_DATABASE_RIVER_MAX_CONNS" \
"AI_GATEWAY_DATABASE_MIN_IDLE_CONNS=$AI_GATEWAY_DATABASE_MIN_IDLE_CONNS" \
"AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS=$AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS"
"${kubectl[@]}" set image "deployment/easyai-api-$site" -n "$NAMESPACE" \
@@ -265,11 +527,36 @@ EOF
-n "$NAMESPACE" --timeout=300s
}
verify_migration_version() {
local expected=$1
[[ $expected =~ ^[0-9]{4}_[a-z0-9][a-z0-9_]*$ ]] || {
echo 'release manifest migration version is invalid' >&2
return 1
}
local primary actual
primary=$("${kubectl[@]}" get cluster easyai-postgres -n "$NAMESPACE" \
-o jsonpath='{.status.currentPrimary}')
[[ $primary =~ ^easyai-postgres-[0-9]+$ ]] || {
echo 'PostgreSQL primary is unavailable for migration verification' >&2
return 1
}
actual=$("${kubectl[@]}" exec -n "$NAMESPACE" "$primary" -c postgres -- \
psql -X -v ON_ERROR_STOP=1 -U postgres -d easyai_ai_gateway -At \
-c 'SELECT version FROM schema_migrations ORDER BY version DESC LIMIT 1')
[[ $actual == "$expected" ]] || {
echo "production migration version mismatch: expected=$expected actual=${actual:-missing}" >&2
return 1
}
}
record_current() {
local manifest=$1
local action=$2
local started_at=$3
local source_sha recorded temporary
local source_sha recorded temporary capacity_controller_enabled=true
if [[ $action == rollback ]]; then
capacity_controller_enabled=false
fi
source_sha=$(manifest_get "$manifest" sourceSha)
recorded=$RELEASES_DIR/$source_sha.json
temporary=$(mktemp "$RELEASES_DIR/.recorded.XXXXXX")
@@ -280,23 +567,102 @@ record_current() {
node "$manifest_tool" record-deployment "$manifest" "$temporary" >/dev/null
install -m 0600 "$temporary" "$recorded"
"${kubectl[@]}" create configmap easyai-gateway-release -n "$NAMESPACE" \
--from-file=manifest.json="$recorded" --dry-run=client -o yaml |
--from-file=manifest.json="$recorded" \
--from-literal=configHash="$(capacity_config_hash)" \
--from-literal=capacityControllerEnabled="$capacity_controller_enabled" \
--from-literal=releaseSha="$source_sha" \
--dry-run=client -o yaml |
"${kubectl[@]}" apply -f - >/dev/null
"${kubectl[@]}" annotate namespace "$NAMESPACE" \
"easyai.io/release=$source_sha" \
"easyai.io/config-hash=$(capacity_config_hash)" \
--overwrite >/dev/null
for deployment in \
easyai-api-ningbo easyai-api-hongkong \
easyai-worker-ningbo easyai-worker-hongkong \
easyai-capacity-controller \
easyai-web-ningbo easyai-web-hongkong; do
"${kubectl[@]}" annotate deployment "$deployment" -n "$NAMESPACE" \
"easyai.io/release=$source_sha" \
"easyai.io/config-hash=$(capacity_config_hash)" \
--overwrite >/dev/null
done
rm -f -- "$temporary"
}
record_certification() {
local run_id=$1
local capacity_profile=$2
local config_hash=$3
local instance_slots=$4
local max_ningbo=$5
local max_hongkong=$6
[[ $run_id =~ ^[0-9a-f-]{36}$ &&
$capacity_profile =~ ^P(24|28|32)$ &&
$config_hash =~ ^[0-9a-f]{64}$ &&
$instance_slots =~ ^[1-9][0-9]*$ &&
$max_ningbo =~ ^[1-9][0-9]*$ &&
$max_hongkong =~ ^[1-9][0-9]*$ ]] || {
echo 'invalid production certification metadata' >&2
return 1
}
local effective_config_hash
effective_config_hash=$("${kubectl[@]}" get configmap easyai-gateway-release -n "$NAMESPACE" \
-o jsonpath='{.data.configHash}' 2>/dev/null || true)
[[ $effective_config_hash == "$config_hash" ]] || {
echo 'production certification config hash does not match the deployed release configuration' >&2
return 1
}
local current_file source_sha recorded temporary capacity_controller_enabled
current_file=$(mktemp "$RELEASES_DIR/.certification-current.XXXXXX")
current_to_file "$current_file"
source_sha=$(manifest_get "$current_file" sourceSha)
recorded=$RELEASES_DIR/$source_sha.json
temporary=$(mktemp "$RELEASES_DIR/.certification.XXXXXX")
rm -f -- "$temporary"
RELEASE_CERTIFIED_AT=$(date -u '+%Y-%m-%dT%H:%M:%SZ') \
RELEASE_ACCEPTANCE_RUN_ID=$run_id \
RELEASE_CONFIG_HASH=$config_hash \
RELEASE_CAPACITY_PROFILE=$capacity_profile \
RELEASE_WORKER_INSTANCE_SLOTS=$instance_slots \
RELEASE_WORKER_MAX_REPLICAS_NINGBO=$max_ningbo \
RELEASE_WORKER_MAX_REPLICAS_HONGKONG=$max_hongkong \
node "$manifest_tool" record-certification "$current_file" "$temporary" >/dev/null
install -m 0600 "$temporary" "$recorded"
capacity_controller_enabled=$("${kubectl[@]}" get configmap easyai-gateway-release -n "$NAMESPACE" \
-o jsonpath='{.data.capacityControllerEnabled}' 2>/dev/null || true)
"${kubectl[@]}" create configmap easyai-gateway-release -n "$NAMESPACE" \
--from-file=manifest.json="$recorded" \
--from-literal=configHash="$config_hash" \
--from-literal=capacityControllerEnabled="${capacity_controller_enabled:-true}" \
--from-literal=releaseSha="$source_sha" \
--from-literal=certifiedRunId="$run_id" \
--from-literal=certifiedCapacityProfile="$capacity_profile" \
--dry-run=client -o yaml |
"${kubectl[@]}" apply -f - >/dev/null
"${kubectl[@]}" annotate namespace "$NAMESPACE" \
"easyai.io/certified-run=$run_id" \
"easyai.io/capacity-profile=$capacity_profile" \
--overwrite >/dev/null
rm -f -- "$temporary" "$current_file"
echo "production_certification=PASS release=$source_sha run_id=$run_id capacity_profile=$capacity_profile config_hash=$config_hash"
}
activate_manifest() {
local manifest=$1
local action=${2:-deploy}
local source_sha base_sha api_image web_image migrations_changed components
local api_changed=false web_changed=false current_file previous_api previous_web started_at
local source_sha base_sha api_image web_image migrations_changed migration_version components
local api_changed=false web_changed=false current_file deployment_snapshot previous_api previous_web started_at
local capacity_controller_replicas
[[ -f $manifest && ! -L $manifest ]] || { echo 'manifest must be a regular file' >&2; return 1; }
node "$manifest_tool" validate "$manifest" >/dev/null
source_sha=$(manifest_get "$manifest" sourceSha)
ACTIVE_RELEASE_SHA=$source_sha
base_sha=$(manifest_get "$manifest" baseReleaseSha)
api_image=$(manifest_get "$manifest" images.api)
web_image=$(manifest_get "$manifest" images.web)
migrations_changed=$(manifest_get "$manifest" migrationsChanged)
migration_version=$(manifest_get "$manifest" migrationVersion)
components=$(manifest_get "$manifest" components)
[[ $api_image =~ @sha256:[0-9a-f]{64}$ && $web_image =~ @sha256:[0-9a-f]{64}$ ]] || {
echo 'release images must be digest-pinned' >&2
@@ -329,35 +695,54 @@ activate_manifest() {
previous_web=
fi
started_at=$(date +%s)
deployment_snapshot=$(mktemp "$RELEASES_DIR/.deployments.XXXXXX")
snapshot_application_deployments "$deployment_snapshot"
if [[ $action == deploy && $migrations_changed == true ]]; then
run_backup "$source_sha"
run_migrator "$source_sha" "$api_image"
fi
if [[ $action == deploy ]]; then
verify_migration_version "$migration_version"
fi
capacity_controller_replicas=2
if [[ $action == rollback ]]; then
capacity_controller_replicas=0
fi
if ! apply_desired_state "$api_image" "$web_image" "$capacity_controller_replicas"; then
echo '[cluster-release] desired state apply failed; restoring exact deployment snapshot' >&2
restore_application_deployments "$deployment_snapshot" || true
rm -f -- "$deployment_snapshot" "$current_file"
return 1
fi
if [[ $action == rollback ]]; then
# Historical API images may predate the capacity-controller process role.
# Rollbacks therefore restore static Worker replicas and leave the
# controller at zero instead of ever starting an incompatible image.
"${kubectl[@]}" scale deployment/easyai-capacity-controller \
-n "$NAMESPACE" --replicas=0
fi
if { [[ $api_changed != true ]] ||
{ rollout_worker_site hongkong "$api_image" &&
rollout_worker_site ningbo "$api_image"; }; } &&
rollout_worker_site ningbo "$api_image" &&
{ [[ $action == rollback ]] ||
rollout_capacity_controller "$api_image"; }; }; } &&
rollout_site hongkong "$api_image" "$web_image" "$api_changed" "$web_changed" &&
rollout_site ningbo "$api_image" "$web_image" "$api_changed" "$web_changed" &&
wait_for_url 'public readiness' "$PUBLIC_BASE_URL/api/v1/readyz" '"ok":true'; then
:
else
if [[ -n $previous_api && -n $previous_web ]]; then
echo '[cluster-release] restoring previous application images' >&2
if [[ $api_changed == true ]]; then
rollout_worker_site hongkong "$previous_api" || true
rollout_worker_site ningbo "$previous_api" || true
fi
rollout_site hongkong "$previous_api" "$previous_web" "$api_changed" "$web_changed" || true
rollout_site ningbo "$previous_api" "$previous_web" "$api_changed" "$web_changed" || true
echo '[cluster-release] restoring exact pre-release Deployment snapshot' >&2
restore_application_deployments "$deployment_snapshot" || true
fi
rm -f -- "$current_file"
rm -f -- "$deployment_snapshot" "$current_file"
return 1
fi
record_current "$manifest" "$action" "$started_at"
rm -f -- "$current_file"
rm -f -- "$deployment_snapshot" "$current_file"
echo "production_release=PASS action=$action release=$source_sha"
}
@@ -378,6 +763,8 @@ case ${1:-} in
[[ $("${kubectl[@]}" get deployment "easyai-worker-$site" -n "$NAMESPACE" \
-o jsonpath='{.spec.template.spec.containers[?(@.name=="worker")].image}') == "$api_image" ]]
done
[[ $("${kubectl[@]}" get deployment easyai-capacity-controller -n "$NAMESPACE" \
-o jsonpath='{.spec.template.spec.containers[?(@.name=="capacity-controller")].image}') == "$api_image" ]]
[[ $("${kubectl[@]}" get deployment easyai-api-ningbo -n "$NAMESPACE" \
-o jsonpath='{.spec.template.spec.containers[?(@.name=="api")].image}') == "$api_image" ]]
[[ $("${kubectl[@]}" get deployment easyai-web-ningbo -n "$NAMESPACE" \
@@ -398,8 +785,12 @@ case ${1:-} in
[[ $# -eq 1 ]] || exit 64
apply_capacity_config
;;
certify)
[[ $# -eq 7 ]] || exit 64
record_certification "$2" "$3" "$4" "$5" "$6" "$7"
;;
*)
echo 'usage: easyai-ai-gateway-cluster-release {status|adopt <manifest>|deploy <manifest>|rollback <sha>|capacity}' >&2
echo 'usage: easyai-ai-gateway-cluster-release {status|adopt <manifest>|deploy <manifest>|rollback <sha>|capacity|certify <run-id> <profile> <config-hash> <slots> <max-ningbo> <max-hongkong>}' >&2
exit 64
;;
esac
@@ -2,6 +2,7 @@ RELEASES_DIR=/var/lib/easyai-ai-gateway-cluster-release/releases
NAMESPACE=easyai
PUBLIC_BASE_URL=https://ai.51easyai.com
K3S_BIN=/usr/local/bin/k3s
DESIRED_STATE_DIR=/usr/local/share/easyai-ai-gateway-release/production
# Worker 容量只需修改本发布配置并执行获授权的 `capacity`,不再修改 Go 代码。
# 副本数属于 Deployment 配置,不能由容器内部环境变量改变;发布工具读取这两个变量并写入 replicas。
@@ -9,9 +10,30 @@ AI_GATEWAY_WORKER_REPLICAS_NINGBO=1
AI_GATEWAY_WORKER_REPLICAS_HONGKONG=1
AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT=24
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=48
AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT=48
AI_GATEWAY_WORKER_MEMORY_REQUEST_MIB=1536
AI_GATEWAY_WORKER_CPU_REQUEST_MILLICORES=500
AI_GATEWAY_DATABASE_MAX_CONNS=32
AI_GATEWAY_API_DATABASE_MAX_CONNS=64
AI_GATEWAY_API_DATABASE_MAX_CONNS=32
AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS=4
AI_GATEWAY_DATABASE_RIVER_MAX_CONNS=8
AI_GATEWAY_API_DATABASE_RIVER_MAX_CONNS=4
AI_GATEWAY_DATABASE_MIN_IDLE_CONNS=4
AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS=30
AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY=24
AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY=24
AI_GATEWAY_WORKER_AUTOSCALING_ENABLED=false
AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO=1
AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG=1
AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO=1
AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG=1
AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA=48
AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS=20
AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS=600
AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS=600
AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT=75
AI_GATEWAY_NODE_MEMORY_HARD_PERCENT=85
AI_GATEWAY_NODE_CPU_TARGET_PERCENT=70
AI_GATEWAY_POSTGRES_CONNECTION_BUDGET=150
# 2 个 API 池(2×32)加 2 个 capacity-controller 池(2×4)。
AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET=72
@@ -0,0 +1,52 @@
apiVersion: postgresql.cnpg.io/v1
kind: Cluster
metadata:
name: easyai-postgres
namespace: easyai
labels:
app.kubernetes.io/name: easyai-postgres
app.kubernetes.io/part-of: easyai-ai-gateway
easyai.io/environment: local-acceptance
spec:
instances: 2
imageName: ghcr.io/cloudnative-pg/postgresql:18.4-standard-trixie@sha256:4587df73024408f5b2be9b4dd6ba2ccee8c9e5dc0c9a87c274c292291cc8a68c
bootstrap:
initdb:
database: easyai_ai_gateway
owner: easyai
secret:
name: easyai-postgres-app
postInitSQL:
- CREATE EXTENSION IF NOT EXISTS pgcrypto;
storage:
storageClass: local-path
size: 8Gi
walStorage:
storageClass: local-path
size: 2Gi
affinity:
enablePodAntiAffinity: true
podAntiAffinityType: required
topologyKey: easyai.io/site
nodeSelector:
easyai.io/database: "true"
resources:
requests:
cpu: 500m
memory: 1Gi
limits:
cpu: "2"
memory: 3Gi
postgresql:
synchronous:
method: any
number: 1
dataDurability: required
parameters:
archive_mode: "off"
max_connections: "200"
shared_buffers: 512MB
wal_compression: "on"
wal_keep_size: 1GB
monitoring:
enablePodMonitor: false
@@ -0,0 +1,8 @@
# Local acceptance dependency lock. Checksums are lowercase SHA-256.
K3D_VERSION=v5.9.0
K3D_DARWIN_ARM64_SHA256=fe106541d5d0a3f18debcd4d432a16f8c0ce3e6ddc06f8fbb6f696a122313e00
K3D_DARWIN_AMD64_SHA256=b4aabc37534f95b9c764e7823f2df923f50d57600837aa60a06266cce47db732
K3S_IMAGE=rancher/k3s:v1.36.2-k3s1
CNPG_VERSION=1.29.1
CNPG_MANIFEST_URL=https://raw.githubusercontent.com/cloudnative-pg/cloudnative-pg/release-1.29/releases/cnpg-1.29.1.yaml
CNPG_MANIFEST_SHA256=e0c5ff41bf5b01c0775bf225929b8423d543cd23e9742e4274fe7de8499cf93a
@@ -0,0 +1,50 @@
apiVersion: k3d.io/v1alpha5
kind: Simple
metadata:
name: easyai-acceptance-local
servers: 3
agents: 0
image: rancher/k3s:v1.36.2-k3s1
kubeAPI:
host: 127.0.0.1
hostPort: "16550"
ports:
- port: 18443:31088
nodeFilters:
- server:0
- port: 19443:31089
nodeFilters:
- server:1
volumes:
- volume: EASYAI_ACCEPTANCE_MEDIA_DIR:/var/lib/easyai-acceptance/media
nodeFilters:
- server:*
options:
k3s:
extraArgs:
- arg: --disable=traefik
nodeFilters:
- server:*
- arg: --disable=servicelb
nodeFilters:
- server:*
- arg: --node-label=easyai.io/site=ningbo
nodeFilters:
- server:0
- arg: --node-label=easyai.io/site=hongkong
nodeFilters:
- server:1
- arg: --node-label=easyai.io/site=los-angeles
nodeFilters:
- server:2
- arg: --node-label=easyai.io/workload=true
nodeFilters:
- server:0
- server:1
- arg: --node-label=easyai.io/database=true
nodeFilters:
- server:0
- server:1
- arg: --node-taint=easyai.io/control-plane-only=true:NoSchedule
nodeFilters:
- server:2
@@ -0,0 +1,54 @@
apiVersion: v1
kind: ConfigMap
metadata:
name: easyai-ai-gateway-config
namespace: easyai
labels:
app.kubernetes.io/part-of: easyai-ai-gateway
easyai.io/environment: local-acceptance
data:
APP_ENV: production
HTTP_ADDR: ":8088"
IDENTITY_MODE: standalone
IDENTITY_SECRET_STORE: file
IDENTITY_SECRET_DIR: /app/data/identity
IDENTITY_SECURITY_EVENTS_HEARTBEAT_INTERVAL_SECONDS: "0"
IDENTITY_SECURITY_EVENTS_STALE_AFTER_SECONDS: "0"
IDENTITY_SECURITY_EVENTS_CLOCK_SKEW_SECONDS: "0"
SERVER_MAIN_BASE_URL: http://easyai-acceptance-emulator.easyai.svc.cluster.local:8090
TASK_PROGRESS_CALLBACK_ENABLED: "true"
TASK_PROGRESS_CALLBACK_URL: http://easyai-acceptance-callback-collector.easyai.svc.cluster.local:8091/callbacks
AI_GATEWAY_PUBLIC_BASE_URL: https://gateway.easyai.local
AI_GATEWAY_WEB_BASE_URL: https://gateway.easyai.local
CORS_ALLOWED_ORIGIN: https://gateway.easyai.local
AI_GATEWAY_GENERATED_STORAGE_DIR: /app/data/static/generated
AI_GATEWAY_UPLOADED_STORAGE_DIR: /app/data/static/uploaded
AI_GATEWAY_LOCAL_TEMP_ASSET_TTL_HOURS: "24"
AI_GATEWAY_LOCAL_RESULT_TTL_HOURS: "24"
AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES: "0"
AI_GATEWAY_LOCAL_RESULT_MAX_BYTES: "268435456"
AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES: "536870912"
AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED: "false"
AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED: "true"
AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT: "24"
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT: "24"
AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT: "48"
AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS: "5"
AI_GATEWAY_DATABASE_IDLE_IN_TRANSACTION_TIMEOUT_SECONDS: "60"
AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS: "30"
AI_GATEWAY_WORKER_AUTOSCALING_ENABLED: "true"
AI_GATEWAY_WORKER_REPLICAS_NINGBO: "1"
AI_GATEWAY_WORKER_REPLICAS_HONGKONG: "1"
AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO: "1"
AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG: "1"
AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO: "2"
AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG: "2"
AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA: "48"
AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS: "20"
AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS: "600"
AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS: "600"
AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT: "75"
AI_GATEWAY_NODE_MEMORY_HARD_PERCENT: "85"
AI_GATEWAY_NODE_CPU_TARGET_PERCENT: "70"
AI_GATEWAY_POSTGRES_CONNECTION_BUDGET: "150"
AI_GATEWAY_WORKER_DATABASE_MAX_CONNS: "32"
@@ -0,0 +1,312 @@
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-acceptance-emulator
namespace: easyai
labels:
app.kubernetes.io/name: easyai-acceptance-emulator
easyai.io/environment: local-acceptance
spec:
replicas: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-acceptance-emulator
template:
metadata:
labels:
app.kubernetes.io/name: easyai-acceptance-emulator
easyai.io/environment: local-acceptance
spec:
nodeSelector:
easyai.io/site: ningbo
containers:
- name: emulator
image: easyai-api
command: ["/app/easyai-ai-gateway-acceptance-emulator"]
env:
- name: HTTP_ADDR
value: ":8090"
ports:
- name: http
containerPort: 8090
readinessProbe:
httpGet:
path: /healthz
port: http
resources:
requests:
cpu: 100m
memory: 128Mi
limits:
cpu: "2"
memory: 1Gi
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop: ["ALL"]
readOnlyRootFilesystem: true
---
apiVersion: v1
kind: Service
metadata:
name: easyai-acceptance-emulator
namespace: easyai
spec:
selector:
app.kubernetes.io/name: easyai-acceptance-emulator
ports:
- name: http
port: 8090
targetPort: http
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-acceptance-callback-collector
namespace: easyai
labels:
app.kubernetes.io/name: easyai-acceptance-callback-collector
easyai.io/environment: local-acceptance
spec:
replicas: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-acceptance-callback-collector
template:
metadata:
labels:
app.kubernetes.io/name: easyai-acceptance-callback-collector
easyai.io/environment: local-acceptance
spec:
nodeSelector:
easyai.io/site: ningbo
containers:
- name: collector
image: easyai-api
command: ["/app/easyai-ai-gateway-acceptance-callback-collector"]
env:
- name: HTTP_ADDR
value: ":8091"
ports:
- name: http
containerPort: 8091
readinessProbe:
httpGet:
path: /healthz
port: http
resources:
requests:
cpu: 20m
memory: 32Mi
limits:
cpu: 250m
memory: 128Mi
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop: ["ALL"]
readOnlyRootFilesystem: true
---
apiVersion: v1
kind: Service
metadata:
name: easyai-acceptance-callback-collector
namespace: easyai
spec:
selector:
app.kubernetes.io/name: easyai-acceptance-callback-collector
ports:
- name: http
port: 8091
targetPort: http
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-acceptance-upstream-proxy
namespace: easyai
labels:
app.kubernetes.io/name: easyai-acceptance-upstream-proxy
easyai.io/environment: local-acceptance
spec:
replicas: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-acceptance-upstream-proxy
template:
metadata:
labels:
app.kubernetes.io/name: easyai-acceptance-upstream-proxy
easyai.io/environment: local-acceptance
spec:
nodeSelector:
easyai.io/site: hongkong
containers:
- name: proxy
image: easyai-acceptance-netem
env:
- name: NETEM_LISTEN_PORT
value: "18090"
- name: NETEM_TARGET_HOST
value: easyai-acceptance-emulator.easyai.svc.cluster.local
- name: NETEM_TARGET_PORT
value: "8090"
ports:
- name: proxy
containerPort: 18090
resources:
requests:
cpu: 20m
memory: 32Mi
limits:
cpu: 250m
memory: 128Mi
securityContext:
allowPrivilegeEscalation: false
capabilities:
add: ["NET_ADMIN"]
drop: ["ALL"]
readOnlyRootFilesystem: true
---
apiVersion: v1
kind: Service
metadata:
name: easyai-acceptance-upstream-proxy
namespace: easyai
spec:
selector:
app.kubernetes.io/name: easyai-acceptance-upstream-proxy
ports:
- name: proxy
port: 8090
targetPort: proxy
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-postgres-proxy-ningbo
namespace: easyai
labels:
app.kubernetes.io/name: easyai-postgres-proxy
easyai.io/site: ningbo
easyai.io/environment: local-acceptance
spec:
replicas: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-postgres-proxy
easyai.io/site: ningbo
template:
metadata:
labels:
app.kubernetes.io/name: easyai-postgres-proxy
easyai.io/site: ningbo
easyai.io/environment: local-acceptance
spec:
nodeSelector:
easyai.io/site: ningbo
containers:
- name: proxy
image: easyai-acceptance-netem
env:
- name: NETEM_LISTEN_PORT
value: "15432"
- name: NETEM_TARGET_HOST
value: easyai-postgres-rw.easyai.svc.cluster.local
- name: NETEM_TARGET_PORT
value: "5432"
ports:
- name: postgres
containerPort: 15432
resources:
requests:
cpu: 20m
memory: 32Mi
limits:
cpu: 250m
memory: 128Mi
securityContext:
allowPrivilegeEscalation: false
capabilities:
add: ["NET_ADMIN"]
drop: ["ALL"]
readOnlyRootFilesystem: true
---
apiVersion: v1
kind: Service
metadata:
name: easyai-postgres-proxy-ningbo
namespace: easyai
spec:
selector:
app.kubernetes.io/name: easyai-postgres-proxy
easyai.io/site: ningbo
ports:
- name: postgres
port: 5432
targetPort: postgres
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-postgres-proxy-hongkong
namespace: easyai
labels:
app.kubernetes.io/name: easyai-postgres-proxy
easyai.io/site: hongkong
easyai.io/environment: local-acceptance
spec:
replicas: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-postgres-proxy
easyai.io/site: hongkong
template:
metadata:
labels:
app.kubernetes.io/name: easyai-postgres-proxy
easyai.io/site: hongkong
easyai.io/environment: local-acceptance
spec:
nodeSelector:
easyai.io/site: hongkong
containers:
- name: proxy
image: easyai-acceptance-netem
env:
- name: NETEM_LISTEN_PORT
value: "15432"
- name: NETEM_TARGET_HOST
value: easyai-postgres-rw.easyai.svc.cluster.local
- name: NETEM_TARGET_PORT
value: "5432"
ports:
- name: postgres
containerPort: 15432
resources:
requests:
cpu: 20m
memory: 32Mi
limits:
cpu: 250m
memory: 128Mi
securityContext:
allowPrivilegeEscalation: false
capabilities:
add: ["NET_ADMIN"]
drop: ["ALL"]
readOnlyRootFilesystem: true
---
apiVersion: v1
kind: Service
metadata:
name: easyai-postgres-proxy-hongkong
namespace: easyai
spec:
selector:
app.kubernetes.io/name: easyai-postgres-proxy
easyai.io/site: hongkong
ports:
- name: postgres
port: 5432
targetPort: postgres
@@ -0,0 +1,172 @@
apiVersion: v1
kind: ConfigMap
metadata:
name: easyai-acceptance-tls-edge
namespace: easyai
data:
default.conf: |
server {
listen 8443 ssl;
server_name gateway.easyai.local;
client_max_body_size 128m;
proxy_request_buffering off;
proxy_read_timeout 900s;
ssl_certificate /etc/acceptance-tls/tls.crt;
ssl_certificate_key /etc/acceptance-tls/tls.key;
location / {
proxy_set_header Host $host;
proxy_set_header X-Forwarded-Proto https;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
proxy_pass http://LOCAL_API_SERVICE.easyai.svc.cluster.local:8088;
}
}
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-acceptance-edge-ningbo
namespace: easyai
labels:
app.kubernetes.io/name: easyai-acceptance-edge
easyai.io/site: ningbo
spec:
replicas: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-acceptance-edge
easyai.io/site: ningbo
template:
metadata:
labels:
app.kubernetes.io/name: easyai-acceptance-edge
easyai.io/site: ningbo
spec:
nodeSelector:
easyai.io/site: ningbo
containers:
- name: edge
image: easyai-web
command:
- /bin/sh
- -ec
- sed 's/LOCAL_API_SERVICE/api-ningbo-edge/g' /etc/acceptance-edge/default.conf > /etc/nginx/conf.d/default.conf && exec nginx -g 'daemon off;'
ports:
- name: https
containerPort: 8443
volumeMounts:
- name: config
mountPath: /etc/acceptance-edge
readOnly: true
- name: nginx-config
mountPath: /etc/nginx/conf.d
- name: tls
mountPath: /etc/acceptance-tls
readOnly: true
resources:
requests:
cpu: 50m
memory: 32Mi
limits:
cpu: 250m
memory: 128Mi
volumes:
- name: config
configMap:
name: easyai-acceptance-tls-edge
- name: tls
secret:
secretName: easyai-acceptance-tls
- name: nginx-config
emptyDir: {}
---
apiVersion: v1
kind: Service
metadata:
name: easyai-acceptance-edge-ningbo
namespace: easyai
spec:
type: NodePort
externalTrafficPolicy: Local
selector:
app.kubernetes.io/name: easyai-acceptance-edge
easyai.io/site: ningbo
ports:
- name: https
port: 8443
targetPort: https
nodePort: 31088
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-acceptance-edge-hongkong
namespace: easyai
labels:
app.kubernetes.io/name: easyai-acceptance-edge
easyai.io/site: hongkong
spec:
replicas: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-acceptance-edge
easyai.io/site: hongkong
template:
metadata:
labels:
app.kubernetes.io/name: easyai-acceptance-edge
easyai.io/site: hongkong
spec:
nodeSelector:
easyai.io/site: hongkong
containers:
- name: edge
image: easyai-web
command:
- /bin/sh
- -ec
- sed 's/LOCAL_API_SERVICE/api-hongkong-edge/g' /etc/acceptance-edge/default.conf > /etc/nginx/conf.d/default.conf && exec nginx -g 'daemon off;'
ports:
- name: https
containerPort: 8443
volumeMounts:
- name: config
mountPath: /etc/acceptance-edge
readOnly: true
- name: nginx-config
mountPath: /etc/nginx/conf.d
- name: tls
mountPath: /etc/acceptance-tls
readOnly: true
resources:
requests:
cpu: 50m
memory: 32Mi
limits:
cpu: 250m
memory: 128Mi
volumes:
- name: config
configMap:
name: easyai-acceptance-tls-edge
- name: tls
secret:
secretName: easyai-acceptance-tls
- name: nginx-config
emptyDir: {}
---
apiVersion: v1
kind: Service
metadata:
name: easyai-acceptance-edge-hongkong
namespace: easyai
spec:
type: NodePort
externalTrafficPolicy: Local
selector:
app.kubernetes.io/name: easyai-acceptance-edge
easyai.io/site: hongkong
ports:
- name: https
port: 8443
targetPort: https
nodePort: 31089
@@ -43,4 +43,21 @@ data:
AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX: easyai-ai-gateway/production/media
AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED: "true"
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT: "48"
AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT: "48"
AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS: "5"
AI_GATEWAY_WORKER_AUTOSCALING_ENABLED: "false"
AI_GATEWAY_WORKER_REPLICAS_NINGBO: "1"
AI_GATEWAY_WORKER_REPLICAS_HONGKONG: "1"
AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO: "1"
AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG: "1"
AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO: "1"
AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG: "1"
AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA: "48"
AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS: "20"
AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS: "600"
AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS: "600"
AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT: "75"
AI_GATEWAY_NODE_MEMORY_HARD_PERCENT: "85"
AI_GATEWAY_NODE_CPU_TARGET_PERCENT: "70"
AI_GATEWAY_POSTGRES_CONNECTION_BUDGET: "150"
AI_GATEWAY_WORKER_DATABASE_MAX_CONNS: "32"
+170 -6
View File
@@ -53,7 +53,11 @@ spec:
- name: AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED
value: "false"
- name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "64"
value: "32"
- name: AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_RIVER_MAX_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_MIN_IDLE_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS
@@ -143,6 +147,128 @@ spec:
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-capacity-controller
labels:
app.kubernetes.io/name: easyai-capacity-controller
app.kubernetes.io/component: capacity-controller
app.kubernetes.io/part-of: easyai-ai-gateway
spec:
replicas: 2
strategy:
type: RollingUpdate
rollingUpdate:
maxSurge: 1
maxUnavailable: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-capacity-controller
template:
metadata:
labels:
app.kubernetes.io/name: easyai-capacity-controller
app.kubernetes.io/component: capacity-controller
app.kubernetes.io/part-of: easyai-ai-gateway
spec:
serviceAccountName: easyai-capacity-controller
terminationGracePeriodSeconds: 30
nodeSelector:
easyai.io/workload: "true"
topologySpreadConstraints:
- maxSkew: 1
topologyKey: easyai.io/site
whenUnsatisfiable: DoNotSchedule
labelSelector:
matchLabels:
app.kubernetes.io/name: easyai-capacity-controller
securityContext:
runAsNonRoot: true
runAsUser: 10001
runAsGroup: 10001
fsGroup: 10001
seccompProfile:
type: RuntimeDefault
containers:
- name: capacity-controller
image: easyai-api
imagePullPolicy: IfNotPresent
envFrom:
- configMapRef:
name: easyai-ai-gateway-config
env:
- name: AI_GATEWAY_PROCESS_ROLE
value: capacity-controller
- name: AI_GATEWAY_DATABASE_URL
valueFrom:
secretKeyRef:
name: easyai-ai-gateway-runtime
key: AI_GATEWAY_DATABASE_URL
- name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS
value: "0"
- name: AI_GATEWAY_DATABASE_RIVER_MAX_CONNS
value: "0"
- name: AI_GATEWAY_DATABASE_MIN_IDLE_CONNS
value: "1"
- name: AI_GATEWAY_WORKER_DATABASE_MAX_CONNS
value: "32"
- name: AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET
value: "72"
- name: POD_NAMESPACE
valueFrom:
fieldRef:
fieldPath: metadata.namespace
- name: POD_NAME
valueFrom:
fieldRef:
fieldPath: metadata.name
ports:
- name: health
containerPort: 8088
startupProbe:
httpGet:
path: /healthz
port: health
periodSeconds: 5
timeoutSeconds: 3
failureThreshold: 24
readinessProbe:
httpGet:
path: /readyz
port: health
periodSeconds: 10
timeoutSeconds: 3
failureThreshold: 6
livenessProbe:
httpGet:
path: /healthz
port: health
periodSeconds: 10
timeoutSeconds: 3
failureThreshold: 6
resources:
requests:
cpu: 50m
memory: 64Mi
limits:
cpu: 200m
memory: 256Mi
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop: ["ALL"]
readOnlyRootFilesystem: true
volumeMounts:
- name: temporary-files
mountPath: /tmp
volumes:
- name: temporary-files
emptyDir:
sizeLimit: 64Mi
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: easyai-api-hongkong
labels:
@@ -196,7 +322,11 @@ spec:
- name: AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED
value: "false"
- name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "64"
value: "32"
- name: AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_RIVER_MAX_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_MIN_IDLE_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS
@@ -313,6 +443,14 @@ spec:
spec:
serviceAccountName: easyai-ai-gateway
terminationGracePeriodSeconds: 90
topologySpreadConstraints:
- maxSkew: 1
topologyKey: kubernetes.io/hostname
whenUnsatisfiable: DoNotSchedule
labelSelector:
matchLabels:
app.kubernetes.io/name: easyai-worker
easyai.io/site: ningbo
nodeSelector:
easyai.io/site: ningbo
easyai.io/workload: "true"
@@ -340,6 +478,10 @@ spec:
value: "true"
- name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "32"
- name: AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_RIVER_MAX_CONNS
value: "8"
- name: AI_GATEWAY_DATABASE_MIN_IDLE_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS
@@ -406,8 +548,8 @@ spec:
command: ["/bin/sh", "-c", "sleep 10"]
resources:
requests:
cpu: 250m
memory: 512Mi
cpu: 500m
memory: 1536Mi
limits:
cpu: "2"
memory: 2Gi
@@ -458,6 +600,14 @@ spec:
spec:
serviceAccountName: easyai-ai-gateway
terminationGracePeriodSeconds: 90
topologySpreadConstraints:
- maxSkew: 1
topologyKey: kubernetes.io/hostname
whenUnsatisfiable: DoNotSchedule
labelSelector:
matchLabels:
app.kubernetes.io/name: easyai-worker
easyai.io/site: hongkong
nodeSelector:
easyai.io/site: hongkong
easyai.io/workload: "true"
@@ -485,6 +635,10 @@ spec:
value: "true"
- name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "32"
- name: AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_RIVER_MAX_CONNS
value: "8"
- name: AI_GATEWAY_DATABASE_MIN_IDLE_CONNS
value: "4"
- name: AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS
@@ -551,8 +705,8 @@ spec:
command: ["/bin/sh", "-c", "sleep 10"]
resources:
requests:
cpu: 250m
memory: 512Mi
cpu: 500m
memory: 1536Mi
limits:
cpu: "2"
memory: 2Gi
@@ -896,6 +1050,16 @@ spec:
---
apiVersion: policy/v1
kind: PodDisruptionBudget
metadata:
name: easyai-capacity-controller
spec:
minAvailable: 1
selector:
matchLabels:
app.kubernetes.io/name: easyai-capacity-controller
---
apiVersion: policy/v1
kind: PodDisruptionBudget
metadata:
name: easyai-worker
spec:
@@ -7,7 +7,7 @@ spec:
matchExpressions:
- key: app.kubernetes.io/name
operator: In
values: [easyai-api, easyai-worker, easyai-web]
values: [easyai-api, easyai-worker, easyai-web, easyai-capacity-controller]
policyTypes: [Ingress]
ingress:
- from:
@@ -29,3 +29,68 @@ roleRef:
apiGroup: rbac.authorization.k8s.io
kind: Role
name: easyai-gateway-security-events
---
apiVersion: v1
kind: ServiceAccount
metadata:
name: easyai-capacity-controller
labels:
app.kubernetes.io/name: easyai-capacity-controller
app.kubernetes.io/part-of: easyai-ai-gateway
automountServiceAccountToken: true
imagePullSecrets:
- name: easyai-registry
---
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
name: easyai-capacity-controller
rules:
- apiGroups: ["apps"]
resources: ["deployments"]
resourceNames: ["easyai-worker-ningbo", "easyai-worker-hongkong"]
verbs: ["get"]
- apiGroups: ["apps"]
resources: ["deployments/scale"]
resourceNames: ["easyai-worker-ningbo", "easyai-worker-hongkong"]
verbs: ["get", "update", "patch"]
- apiGroups: [""]
resources: ["pods"]
verbs: ["get", "list", "patch"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
name: easyai-capacity-controller
subjects:
- kind: ServiceAccount
name: easyai-capacity-controller
roleRef:
apiGroup: rbac.authorization.k8s.io
kind: Role
name: easyai-capacity-controller
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
name: easyai-capacity-controller
rules:
- apiGroups: [""]
resources: ["nodes"]
verbs: ["get", "list"]
- apiGroups: ["metrics.k8s.io"]
resources: ["nodes", "pods"]
verbs: ["get", "list"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
name: easyai-capacity-controller
subjects:
- kind: ServiceAccount
name: easyai-capacity-controller
namespace: easyai
roleRef:
apiGroup: rbac.authorization.k8s.io
kind: ClusterRole
name: easyai-capacity-controller
@@ -0,0 +1,106 @@
# 本地三节点同构验收
## 链路边界
本地验收不使用旧 `simulation` 请求模式。压测端发送正式 Gemini `generateContent` 和正式
多参考图视频请求,完整经过:
```text
双 TLS 入口 → API 鉴权与 validation 门禁 → 生产候选路由
→ Base64 解码及媒体物化 → PostgreSQL/River → Worker/租约
→ 真实 Gemini/Volces 客户端 → 协议模拟上游
→ 结果持久化 → 独立钱包结算 → 独立回调收集器 → 客户端响应
```
只有候选已经确定后的 Base URL 和凭据被验收 Run 替换。模拟上游实现真实协议、延迟、
异步轮询及示例媒体返回;API、Worker、候选选择、账务和恢复代码不走测试旁路。
## 资源和工具
本地集群名固定为 `easyai-acceptance-local`,包含三个 k3d/K3s server
| 节点 | Docker 上限 | 调度职责 |
|---|---:|---|
| server-0 / ningbo | 4 CPU / 8 GiB | API、Worker、PostgreSQL |
| server-1 / hongkong | 4 CPU / 8 GiB | API、Worker、PostgreSQL |
| server-2 / los-angeles | 2 CPU / 4 GiB | 控制面和 etcd`NoSchedule` |
Docker Desktop 必须配置至少 24 GiB 内存;脚本只检查,不修改 Docker 设置。k3d、K3s 和
CNPG 版本及 SHA-256 位于
`deploy/kubernetes/local-acceptance/dependencies.lock`。高并发阶段使用宿主原生架构镜像,
最后才导入 release manifest 中精确的 `linux/amd64@sha256` 制品做启动、迁移和媒体冒烟。
```bash
scripts/acceptance/local-cluster.sh install-tools
scripts/acceptance/local-cluster.sh preflight
```
失败的 `up` 默认保留集群。只有下面的显式命令会销毁:
```bash
scripts/acceptance/local-cluster.sh down --confirm
```
## 脱敏生产快照
先使用只有 `SELECT` 权限的数据库角色导出 `acceptance-snapshot/v1`
```bash
AI_GATEWAY_ACCEPTANCE_SNAPSHOT_DATABASE_URL='<只读连接>' \
scripts/acceptance/export-production-snapshot.sh \
--release-sha <当前线上完整SHA> \
--output .local-secrets/acceptance/production-snapshot.json
```
快照只包含一个 Gemini 图片候选和一个 `omni_video.max_images >= 9` 视频候选所需的能力、
协议、路由、限流、价格、重试和运行策略。导出器拒绝密码、Secret、Token、认证、代理和
连接串字段,文件权限固定为 `0600`,并记录配置哈希和快照 SHA-256。
本地导入必须先匹配 `acceptance_local_cluster_id` 数据库标记,否则拒绝写入。集群内创建
32 个隔离身份/API Key、独立钱包和专属 RunKey 与 Run Token 只保存在 `0600` 临时文件,
不进入 Pod Spec、报告或日志。
## 创建与验收
源码必须已提交且工作区干净,确保本地镜像和报告中的完整 Git SHA 可复现:
```bash
scripts/acceptance/local-cluster.sh up \
--snapshot .local-secrets/acceptance/production-snapshot.json
scripts/acceptance/run-local-acceptance.sh quick
scripts/acceptance/run-local-acceptance.sh full \
--release-manifest dist/releases/<完整SHA>.json
```
`full` 依次执行:
1. Gemini、3/6/9 图视频、账务和回调快速验收。
2. 固定 `1+1` Worker 的 P24、P28、P32,每档完整运行三次。
3. 1000×256 KiB、128×2 MiB、32×8 MiB Gemini 与 1200/96 视频负载。
4. 40±20 ms/0.5% 丢包、上游断链 10 秒、Worker—数据库断链 30 秒。
5. 确实持有远程任务的 Worker 强杀,以及容量控制器 Leader 切换。
6. `1+1 → 1+2/2+2 → drain → 1+1` 弹性验证。
7. 80% 两小时混合流量和 120% 十分钟主动限流。
8. 精确 amd64 release 镜像的迁移、启动与媒体冒烟。
网络故障只能作用于带 `easyai.io/environment=local-acceptance` 标签的代理 Pod
```bash
scripts/acceptance/network-fault.sh baseline
scripts/acceptance/network-fault.sh weak-link
scripts/acceptance/network-fault.sh upstream-outage
scripts/acceptance/network-fault.sh database-outage hongkong
scripts/acceptance/network-fault.sh reset
```
统一报告为 `acceptance-report/v1`,位于
`dist/acceptance/local/<run-id>/acceptance-report.json`。本地吞吐只用于筛选候选档位;
生产 certified profile 必须在线上模拟验收中重新确定。
## 不能被本地替代的门禁
本地环境不能认证真实 WireGuard 六向链路、生产磁盘、跨地域 CNPG 同步、供应商配额或
生产节点混部资源。线上 `--execute` 会重新导出生产候选配置;配置哈希与本地快照不一致时
立即停止,必须基于新快照重跑本地阻断验收。
+23 -6
View File
@@ -8,7 +8,7 @@
验收脚本是显式生产写操作,会暂停新正式任务、调整 Worker 容量并在恢复场景删除一个
Worker Pod。它不会自动发布新版本;只能针对已经发布、完整 SHA 和 digest 固定的 manifest
执行。
执行。进入线上前必须先通过[本地三节点同构验收](local-isomorphic-acceptance.md)。
## 隔离与流量门禁
@@ -79,8 +79,9 @@ AI_GATEWAY_WORKER_REPLICAS_NINGBO
AI_GATEWAY_WORKER_REPLICAS_HONGKONG
```
`AI_GATEWAY_WORKER_REPLICAS_*` 是发布工具配置,不会传入容器。当前验收固定宁波、香港
一个 2 GiB Worker不创建额外 Pod
`AI_GATEWAY_WORKER_REPLICAS_*` 是发布工具配置,不会传入容器。容量档位先固定宁波、香港
一个 2 GiB Worker再按实时内存、CPU 和 PostgreSQL 预算验证 `1+2`、条件允许时的
`2+2`,以及安全 drain 回到 `1+1`
| 档位 | 单实例执行槽 | 数据库池 | 媒体并发 | 全局执行槽 |
|---|---:|---:|---:|---:|
@@ -120,7 +121,8 @@ easyai-ai-gateway-cluster-release capacity
```bash
scripts/cluster/run-production-acceptance.sh \
--execute dist/releases/<完整SHA>.json
--execute dist/releases/<完整SHA>.json \
--local-report dist/acceptance/local/<本地run-id>/acceptance-report.json
```
脚本依次执行:
@@ -133,11 +135,25 @@ scripts/cluster/run-production-acceptance.sh \
6. 执行两条真实小流量请求。
7. 校验任务、账务、回调、队列、连接池、RSS、节点内存、数据库同步、六向 WireGuard 和
公网健康。
8. 重新校验 SHA/digest/revision 后切回 `live`
8. 完成 Run 并输出 `acceptance-report/v1`,保持 `validation`,等待人工确认
报告保存在 `dist/acceptance/<run-id>/`,包含每阶段吞吐、p50/p95/p99、队列和资源 CSV、
协议模拟器统计及最终摘要。报告不包含图片原文、Base64、Token、密码或连接串。
确认整体报告后单独执行人工开闸:
```bash
scripts/cluster/run-production-acceptance.sh \
--promote dist/releases/<完整SHA>.json \
--run-id <线上run-id>
```
`--promote` 会重新读取 Run、流量 revision、release SHA、镜像 digest、配置哈希及暂存容量
配置,任一 CAS 不一致即拒绝。成功切到 `live` 后同一命令继续运行有限期监控:前 2 小时
每 10 秒采样,之后每 60 秒采样至 24 小时,不恢复已经停用的周期性巡检。队列、RSS、
连接池、同步复制、租约、重复结算、WireGuard 或公网健康触发硬门禁时,监控通过 CAS
自动切回 `validation`,并恢复 Run 开始前保存的精确容量快照;不会自动回滚数据库。
## 通过和失败处理
通过要求包括:任务全部成功、Gemini 输出可解码且哈希/字节数正确、视频参考图校验全部
@@ -145,6 +161,7 @@ scripts/cluster/run-production-acceptance.sh \
超过当前档位、无 OOM/重启、Pod RSS 小于 1.5 GiB、节点内存低于 80%、PostgreSQL 连接
低于 75%、连接池不连续满载、双副本同步、六向链路和公网健康通过。
失败 Run 会记录原因、恢复上一稳定容量并保持 `validation`排除问题后调用 Run
失败 Run 会记录原因、恢复上一稳定容量并保持 `validation`成功 Run 也保持
`validation`,只有上述人工 `--promote` 才能开闸。排除失败问题后调用 Run 的
`retry` 接口重新执行。只有人工决定放弃验收时,才使用带相同 CAS 字段的 `abort` 接口
恢复 `live`;该操作不会把失败验收标记为通过。
@@ -0,0 +1,73 @@
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"$id": "https://easyai.local/schemas/acceptance-report-v1.schema.json",
"title": "EasyAI acceptance-report/v1",
"type": "object",
"additionalProperties": false,
"required": ["schemaVersion", "runId", "release", "snapshot", "stages", "gates", "promotion", "createdAt", "secretSafe"],
"properties": {
"schemaVersion": {"const": "acceptance-report/v1"},
"runId": {"type": "string", "minLength": 1},
"release": {
"type": "object",
"additionalProperties": false,
"required": ["sourceSha", "apiImageDigest", "workerImageDigest"],
"properties": {
"sourceSha": {"type": "string", "pattern": "^[0-9a-f]{40}$"},
"apiImageDigest": {"type": "string", "pattern": "^sha256:[0-9a-f]{64}$"},
"workerImageDigest": {"type": "string", "pattern": "^sha256:[0-9a-f]{64}$"}
}
},
"snapshot": {
"type": "object",
"additionalProperties": false,
"required": ["sourceReleaseSha", "configHash", "sha256"],
"properties": {
"sourceReleaseSha": {"type": "string", "pattern": "^[0-9a-f]{40}$"},
"configHash": {"type": "string", "pattern": "^[0-9a-f]{64}$"},
"sha256": {"type": "string", "pattern": "^[0-9a-f]{64}$"}
}
},
"stages": {
"type": "object",
"additionalProperties": {"$ref": "#/$defs/stage"}
},
"certifiedProfile": {"type": ["object", "null"]},
"gates": {
"type": "array",
"items": {
"type": "object",
"additionalProperties": false,
"required": ["id", "passed"],
"properties": {
"id": {"type": "string", "minLength": 1},
"passed": {"type": "boolean"},
"detail": {"type": "string"}
}
}
},
"promotion": {
"type": "object",
"additionalProperties": false,
"required": ["ready", "requiresManualConfirmation", "trafficMode"],
"properties": {
"ready": {"type": "boolean"},
"requiresManualConfirmation": {"const": true},
"trafficMode": {"enum": ["validation", "live", "not-applicable"]},
"promotedAt": {"type": "string", "format": "date-time"}
}
},
"createdAt": {"type": "string", "format": "date-time"},
"secretSafe": {"const": true}
},
"$defs": {
"stage": {
"type": "object",
"additionalProperties": true,
"required": ["status"],
"properties": {
"status": {"enum": ["passed", "failed", "not-run"]}
}
}
}
}
@@ -0,0 +1,35 @@
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"$id": "https://easyai.local/schemas/acceptance-snapshot-v1.schema.json",
"title": "EasyAI acceptance-snapshot/v1",
"type": "object",
"additionalProperties": false,
"required": ["schemaVersion", "source", "candidates", "snapshotSha256", "secretSafe"],
"properties": {
"schemaVersion": {"const": "acceptance-snapshot/v1"},
"source": {
"type": "object",
"additionalProperties": false,
"required": ["releaseSha", "configHash", "createdAt"],
"properties": {
"releaseSha": {"type": "string", "pattern": "^[0-9a-f]{40}$"},
"configHash": {"type": "string", "pattern": "^[0-9a-f]{64}$"},
"createdAt": {"type": "string", "format": "date-time"}
}
},
"candidates": {
"type": "array",
"minItems": 2,
"maxItems": 2,
"items": {
"type": "object",
"required": ["workload", "provider", "baseModel", "platform", "platformModel", "runtimePolicy", "pricingRules", "metadata"],
"properties": {
"workload": {"enum": ["gemini_image_edit", "multi_reference_video"]}
}
}
},
"snapshotSha256": {"type": "string", "pattern": "^[0-9a-f]{64}$"},
"secretSafe": {"const": true}
}
}
+43
View File
@@ -0,0 +1,43 @@
#!/usr/bin/env bash
set -euo pipefail
usage() {
cat <<'EOF'
Usage:
AI_GATEWAY_ACCEPTANCE_SNAPSHOT_DATABASE_URL='postgresql://...' \
scripts/acceptance/export-production-snapshot.sh \
--release-sha <full-sha> --output <snapshot.json>
This command is database read-only. The supplied role should have SELECT access
only. It never prints the database URL or candidate credentials.
EOF
}
script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
repository_root=$(cd "$script_dir/../.." && pwd)
[[ ${1:-} == --release-sha && ${3:-} == --output && $# -eq 4 ]] || {
usage >&2
exit 64
}
release_sha=$2
output=$4
database_url=${AI_GATEWAY_ACCEPTANCE_SNAPSHOT_DATABASE_URL:-}
[[ $release_sha =~ ^[0-9a-f]{40}$ ]] || {
echo 'release SHA must be a full lowercase Git SHA' >&2
exit 64
}
[[ -n $database_url ]] || {
echo 'AI_GATEWAY_ACCEPTANCE_SNAPSHOT_DATABASE_URL is required' >&2
exit 64
}
umask 077
(
cd "$repository_root/apps/api"
AI_GATEWAY_DATABASE_URL="$database_url" \
go run ./cmd/acceptance-snapshot export \
--release-sha "$release_sha" \
--output "$output"
)
chmod 0600 "$output"
+556
View File
@@ -0,0 +1,556 @@
#!/usr/bin/env bash
set -euo pipefail
script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
repository_root=$(cd "$script_dir/../.." && pwd)
manifest_root="$repository_root/deploy/kubernetes/local-acceptance"
lock_file="$manifest_root/dependencies.lock"
private_root="$repository_root/.local-secrets/acceptance"
tools_root="$private_root/tools"
state_root="$private_root/state"
media_root="$private_root/media"
cluster_name=easyai-acceptance-local
context="k3d-$cluster_name"
namespace=easyai
usage() {
cat <<'EOF'
Usage:
scripts/acceptance/local-cluster.sh install-tools
scripts/acceptance/local-cluster.sh preflight
scripts/acceptance/local-cluster.sh up --snapshot <acceptance-snapshot-v1.json>
scripts/acceptance/local-cluster.sh status
scripts/acceptance/local-cluster.sh down --confirm
The full cluster requires Docker Desktop memory >= 24 GiB. Failed `up` runs are
kept intact for diagnosis. `down` is the only destructive lifecycle action.
EOF
}
load_lock() {
[[ -f $lock_file && ! -L $lock_file ]] || {
echo "missing dependency lock: $lock_file" >&2
exit 1
}
# shellcheck source=/dev/null
source "$lock_file"
: "${K3D_VERSION:?}" "${K3S_IMAGE:?}" "${CNPG_MANIFEST_URL:?}" "${CNPG_MANIFEST_SHA256:?}"
}
require_command() {
command -v "$1" >/dev/null 2>&1 || {
echo "$1 is required" >&2
exit 1
}
}
private_file() {
local path=$1
[[ -f $path && ! -L $path ]] || return 1
local mode
if [[ $(uname -s) == Darwin ]]; then
mode=$(stat -f '%Lp' "$path")
else
mode=$(stat -c '%a' "$path")
fi
[[ $mode == 600 ]]
}
install_tools() {
load_lock
require_command curl
require_command shasum
install -d -m 0700 "$tools_root"
local os arch checksum url temporary
os=$(uname -s | tr '[:upper:]' '[:lower:]')
arch=$(uname -m)
[[ $os == darwin ]] || {
echo 'the pinned local acceptance installer currently supports macOS only' >&2
exit 1
}
case $arch in
arm64)
arch=arm64
checksum=$K3D_DARWIN_ARM64_SHA256
;;
x86_64)
arch=amd64
checksum=$K3D_DARWIN_AMD64_SHA256
;;
*)
echo "unsupported host architecture: $arch" >&2
exit 1
;;
esac
url="https://github.com/k3d-io/k3d/releases/download/$K3D_VERSION/k3d-darwin-$arch"
temporary="$tools_root/k3d.download"
curl -fsSL "$url" -o "$temporary"
[[ $(shasum -a 256 "$temporary" | awk '{print $1}') == "$checksum" ]] || {
echo 'k3d checksum mismatch' >&2
exit 1
}
chmod 0555 "$temporary"
mv "$temporary" "$tools_root/k3d"
echo "local_acceptance_tools=PASS k3d=$K3D_VERSION path=$tools_root/k3d"
}
k3d_binary() {
if [[ -x $tools_root/k3d ]]; then
printf '%s\n' "$tools_root/k3d"
return
fi
command -v k3d
}
docker_memory_bytes() {
docker info --format '{{.MemTotal}}'
}
preflight() {
load_lock
require_command docker
require_command kubectl
require_command jq
require_command openssl
require_command sed
require_command shasum
require_command go
require_command node
require_command curl
local k3d memory_bytes architecture
k3d=$(k3d_binary) || {
echo 'k3d is missing; run install-tools' >&2
exit 1
}
"$k3d" version | grep -Fq "$K3D_VERSION" || {
echo "k3d must be pinned to $K3D_VERSION" >&2
exit 1
}
docker info >/dev/null
memory_bytes=$(docker_memory_bytes)
[[ $memory_bytes =~ ^[0-9]+$ && $memory_bytes -ge 25769803776 ]] || {
memory_gib=$(awk -v bytes="${memory_bytes:-0}" 'BEGIN {printf "%.1f", bytes/1024/1024/1024}')
echo "Docker Desktop memory is ${memory_gib} GiB; local acceptance requires at least 24 GiB" >&2
exit 1
}
architecture=$(docker info --format '{{.Architecture}}')
[[ $architecture == aarch64 || $architecture == arm64 || $architecture == x86_64 ]] || {
echo "unsupported Docker architecture: $architecture" >&2
exit 1
}
docker buildx version >/dev/null
echo "local_acceptance_preflight=PASS docker_memory_bytes=$memory_bytes architecture=$architecture k3d=$K3D_VERSION"
}
ensure_private_material() {
install -d -m 0700 "$private_root" "$state_root" "$media_root"
local env_file="$state_root/local.env"
if [[ ! -e $env_file ]]; then
umask 077
{
printf 'LOCAL_CLUSTER_ID=%s\n' "easyai-local-$(openssl rand -hex 12)"
printf 'POSTGRES_PASSWORD=%s\n' "$(openssl rand -hex 32)"
printf 'CONFIG_JWT_SECRET=%s\n' "$(openssl rand -hex 32)"
printf 'SERVER_MAIN_INTERNAL_TOKEN=%s\n' "$(openssl rand -hex 32)"
printf 'SERVER_MAIN_INTERNAL_KEY=local-acceptance\n'
printf 'SERVER_MAIN_INTERNAL_SECRET=%s\n' "$(openssl rand -hex 32)"
} >"$env_file"
chmod 0600 "$env_file"
fi
private_file "$env_file" || {
echo "local acceptance state must be a 0600 regular file: $env_file" >&2
exit 1
}
}
generate_tls() {
local ca_key="$state_root/ca.key" ca_crt="$state_root/ca.crt"
local tls_key="$state_root/tls.key" tls_csr="$state_root/tls.csr"
local tls_crt="$state_root/tls.crt" extensions="$state_root/tls.ext"
if [[ ! -e $ca_key ]]; then
umask 077
openssl genrsa -out "$ca_key" 3072 >/dev/null 2>&1
openssl req -x509 -new -key "$ca_key" -sha256 -days 30 \
-subj '/CN=EasyAI local acceptance CA' -out "$ca_crt" >/dev/null 2>&1
openssl genrsa -out "$tls_key" 2048 >/dev/null 2>&1
openssl req -new -key "$tls_key" -subj '/CN=gateway.easyai.local' \
-out "$tls_csr" >/dev/null 2>&1
printf 'subjectAltName=DNS:gateway.easyai.local\nextendedKeyUsage=serverAuth\n' >"$extensions"
openssl x509 -req -in "$tls_csr" -CA "$ca_crt" -CAkey "$ca_key" \
-CAcreateserial -out "$tls_crt" -days 30 -sha256 -extfile "$extensions" \
>/dev/null 2>&1
chmod 0600 "$ca_key" "$ca_crt" "$tls_key" "$tls_csr" "$tls_crt" "$extensions"
fi
private_file "$ca_crt" && private_file "$tls_key" && private_file "$tls_crt"
}
render_k3d_config() {
local output=$1
[[ $media_root != *'|'* ]]
sed "s|EASYAI_ACCEPTANCE_MEDIA_DIR|$media_root|g" "$manifest_root/k3d.yaml" >"$output"
}
wait_rollout() {
kubectl --context "$context" -n "$namespace" rollout status "$1" --timeout="${2:-10m}"
}
apply_runtime_secrets() {
# shellcheck source=/dev/null
source "$state_root/local.env"
local database_ningbo database_hongkong database_direct
database_ningbo="postgresql://easyai:${POSTGRES_PASSWORD}@easyai-postgres-proxy-ningbo.easyai.svc.cluster.local:5432/easyai_ai_gateway?sslmode=require"
database_hongkong="postgresql://easyai:${POSTGRES_PASSWORD}@easyai-postgres-proxy-hongkong.easyai.svc.cluster.local:5432/easyai_ai_gateway?sslmode=require"
database_direct="postgresql://easyai:${POSTGRES_PASSWORD}@easyai-postgres-rw.easyai.svc.cluster.local:5432/easyai_ai_gateway?sslmode=require"
kubectl --context "$context" create namespace "$namespace" --dry-run=client -o yaml |
kubectl --context "$context" apply -f - >/dev/null
kubectl --context "$context" -n "$namespace" create secret generic easyai-postgres-app \
--type=kubernetes.io/basic-auth \
--from-literal=username=easyai --from-literal=password="$POSTGRES_PASSWORD" \
--dry-run=client -o yaml | kubectl --context "$context" apply -f - >/dev/null
kubectl --context "$context" -n "$namespace" create secret generic easyai-ai-gateway-runtime \
--from-literal=AI_GATEWAY_DATABASE_URL="$database_direct" \
--from-literal=CONFIG_JWT_SECRET="$CONFIG_JWT_SECRET" \
--from-literal=SERVER_MAIN_INTERNAL_TOKEN="$SERVER_MAIN_INTERNAL_TOKEN" \
--from-literal=SERVER_MAIN_INTERNAL_KEY="$SERVER_MAIN_INTERNAL_KEY" \
--from-literal=SERVER_MAIN_INTERNAL_SECRET="$SERVER_MAIN_INTERNAL_SECRET" \
--dry-run=client -o yaml | kubectl --context "$context" apply -f - >/dev/null
kubectl --context "$context" -n "$namespace" create secret generic easyai-database-ningbo \
--from-literal=AI_GATEWAY_DATABASE_URL="$database_ningbo" \
--dry-run=client -o yaml | kubectl --context "$context" apply -f - >/dev/null
kubectl --context "$context" -n "$namespace" create secret generic easyai-database-hongkong \
--from-literal=AI_GATEWAY_DATABASE_URL="$database_hongkong" \
--dry-run=client -o yaml | kubectl --context "$context" apply -f - >/dev/null
kubectl --context "$context" -n "$namespace" create secret generic easyai-oss-backup \
--from-literal=ACCESS_KEY_ID=local-disabled \
--from-literal=ACCESS_SECRET_KEY=local-disabled \
--dry-run=client -o yaml | kubectl --context "$context" apply -f - >/dev/null
kubectl --context "$context" -n "$namespace" create secret tls easyai-acceptance-tls \
--cert="$state_root/tls.crt" --key="$state_root/tls.key" \
--dry-run=client -o yaml | kubectl --context "$context" apply -f - >/dev/null
}
run_migrations() {
local api_image=$1
kubectl --context "$context" -n "$namespace" delete job easyai-local-migrate \
--ignore-not-found --wait=true >/dev/null
cat <<EOF | kubectl --context "$context" apply -f - >/dev/null
apiVersion: batch/v1
kind: Job
metadata:
name: easyai-local-migrate
namespace: easyai
spec:
backoffLimit: 0
template:
spec:
restartPolicy: Never
containers:
- name: migrate
image: $api_image
command: ["/bin/sh", "-ec", "cd /app && exec /app/easyai-ai-gateway-migrate"]
envFrom:
- secretRef:
name: easyai-ai-gateway-runtime
EOF
kubectl --context "$context" -n "$namespace" wait \
--for=condition=complete job/easyai-local-migrate --timeout=10m >/dev/null
}
mark_local_database() {
# shellcheck source=/dev/null
source "$state_root/local.env"
local primary
primary=$(kubectl --context "$context" -n "$namespace" get cluster easyai-postgres \
-o 'jsonpath={.status.currentPrimary}')
kubectl --context "$context" -n "$namespace" exec "$primary" -c postgres -- \
psql -X -v ON_ERROR_STOP=1 -U postgres -d easyai_ai_gateway \
-v cluster_id="$LOCAL_CLUSTER_ID" -c \
"INSERT INTO system_settings(setting_key,value)
VALUES('acceptance_local_cluster_id',jsonb_build_object('clusterId', :'cluster_id'))
ON CONFLICT(setting_key) DO UPDATE SET value=excluded.value,updated_at=now();" \
>/dev/null
}
import_snapshot() {
local api_image=$1 snapshot=$2
# shellcheck source=/dev/null
source "$state_root/local.env"
kubectl --context "$context" -n "$namespace" create configmap easyai-acceptance-snapshot \
--from-file=snapshot.json="$snapshot" --dry-run=client -o yaml |
kubectl --context "$context" apply -f - >/dev/null
kubectl --context "$context" -n "$namespace" delete job easyai-local-snapshot-import \
--ignore-not-found --wait=true >/dev/null
cat <<EOF | kubectl --context "$context" apply -f - >/dev/null
apiVersion: batch/v1
kind: Job
metadata:
name: easyai-local-snapshot-import
namespace: easyai
spec:
backoffLimit: 0
template:
spec:
restartPolicy: Never
containers:
- name: import
image: $api_image
command:
- /app/easyai-ai-gateway-acceptance-snapshot
- import
- --input
- /snapshot/snapshot.json
- --local-cluster-id
- $LOCAL_CLUSTER_ID
envFrom:
- secretRef:
name: easyai-ai-gateway-runtime
volumeMounts:
- name: snapshot
mountPath: /snapshot
readOnly: true
volumes:
- name: snapshot
configMap:
name: easyai-acceptance-snapshot
EOF
kubectl --context "$context" -n "$namespace" wait \
--for=condition=complete job/easyai-local-snapshot-import --timeout=5m >/dev/null
}
render_and_apply_application() {
local api_image=$1 web_image=$2 rendered=$state_root/application.rendered.yaml
sed \
-e "s|image: easyai-api|image: $api_image|g" \
-e "s|image: easyai-web|image: $web_image|g" \
-e 's/nodePort: 31088/nodePort: 32088/g' \
-e 's/nodePort: 31089/nodePort: 32089/g' \
"$repository_root/deploy/kubernetes/production/application.yaml" >"$rendered"
chmod 0600 "$rendered"
kubectl --context "$context" -n "$namespace" apply \
-f "$repository_root/deploy/kubernetes/production/service-account-rbac.yaml" >/dev/null
kubectl --context "$context" -n "$namespace" apply -f "$manifest_root/local-config.yaml" >/dev/null
kubectl --context "$context" -n "$namespace" apply -f "$rendered" >/dev/null
for workload in easyai-api-ningbo easyai-worker-ningbo; do
kubectl --context "$context" -n "$namespace" set env deployment/"$workload" \
--from=secret/easyai-database-ningbo >/dev/null
done
for workload in easyai-api-hongkong easyai-worker-hongkong; do
kubectl --context "$context" -n "$namespace" set env deployment/"$workload" \
--from=secret/easyai-database-hongkong >/dev/null
done
for workload in easyai-api-ningbo easyai-api-hongkong easyai-worker-ningbo easyai-worker-hongkong; do
kubectl --context "$context" -n "$namespace" patch deployment "$workload" --type=json \
-p='[{"op":"replace","path":"/spec/template/spec/volumes/0","value":{"name":"application-data","hostPath":{"path":"/var/lib/easyai-acceptance/media","type":"DirectoryOrCreate"}}}]' \
>/dev/null
done
kubectl --context "$context" -n "$namespace" scale \
deployment/easyai-web-ningbo deployment/easyai-web-hongkong --replicas=0 >/dev/null
}
bootstrap_runtime() {
local api_image=$1 api_digest=$2 snapshot=$3
# shellcheck source=/dev/null
source "$state_root/local.env"
local snapshot_hash release_sha runtime_file
snapshot_hash=$(jq -r '.snapshotSha256' "$snapshot")
release_sha=$(git -C "$repository_root" rev-parse HEAD)
runtime_file="$state_root/runtime-$snapshot_hash.json"
kubectl --context "$context" -n "$namespace" delete pod easyai-local-bootstrap \
--ignore-not-found --wait=true >/dev/null
cat <<EOF | kubectl --context "$context" apply -f - >/dev/null
apiVersion: v1
kind: Pod
metadata:
name: easyai-local-bootstrap
namespace: easyai
labels:
easyai.io/environment: local-acceptance
spec:
restartPolicy: Never
containers:
- name: bootstrap
image: $api_image
command:
- /app/easyai-ai-gateway-acceptance-bootstrap
- --local-cluster-id
- $LOCAL_CLUSTER_ID
- --release-sha
- $release_sha
- --api-image-digest
- $api_digest
- --worker-image-digest
- $api_digest
- --emulator-base-url
- http://easyai-acceptance-upstream-proxy.easyai.svc.cluster.local:8090
- --callback-url
- http://easyai-acceptance-callback-collector.easyai.svc.cluster.local:8091/callbacks
- --identity-shards
- "32"
- --output
- /tmp/runtime.json
envFrom:
- secretRef:
name: easyai-ai-gateway-runtime
EOF
kubectl --context "$context" -n "$namespace" wait \
--for=jsonpath='{.status.phase}'=Succeeded pod/easyai-local-bootstrap --timeout=5m >/dev/null
umask 077
kubectl --context "$context" -n "$namespace" cp \
easyai-local-bootstrap:/tmp/runtime.json "$runtime_file" >/dev/null
chmod 0600 "$runtime_file"
kubectl --context "$context" -n "$namespace" delete pod easyai-local-bootstrap \
--wait=true >/dev/null
printf '%s\n' "$runtime_file" >"$state_root/current-runtime-path"
chmod 0600 "$state_root/current-runtime-path"
}
up_cluster() {
[[ ${1:-} == --snapshot && $# -eq 2 ]] || {
usage >&2
exit 64
}
local snapshot=$2
[[ -f $snapshot && ! -L $snapshot ]] || {
echo 'snapshot must be a regular non-symlink file' >&2
exit 1
}
[[ -z $(git -C "$repository_root" status --short) ]] || {
echo 'local acceptance requires a clean source tree so report and image source are reproducible' >&2
exit 1
}
preflight
ensure_private_material
generate_tls
(
cd "$repository_root/apps/api"
go run ./cmd/acceptance-snapshot validate --input "$snapshot"
)
cp "$snapshot" "$state_root/snapshot.json"
chmod 0600 "$state_root/snapshot.json"
local k3d config native_arch release_sha api_image web_image netem_image api_digest
k3d=$(k3d_binary)
if "$k3d" cluster list --no-headers | awk '{print $1}' | grep -Fxq "$cluster_name"; then
echo "cluster $cluster_name already exists; use status or explicit down --confirm" >&2
exit 1
fi
native_arch=$(docker info --format '{{.Architecture}}')
case $native_arch in
aarch64|arm64) native_arch=arm64 ;;
x86_64) native_arch=amd64 ;;
esac
release_sha=$(git -C "$repository_root" rev-parse HEAD)
api_image="easyai-acceptance-api:$release_sha"
web_image="easyai-acceptance-web:$release_sha"
netem_image="easyai-acceptance-netem:$release_sha"
docker buildx build --load --platform "linux/$native_arch" --target api \
-t "$api_image" "$repository_root"
docker buildx build --load --platform "linux/$native_arch" --target web \
-t "$web_image" "$repository_root"
docker buildx build --load --platform "linux/$native_arch" --target acceptance-netem \
-t "$netem_image" "$repository_root"
api_digest=$(docker image inspect "$api_image" --format '{{.Id}}')
[[ $api_digest =~ ^sha256:[0-9a-f]{64}$ ]]
config="$state_root/k3d.rendered.yaml"
render_k3d_config "$config"
EASYAI_ACCEPTANCE_MEDIA_DIR="$media_root" "$k3d" cluster create --config "$config"
docker update --cpus 4 --memory 8g "${cluster_name}-server-0" >/dev/null
docker update --cpus 4 --memory 8g "${cluster_name}-server-1" >/dev/null
docker update --cpus 2 --memory 4g "${cluster_name}-server-2" >/dev/null
"$k3d" image import -c "$cluster_name" "$api_image" "$web_image" "$netem_image"
local cnpg_manifest="$state_root/cnpg.yaml"
curl -fsSL "$CNPG_MANIFEST_URL" -o "$cnpg_manifest"
[[ $(shasum -a 256 "$cnpg_manifest" | awk '{print $1}') == "$CNPG_MANIFEST_SHA256" ]] || {
echo 'CNPG manifest checksum mismatch' >&2
exit 1
}
chmod 0600 "$cnpg_manifest"
kubectl --context "$context" apply --server-side --force-conflicts -f "$cnpg_manifest" >/dev/null
kubectl --context "$context" -n cnpg-system rollout status \
deployment/cnpg-controller-manager --timeout=10m
apply_runtime_secrets
kubectl --context "$context" -n "$namespace" apply -f "$manifest_root/database.yaml" >/dev/null
kubectl --context "$context" -n "$namespace" wait --for=condition=Ready \
cluster/easyai-postgres --timeout=15m >/dev/null
run_migrations "$api_image"
mark_local_database
import_snapshot "$api_image" "$snapshot"
sed "s|image: easyai-api|image: $api_image|g; s|image: easyai-acceptance-netem|image: $netem_image|g" \
"$manifest_root/support-services.yaml" |
kubectl --context "$context" -n "$namespace" apply -f - >/dev/null
wait_rollout deployment/easyai-acceptance-emulator
wait_rollout deployment/easyai-acceptance-callback-collector
wait_rollout deployment/easyai-acceptance-upstream-proxy
wait_rollout deployment/easyai-postgres-proxy-ningbo
wait_rollout deployment/easyai-postgres-proxy-hongkong
render_and_apply_application "$api_image" "$web_image"
sed "s|image: easyai-web|image: $web_image|g" "$manifest_root/tls-edges.yaml" |
kubectl --context "$context" -n "$namespace" apply -f - >/dev/null
wait_rollout deployment/easyai-api-ningbo
wait_rollout deployment/easyai-api-hongkong
wait_rollout deployment/easyai-worker-ningbo
wait_rollout deployment/easyai-worker-hongkong
wait_rollout deployment/easyai-capacity-controller
wait_rollout deployment/easyai-acceptance-edge-ningbo
wait_rollout deployment/easyai-acceptance-edge-hongkong
bootstrap_runtime "$api_image" "$api_digest" "$snapshot"
"$script_dir/network-fault.sh" baseline
echo "local_acceptance_cluster=PASS context=$context gateways=https://127.0.0.1:18443,https://127.0.0.1:19443 ca_file=$state_root/ca.crt"
}
status_cluster() {
local k3d
k3d=$(k3d_binary) || {
echo 'k3d is unavailable' >&2
exit 1
}
"$k3d" cluster list --no-headers | awk '{print $1}' | grep -Fxq "$cluster_name" || {
echo "cluster $cluster_name does not exist" >&2
exit 1
}
kubectl --context "$context" get nodes -o wide
kubectl --context "$context" -n "$namespace" get cluster,pod,deployment
echo "local_acceptance_status=PASS context=$context"
}
down_cluster() {
[[ ${1:-} == --confirm && $# -eq 1 ]] || {
echo 'down requires the literal --confirm flag' >&2
exit 64
}
local k3d
k3d=$(k3d_binary)
"$k3d" cluster delete "$cluster_name"
echo "local_acceptance_down=PASS preserved_private_root=$private_root"
}
case ${1:-} in
install-tools)
[[ $# -eq 1 ]] || { usage >&2; exit 64; }
install_tools
;;
preflight)
[[ $# -eq 1 ]] || { usage >&2; exit 64; }
preflight
;;
up)
shift
up_cluster "$@"
;;
status)
[[ $# -eq 1 ]] || { usage >&2; exit 64; }
status_cluster
;;
down)
shift
down_cluster "$@"
;;
*)
usage >&2
exit 64
;;
esac
+17
View File
@@ -0,0 +1,17 @@
#!/bin/sh
set -eu
listen_port=${NETEM_LISTEN_PORT:-15432}
target_host=${NETEM_TARGET_HOST:-}
target_port=${NETEM_TARGET_PORT:-}
case "$listen_port:$target_port" in
*[!0-9:]*|:*|*:) echo "NETEM_LISTEN_PORT and NETEM_TARGET_PORT must be numeric" >&2; exit 64 ;;
esac
case "$target_host" in
""|*[!A-Za-z0-9.-]*) echo "NETEM_TARGET_HOST must be a DNS name or address" >&2; exit 64 ;;
esac
exec socat \
"TCP-LISTEN:${listen_port},fork,reuseaddr,keepalive" \
"TCP:${target_host}:${target_port},connect-timeout=10"
+96
View File
@@ -0,0 +1,96 @@
#!/usr/bin/env bash
set -euo pipefail
usage() {
cat <<'EOF'
Usage:
scripts/acceptance/network-fault.sh baseline
scripts/acceptance/network-fault.sh weak-link
scripts/acceptance/network-fault.sh upstream-outage
scripts/acceptance/network-fault.sh database-outage <ningbo|hongkong>
scripts/acceptance/network-fault.sh reset
Only operates on Pods carrying easyai.io/environment=local-acceptance in the
easyai-acceptance-local k3d context.
EOF
}
profile=${1:-}
site=${2:-}
namespace=${AI_GATEWAY_LOCAL_ACCEPTANCE_NAMESPACE:-easyai}
context=${AI_GATEWAY_LOCAL_ACCEPTANCE_CONTEXT:-k3d-easyai-acceptance-local}
command -v kubectl >/dev/null 2>&1 || {
echo 'kubectl is required' >&2
exit 1
}
[[ $(kubectl config get-contexts "$context" -o name) == "$context" ]] || {
echo "refusing netem because context $context does not exist" >&2
exit 1
}
pod_for() {
local name=$1
kubectl --context "$context" -n "$namespace" get pod \
-l "app.kubernetes.io/name=$name,easyai.io/environment=local-acceptance" \
-o 'jsonpath={.items[0].metadata.name}'
}
apply_qdisc() {
local pod=$1
shift
[[ -n $pod ]] || {
echo 'local acceptance proxy Pod was not found' >&2
exit 1
}
kubectl --context "$context" -n "$namespace" exec "$pod" -- \
tc qdisc replace dev eth0 root netem "$@"
}
reset_qdisc() {
local pod=$1
[[ -z $pod ]] || kubectl --context "$context" -n "$namespace" exec "$pod" -- \
tc qdisc del dev eth0 root >/dev/null 2>&1 || true
}
upstream_pod=$(pod_for easyai-acceptance-upstream-proxy)
ningbo_db_pod=$(kubectl --context "$context" -n "$namespace" get pod \
-l 'app.kubernetes.io/name=easyai-postgres-proxy,easyai.io/site=ningbo,easyai.io/environment=local-acceptance' \
-o 'jsonpath={.items[0].metadata.name}')
hongkong_db_pod=$(kubectl --context "$context" -n "$namespace" get pod \
-l 'app.kubernetes.io/name=easyai-postgres-proxy,easyai.io/site=hongkong,easyai.io/environment=local-acceptance' \
-o 'jsonpath={.items[0].metadata.name}')
case $profile in
baseline)
apply_qdisc "$upstream_pod" delay 20ms
apply_qdisc "$ningbo_db_pod" delay 20ms
apply_qdisc "$hongkong_db_pod" delay 20ms
;;
weak-link)
apply_qdisc "$upstream_pod" delay 20ms 10ms distribution normal loss 0.5%
apply_qdisc "$ningbo_db_pod" delay 20ms 10ms distribution normal loss 0.5%
apply_qdisc "$hongkong_db_pod" delay 20ms 10ms distribution normal loss 0.5%
;;
upstream-outage)
apply_qdisc "$upstream_pod" loss 100%
;;
database-outage)
case $site in
ningbo) apply_qdisc "$ningbo_db_pod" loss 100% ;;
hongkong) apply_qdisc "$hongkong_db_pod" loss 100% ;;
*) usage >&2; exit 64 ;;
esac
;;
reset)
reset_qdisc "$upstream_pod"
reset_qdisc "$ningbo_db_pod"
reset_qdisc "$hongkong_db_pod"
;;
*)
usage >&2
exit 64
;;
esac
echo "local_acceptance_netem=PASS profile=$profile${site:+ site=$site}"
+316
View File
@@ -0,0 +1,316 @@
#!/usr/bin/env node
import { createHash } from 'node:crypto';
import { lstat, mkdir, readFile, readdir, writeFile } from 'node:fs/promises';
import { dirname, resolve } from 'node:path';
const shaPattern = /^[0-9a-f]{40}$/;
const digestPattern = /^sha256:[0-9a-f]{64}$/;
const hashPattern = /^[0-9a-f]{64}$/;
const unsafeKeyPattern = /(password|secret|token|credential|authorization|api.?key|access.?key|private.?key|connection.?string|database.?url|proxy)/i;
function parseArgs(argv) {
const command = argv.shift();
const values = {};
while (argv.length > 0) {
const key = argv.shift();
if (!key?.startsWith('--') || argv.length === 0) throw new Error('invalid arguments');
values[key.slice(2)] = argv.shift();
}
return { command, values };
}
async function regularJSON(path) {
const absolute = resolve(path);
const info = await lstat(absolute);
if (!info.isFile() || info.isSymbolicLink()) throw new Error(`${path} must be a regular non-symlink file`);
return JSON.parse(await readFile(absolute, 'utf8'));
}
function assertSecretSafe(value, path = '$') {
if (Array.isArray(value)) {
value.forEach((item, index) => assertSecretSafe(item, `${path}[${index}]`));
return;
}
if (value && typeof value === 'object') {
for (const [key, nested] of Object.entries(value)) {
if (key !== 'secretSafe' && unsafeKeyPattern.test(key)) {
throw new Error(`unsafe report key at ${path}.${key}`);
}
assertSecretSafe(nested, `${path}.${key}`);
}
return;
}
if (typeof value === 'string' && /(?:postgres(?:ql)?|https?):\/\/[^\s]+/i.test(value)) {
throw new Error(`URL-like value is forbidden at ${path}`);
}
}
function validate(report) {
if (report.schemaVersion !== 'acceptance-report/v1' || report.secretSafe !== true) {
throw new Error('unsupported or non-secret-safe acceptance report');
}
if (!report.runId || !shaPattern.test(report.release?.sourceSha ?? '') ||
!digestPattern.test(report.release?.apiImageDigest ?? '') ||
!digestPattern.test(report.release?.workerImageDigest ?? '') ||
!shaPattern.test(report.snapshot?.sourceReleaseSha ?? '') ||
!hashPattern.test(report.snapshot?.configHash ?? '') ||
!hashPattern.test(report.snapshot?.sha256 ?? '')) {
throw new Error('acceptance report identity fields are invalid');
}
if (!report.stages || !Array.isArray(report.gates) ||
report.promotion?.requiresManualConfirmation !== true ||
!['validation', 'live', 'not-applicable'].includes(report.promotion?.trafficMode)) {
throw new Error('acceptance report stages, gates, or promotion are invalid');
}
for (const [name, stage] of Object.entries(report.stages)) {
if (!['passed', 'failed', 'not-run'].includes(stage?.status)) {
throw new Error(`invalid stage status for ${name}`);
}
}
assertSecretSafe(report);
}
async function loadReports(directory) {
const entries = (await readdir(resolve(directory), { withFileTypes: true }))
.filter((entry) => entry.isFile() && entry.name.endsWith('.json'))
.map((entry) => entry.name)
.sort();
const reports = [];
for (const name of entries) {
const report = await regularJSON(resolve(directory, name));
if (report.schemaVersion === 'acceptance-load-report/v1') {
if (report.secretSafe !== true) throw new Error(`load report ${name} is not secret-safe`);
reports.push({
file: name,
profile: report.profile,
passed: report.passed,
phases: report.phases,
startedAt: report.startedAt,
finishedAt: report.finishedAt
});
}
}
return reports;
}
async function buildLocal(values) {
for (const key of ['runtime', 'snapshot', 'reports', 'artifact', 'output']) {
if (!values[key]) throw new Error(`--${key} is required`);
}
const runtime = await regularJSON(values.runtime);
const snapshot = await regularJSON(values.snapshot);
const artifact = await regularJSON(values.artifact);
const loads = await loadReports(values.reports);
if (runtime.schemaVersion !== 'acceptance-runtime/v1' ||
snapshot.schemaVersion !== 'acceptance-snapshot/v1' ||
artifact.schemaVersion !== 'acceptance-artifact-smoke/v1') {
throw new Error('local report inputs have incompatible schema versions');
}
if (runtime.releaseSha !== artifact.releaseSha ||
artifact.apiImageDigest !== values['api-digest'] ||
runtime.snapshotConfigHash !== snapshot.source.configHash ||
runtime.snapshotSha256 !== snapshot.snapshotSha256) {
throw new Error('local report CAS identity mismatch');
}
const localPassed = loads.length > 0 && loads.every((item) => item.passed === true);
const gates = [
{ id: 'local_load_reports', passed: localPassed, detail: `${loads.length} load reports` },
{ id: 'amd64_artifact_smoke', passed: artifact.passed === true }
];
const report = {
schemaVersion: 'acceptance-report/v1',
runId: runtime.runId,
release: {
sourceSha: artifact.releaseSha,
apiImageDigest: artifact.apiImageDigest,
workerImageDigest: artifact.apiImageDigest
},
snapshot: {
sourceReleaseSha: snapshot.source.releaseSha,
configHash: snapshot.source.configHash,
sha256: snapshot.snapshotSha256
},
stages: {
localNative: { status: localPassed ? 'passed' : 'failed', loadReports: loads },
amd64Artifact: { status: artifact.passed ? 'passed' : 'failed', ...artifact },
onlineSimulation: { status: 'not-run' },
realCanary: { status: 'not-run' }
},
certifiedProfile: values['certified-profile'] ? {
name: values['certified-profile'],
source: 'local-candidate-only'
} : null,
gates,
promotion: {
ready: false,
requiresManualConfirmation: true,
trafficMode: 'not-applicable'
},
createdAt: new Date().toISOString(),
secretSafe: true
};
validate(report);
await mkdir(dirname(resolve(values.output)), { recursive: true, mode: 0o700 });
await writeFile(resolve(values.output), `${JSON.stringify(report, null, 2)}\n`, { mode: 0o600 });
process.stdout.write(`acceptance_report_local=PASS sha256=${createHash('sha256').update(JSON.stringify(report)).digest('hex')}\n`);
}
async function buildLocalPartial(values) {
for (const key of ['runtime', 'snapshot', 'reports', 'output', 'failure-phase']) {
if (!values[key]) throw new Error(`--${key} is required`);
}
const runtime = await regularJSON(values.runtime);
const snapshot = await regularJSON(values.snapshot);
const loads = await loadReports(values.reports);
if (runtime.schemaVersion !== 'acceptance-runtime/v1' ||
snapshot.schemaVersion !== 'acceptance-snapshot/v1' ||
!shaPattern.test(runtime.releaseSha ?? '') ||
!digestPattern.test(runtime.apiImageDigest ?? '') ||
!digestPattern.test(runtime.workerImageDigest ?? '') ||
runtime.snapshotConfigHash !== snapshot.source.configHash ||
runtime.snapshotSha256 !== snapshot.snapshotSha256) {
throw new Error('partial local report inputs have incompatible CAS identity');
}
const completedLoadsPassed = loads.every((item) => item.passed === true);
const report = {
schemaVersion: 'acceptance-report/v1',
runId: runtime.runId,
release: {
sourceSha: runtime.releaseSha,
apiImageDigest: runtime.apiImageDigest,
workerImageDigest: runtime.workerImageDigest
},
snapshot: {
sourceReleaseSha: snapshot.source.releaseSha,
configHash: snapshot.source.configHash,
sha256: snapshot.snapshotSha256
},
stages: {
localNative: {
status: 'failed',
failurePhase: values['failure-phase'],
completedLoadReportsPassed: completedLoadsPassed,
loadReports: loads
},
amd64Artifact: { status: 'not-run' },
onlineSimulation: { status: 'not-run' },
realCanary: { status: 'not-run' }
},
certifiedProfile: values['certified-profile'] ? {
name: values['certified-profile'],
source: 'last-local-stable-candidate'
} : null,
gates: [
{
id: 'local_execution_complete',
passed: false,
detail: `stopped during ${values['failure-phase']}`
},
{
id: 'completed_load_reports',
passed: completedLoadsPassed,
detail: `${loads.length} load reports`
}
],
promotion: {
ready: false,
requiresManualConfirmation: true,
trafficMode: 'not-applicable'
},
createdAt: new Date().toISOString(),
secretSafe: true
};
validate(report);
await mkdir(dirname(resolve(values.output)), { recursive: true, mode: 0o700 });
await writeFile(resolve(values.output), `${JSON.stringify(report, null, 2)}\n`, { mode: 0o600 });
process.stdout.write('acceptance_report_local_partial=PASS\n');
}
async function mergeProduction(values) {
for (const key of ['local-report', 'production-summary', 'output']) {
if (!values[key]) throw new Error(`--${key} is required`);
}
const local = await regularJSON(values['local-report']);
const production = await regularJSON(values['production-summary']);
validate(local);
if (local.release.sourceSha !== production.releaseSha ||
local.release.apiImageDigest !== production.apiImageDigest ||
local.snapshot.configHash !== production.snapshotConfigHash ||
local.snapshot.sha256 !== production.snapshotSha256) {
throw new Error('production summary does not match local acceptance CAS fields');
}
local.stages.onlineSimulation = {
status: production.passed === true ? 'passed' : 'failed',
stableCapacityProfile: production.stableCapacityProfile,
tasks: production.tasks,
geminiTasks: production.geminiTasks,
videoTasks: production.videoTasks
};
local.stages.realCanary = {
status: production.realCanaryPassed === true ? 'passed' : 'failed'
};
local.certifiedProfile = production.certifiedProfile ?? local.certifiedProfile;
local.gates.push(...(production.gates ?? []));
local.promotion = {
ready: production.passed === true && production.realCanaryPassed === true &&
local.gates.every((gate) => gate.passed === true),
requiresManualConfirmation: true,
trafficMode: 'validation'
};
local.createdAt = new Date().toISOString();
validate(local);
await mkdir(dirname(resolve(values.output)), { recursive: true, mode: 0o700 });
await writeFile(resolve(values.output), `${JSON.stringify(local, null, 2)}\n`, { mode: 0o600 });
process.stdout.write(`acceptance_report_production=PASS promotion_ready=${local.promotion.ready}\n`);
}
const { command, values } = parseArgs(process.argv.slice(2));
if (command === 'validate') {
const report = await regularJSON(values.input ?? '');
validate(report);
process.stdout.write('acceptance_report_validate=PASS\n');
} else if (command === 'build-local') {
await buildLocal(values);
} else if (command === 'build-local-partial') {
await buildLocalPartial(values);
} else if (command === 'merge-production') {
await mergeProduction(values);
} else if (command === 'mark-promoted') {
const input = values.input ?? '';
const output = values.output ?? input;
const report = await regularJSON(input);
validate(report);
if (report.promotion.ready !== true || report.promotion.trafficMode !== 'validation') {
throw new Error('only a ready validation report can be marked promoted');
}
report.promotion.trafficMode = 'live';
report.promotion.promotedAt = new Date().toISOString();
await writeFile(resolve(output), `${JSON.stringify(report, null, 2)}\n`, { mode: 0o600 });
process.stdout.write('acceptance_report_promoted=PASS\n');
} else if (command === 'attach-monitor') {
const input = values.input ?? '';
const output = values.output ?? input;
const report = await regularJSON(input);
const monitor = await regularJSON(values.monitor ?? '');
validate(report);
if (monitor.schemaVersion !== 'acceptance-monitor-report/v1' ||
monitor.runId !== report.runId || monitor.releaseSha !== report.release.sourceSha ||
monitor.secretSafe !== true) {
throw new Error('monitor report does not match overall acceptance report');
}
report.stages.postPromotionMonitor = {
status: monitor.passed === true ? 'passed' : 'failed',
startedAt: monitor.startedAt,
finishedAt: monitor.finishedAt,
samples: monitor.samples,
failureGateId: monitor.failureGateId
};
if (monitor.passed !== true) {
report.promotion.ready = false;
report.promotion.trafficMode = 'validation';
}
await writeFile(resolve(output), `${JSON.stringify(report, null, 2)}\n`, { mode: 0o600 });
process.stdout.write(`acceptance_report_monitor=PASS monitor_passed=${monitor.passed === true}\n`);
} else {
throw new Error('usage: report.mjs {validate|build-local|build-local-partial|merge-production|mark-promoted|attach-monitor} [options]');
}
+536
View File
@@ -0,0 +1,536 @@
#!/usr/bin/env bash
set -euo pipefail
script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
repository_root=$(cd "$script_dir/../.." && pwd)
private_root="$repository_root/.local-secrets/acceptance"
state_root="$private_root/state"
context=k3d-easyai-acceptance-local
namespace=easyai
runtime_path_file="$state_root/current-runtime-path"
snapshot="$state_root/snapshot.json"
load_binary="$state_root/easyai-ai-gateway-acceptance-load"
active_load_pid=
netem_active=false
current_phase=initialization
stable_profile=
usage() {
cat <<'EOF'
Usage:
scripts/acceptance/run-local-acceptance.sh quick
scripts/acceptance/run-local-acceptance.sh full --release-manifest dist/releases/<SHA>.json
scripts/acceptance/run-local-acceptance.sh artifact-smoke --release-manifest dist/releases/<SHA>.json
`full` executes P24/P28/P32 three times, the fault matrix, autoscaling/drain,
80% soak, 120% overload, and exact linux/amd64 artifact smoke. The load process
runs outside K3s and splits requests 50/50 across both TLS entrances.
EOF
}
private_file() {
local path=$1 mode
[[ -f $path && ! -L $path ]] || return 1
if [[ $(uname -s) == Darwin ]]; then
mode=$(stat -f '%Lp' "$path")
else
mode=$(stat -c '%a' "$path")
fi
[[ $mode == 600 ]]
}
cleanup() {
local status=$?
trap - EXIT
if [[ -n $active_load_pid ]]; then
kill "$active_load_pid" >/dev/null 2>&1 || true
wait "$active_load_pid" >/dev/null 2>&1 || true
fi
if [[ $netem_active == true ]]; then
"$script_dir/network-fault.sh" reset >/dev/null 2>&1 || true
fi
if [[ $status -ne 0 && -n ${runtime:-} && -n ${report_root:-} && -f ${runtime:-} && -f $snapshot ]]; then
restore_profile_best_effort "${stable_profile:-P24}"
local partial_args=(
"$script_dir/report.mjs" build-local-partial
--runtime "$runtime" \
--snapshot "$snapshot" \
--reports "$report_root" \
--failure-phase "$current_phase" \
--output "$report_root/acceptance-report.partial.json"
)
if [[ -n $stable_profile ]]; then
partial_args+=(--certified-profile "$stable_profile")
fi
if node "${partial_args[@]}" >/dev/null 2>&1; then
echo "local_acceptance=FAILED phase=$current_phase partial_report=$report_root/acceptance-report.partial.json" >&2
else
echo "local_acceptance=FAILED phase=$current_phase partial_report=unavailable" >&2
fi
fi
exit "$status"
}
trap cleanup EXIT
trap 'current_phase=signal_interrupted; exit 130' HUP INT TERM
require_local_cluster() {
command -v kubectl >/dev/null 2>&1
command -v jq >/dev/null 2>&1
command -v node >/dev/null 2>&1
[[ $(kubectl config get-contexts "$context" -o name) == "$context" ]]
private_file "$runtime_path_file" && private_file "$snapshot"
runtime=$(<"$runtime_path_file")
private_file "$runtime" || {
echo 'local acceptance runtime file is missing or not mode 0600' >&2
exit 1
}
jq -e '.schemaVersion == "acceptance-runtime/v1"' "$runtime" >/dev/null
run_id=$(jq -r '.runId' "$runtime")
release_sha=$(jq -r '.releaseSha' "$runtime")
report_root="$repository_root/dist/acceptance/local/$run_id"
install -d -m 0700 "$report_root"
(
cd "$repository_root/apps/api"
env -u AI_GATEWAY_TEST_DATABASE_URL go build -trimpath \
-o "$load_binary" ./cmd/acceptance-load
)
}
database_query() {
local sql=$1 primary
primary=$(kubectl --context "$context" -n "$namespace" get cluster easyai-postgres \
-o 'jsonpath={.status.currentPrimary}')
kubectl --context "$context" -n "$namespace" exec "$primary" -c postgres -- \
psql -X -v ON_ERROR_STOP=1 -U postgres -d easyai_ai_gateway -At -c "$sql"
}
load_environment() {
acceptance_api_keys=$(jq -r '.apiKeys | join(",")' "$runtime")
acceptance_run_token=$(jq -r '.runToken' "$runtime")
acceptance_gemini_model=$(jq -r '.geminiModel' "$runtime")
acceptance_video_model=$(jq -r '.videoModel' "$runtime")
acceptance_emulator_url=$(jq -r '.emulatorBaseUrl' "$runtime")
}
run_load() {
local profile=$1 report_path=$2
shift 2
AI_GATEWAY_ACCEPTANCE_GATEWAYS='https://127.0.0.1:18443,https://127.0.0.1:19443' \
AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME=gateway.easyai.local \
AI_GATEWAY_ACCEPTANCE_GATEWAY_CA_FILE="$state_root/ca.crt" \
AI_GATEWAY_ACCEPTANCE_EMULATOR_URL="$acceptance_emulator_url" \
AI_GATEWAY_ACCEPTANCE_API_KEYS="$acceptance_api_keys" \
AI_GATEWAY_ACCEPTANCE_RUN_ID="$run_id" \
AI_GATEWAY_ACCEPTANCE_RUN_TOKEN="$acceptance_run_token" \
AI_GATEWAY_ACCEPTANCE_GEMINI_MODEL="$acceptance_gemini_model" \
AI_GATEWAY_ACCEPTANCE_VIDEO_MODEL="$acceptance_video_model" \
"$load_binary" -profile "$profile" -report "$report_path" "$@" \
>"$report_path.stdout"
chmod 0600 "$report_path" "$report_path.stdout"
jq -e '.schemaVersion == "acceptance-load-report/v1" and .secretSafe == true and .passed == true' \
"$report_path" >/dev/null
}
sample_resources() {
local output=$1 stop_file=$2
printf 'timestamp,scope,name,cpu,memory\n' >"$output"
while [[ ! -e $stop_file ]]; do
kubectl --context "$context" top nodes --no-headers 2>/dev/null |
awk -v timestamp="$(date -u '+%Y-%m-%dT%H:%M:%SZ')" \
'{print timestamp ",node," $1 "," $3 "," $5}' >>"$output" || true
kubectl --context "$context" -n "$namespace" top pods --no-headers 2>/dev/null |
awk -v timestamp="$(date -u '+%Y-%m-%dT%H:%M:%SZ')" \
'{print timestamp ",pod," $1 "," $2 "," $3}' >>"$output" || true
sleep 5
done
}
verify_hard_gates() {
local unhealthy ready sync_state connections max_connections queue duplicate_remote duplicate_billing
ready=$(kubectl --context "$context" get nodes -o json |
jq '[.items[] | select(any(.status.conditions[]; .type=="Ready" and .status=="True"))] | length')
[[ $ready == 3 ]]
[[ $(kubectl --context "$context" get nodes -o json |
jq '[.items[] | select(any(.status.conditions[]; .type=="MemoryPressure" and .status=="True"))] | length') == 0 ]]
[[ $(kubectl --context "$context" -n "$namespace" get cluster easyai-postgres \
-o 'jsonpath={.status.readyInstances}') == 2 ]]
sync_state=$(database_query "SELECT COALESCE(string_agg(sync_state,','),'') FROM pg_stat_replication;")
[[ ",$sync_state," == *,sync,* || ",$sync_state," == *,quorum,* ]]
unhealthy=$(kubectl --context "$context" -n "$namespace" get pods -o json |
jq '[.items[]
| select(.metadata.labels["app.kubernetes.io/name"] == "easyai-api" or
.metadata.labels["app.kubernetes.io/name"] == "easyai-worker")
| .status.containerStatuses[]?
| select(.restartCount > 0 or .lastState.terminated.reason == "OOMKilled" or
.state.terminated.reason == "OOMKilled")] | length')
[[ $unhealthy == 0 ]]
while read -r _node _cpu _cpu_percent _memory memory_percent; do
memory_percent=${memory_percent%\%}
[[ $memory_percent =~ ^[0-9]+$ ]]
(( memory_percent < 80 ))
done < <(kubectl --context "$context" top nodes --no-headers)
while read -r _pod _cpu memory; do
case $memory in
*Gi) memory=$(awk -v value="${memory%Gi}" 'BEGIN {printf "%.0f", value*1024}') ;;
*Mi) memory=${memory%Mi} ;;
*Ki) memory=$(awk -v value="${memory%Ki}" 'BEGIN {printf "%.0f", value/1024}') ;;
*) return 1 ;;
esac
(( memory < 1536 ))
done < <(kubectl --context "$context" -n "$namespace" top pods \
-l 'app.kubernetes.io/part-of=easyai-ai-gateway' --no-headers)
connections=$(database_query "SELECT count(*) FROM pg_stat_activity WHERE backend_type='client backend';")
max_connections=$(database_query "SELECT setting::int FROM pg_settings WHERE name='max_connections';")
(( connections < 150 && connections * 4 < max_connections * 3 ))
queue=$(database_query "SELECT count(*) FILTER (WHERE status='queued')||':'||count(*) FILTER (WHERE status='running') FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid;")
[[ $queue == 0:0 ]]
duplicate_remote=$(database_query "SELECT count(*) FROM (SELECT remote_task_id FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid AND remote_task_id IS NOT NULL GROUP BY remote_task_id HAVING count(*)>1) d;")
duplicate_billing=$(database_query "SELECT count(*) FROM (SELECT reference_id,transaction_type FROM gateway_wallet_transactions WHERE reference_type='gateway_task' AND reference_id IN (SELECT id::text FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid) GROUP BY reference_id,transaction_type HAVING count(*)>1) d;")
[[ $duplicate_remote == 0 && $duplicate_billing == 0 ]]
kubectl --context "$context" -n "$namespace" exec deployment/easyai-acceptance-callback-collector -- \
wget -qO- http://127.0.0.1:8091/report |
jq -e '.duplicates == 0 and .invalid == 0' >/dev/null
}
apply_profile() {
local profile=$1 slots pool
case $profile in
P24) slots=24; pool=32 ;;
P28) slots=28; pool=36 ;;
P32) slots=32; pool=40 ;;
*) return 64 ;;
esac
local global=$((slots * 2))
kubectl --context "$context" -n "$namespace" patch configmap easyai-ai-gateway-config \
--type=merge -p "$(jq -cn \
--arg slots "$slots" --arg global "$global" \
'{data:{
AI_GATEWAY_WORKER_AUTOSCALING_ENABLED:"false",
AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT:$slots,
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT:$slots,
AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT:$global,
AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA:$global
}}')" >/dev/null
kubectl --context "$context" -n "$namespace" scale \
deployment/easyai-worker-ningbo deployment/easyai-worker-hongkong --replicas=1 >/dev/null
for deployment in easyai-worker-ningbo easyai-worker-hongkong; do
kubectl --context "$context" -n "$namespace" set env deployment/"$deployment" \
AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT="$slots" \
AI_GATEWAY_DATABASE_MAX_CONNS="$pool" \
AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY="$slots" \
AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY="$slots" >/dev/null
kubectl --context "$context" -n "$namespace" rollout status deployment/"$deployment" --timeout=10m
done
for deployment in easyai-api-ningbo easyai-api-hongkong easyai-capacity-controller; do
kubectl --context "$context" -n "$namespace" rollout restart deployment/"$deployment" >/dev/null
kubectl --context "$context" -n "$namespace" rollout status deployment/"$deployment" --timeout=10m
done
}
restore_profile_best_effort() {
local profile=$1 slots pool global
case $profile in
P24) slots=24; pool=32 ;;
P28) slots=28; pool=36 ;;
P32) slots=32; pool=40 ;;
*) return ;;
esac
global=$((slots * 2))
kubectl --context "$context" -n "$namespace" patch configmap easyai-ai-gateway-config \
--type=merge -p "$(jq -cn \
--arg slots "$slots" --arg global "$global" \
'{data:{
AI_GATEWAY_WORKER_AUTOSCALING_ENABLED:"false",
AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT:$slots,
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT:$slots,
AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT:$global,
AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA:$global
}}')" >/dev/null 2>&1 || true
kubectl --context "$context" -n "$namespace" scale \
deployment/easyai-worker-ningbo deployment/easyai-worker-hongkong \
--replicas=1 >/dev/null 2>&1 || true
for deployment in easyai-worker-ningbo easyai-worker-hongkong; do
kubectl --context "$context" -n "$namespace" set env deployment/"$deployment" \
AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT="$slots" \
AI_GATEWAY_DATABASE_MAX_CONNS="$pool" \
AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY="$slots" \
AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY="$slots" >/dev/null 2>&1 || true
done
}
run_profile_round() {
local profile=$1 repetition=$2 prefix
prefix="$report_root/${profile,,}-$repetition"
local stop_file="$prefix.resources.stop"
sample_resources "$prefix.resources.csv" "$stop_file" &
local sampler=$!
run_load gemini-baseline "$prefix-gemini-baseline.json"
run_load gemini-large "$prefix-gemini-large.json"
run_load gemini-peak "$prefix-gemini-peak.json"
run_load video-throughput "$prefix-video-throughput.json"
run_recovery "$prefix-video-recovery.json"
touch "$stop_file"
wait "$sampler"
verify_hard_gates
}
wait_for_recovery_owner() {
local deadline=$((SECONDS + 120)) owner
while (( SECONDS < deadline )); do
owner=$(database_query "
SELECT task.id::text||','||worker.pod_name||','||worker.site
FROM gateway_tasks task
JOIN gateway_worker_instances worker ON worker.instance_id=task.locked_by
WHERE task.acceptance_run_id='$run_id'::uuid
AND task.status='running'
AND task.request::text LIKE '%acceptance-long-recovery%'
AND worker.status='active'
ORDER BY task.updated_at
LIMIT 1;")
[[ -z $owner ]] || { printf '%s\n' "$owner"; return; }
sleep 1
done
return 1
}
run_recovery() {
local report_path=$1 owner task_id pod site
run_load video-recovery "$report_path" &
active_load_pid=$!
owner=$(wait_for_recovery_owner)
IFS=',' read -r task_id pod site <<<"$owner"
[[ $task_id =~ ^[0-9a-f-]{36}$ && $pod == easyai-worker-"$site"-* ]]
kubectl --context "$context" -n "$namespace" delete pod "$pod" \
--force --grace-period=0 --wait=false >/dev/null
kubectl --context "$context" -n "$namespace" rollout status \
deployment/easyai-worker-"$site" --timeout=10m
wait "$active_load_pid"
active_load_pid=
}
run_fault_matrix() {
local duration=${AI_GATEWAY_LOCAL_WEAK_LINK_DURATION:-5m}
"$script_dir/network-fault.sh" weak-link
netem_active=true
run_load mixed-soak "$report_root/fault-weak-link.json" \
-duration "$duration" -image-rate 1 -video-rate 1
"$script_dir/network-fault.sh" reset
netem_active=false
run_load video-throughput "$report_root/fault-upstream-outage.json" &
active_load_pid=$!
sleep 5
"$script_dir/network-fault.sh" upstream-outage
netem_active=true
sleep 10
"$script_dir/network-fault.sh" reset
netem_active=false
wait "$active_load_pid"
active_load_pid=
run_load video-recovery "$report_root/fault-database-outage.json" &
active_load_pid=$!
sleep 10
"$script_dir/network-fault.sh" database-outage hongkong
netem_active=true
sleep 30
"$script_dir/network-fault.sh" reset
netem_active=false
wait "$active_load_pid"
active_load_pid=
local leader
leader=$(kubectl --context "$context" -n "$namespace" get pods \
-l app.kubernetes.io/name=easyai-capacity-controller -o name |
while read -r pod; do
status=$(kubectl --context "$context" -n "$namespace" exec "$pod" -- \
wget -qO- http://127.0.0.1:8088/status)
[[ $(jq -r '.leader' <<<"$status") == true ]] && { printf '%s\n' "${pod#pod/}"; break; }
done)
[[ -n $leader ]]
kubectl --context "$context" -n "$namespace" delete pod "$leader" --wait=false >/dev/null
kubectl --context "$context" -n "$namespace" rollout status \
deployment/easyai-capacity-controller --timeout=5m
verify_hard_gates
}
run_autoscaling() {
local profile=$1 slots
slots=${profile#P}
kubectl --context "$context" -n "$namespace" patch configmap easyai-ai-gateway-config \
--type=merge -p "$(jq -cn --arg global "$((slots * 4))" \
'{data:{
AI_GATEWAY_WORKER_AUTOSCALING_ENABLED:"true",
AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO:"1",
AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG:"1",
AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO:"2",
AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG:"2",
AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT:$global
}}')" >/dev/null
kubectl --context "$context" -n "$namespace" rollout restart \
deployment/easyai-capacity-controller >/dev/null
kubectl --context "$context" -n "$namespace" rollout status \
deployment/easyai-capacity-controller --timeout=5m
run_load video-throughput "$report_root/autoscaling-load.json" &
active_load_pid=$!
local deadline=$((SECONDS + 420)) total=2
while (( SECONDS < deadline )); do
total=$(kubectl --context "$context" -n "$namespace" get deployment \
easyai-worker-ningbo easyai-worker-hongkong \
-o json | jq '[.items[].spec.replicas] | add')
(( total >= 3 )) && break
sleep 5
done
(( total >= 3 ))
wait "$active_load_pid"
active_load_pid=
deadline=$((SECONDS + 780))
while (( SECONDS < deadline )); do
total=$(kubectl --context "$context" -n "$namespace" get deployment \
easyai-worker-ningbo easyai-worker-hongkong \
-o json | jq '[.items[].spec.replicas] | add')
(( total == 2 )) && break
sleep 10
done
(( total == 2 ))
verify_hard_gates
}
artifact_smoke() {
local manifest=$1
node "$repository_root/scripts/release-manifest.mjs" validate "$manifest" >/dev/null
local manifest_sha api_image web_image api_digest artifact_report
manifest_sha=$(node "$repository_root/scripts/release-manifest.mjs" get "$manifest" sourceSha)
api_image=$(node "$repository_root/scripts/release-manifest.mjs" get "$manifest" images.api)
web_image=$(node "$repository_root/scripts/release-manifest.mjs" get "$manifest" images.web)
api_digest=${api_image##*@}
[[ $manifest_sha == "$release_sha" && $api_digest =~ ^sha256:[0-9a-f]{64}$ ]]
docker pull --platform linux/amd64 "$api_image" >/dev/null
docker pull --platform linux/amd64 "$web_image" >/dev/null
"$private_root/tools/k3d" image import -c easyai-acceptance-local "$api_image" "$web_image"
for deployment in easyai-api-ningbo easyai-api-hongkong easyai-worker-ningbo \
easyai-worker-hongkong easyai-capacity-controller easyai-acceptance-emulator \
easyai-acceptance-callback-collector; do
kubectl --context "$context" -n "$namespace" set image deployment/"$deployment" \
"*=$api_image" >/dev/null
kubectl --context "$context" -n "$namespace" rollout status deployment/"$deployment" --timeout=15m
done
for deployment in easyai-acceptance-edge-ningbo easyai-acceptance-edge-hongkong; do
kubectl --context "$context" -n "$namespace" set image deployment/"$deployment" \
"*=$web_image" >/dev/null
kubectl --context "$context" -n "$namespace" rollout status deployment/"$deployment" --timeout=10m
done
kubectl --context "$context" -n "$namespace" delete job easyai-artifact-migrate \
--ignore-not-found --wait=true >/dev/null
cat <<EOF | kubectl --context "$context" apply -f - >/dev/null
apiVersion: batch/v1
kind: Job
metadata:
name: easyai-artifact-migrate
namespace: easyai
spec:
backoffLimit: 0
template:
spec:
restartPolicy: Never
containers:
- name: migrate
image: $api_image
command: ["/bin/sh", "-ec", "cd /app && exec /app/easyai-ai-gateway-migrate"]
envFrom:
- secretRef:
name: easyai-ai-gateway-runtime
EOF
kubectl --context "$context" -n "$namespace" wait \
--for=condition=complete job/easyai-artifact-migrate --timeout=10m >/dev/null
run_load simulated-smoke "$report_root/artifact-simulated-smoke.json"
verify_hard_gates
artifact_report="$report_root/artifact-smoke.json"
jq -n \
--arg releaseSha "$manifest_sha" \
--arg apiImageDigest "$api_digest" \
--arg webImageDigest "${web_image##*@}" \
'{
schemaVersion:"acceptance-artifact-smoke/v1",
releaseSha:$releaseSha,
apiImageDigest:$apiImageDigest,
webImageDigest:$webImageDigest,
architecture:"linux/amd64",
migrationSmoke:true,
startupSmoke:true,
mediaSmoke:true,
passed:true,
secretSafe:true
}' >"$artifact_report"
chmod 0600 "$artifact_report"
}
quick() {
current_phase=quick
run_load simulated-smoke "$report_root/quick.json"
verify_hard_gates
echo "local_acceptance_quick=PASS run_id=$run_id report=$report_root/quick.json"
}
full() {
local manifest=$1 profile repetition
quick
for profile in P24 P28 P32; do
current_phase="capacity_${profile,,}_apply"
apply_profile "$profile"
for repetition in 1 2 3; do
current_phase="capacity_${profile,,}_round_$repetition"
run_profile_round "$profile" "$repetition"
done
stable_profile=$profile
done
current_phase=fault_matrix
run_fault_matrix
current_phase=autoscaling_and_drain
run_autoscaling "$stable_profile"
current_phase=certified_soak
run_load mixed-soak "$report_root/mixed-soak.json" \
-duration "${AI_GATEWAY_LOCAL_SOAK_DURATION:-2h}" \
-image-rate "${AI_GATEWAY_LOCAL_CERTIFIED_IMAGE_RATE:-1}" \
-video-rate "${AI_GATEWAY_LOCAL_CERTIFIED_VIDEO_RATE:-1}"
current_phase=overload_shedding
run_load mixed-overload "$report_root/mixed-overload.json" \
-duration "${AI_GATEWAY_LOCAL_OVERLOAD_DURATION:-10m}" \
-image-rate "${AI_GATEWAY_LOCAL_OVERLOAD_IMAGE_RATE:-2}" \
-video-rate "${AI_GATEWAY_LOCAL_OVERLOAD_VIDEO_RATE:-2}"
current_phase=amd64_artifact_smoke
artifact_smoke "$manifest"
current_phase=final_report
node "$script_dir/report.mjs" build-local \
--runtime "$runtime" \
--snapshot "$snapshot" \
--reports "$report_root" \
--artifact "$report_root/artifact-smoke.json" \
--api-digest "$(node "$repository_root/scripts/release-manifest.mjs" get "$manifest" images.api | sed 's/.*@//')" \
--certified-profile "$stable_profile" \
--output "$report_root/acceptance-report.json"
echo "local_acceptance_full=PASS run_id=$run_id certified_profile=$stable_profile report=$report_root/acceptance-report.json"
}
command=${1:-}
shift || true
require_local_cluster
load_environment
case $command in
quick)
[[ $# -eq 0 ]] || { usage >&2; exit 64; }
quick
;;
artifact-smoke)
[[ ${1:-} == --release-manifest && $# -eq 2 ]] || { usage >&2; exit 64; }
artifact_smoke "$2"
;;
full)
[[ ${1:-} == --release-manifest && $# -eq 2 ]] || { usage >&2; exit 64; }
full "$2"
;;
*)
usage >&2
exit 64
;;
esac
+8 -1
View File
@@ -7,13 +7,19 @@ source "$script_dir/common.sh"
load_cluster_env
cluster_ssh "$CLUSTER_NINGBO_HOST" \
'install -d -m 0755 /usr/local/lib/easyai-ai-gateway-release /usr/local/share/easyai-ai-gateway-release'
'install -d -m 0755 /usr/local/lib/easyai-ai-gateway-release /usr/local/share/easyai-ai-gateway-release/production'
cluster_scp "$cluster_root/deploy/kubernetes/easyai-ai-gateway-cluster-release" \
"$CLUSTER_NINGBO_HOST:/usr/local/sbin/easyai-ai-gateway-cluster-release"
cluster_scp "$cluster_root/deploy/kubernetes/easyai-ai-gateway-cluster-release.conf.example" \
"$CLUSTER_NINGBO_HOST:/usr/local/share/easyai-ai-gateway-release/easyai-ai-gateway-cluster-release.conf.example"
cluster_scp "$cluster_root/scripts/release-manifest.mjs" \
"$CLUSTER_NINGBO_HOST:/usr/local/lib/easyai-ai-gateway-release/release-manifest.mjs"
for desired_state_file in \
namespace.yaml service-account-rbac.yaml application-config.yaml application.yaml \
database.yaml network-policy.yaml kustomization.yaml; do
cluster_scp "$cluster_root/deploy/kubernetes/production/$desired_state_file" \
"$CLUSTER_NINGBO_HOST:/usr/local/share/easyai-ai-gateway-release/production/$desired_state_file"
done
cluster_ssh "$CLUSTER_NINGBO_HOST" 'bash -s' <<'REMOTE'
set -euo pipefail
if [[ ! -e /etc/easyai-ai-gateway-cluster-release.conf ]]; then
@@ -24,5 +30,6 @@ fi
chmod 0750 /usr/local/sbin/easyai-ai-gateway-cluster-release
chmod 0640 /etc/easyai-ai-gateway-cluster-release.conf
chmod 0755 /usr/local/lib/easyai-ai-gateway-release/release-manifest.mjs
find /usr/local/share/easyai-ai-gateway-release/production -type f -exec chmod 0644 {} +
REMOTE
echo 'cluster_release_helper_install=PASS'
+352
View File
@@ -0,0 +1,352 @@
#!/usr/bin/env bash
set -euo pipefail
script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck source=scripts/cluster/common.sh
source "$script_dir/common.sh"
cluster_root=${cluster_root:?}
usage() {
cat <<'EOF'
Usage:
scripts/cluster/monitor-production-release.sh \
--run-id <production-run-id> --release-manifest dist/releases/<SHA>.json
Runs a finite 24-hour post-promotion monitor: 10-second samples for the first
2 hours, then 60-second samples. Test-only duration overrides are available via
AI_GATEWAY_RELEASE_MONITOR_DENSE_SECONDS and
AI_GATEWAY_RELEASE_MONITOR_TOTAL_SECONDS.
EOF
}
[[ ${1:-} == --run-id && ${3:-} == --release-manifest && $# -eq 4 ]] || {
usage >&2
exit 64
}
run_id=$2
release_manifest=$4
[[ $run_id =~ ^[0-9a-f-]{36}$ && -f $release_manifest && ! -L $release_manifest ]] || {
echo 'invalid monitor Run ID or release manifest' >&2
exit 64
}
load_cluster_env
require_commands curl jq node openssl
: "${AI_GATEWAY_ACCEPTANCE_ADMIN_TOKEN:?}"
: "${AI_GATEWAY_RELEASE_MONITOR_DENSE_SECONDS:=7200}"
: "${AI_GATEWAY_RELEASE_MONITOR_TOTAL_SECONDS:=86400}"
: "${AI_GATEWAY_RELEASE_MONITOR_DENSE_INTERVAL_SECONDS:=10}"
: "${AI_GATEWAY_RELEASE_MONITOR_SPARSE_INTERVAL_SECONDS:=60}"
[[ $AI_GATEWAY_RELEASE_MONITOR_DENSE_SECONDS =~ ^[0-9]+$ &&
$AI_GATEWAY_RELEASE_MONITOR_TOTAL_SECONDS =~ ^[1-9][0-9]*$ &&
$AI_GATEWAY_RELEASE_MONITOR_DENSE_INTERVAL_SECONDS =~ ^[1-9][0-9]*$ &&
$AI_GATEWAY_RELEASE_MONITOR_SPARSE_INTERVAL_SECONDS =~ ^[1-9][0-9]*$ &&
$AI_GATEWAY_RELEASE_MONITOR_DENSE_SECONDS -le $AI_GATEWAY_RELEASE_MONITOR_TOTAL_SECONDS ]] || {
echo 'invalid monitor duration configuration' >&2
exit 64
}
node "$cluster_root/scripts/release-manifest.mjs" validate "$release_manifest" >/dev/null
release_sha=$(node "$cluster_root/scripts/release-manifest.mjs" get "$release_manifest" sourceSha)
api_image=$(node "$cluster_root/scripts/release-manifest.mjs" get "$release_manifest" images.api)
api_digest=${api_image##*@}
namespace=${AI_GATEWAY_K3S_NAMESPACE:-easyai}
report_root=${AI_GATEWAY_ACCEPTANCE_REPORT_DIR:-"$cluster_root/dist/acceptance/$run_id"}
overall_report=$report_root/acceptance-report.json
monitor_report=$report_root/post-promotion-monitor.json
observations=$report_root/post-promotion-observations.csv
capacity_snapshot=$report_root/pre-acceptance-capacity.json
[[ -f $overall_report && ! -L $overall_report &&
-f $capacity_snapshot && ! -L $capacity_snapshot ]] || {
echo 'monitor requires the overall report and pre-acceptance capacity snapshot' >&2
exit 1
}
temporary_root=$(mktemp -d)
chmod 0700 "$temporary_root"
trap '[[ $temporary_root == /tmp/* || $temporary_root == /var/folders/* ]] && rm -rf -- "$temporary_root"' EXIT
remote_kubectl() {
local command_text
printf -v command_text '%q ' "$@"
cluster_ssh "$CLUSTER_NINGBO_HOST" "k3s kubectl $command_text"
}
database_query() {
local sql=$1 primary
primary=$(remote_kubectl get cluster easyai-postgres -n "$namespace" \
-o 'jsonpath={.status.currentPrimary}')
remote_kubectl exec -n "$namespace" "$primary" -c postgres -- \
psql -X -v ON_ERROR_STOP=1 -U postgres -d easyai_ai_gateway -At -c "$sql"
}
admin_request() {
local method=$1 path=$2 body=${3:-} output=$4
local service_ip service_port token_encoded body_encoded response status payload script command
service_ip=$(remote_kubectl get service api-ningbo-edge -n "$namespace" -o 'jsonpath={.spec.clusterIP}')
service_port=$(remote_kubectl get service api-ningbo-edge -n "$namespace" -o 'jsonpath={.spec.ports[0].port}')
token_encoded=$(printf '%s' "$AI_GATEWAY_ACCEPTANCE_ADMIN_TOKEN" | openssl base64 -A)
body_encoded=$(printf '%s' "$body" | openssl base64 -A)
# shellcheck disable=SC2016
script='
set -euo pipefail
token=$(printf "%s" "$4" | base64 -d)
body=$(printf "%s" "$5" | base64 -d)
args=(-sS --max-time 30 -w "\n%{http_code}" -X "$1" -H "Authorization: Bearer $token")
if [[ -n $body ]]; then args+=(-H "Content-Type: application/json" --data-binary "$body"); fi
curl "${args[@]}" "http://$2:$3$6"
'
printf -v command 'bash -c %q -- %q %q %q %q %q %q' \
"$script" "$method" "$service_ip" "$service_port" "$token_encoded" "$body_encoded" "$path"
response=$(cluster_ssh "$CLUSTER_NINGBO_HOST" "$command")
status=${response##*$'\n'}
payload=${response%$'\n'*}
printf '%s' "$payload" >"$output"
[[ $status =~ ^2[0-9][0-9]$ ]]
}
metric_sum() {
local pattern=$1 sum=0 pod value metrics
while read -r pod; do
[[ -n $pod ]] || continue
metrics=$(remote_kubectl exec -n "$namespace" "$pod" -- \
wget -qO- http://127.0.0.1:8088/metrics)
value=$(awk -v pattern="$pattern" '$0 ~ pattern {sum += $NF} END {print sum+0}' <<<"$metrics")
sum=$(awk -v left="$sum" -v right="$value" 'BEGIN {print left+right}')
done < <(remote_kubectl get pods -n "$namespace" \
-l app.kubernetes.io/name=easyai-worker \
-o 'jsonpath={range .items[*]}{.metadata.name}{"\n"}{end}')
printf '%s\n' "$sum"
}
pool_saturated_pods() {
local pod metrics max acquired idle saturated=0
while read -r pod; do
[[ -n $pod ]] || continue
metrics=$(remote_kubectl exec -n "$namespace" "$pod" -- \
wget -qO- http://127.0.0.1:8088/metrics)
max=$(awk '$1=="easyai_gateway_postgres_pool_max_connections" {print int($2)}' <<<"$metrics")
acquired=$(awk '$1=="easyai_gateway_postgres_pool_acquired_connections" {print int($2)}' <<<"$metrics")
idle=$(awk '$1=="easyai_gateway_postgres_pool_idle_connections" {print int($2)}' <<<"$metrics")
[[ -z $max || -z $acquired || -z $idle ]] || {
(( acquired == max && idle == 0 )) && saturated=$((saturated + 1))
}
done < <(remote_kubectl get pods -n "$namespace" \
-l 'app.kubernetes.io/part-of=easyai-ai-gateway' \
-o 'jsonpath={range .items[*]}{.metadata.name}{"\n"}{end}')
printf '%s\n' "$saturated"
}
check_wireguard() {
local -a hosts=("$CLUSTER_NINGBO_HOST" "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST")
local -a ips=(10.77.0.1 10.77.0.2 10.77.0.3)
local source target output loss average handshakes
for source in 0 1 2; do
handshakes=$(cluster_ssh "${hosts[$source]}" \
"wg show wg0 latest-handshakes | awk '\$2 > 0 && systime()-\$2 < 180 {count++} END {print count+0}'")
(( handshakes == 2 )) || return 1
for target in 0 1 2; do
(( source == target )) && continue
output=$(cluster_ssh "${hosts[$source]}" "ping -q -c 10 -W 2 ${ips[$target]}") || return 1
loss=$(awk -F',' '/packet loss/ {gsub(/[^0-9.]/,"",$3); print $3}' <<<"$output")
average=$(awk -F'=' '/min\\/avg\\/max/ {split($2,v,"/"); gsub(/ /,"",v[2]); print v[2]}' <<<"$output")
if (( source == 2 || target == 2 )); then
awk -v loss="$loss" -v average="$average" 'BEGIN {exit !(loss < 1 && average < 300)}' || return 1
else
awk -v loss="$loss" -v average="$average" 'BEGIN {exit !(loss < 1 && average < 80)}' || return 1
fi
done
done
}
restore_previous_capacity() {
local remote_file="/root/easyai-pre-acceptance-$run_id.json"
cluster_scp "$capacity_snapshot" "$CLUSTER_NINGBO_HOST:$remote_file" >/dev/null
cluster_ssh "$CLUSTER_NINGBO_HOST" \
"k3s kubectl apply -f '$remote_file' >/dev/null && unlink '$remote_file'"
for deployment in easyai-api-ningbo easyai-api-hongkong easyai-worker-ningbo \
easyai-worker-hongkong easyai-capacity-controller; do
remote_kubectl rollout status deployment/"$deployment" -n "$namespace" --timeout=600s
done
}
pause_and_restore() {
local gate_id=$1 mode_file=$temporary_root/mode.json pause_file=$temporary_root/pause.json body
admin_request GET /api/admin/system/acceptance/traffic-mode '' "$mode_file"
body=$(jq -cn \
--argjson revision "$(jq -r '.revision' "$mode_file")" \
--arg releaseSha "$release_sha" \
--arg apiDigest "$api_digest" \
--arg reason "$gate_id" \
'{
revision:$revision,
releaseSha:$releaseSha,
apiImageDigest:$apiDigest,
workerImageDigest:$apiDigest,
reason:$reason
}')
admin_request POST /api/admin/system/acceptance/traffic-mode/pause "$body" "$pause_file"
[[ $(jq -r '.mode' "$pause_file") == validation ]]
restore_previous_capacity
}
started_epoch=$(date +%s)
started_at=$(date -u '+%Y-%m-%dT%H:%M:%SZ')
monitor_since=$(date -u '+%Y-%m-%d %H:%M:%S+00')
baseline_lease_failure=$(metric_sum 'outcome="failure"')
baseline_lease_lost=$(metric_sum 'outcome="lost"')
samples=0
queue_growth_streak=0
pool_saturation_streak=0
previous_queue=0
failure_gate_id=
last_wireguard_epoch=0
printf 'timestamp,queue_depth,oldest_wait_seconds,connections,max_connections,saturated_pods,node_max_memory_percent,lease_failure,lease_lost\n' >"$observations"
while (( $(date +%s) - started_epoch < AI_GATEWAY_RELEASE_MONITOR_TOTAL_SECONDS )); do
now_epoch=$(date +%s)
elapsed=$((now_epoch - started_epoch))
interval=$AI_GATEWAY_RELEASE_MONITOR_SPARSE_INTERVAL_SECONDS
(( elapsed < AI_GATEWAY_RELEASE_MONITOR_DENSE_SECONDS )) &&
interval=$AI_GATEWAY_RELEASE_MONITOR_DENSE_INTERVAL_SECONDS
samples=$((samples + 1))
if ! curl -fsS --max-time 10 https://ai.51easyai.com/api/v1/healthz >/dev/null ||
! curl -fsS --max-time 10 https://ai.51easyai.com/api/v1/readyz >/dev/null; then
failure_gate_id=public_health
fi
ready_nodes=$(remote_kubectl get nodes -o json |
jq '[.items[] | select(any(.status.conditions[]; .type=="Ready" and .status=="True"))] | length')
pressure_nodes=$(remote_kubectl get nodes -o json |
jq '[.items[] | select(any(.status.conditions[]; .type=="MemoryPressure" and .status=="True"))] | length')
[[ $ready_nodes == 3 && $pressure_nodes == 0 ]] || failure_gate_id=${failure_gate_id:-node_health}
node_max_memory=$(remote_kubectl top nodes --no-headers |
awk '{gsub(/%/,"",$5); if ($5>max) max=$5} END {print max+0}')
(( node_max_memory < 85 )) || failure_gate_id=${failure_gate_id:-node_memory_hard}
unhealthy=$(remote_kubectl get pods -n "$namespace" -o json |
jq '[.items[]
| select(.metadata.labels["app.kubernetes.io/name"] == "easyai-api" or
.metadata.labels["app.kubernetes.io/name"] == "easyai-worker")
| .status.containerStatuses[]?
| select(.restartCount > 0 or .lastState.terminated.reason == "OOMKilled" or
.state.terminated.reason == "OOMKilled")] | length')
[[ $unhealthy == 0 ]] || failure_gate_id=${failure_gate_id:-pod_restart_or_oom}
pod_max_memory=$(remote_kubectl top pods -n "$namespace" \
-l 'app.kubernetes.io/part-of=easyai-ai-gateway' --no-headers |
awk '{
value=$3
if (value ~ /Gi$/) {sub(/Gi$/,"",value); value*=1024}
else if (value ~ /Mi$/) {sub(/Mi$/,"",value)}
else if (value ~ /Ki$/) {sub(/Ki$/,"",value); value/=1024}
if (value>max) max=value
} END {printf "%.0f",max}')
(( pod_max_memory < 1536 )) || failure_gate_id=${failure_gate_id:-pod_memory_hard}
[[ $(remote_kubectl get cluster easyai-postgres -n "$namespace" \
-o 'jsonpath={.status.readyInstances}') == 2 ]] || failure_gate_id=${failure_gate_id:-postgres_ready}
sync_state=$(database_query "SELECT COALESCE(string_agg(sync_state,','),'') FROM pg_stat_replication;")
[[ ",$sync_state," == *,sync,* || ",$sync_state," == *,quorum,* ]] ||
failure_gate_id=${failure_gate_id:-postgres_replication}
connections=$(database_query "SELECT count(*) FROM pg_stat_activity WHERE backend_type='client backend';")
max_connections=$(database_query "SELECT setting::int FROM pg_settings WHERE name='max_connections';")
(( connections < 150 && connections * 4 < max_connections * 3 )) ||
failure_gate_id=${failure_gate_id:-postgres_connections}
queue_state=$(database_query "
SELECT count(*) FILTER (WHERE status='queued')||':'||
COALESCE(EXTRACT(EPOCH FROM now()-MIN(created_at)) FILTER (WHERE status='queued'),0)::bigint
FROM gateway_tasks
WHERE run_mode='production';")
IFS=: read -r queue_depth oldest_wait <<<"$queue_state"
if (( queue_depth > previous_queue && queue_depth > 0 )); then
queue_growth_streak=$((queue_growth_streak + 1))
else
queue_growth_streak=0
fi
previous_queue=$queue_depth
(( oldest_wait <= 900 && queue_growth_streak < 4 )) ||
failure_gate_id=${failure_gate_id:-queue_growth_or_oldest_wait}
saturated=$(pool_saturated_pods)
if (( saturated > 0 )); then
pool_saturation_streak=$((pool_saturation_streak + 1))
else
pool_saturation_streak=0
fi
(( pool_saturation_streak * interval < 30 )) ||
failure_gate_id=${failure_gate_id:-postgres_pool_saturation}
lease_failure=$(metric_sum 'outcome="failure"')
lease_lost=$(metric_sum 'outcome="lost"')
awk -v current="$lease_failure" -v baseline="$baseline_lease_failure" 'BEGIN {exit !(current<=baseline)}' ||
failure_gate_id=${failure_gate_id:-lease_renewal_failure}
awk -v current="$lease_lost" -v baseline="$baseline_lease_lost" 'BEGIN {exit !(current<=baseline)}' ||
failure_gate_id=${failure_gate_id:-lease_lost}
duplicates=$(database_query "
SELECT
(SELECT count(*) FROM (
SELECT remote_task_id
FROM gateway_tasks
WHERE run_mode='production' AND created_at >= '$monitor_since'::timestamptz
AND remote_task_id IS NOT NULL
GROUP BY remote_task_id HAVING count(*)>1
) d)
+
(SELECT count(*) FROM (
SELECT reference_id,transaction_type
FROM gateway_wallet_transactions
WHERE created_at >= '$monitor_since'::timestamptz
AND reference_type='gateway_task'
GROUP BY reference_id,transaction_type HAVING count(*)>1
) d);")
[[ $duplicates == 0 ]] || failure_gate_id=${failure_gate_id:-duplicate_execution_or_billing}
if (( now_epoch - last_wireguard_epoch >= 300 )); then
check_wireguard || failure_gate_id=${failure_gate_id:-wireguard}
last_wireguard_epoch=$now_epoch
fi
printf '%s,%s,%s,%s,%s,%s,%s,%s,%s\n' \
"$(date -u '+%Y-%m-%dT%H:%M:%SZ')" "$queue_depth" "$oldest_wait" \
"$connections" "$max_connections" "$saturated" "$node_max_memory" \
"$lease_failure" "$lease_lost" >>"$observations"
[[ -z $failure_gate_id ]] || break
sleep "$interval"
done
finished_at=$(date -u '+%Y-%m-%dT%H:%M:%SZ')
passed=true
traffic_mode=live
restored=false
if [[ -n $failure_gate_id ]]; then
passed=false
pause_and_restore "$failure_gate_id"
traffic_mode=validation
restored=true
fi
jq -n \
--arg runId "$run_id" \
--arg releaseSha "$release_sha" \
--arg startedAt "$started_at" \
--arg finishedAt "$finished_at" \
--arg failureGateId "$failure_gate_id" \
--arg trafficMode "$traffic_mode" \
--argjson samples "$samples" \
--argjson passed "$passed" \
--argjson restored "$restored" \
'{
schemaVersion:"acceptance-monitor-report/v1",
runId:$runId,
releaseSha:$releaseSha,
startedAt:$startedAt,
finishedAt:$finishedAt,
samples:$samples,
passed:$passed,
failureGateId:($failureGateId | if length>0 then . else null end),
trafficMode:$trafficMode,
previousCapacityRestored:$restored,
secretSafe:true
}' >"$monitor_report"
chmod 0600 "$monitor_report" "$observations"
node "$cluster_root/scripts/acceptance/report.mjs" attach-monitor \
--input "$overall_report" --monitor "$monitor_report" --output "$overall_report" >/dev/null
if [[ $passed == true ]]; then
echo "production_release_monitor=PASS run_id=$run_id samples=$samples duration_seconds=$AI_GATEWAY_RELEASE_MONITOR_TOTAL_SECONDS traffic_mode=live"
else
echo "production_release_monitor=FAIL run_id=$run_id gate_id=$failure_gate_id traffic_mode=validation previous_capacity_restored=true" >&2
exit 1
fi
+170
View File
@@ -0,0 +1,170 @@
#!/usr/bin/env bash
set -euo pipefail
script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck source=scripts/cluster/common.sh
source "$script_dir/common.sh"
cluster_root=${cluster_root:?}
usage() {
cat <<'EOF'
Usage:
scripts/cluster/run-production-acceptance.sh \
--promote dist/releases/<SHA>.json --run-id <production-run-id>
This is the only acceptance action that switches production traffic from
validation to live. It requires an already-passed overall acceptance report and
rechecks the release, digest, config, Run, and traffic-mode CAS fields.
EOF
}
[[ ${1:-} == --promote && ${3:-} == --run-id && $# -eq 4 ]] || {
usage >&2
exit 64
}
release_manifest=$2
run_id=$4
[[ -f $release_manifest && ! -L $release_manifest ]] || {
echo 'release manifest must be a regular non-symlink file' >&2
exit 1
}
[[ $run_id =~ ^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$ ]] || {
echo 'run ID must be a lowercase UUID' >&2
exit 64
}
load_cluster_env
require_commands curl git jq node openssl
: "${AI_GATEWAY_ACCEPTANCE_ADMIN_TOKEN:?}"
node "$cluster_root/scripts/release-manifest.mjs" validate "$release_manifest" >/dev/null
release_sha=$(node "$cluster_root/scripts/release-manifest.mjs" get "$release_manifest" sourceSha)
api_image=$(node "$cluster_root/scripts/release-manifest.mjs" get "$release_manifest" images.api)
api_digest=${api_image##*@}
[[ $(git -C "$cluster_root" rev-parse HEAD) == "$release_sha" &&
-z $(git -C "$cluster_root" status --short) ]] || {
echo 'manual promotion requires the clean release source at HEAD' >&2
exit 1
}
report=${AI_GATEWAY_ACCEPTANCE_REPORT_DIR:-"$cluster_root/dist/acceptance/$run_id"}/acceptance-report.json
[[ -f $report && ! -L $report ]] || {
echo "passed overall acceptance report is missing: $report" >&2
exit 1
}
node "$cluster_root/scripts/acceptance/report.mjs" validate --input "$report" >/dev/null
[[ $(jq -r '.runId' "$report") == "$run_id" &&
$(jq -r '.release.sourceSha' "$report") == "$release_sha" &&
$(jq -r '.release.apiImageDigest' "$report") == "$api_digest" &&
$(jq -r '.stages.onlineSimulation.status' "$report") == passed &&
$(jq -r '.stages.realCanary.status' "$report") == passed &&
$(jq -r '.promotion.ready' "$report") == true &&
$(jq -r '.promotion.trafficMode' "$report") == validation ]] || {
echo 'overall acceptance report is not ready for this manual promotion' >&2
exit 1
}
namespace=${AI_GATEWAY_K3S_NAMESPACE:-easyai}
remote_release_helper=${AI_GATEWAY_REMOTE_RELEASE_HELPER:-/usr/local/sbin/easyai-ai-gateway-cluster-release}
temporary_root=$(mktemp -d)
chmod 0700 "$temporary_root"
trap '[[ $temporary_root == /tmp/* || $temporary_root == /var/folders/* ]] && rm -rf -- "$temporary_root"' EXIT
remote_kubectl() {
local command_text
printf -v command_text '%q ' "$@"
cluster_ssh "$CLUSTER_NINGBO_HOST" "k3s kubectl $command_text"
}
admin_request() {
local method=$1 path=$2 body=${3:-} output=$4
local service_ip service_port token_encoded body_encoded response status payload remote_script command
service_ip=$(remote_kubectl get service api-ningbo-edge -n "$namespace" \
-o 'jsonpath={.spec.clusterIP}')
service_port=$(remote_kubectl get service api-ningbo-edge -n "$namespace" \
-o 'jsonpath={.spec.ports[0].port}')
token_encoded=$(printf '%s' "$AI_GATEWAY_ACCEPTANCE_ADMIN_TOKEN" | openssl base64 -A)
body_encoded=$(printf '%s' "$body" | openssl base64 -A)
# shellcheck disable=SC2016
remote_script='
set -euo pipefail
token=$(printf "%s" "$4" | base64 -d)
body=$(printf "%s" "$5" | base64 -d)
args=(-sS --max-time 30 -w "\n%{http_code}" -X "$1" -H "Authorization: Bearer $token")
if [[ -n $body ]]; then
args+=(-H "Content-Type: application/json" --data-binary "$body")
fi
curl "${args[@]}" "http://$2:$3$6"
'
printf -v command 'bash -c %q -- %q %q %q %q %q %q' \
"$remote_script" "$method" "$service_ip" "$service_port" \
"$token_encoded" "$body_encoded" "$path"
response=$(cluster_ssh "$CLUSTER_NINGBO_HOST" "$command")
status=${response##*$'\n'}
payload=${response%$'\n'*}
printf '%s' "$payload" >"$output"
[[ $status =~ ^2[0-9][0-9]$ ]] || {
echo "acceptance control API failed: method=$method path=$path status=$status" >&2
return 1
}
}
current_manifest=$temporary_root/current.json
cluster_ssh "$CLUSTER_NINGBO_HOST" "$remote_release_helper status" >"$current_manifest"
node "$cluster_root/scripts/release-manifest.mjs" validate "$current_manifest" >/dev/null
[[ $(node "$cluster_root/scripts/release-manifest.mjs" get "$current_manifest" sourceSha) == "$release_sha" &&
$(node "$cluster_root/scripts/release-manifest.mjs" get "$current_manifest" images.api) == "$api_image" ]] || {
echo 'production release changed after acceptance' >&2
exit 1
}
for deployment in easyai-api-ningbo easyai-api-hongkong easyai-worker-ningbo easyai-worker-hongkong; do
[[ $(remote_kubectl get deployment "$deployment" -n "$namespace" \
-o 'jsonpath={.spec.template.spec.containers[0].image}') == "$api_image" ]]
done
mode_file=$temporary_root/mode.json
run_file=$temporary_root/run.json
promotion_file=$temporary_root/promotion.json
admin_request GET /api/admin/system/acceptance/traffic-mode '' "$mode_file"
admin_request GET "/api/admin/system/acceptance/runs/$run_id" '' "$run_file"
[[ $(jq -r '.mode' "$mode_file") == validation &&
$(jq -r '.runId' "$mode_file") == "$run_id" &&
$(jq -r '.releaseSha' "$mode_file") == "$release_sha" &&
$(jq -r '.apiImageDigest' "$mode_file") == "$api_digest" &&
$(jq -r '.status' "$run_file") == passed ]] || {
echo 'production acceptance Run or traffic-mode CAS changed before promotion' >&2
exit 1
}
config_hash=$(jq -r '.config.validationConfigHash' "$run_file")
profiles=$(jq -c '.config.validationCapacityProfiles' "$run_file")
[[ $config_hash =~ ^[0-9a-f]{64}$ && $(jq 'type == "array" and length > 0' <<<"$profiles") == true ]]
[[ $(remote_kubectl get configmap easyai-gateway-release -n "$namespace" \
-o 'jsonpath={.data.configHash}') == "$config_hash" ]] || {
echo 'production configuration changed after acceptance' >&2
exit 1
}
body=$(jq -cn \
--argjson revision "$(jq -r '.revision' "$mode_file")" \
--arg releaseSha "$release_sha" \
--arg apiDigest "$api_digest" \
--arg configHash "$config_hash" \
--argjson capacityProfiles "$profiles" \
'{
revision:$revision,
releaseSha:$releaseSha,
apiImageDigest:$apiDigest,
workerImageDigest:$apiDigest,
configHash:$configHash,
capacityProfiles:$capacityProfiles
}')
admin_request POST "/api/admin/system/acceptance/runs/$run_id/promote" "$body" "$promotion_file"
[[ $(jq -r '.mode' "$promotion_file") == live &&
$(jq -r '.runId' "$promotion_file") == "$run_id" ]] || {
echo 'live promotion CAS did not complete' >&2
exit 1
}
node "$cluster_root/scripts/acceptance/report.mjs" mark-promoted \
--input "$report" --output "$report" >/dev/null
echo "production_acceptance_promotion=PASS run_id=$run_id release=$release_sha traffic_mode=live"
exec "$script_dir/monitor-production-release.sh" \
--run-id "$run_id" --release-manifest "$release_manifest"
File diff suppressed because it is too large Load Diff
+3
View File
@@ -114,6 +114,9 @@ else
fi
scp "${scp_options[@]}" "$manifest" "$production_host:$remote_manifest"
if [[ $deploy_mode == kubernetes ]]; then
"$root/scripts/cluster/install-release-helper.sh"
fi
# shellcheck disable=SC2029 # remote_manifest is derived from a validated full SHA.
ssh "${ssh_options[@]}" "$production_host" "$remote_helper" deploy "$remote_manifest"
echo "production_deploy=PASS release=$source_sha"
+15
View File
@@ -15,6 +15,21 @@
"non-null column addition"
],
"reason": "The table stores ephemeral Worker heartbeats and currently has only two production rows. PostgreSQL adds the integer column with a constant default, then the migration backfills each live row from its existing allocation so a rolling deployment cannot assign unsafe failover capacity."
},
"apps/api/migrations/0095_oidc_session_context_identity.sql": {
"sha256": "3e6d17b5b564047c4a77ef3e4cf4ed88caba375a34a33bc7ecfbd7c59d6ab479",
"allowedViolations": [
"destructive DROP operation",
"non-null column addition"
],
"reason": "The migration adds only nullable session identity columns, backfills bound sessions, and replaces one CHECK constraint without dropping data. Production preflight on 2026-07-31 found zero gateway_oidc_sessions rows and a 48 KiB relation, so the constraint validation and metadata lock are bounded; the non-null finding comes from the new CHECK expression rather than a NOT NULL column addition."
},
"apps/api/migrations/0096_worker_capacity_controller.sql": {
"sha256": "52c2c5f31eaa61b96508ae9b2c9fa164e0f1544184c55a1ba85f64b9aca442cb",
"allowedViolations": [
"destructive DROP operation"
],
"reason": "The migration replaces only the existing Worker status CHECK constraint so a rolling deployment can mark lost instances stale. It does not drop data, columns, tables, or indexes; the old and new binaries both continue to read active and draining rows safely."
}
}
}
+6
View File
@@ -373,10 +373,16 @@ esac
mkdir -p dist/releases
manifest=$root/dist/releases/$source_sha.json
[[ ! -e $manifest ]] || fail "release manifest already exists: $manifest"
migration_file=$(git ls-files 'apps/api/migrations/*.sql' | LC_ALL=C sort | tail -n 1)
migration_version=${migration_file##*/}
migration_version=${migration_version%.sql}
[[ $migration_version =~ ^[0-9]{4}_[a-z0-9][a-z0-9_]*$ ]] ||
fail 'could not resolve the latest database migration version'
RELEASE_SOURCE_SHA=$source_sha \
RELEASE_BASE_SHA=$base_sha \
RELEASE_COMPONENTS=$component_list \
RELEASE_MIGRATIONS_CHANGED=$migrations_changed \
RELEASE_MIGRATION_VERSION=$migration_version \
RELEASE_API_IMAGE=$api_image \
RELEASE_WEB_IMAGE=$web_image \
RELEASE_BUILD_COMPLETED_AT=$build_completed_at \
+55
View File
@@ -32,6 +32,7 @@ function releaseIntegrity(manifest) {
const payload = structuredClone(manifest)
delete payload.integrity
delete payload.deployment
delete payload.certification
return createHash('sha256').update(JSON.stringify(payload)).digest('hex')
}
@@ -58,6 +59,8 @@ function validate(manifest) {
'migrationsChanged', 'images', 'smoke', 'integrity',
]
if (manifest.deployment !== undefined) rootKeys.push('deployment')
if (manifest.certification !== undefined) rootKeys.push('certification')
if (manifest.migrationVersion !== undefined) rootKeys.push('migrationVersion')
requireExactKeys(manifest, rootKeys, 'manifest')
if (manifest.schemaVersion !== 1) fail('schemaVersion must be 1')
if (!shaPattern.test(manifest.sourceSha || '')) fail('sourceSha must be a full lowercase Git SHA')
@@ -84,6 +87,10 @@ function validate(manifest) {
if (manifest.migrationsChanged && !manifest.components.includes('api')) {
fail('migration changes require the api component')
}
if (manifest.migrationVersion !== undefined &&
!/^\d{4}_[a-z0-9][a-z0-9_]*$/.test(manifest.migrationVersion)) {
fail('migrationVersion must identify the latest migration without .sql')
}
if (!manifest.images || typeof manifest.images !== 'object') fail('images must be an object')
requireExactKeys(manifest.images, ['api', 'web'], 'images')
for (const component of allowedComponents) {
@@ -126,6 +133,28 @@ function validate(manifest) {
fail('deployment record is invalid')
}
}
if (manifest.certification !== undefined) {
if (!manifest.certification || typeof manifest.certification !== 'object') {
fail('certification must be an object')
}
requireExactKeys(manifest.certification, [
'status', 'runId', 'certifiedAt', 'configHash', 'capacityProfile',
'workerInstanceSlots', 'workerMaxReplicasNingbo', 'workerMaxReplicasHongkong',
], 'certification')
if (manifest.certification.status !== 'passed' ||
!/^[0-9a-f-]{36}$/.test(manifest.certification.runId || '') ||
!Number.isFinite(Date.parse(manifest.certification.certifiedAt)) ||
!/^[0-9a-f]{64}$/.test(manifest.certification.configHash || '') ||
!['P24', 'P28', 'P32'].includes(manifest.certification.capacityProfile) ||
!Number.isInteger(manifest.certification.workerInstanceSlots) ||
manifest.certification.workerInstanceSlots < 1 ||
!Number.isInteger(manifest.certification.workerMaxReplicasNingbo) ||
manifest.certification.workerMaxReplicasNingbo < 1 ||
!Number.isInteger(manifest.certification.workerMaxReplicasHongkong) ||
manifest.certification.workerMaxReplicasHongkong < 1) {
fail('certification record is invalid')
}
}
const forbiddenKey = /(password|secret|token|credential|private.?key)/i
const visit = (value, path = '') => {
@@ -166,6 +195,7 @@ if (command === 'create') {
targetPlatform: 'linux/amd64',
components,
migrationsChanged,
migrationVersion: requireEnv('RELEASE_MIGRATION_VERSION', /^\d{4}_[a-z0-9][a-z0-9_]*$/),
images: {
api: requireEnv('RELEASE_API_IMAGE', digestImagePattern),
web: requireEnv('RELEASE_WEB_IMAGE', digestImagePattern),
@@ -206,6 +236,31 @@ if (command === 'record-deployment') {
process.exit(0)
}
if (command === 'record-certification') {
if (!field) fail('record-certification requires an output path')
const certifiedAt = requireEnv('RELEASE_CERTIFIED_AT')
const runId = requireEnv('RELEASE_ACCEPTANCE_RUN_ID', /^[0-9a-f-]{36}$/)
const configHash = requireEnv('RELEASE_CONFIG_HASH', /^[0-9a-f]{64}$/)
const capacityProfile = requireEnv('RELEASE_CAPACITY_PROFILE', /^P(?:24|28|32)$/)
const workerInstanceSlots = Number(process.env.RELEASE_WORKER_INSTANCE_SLOTS || '')
const workerMaxReplicasNingbo = Number(process.env.RELEASE_WORKER_MAX_REPLICAS_NINGBO || '')
const workerMaxReplicasHongkong = Number(process.env.RELEASE_WORKER_MAX_REPLICAS_HONGKONG || '')
manifest.certification = {
status: 'passed',
runId,
certifiedAt,
configHash,
capacityProfile,
workerInstanceSlots,
workerMaxReplicasNingbo,
workerMaxReplicasHongkong,
}
validate(manifest)
writeFileSync(field, `${JSON.stringify(manifest, null, 2)}\n`, { flag: 'wx', mode: 0o600 })
console.log(`release_manifest=CERTIFIED release=${manifest.releaseId} run=${runId}`)
process.exit(0)
}
if (command === 'validate') {
const digest = createHash('sha256').update(readFileSync(path)).digest('hex')
console.log(`release_manifest=PASS release=${manifest.releaseId} sha256=${digest}`)
+124
View File
@@ -0,0 +1,124 @@
import assert from 'node:assert/strict';
import { execFileSync } from 'node:child_process';
import { mkdtempSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { join, resolve } from 'node:path';
const root = resolve(import.meta.dirname, '../..');
const temporary = mkdtempSync(join(tmpdir(), 'easyai-acceptance-report-'));
const reports = join(temporary, 'reports');
mkdirSync(reports);
const releaseSha = 'a'.repeat(40);
const digest = `sha256:${'b'.repeat(64)}`;
const configHash = 'c'.repeat(64);
const snapshotHash = 'd'.repeat(64);
const runID = '00000000-0000-4000-8000-000000000001';
const writeJSON = (path, value) => writeFileSync(path, `${JSON.stringify(value)}\n`, { mode: 0o600 });
const runtime = join(temporary, 'runtime.json');
const snapshot = join(temporary, 'snapshot.json');
const artifact = join(temporary, 'artifact.json');
const output = join(temporary, 'acceptance-report.json');
const partialOutput = join(temporary, 'acceptance-report.partial.json');
const productionPartial = join(temporary, 'production-summary.partial.json');
const mergedPartialOutput = join(temporary, 'acceptance-report.production-partial.json');
writeJSON(runtime, {
schemaVersion: 'acceptance-runtime/v1',
runId: runID,
releaseSha,
apiImageDigest: digest,
workerImageDigest: digest,
snapshotConfigHash: configHash,
snapshotSha256: snapshotHash,
apiKeys: ['must-not-be-copied'],
runToken: 'must-not-be-copied'
});
writeJSON(snapshot, {
schemaVersion: 'acceptance-snapshot/v1',
source: { releaseSha: 'e'.repeat(40), configHash },
snapshotSha256: snapshotHash
});
writeJSON(artifact, {
schemaVersion: 'acceptance-artifact-smoke/v1',
releaseSha,
apiImageDigest: digest,
passed: true,
secretSafe: true
});
writeJSON(join(reports, 'load.json'), {
schemaVersion: 'acceptance-load-report/v1',
profile: 'simulated-smoke',
passed: true,
secretSafe: true,
phases: [],
startedAt: '2026-01-01T00:00:00Z',
finishedAt: '2026-01-01T00:01:00Z'
});
execFileSync('node', [
join(root, 'scripts/acceptance/report.mjs'),
'build-local',
'--runtime', runtime,
'--snapshot', snapshot,
'--reports', reports,
'--artifact', artifact,
'--api-digest', digest,
'--certified-profile', 'P24',
'--output', output
]);
const report = JSON.parse(readFileSync(output, 'utf8'));
assert.equal(report.schemaVersion, 'acceptance-report/v1');
assert.equal(report.stages.localNative.status, 'passed');
assert.equal(report.stages.amd64Artifact.status, 'passed');
assert.equal(report.promotion.ready, false);
assert.equal(readFileSync(output, 'utf8').includes('must-not-be-copied'), false);
execFileSync('node', [join(root, 'scripts/acceptance/report.mjs'), 'validate', '--input', output]);
execFileSync('node', [
join(root, 'scripts/acceptance/report.mjs'),
'build-local-partial',
'--runtime', runtime,
'--snapshot', snapshot,
'--reports', reports,
'--certified-profile', 'P24',
'--failure-phase', 'capacity_p28_round_2',
'--output', partialOutput
]);
const partialReport = JSON.parse(readFileSync(partialOutput, 'utf8'));
assert.equal(partialReport.stages.localNative.status, 'failed');
assert.equal(partialReport.stages.localNative.failurePhase, 'capacity_p28_round_2');
assert.equal(partialReport.stages.amd64Artifact.status, 'not-run');
assert.equal(partialReport.promotion.ready, false);
assert.equal(readFileSync(partialOutput, 'utf8').includes('must-not-be-copied'), false);
execFileSync('node', [join(root, 'scripts/acceptance/report.mjs'), 'validate', '--input', partialOutput]);
writeJSON(productionPartial, {
releaseSha,
apiImageDigest: digest,
snapshotConfigHash: configHash,
snapshotSha256: snapshotHash,
stableCapacityProfile: 'P24',
realCanaryPassed: false,
gates: [
{
id: 'workflow_interrupted',
passed: false,
detail: 'production acceptance interrupted by signal'
}
],
passed: false
});
execFileSync('node', [
join(root, 'scripts/acceptance/report.mjs'),
'merge-production',
'--local-report', output,
'--production-summary', productionPartial,
'--output', mergedPartialOutput
]);
const mergedPartial = JSON.parse(readFileSync(mergedPartialOutput, 'utf8'));
assert.equal(mergedPartial.stages.onlineSimulation.status, 'failed');
assert.equal(mergedPartial.stages.realCanary.status, 'failed');
assert.equal(mergedPartial.promotion.ready, false);
assert.equal(mergedPartial.promotion.trafficMode, 'validation');
assert.equal(mergedPartial.gates.at(-1).id, 'workflow_interrupted');
execFileSync('node', [join(root, 'scripts/acceptance/report.mjs'), 'validate', '--input', mergedPartialOutput]);
+14
View File
@@ -117,6 +117,7 @@ RELEASE_SOURCE_SHA=$source_sha \
RELEASE_BASE_SHA=$base_sha \
RELEASE_COMPONENTS=api \
RELEASE_MIGRATIONS_CHANGED=false \
RELEASE_MIGRATION_VERSION=0097_worker_runtime_owner_index \
RELEASE_API_IMAGE=$api_image \
RELEASE_WEB_IMAGE=$web_image \
RELEASE_BUILD_COMPLETED_AT=2026-07-22T00:00:00Z \
@@ -129,6 +130,7 @@ RELEASE_SOURCE_SHA=$source_sha \
RELEASE_BASE_SHA='' \
RELEASE_COMPONENTS=api \
RELEASE_MIGRATIONS_CHANGED=true \
RELEASE_MIGRATION_VERSION=0097_worker_runtime_owner_index \
RELEASE_API_IMAGE=$api_image \
RELEASE_WEB_IMAGE=$web_image \
RELEASE_BUILD_COMPLETED_AT=2026-07-22T00:00:00Z \
@@ -143,6 +145,17 @@ RELEASE_DEPLOY_DURATION_SECONDS=17 \
RELEASE_DEPLOY_ACTION=deploy \
node "$root/scripts/release-manifest.mjs" record-deployment "$manifest" "$recorded" >/dev/null
node "$root/scripts/release-manifest.mjs" validate "$recorded" >/dev/null
certified=$tmp/certified.json
RELEASE_CERTIFIED_AT=2026-07-22T00:02:00Z \
RELEASE_ACCEPTANCE_RUN_ID=11111111-1111-4111-8111-111111111111 \
RELEASE_CONFIG_HASH=$(printf '3%.0s' {1..64}) \
RELEASE_CAPACITY_PROFILE=P24 \
RELEASE_WORKER_INSTANCE_SLOTS=24 \
RELEASE_WORKER_MAX_REPLICAS_NINGBO=1 \
RELEASE_WORKER_MAX_REPLICAS_HONGKONG=2 \
node "$root/scripts/release-manifest.mjs" record-certification "$recorded" "$certified" >/dev/null
node "$root/scripts/release-manifest.mjs" validate "$certified" >/dev/null
[[ $(node "$root/scripts/release-manifest.mjs" get "$certified" certification.capacityProfile) == P24 ]]
sensitive=$tmp/sensitive.json
node -e 'const fs=require("fs"); const value=JSON.parse(fs.readFileSync(process.argv[1])); value.accessToken="forbidden"; fs.writeFileSync(process.argv[2], JSON.stringify(value))' "$manifest" "$sensitive"
@@ -234,6 +247,7 @@ RELEASE_SOURCE_SHA=cccccccccccccccccccccccccccccccccccccccc \
RELEASE_BASE_SHA=dddddddddddddddddddddddddddddddddddddddd \
RELEASE_COMPONENTS=web \
RELEASE_MIGRATIONS_CHANGED=false \
RELEASE_MIGRATION_VERSION=0097_worker_runtime_owner_index \
RELEASE_API_IMAGE=$api_image \
RELEASE_WEB_IMAGE=$web_image \
RELEASE_BUILD_COMPLETED_AT=2026-07-22T00:02:00Z \