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/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 (