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

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

Browse Source 
实现本地三节点 K3s 同构环境、脱敏生产快照、Gemini 图片和多参考图视频协议模拟、统一验收报告及故障注入。\n\n新增 Worker 容量控制器、资源与连接预算、任务恢复保护，并将生产验收拆分为 validation 执行和人工 CAS 放量。\n\n验证包括 Go 全量测试、PostgreSQL HTTP 集成测试、go vet、OpenAPI、ShellCheck、前端检查、迁移及发布脚本测试。
This commit is contained in:
wangbo
2026-07-31 18:02:24 +08:00
parent 93d36d0e55
commit e05922b0f4
94 changed files with 12379 additions and 826 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/cmd/acceptance-bootstrap/main.go
+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
}
apps/api/cmd/acceptance-callback-collector/main.go
+120
View File
@@ -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)
}
apps/api/cmd/acceptance-load/main.go
+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
}
apps/api/cmd/acceptance-load/main_test.go
+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")
}
}
apps/api/cmd/acceptance-snapshot/main.go
+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
}
apps/api/cmd/gateway/main.go
+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,
}
apps/api/docs/swagger.json
+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": {
apps/api/docs/swagger.yaml
+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: 返回客户端名称、英文名称、平台名和图标路径；数据库对象尚未创建时返回默认设置。
apps/api/internal/acceptanceemulator/server.go
+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 {
apps/api/internal/acceptanceemulator/server_test.go
+25 -3
View File
@@ -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)
}
apps/api/internal/acceptancesnapshot/snapshot.go
+1101
View File
File diff suppressed because it is too large Load Diff
apps/api/internal/acceptancesnapshot/snapshot_test.go
+180
View File
@@ -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
}
apps/api/internal/capacitycontroller/controller.go
+413
View File
@@ -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
}
}
apps/api/internal/capacitycontroller/kubernetes.go
+440
View File
@@ -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
}
apps/api/internal/capacitycontroller/kubernetes_test.go
+94
View File
@@ -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)
}
}
apps/api/internal/capacitycontroller/planner.go
+281
View File
@@ -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)
}
apps/api/internal/capacitycontroller/planner_test.go
+164
View File
@@ -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)
}
}
}
apps/api/internal/clients/gemini.go
+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
apps/api/internal/clients/types.go
+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.
apps/api/internal/clients/volces.go
+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
}
apps/api/internal/config/config.go
+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 {
apps/api/internal/config/config_test.go
+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)
apps/api/internal/httpapi/acceptance_handlers.go
+85 -9
View File
@@ -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
apps/api/internal/httpapi/acceptance_state_integration_test.go
+165
View File
@@ -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
}
apps/api/internal/httpapi/async_worker_acceptance_integration_test.go
+24 -14
View File
@@ -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()
apps/api/internal/httpapi/core_flow_integration_test.go
+160 -125
View File
@@ -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) {
apps/api/internal/httpapi/failure_policy_http_acceptance_integration_test.go
+11 -7
View File
@@ -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{
apps/api/internal/httpapi/gemini_base64_stress_integration_test.go
+2 -2
View File
@@ -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},
apps/api/internal/httpapi/gemini_compat.go
+5
View File
@@ -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
apps/api/internal/httpapi/handlers.go
+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)
apps/api/internal/httpapi/keling_compat_integration_test.go
+11 -7
View File
@@ -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)
apps/api/internal/httpapi/keling_simulation_integration_test.go
+11 -7
View File
@@ -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,
apps/api/internal/httpapi/kling_compat.go
+3
View File
@@ -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):
apps/api/internal/httpapi/kling_compat_integration_test.go
+11 -7
View File
@@ -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 {
apps/api/internal/httpapi/oidc_jit_integration_test.go
+22 -2
View File
@@ -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(),
apps/api/internal/httpapi/server.go
+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,
apps/api/internal/identity/configuration.go
+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...)
apps/api/internal/identity/configuration_test.go
+24
View File
@@ -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)
}
}
apps/api/internal/runner/admission.go
+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
}
apps/api/internal/runner/execution_lease.go
+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
}
apps/api/internal/runner/queue_worker.go
+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,
apps/api/internal/runner/service.go
+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
apps/api/internal/securityevents/metrics.go
+24
View File
@@ -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()},
apps/api/internal/store/acceptance.go
+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 {
apps/api/internal/store/acceptance_integration_test.go
+58 -3
View File
@@ -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)
apps/api/internal/store/admission_queue_integration_test.go
+47
View File
@@ -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, `
apps/api/internal/store/billing_v2_integration_test.go
+153 -1
View File
@@ -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()
apps/api/internal/store/capacity_profiles.go
+532
View File
@@ -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
}
apps/api/internal/store/capacity_profiles_test.go
+51
View File
@@ -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)
}
}
apps/api/internal/store/leadership.go
+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 {
apps/api/internal/store/oidc_users_integration_test.go
+10
View File
@@ -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)
}
}
apps/api/internal/store/postgres.go
+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
}
apps/api/internal/store/rate_limits.go
+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 {
apps/api/internal/store/runtime_types.go
+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 {
apps/api/internal/store/task_history.go
+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
apps/api/internal/store/tasks_runtime.go
+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
}
apps/api/internal/store/worker_registry.go
+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 (
apps/api/migrations/0096_worker_capacity_controller.sql
+39
View File
@@ -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;
apps/api/migrations/0097_worker_runtime_owner_index.sql
+5
View File
@@ -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;