fix(worker): 防止事务泄漏并恢复滞留队列

为 PostgreSQL 连接增加可配置的事务空闲与锁等待超时,并在请求取消或 Worker 退出后使用独立有界上下文回滚事务。\n\n恢复过期任务时按批次使用 SKIP LOCKED,周期重建缺失或终态 River Job;锁等待超时只在确认尚未提交上游时安全重排,避免重复执行和重复结算。\n\n验证:完整 Go 测试、go vet、生产 Kustomize 渲染、gofmt 与 git diff --check 均通过。
This commit is contained in:
2026-07-30 17:27:11 +08:00
parent bc44af751e
commit 6bab0f0749
32 changed files with 545 additions and 106 deletions
+5 -1
View File
@@ -37,7 +37,11 @@ func main() {
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
db, err := store.ConnectWithMaxConns(ctx, cfg.DatabaseURL, cfg.DatabaseMaxConns)
db, err := store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
MaxConns: cfg.DatabaseMaxConns,
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
})
if err != nil {
stop()
logger.Error("connect postgres failed", "error", err)
+41 -14
View File
@@ -59,6 +59,8 @@ type Config struct {
BillingEngineMode string
ProcessRole string
DatabaseMaxConns int
DatabaseIdleInTransactionTimeoutSeconds int
DatabaseLockTimeoutSeconds int
MediaRequestConcurrency int
MediaMaterializationConcurrency int
AsyncQueueWorkerEnabled bool
@@ -106,20 +108,27 @@ func Load() Config {
TaskProgressCallbackURL: env("TASK_PROGRESS_CALLBACK_URL",
strings.TrimRight(env("SERVER_MAIN_BASE_URL", "http://localhost:3000"), "/")+"/internal/platform/task-progress-callbacks",
),
TaskProgressCallbackTimeoutMS: envIntValidated("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", 5000),
TaskProgressCallbackMaxAttempts: envIntValidated("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", 10),
TaskCleanupEnabled: env("AI_GATEWAY_TASK_CLEANUP_ENABLED", "false") == "true",
TaskRetentionDays: envIntValidated("AI_GATEWAY_TASK_RETENTION_DAYS", 30),
TaskAnalysisRetentionDays: envIntValidated("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS", 7),
TaskCleanupIntervalSeconds: envIntValidated("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS", 300),
TaskCleanupBatchSize: envIntValidated("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE", 1000),
CORSAllowedOrigin: env("CORS_ALLOWED_ORIGIN", "http://localhost:5178,http://127.0.0.1:5178"),
GlobalHTTPProxy: globalProxy.HTTPProxy,
GlobalHTTPProxySource: globalProxy.Source,
LogLevel: logLevel(env("LOG_LEVEL", "info")),
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),
TaskProgressCallbackTimeoutMS: envIntValidated("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", 5000),
TaskProgressCallbackMaxAttempts: envIntValidated("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", 10),
TaskCleanupEnabled: env("AI_GATEWAY_TASK_CLEANUP_ENABLED", "false") == "true",
TaskRetentionDays: envIntValidated("AI_GATEWAY_TASK_RETENTION_DAYS", 30),
TaskAnalysisRetentionDays: envIntValidated("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS", 7),
TaskCleanupIntervalSeconds: envIntValidated("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS", 300),
TaskCleanupBatchSize: envIntValidated("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE", 1000),
CORSAllowedOrigin: env("CORS_ALLOWED_ORIGIN", "http://localhost:5178,http://127.0.0.1:5178"),
GlobalHTTPProxy: globalProxy.HTTPProxy,
GlobalHTTPProxySource: globalProxy.Source,
LogLevel: logLevel(env("LOG_LEVEL", "info")),
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),
DatabaseIdleInTransactionTimeoutSeconds: envOptionalIntValidated(
"AI_GATEWAY_DATABASE_IDLE_IN_TRANSACTION_TIMEOUT_SECONDS",
60,
),
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),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
@@ -138,6 +147,12 @@ func (c Config) Validate() error {
if c.DatabaseMaxConns < 0 || c.DatabaseMaxConns > 1000 {
return errors.New("AI_GATEWAY_DATABASE_MAX_CONNS must be between 1 and 1000 when configured")
}
if c.DatabaseIdleInTransactionTimeoutSeconds < 0 || c.DatabaseIdleInTransactionTimeoutSeconds > 3600 {
return errors.New("AI_GATEWAY_DATABASE_IDLE_IN_TRANSACTION_TIMEOUT_SECONDS must be between 0 and 3600")
}
if c.DatabaseLockTimeoutSeconds < 0 || c.DatabaseLockTimeoutSeconds > 3600 {
return errors.New("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS must be between 0 and 3600")
}
if c.MediaMaterializationConcurrency != 0 && (c.MediaMaterializationConcurrency < 1 || c.MediaMaterializationConcurrency > 256) {
return errors.New("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY must be between 1 and 256")
}
@@ -336,6 +351,18 @@ func envIntValidated(key string, fallback int) int {
return parsed
}
func envOptionalIntValidated(key string, fallback int) int {
value := envValue(key)
if value == "" {
return fallback
}
parsed, err := strconv.Atoi(value)
if err != nil {
return -1
}
return parsed
}
func envInt64Validated(key string, fallback int64) int64 {
value := envValue(key)
if value == "" {
+22
View File
@@ -144,6 +144,13 @@ func TestProcessRolePrecedenceAndCompatibility(t *testing.T) {
func TestValidateProcessRoleAndDatabasePool(t *testing.T) {
cfg := Load()
if cfg.DatabaseIdleInTransactionTimeoutSeconds != 60 || cfg.DatabaseLockTimeoutSeconds != 30 {
t.Fatalf(
"database transaction timeouts = %d/%d, want 60/30",
cfg.DatabaseIdleInTransactionTimeoutSeconds,
cfg.DatabaseLockTimeoutSeconds,
)
}
cfg.ProcessRole = "invalid"
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "PROCESS_ROLE") {
t.Fatalf("Validate() error = %v, want invalid process role", err)
@@ -154,6 +161,21 @@ func TestValidateProcessRoleAndDatabasePool(t *testing.T) {
t.Fatalf("Validate() error = %v, want invalid database max conns", err)
}
cfg.DatabaseMaxConns = 16
cfg.DatabaseIdleInTransactionTimeoutSeconds = 3601
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDLE_IN_TRANSACTION") {
t.Fatalf("Validate() error = %v, want invalid idle transaction timeout", err)
}
cfg.DatabaseIdleInTransactionTimeoutSeconds = 60
cfg.DatabaseLockTimeoutSeconds = 3601
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "LOCK_TIMEOUT") {
t.Fatalf("Validate() error = %v, want invalid lock timeout", err)
}
cfg.DatabaseLockTimeoutSeconds = 30
t.Setenv("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", "not-an-integer")
if err := Load().Validate(); err == nil || !strings.Contains(err.Error(), "LOCK_TIMEOUT") {
t.Fatalf("Validate() error = %v, want invalid non-integer lock timeout", err)
}
t.Setenv("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", "30")
cfg.MediaMaterializationConcurrency = 257
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_MATERIALIZATION_CONCURRENCY") {
t.Fatalf("Validate() error = %v, want invalid media materialization concurrency", err)
+38
View File
@@ -80,6 +80,38 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs
w.service.logger.Debug("river async task interrupted and requeued", "taskID", task.ID, "status", queued.Status, "riverJobID", job.ID)
return river.JobSnooze(0)
}
if store.IsPostgresLockTimeout(runErr) {
queued, changed, queueErr := w.service.store.RequeueTaskBeforeUpstreamSubmission(
context.WithoutCancel(ctx),
task.ID,
executionToken,
time.Second,
)
if queueErr != nil {
return queueErr
}
if changed {
if eventErr := w.service.emit(
context.WithoutCancel(ctx),
task.ID,
"task.queued",
"queued",
"database_lock_timeout",
0.2,
"async task queued after database lock timeout",
map[string]any{"code": "database_lock_timeout"},
task.RunMode == "simulation",
); eventErr != nil {
w.service.logger.Warn("record database lock timeout requeue event failed", "taskID", task.ID, "error", eventErr)
}
w.service.logger.Warn("river async task requeued after database lock timeout",
"taskID", task.ID,
"status", queued.Status,
"riverJobID", job.ID,
)
return river.JobSnooze(time.Second)
}
}
w.service.logger.Warn("river async task completed with failure", "taskID", task.ID, "error", runErr, "riverJobID", job.ID)
return nil
}
@@ -570,6 +602,12 @@ func (s *Service) recoverOrphanedAsyncRiverJobs(ctx context.Context) {
if recovered > 0 {
s.logger.Warn("orphaned river jobs recovered", "count", recovered)
}
if err := s.recoverAsyncRiverJobs(ctx); err != nil {
if ctx.Err() == nil {
s.logger.Warn("recover queued tasks without active river jobs failed", "error", err)
}
return
}
}
recoverJobs()
ticker := time.NewTicker(orphanedRiverJobScanInterval)
+1 -1
View File
@@ -153,7 +153,7 @@ func (s *Store) BatchAccessRules(ctx context.Context, input AccessRuleBatchInput
if err != nil {
return nil, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
for _, resource := range dedupeAccessRuleResources(input.DeleteResources) {
resource = normalizeAccessRuleResource(resource, input.Effect)
+5 -5
View File
@@ -134,7 +134,7 @@ func (s *Store) QueueTaskAdmissionWithHook(
if err != nil {
return TaskAdmission{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil {
return TaskAdmission{}, err
@@ -232,7 +232,7 @@ func (s *Store) tryTaskAdmission(
if err != nil {
return TaskAdmissionResult{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil {
return TaskAdmissionResult{}, err
@@ -456,7 +456,7 @@ func (s *Store) RebindWaitingTaskAdmission(ctx context.Context, input TaskAdmiss
if err != nil {
return TaskAdmission{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil {
return TaskAdmission{}, err
}
@@ -633,7 +633,7 @@ func (s *Store) DeleteTaskAdmission(ctx context.Context, taskID string) error {
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil {
return err
}
@@ -794,7 +794,7 @@ func (s *Store) ReapExpiredTaskAdmissions(ctx context.Context, limit int) (TaskA
if err != nil {
return TaskAdmissionReapResult{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
rows, err := tx.Query(ctx, `
SELECT task_id::text
FROM gateway_task_admissions
@@ -543,6 +543,160 @@ SELECT
t.Fatalf("recovered task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases)
}
lockedRecoveryTask := createTask(true)
unlockedRecoveryTask := createTask(true)
if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
execution_token = gen_random_uuid(),
execution_lease_expires_at = now() - interval '1 second'
WHERE id = ANY($1::uuid[])`, []string{lockedRecoveryTask.ID, unlockedRecoveryTask.ID}); err != nil {
t.Fatalf("prepare locked recovery tasks: %v", err)
}
lockTx, err := second.pool.Begin(ctx)
if err != nil {
t.Fatalf("begin task row lock: %v", err)
}
var lockedTaskID string
if err := lockTx.QueryRow(ctx, `
SELECT id::text
FROM gateway_tasks
WHERE id = $1::uuid
FOR UPDATE`, lockedRecoveryTask.ID).Scan(&lockedTaskID); err != nil {
rollbackTransaction(lockTx)
t.Fatalf("lock interrupted task row: %v", err)
}
recoveryCtx, recoveryCancel := context.WithTimeout(ctx, 3*time.Second)
recovery, err = first.RecoverInterruptedRuntimeState(recoveryCtx)
recoveryCancel()
if err != nil {
rollbackTransaction(lockTx)
t.Fatalf("recover unlocked task while another row is locked: %v", err)
}
if recovery.RequeuedAsyncTasks < 1 {
rollbackTransaction(lockTx)
t.Fatalf("SKIP LOCKED recovery = %+v, want at least one requeued task", recovery)
}
var lockedStatus, unlockedStatus string
if err := first.pool.QueryRow(ctx, `
SELECT
(SELECT status FROM gateway_tasks WHERE id = $1::uuid),
(SELECT status FROM gateway_tasks WHERE id = $2::uuid)`,
lockedRecoveryTask.ID,
unlockedRecoveryTask.ID,
).Scan(&lockedStatus, &unlockedStatus); err != nil {
rollbackTransaction(lockTx)
t.Fatalf("read SKIP LOCKED recovery states: %v", err)
}
if lockedStatus != "running" || unlockedStatus != "queued" {
rollbackTransaction(lockTx)
t.Fatalf("SKIP LOCKED recovery statuses = %s/%s, want running/queued", lockedStatus, unlockedStatus)
}
rollbackTransaction(lockTx)
recovery, err = first.RecoverInterruptedRuntimeState(ctx)
if err != nil {
t.Fatalf("recover previously locked task: %v", err)
}
if recovery.RequeuedAsyncTasks < 1 {
t.Fatalf("unlocked recovery = %+v, want requeued task", recovery)
}
if err := first.pool.QueryRow(ctx, `
SELECT status
FROM gateway_tasks
WHERE id = $1::uuid`, lockedRecoveryTask.ID).Scan(&lockedStatus); err != nil {
t.Fatalf("read unlocked recovery status: %v", err)
}
if lockedStatus != "queued" {
t.Fatalf("unlocked recovery status = %s, want queued", lockedStatus)
}
preSubmissionTask := createTask(true)
preSubmissionToken := uuid.NewString()
if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
execution_token = $2::uuid,
execution_lease_expires_at = now() + interval '5 minutes'
WHERE id = $1::uuid`,
preSubmissionTask.ID,
preSubmissionToken,
); err != nil {
t.Fatalf("prepare pre-submission lock timeout 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', 'not_submitted')`,
preSubmissionTask.ID,
); err != nil {
t.Fatalf("create pre-submission attempt: %v", err)
}
requeuedTask, changed, err := first.RequeueTaskBeforeUpstreamSubmission(
ctx,
preSubmissionTask.ID,
preSubmissionToken,
time.Second,
)
if err != nil || !changed || requeuedTask.Status != "queued" {
t.Fatalf("pre-submission lock timeout requeue task=%+v changed=%v err=%v", requeuedTask, changed, err)
}
var preSubmissionAttempts int
if err := first.pool.QueryRow(ctx, `
SELECT count(*)
FROM gateway_task_attempts
WHERE task_id = $1::uuid`, preSubmissionTask.ID).Scan(&preSubmissionAttempts); err != nil {
t.Fatalf("count cleaned pre-submission attempts: %v", err)
}
if preSubmissionAttempts != 0 {
t.Fatalf("pre-submission lock timeout left %d attempts, want 0", preSubmissionAttempts)
}
ambiguousTask := createTask(true)
ambiguousToken := uuid.NewString()
if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
execution_token = $2::uuid,
execution_lease_expires_at = now() + interval '5 minutes'
WHERE id = $1::uuid`,
ambiguousTask.ID,
ambiguousToken,
); err != nil {
t.Fatalf("prepare ambiguous submission 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')`,
ambiguousTask.ID,
); err != nil {
t.Fatalf("create ambiguous submission attempt: %v", err)
}
_, changed, err = first.RequeueTaskBeforeUpstreamSubmission(
ctx,
ambiguousTask.ID,
ambiguousToken,
time.Second,
)
if err != nil {
t.Fatalf("guard ambiguous submission task requeue: %v", err)
}
if changed {
t.Fatal("ambiguous submission task was requeued")
}
var ambiguousStatus string
if err := first.pool.QueryRow(ctx, `
SELECT status
FROM gateway_tasks
WHERE id = $1::uuid`, ambiguousTask.ID).Scan(&ambiguousStatus); err != nil {
t.Fatalf("read ambiguous submission task: %v", err)
}
if ambiguousStatus != "running" {
t.Fatalf("ambiguous submission task status = %s, want running", ambiguousStatus)
}
terminalResidue := createTask(true)
result, err = first.TryTaskAdmission(ctx, inputFor(terminalResidue, 100, ""))
if err != nil || !result.Admitted {
+2 -2
View File
@@ -101,7 +101,7 @@ func (s *Store) CreateBaseModel(ctx context.Context, input BaseModelInput) (Base
if err != nil {
return BaseModel{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
item, err := scanBaseModel(tx.QueryRow(ctx, `
INSERT INTO base_model_catalog (
@@ -162,7 +162,7 @@ func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelI
if err != nil {
return BaseModel{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
item, err := scanBaseModel(tx.QueryRow(ctx, `
UPDATE base_model_catalog
@@ -157,7 +157,7 @@ RETURNING outbox.id::text, outbox.task_id::text, outbox.action, outbox.amount::t
}
func (s *Store) ProcessBillingSettlement(ctx context.Context, settlement BillingSettlement) error {
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
return s.beginTransaction(ctx, func(tx pgx.Tx) error {
var action string
var amount string
var currency string
@@ -479,7 +479,7 @@ func (s *Store) MarkBillingSettlementFailed(ctx context.Context, settlement Bill
taskStatus = "manual_review"
manualReason = "maximum settlement attempts exceeded"
}
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
return s.beginTransaction(ctx, func(tx pgx.Tx) error {
tag, err := tx.Exec(ctx, `
UPDATE settlement_outbox
SET status = $3,
@@ -19,7 +19,7 @@ func (s *Store) ListTaskBinaryResultBackfillBatch(ctx context.Context, afterID s
batchSize = 100
}
var items []TaskBinaryResultBackfillItem
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil {
return err
}
@@ -64,7 +64,7 @@ func (s *Store) UpdateTaskBinaryResultBackfill(ctx context.Context, taskID strin
return false, err
}
updated := false
err = pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err = s.beginTransaction(ctx, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil {
return err
}
+1 -1
View File
@@ -62,7 +62,7 @@ func (s *Store) UpsertConversationMessages(ctx context.Context, conversationID s
if err != nil {
return nil, 0, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
refs := make([]TaskMessageRefInput, 0, len(messages))
newCount := 0
+1 -1
View File
@@ -272,7 +272,7 @@ func (s *Store) DeleteUserGroup(ctx context.Context, id string) error {
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `
DELETE FROM gateway_access_rules
@@ -59,7 +59,7 @@ func (s *Store) beginIdentityConfigurationLifecycleTx(ctx context.Context) (pgx.
return nil, err
}
if err := lockIdentityConfigurationLifecycle(ctx, tx); err != nil {
_ = tx.Rollback(ctx)
rollbackTransaction(tx)
return nil, err
}
return tx, nil
@@ -171,7 +171,7 @@ func (s *Store) ApplyIdentityManifest(ctx context.Context, id string, expectedVe
if err != nil {
return identity.Revision{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
current, err := scanIdentityRevision(tx.QueryRow(ctx, `SELECT `+identityRevisionColumns+`
FROM gateway_identity_configuration_revisions WHERE id=$1::uuid FOR UPDATE`, id))
if err != nil {
@@ -282,7 +282,7 @@ func (s *Store) ActivateIdentityRevision(ctx context.Context, id string, expecte
if err != nil {
return identity.Revision{}, false, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var foreignPairingReservation bool
if err := tx.QueryRow(ctx, `SELECT EXISTS(
SELECT 1 FROM gateway_identity_pairing_start_reservation
@@ -364,7 +364,7 @@ func (s *Store) DisableActiveIdentityRevision(ctx context.Context, expectedVersi
if err != nil {
return identity.Revision{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if ok, err := hasBreakGlassManager(ctx, tx); err != nil {
return identity.Revision{}, err
} else if !ok {
+5 -5
View File
@@ -20,7 +20,7 @@ func (s *Store) ReserveIdentityPairingStart(ctx context.Context, attemptID strin
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `DELETE FROM gateway_identity_pairing_start_reservation
WHERE state='starting' AND expires_at <= now()`); err != nil {
return err
@@ -68,7 +68,7 @@ func (s *Store) CommitIdentityPairingStart(ctx context.Context, revision identit
if err != nil {
return identity.PairingExchange{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var reserved bool
if err := tx.QueryRow(ctx, `SELECT true FROM gateway_identity_pairing_start_reservation
WHERE singleton=true AND attempt_id=$1::uuid AND state='starting' FOR UPDATE`, exchange.ID).Scan(&reserved); errors.Is(err, pgx.ErrNoRows) {
@@ -256,7 +256,7 @@ func (s *Store) UpdateIdentityPairingExchange(ctx context.Context, id string, ex
if err != nil {
return identity.PairingExchange{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
exchange, err := scanIdentityPairing(tx.QueryRow(ctx, `
UPDATE gateway_identity_onboarding_exchanges SET status=$3,remote_version=$4,last_error_category=NULLIF($5,''),
auth_center_audit_id=COALESCE(NULLIF($6,''),auth_center_audit_id),version=version+1,updated_at=now()
@@ -297,7 +297,7 @@ func (s *Store) CancelIdentityPairingExchange(ctx context.Context, id string, ex
if err != nil {
return identity.PairingExchange{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
exchange, err := scanIdentityPairing(tx.QueryRow(ctx, `SELECT `+identityPairingColumns+`
FROM gateway_identity_onboarding_exchanges WHERE id=$1::uuid FOR UPDATE`, id))
if errors.Is(err, pgx.ErrNoRows) {
@@ -352,7 +352,7 @@ func (s *Store) CompleteIdentityPairingCleanup(ctx context.Context, id string, e
if err != nil {
return identity.PairingExchange{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var revisionID string
if err := tx.QueryRow(ctx, `SELECT revision_id::text FROM gateway_identity_onboarding_exchanges
WHERE id=$1::uuid AND version=$2 AND status='cancelled' AND cleanup_status='pending' FOR UPDATE`, id, expectedVersion).Scan(&revisionID); errors.Is(err, pgx.ErrNoRows) {
@@ -24,7 +24,7 @@ func (s *Store) resolveOrProvisionOIDCMultiTenantUser(ctx context.Context, input
if err != nil {
return ResolveOrProvisionOIDCUserResult{}, err
}
defer func() { _ = tx.Rollback(ctx) }()
defer rollbackTransaction(tx)
bindingID, gatewayTenantID, tenantKey, groupID, groupKey, accessStatus, err :=
loadOIDCTenantBinding(ctx, tx, input)
@@ -78,7 +78,7 @@ func (s *Store) ApplyOIDCTenantBindingSync(ctx context.Context, bindingID string
if err != nil {
return err
}
defer func() { _ = tx.Rollback(ctx) }()
defer rollbackTransaction(tx)
if unchanged {
tag, err := tx.Exec(ctx, `UPDATE gateway_oidc_tenant_bindings
SET last_sync_at=$2::timestamptz,next_sync_at=$2::timestamptz+interval '15 minutes',sync_failure_count=0,
@@ -129,7 +129,7 @@ func (s *Store) RejectOIDCTenantBinding(ctx context.Context, bindingID, category
if err != nil {
return err
}
defer func() { _ = tx.Rollback(ctx) }()
defer rollbackTransaction(tx)
tag, err := tx.Exec(ctx, `UPDATE gateway_oidc_tenant_bindings
SET access_status='disabled',metadata_status='rejected',last_sync_at=$2::timestamptz,
next_sync_at=$2::timestamptz+interval '15 minutes',
+1 -1
View File
@@ -69,7 +69,7 @@ func (s *Store) ResolveOrProvisionOIDCUser(ctx context.Context, input ResolveOrP
if err != nil {
return ResolveOrProvisionOIDCUserResult{}, err
}
defer func() { _ = tx.Rollback(ctx) }()
defer rollbackTransaction(tx)
projection, err := loadOIDCUserProjection(ctx, tx, input.Subject)
if err == nil {
+2 -2
View File
@@ -47,7 +47,7 @@ func (s *Store) ReplacePlatformModels(ctx context.Context, platformID string, in
if err != nil {
return nil, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
keptIDs := make([]string, 0, len(inputs))
for _, input := range inputs {
@@ -372,7 +372,7 @@ func (s *Store) DeletePlatformModel(ctx context.Context, id string) error {
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
result, err := tx.Exec(ctx, `DELETE FROM platform_models WHERE id = $1::uuid`, id)
if err != nil {
+42 -8
View File
@@ -8,6 +8,7 @@ import (
"encoding/json"
"errors"
"net"
"strconv"
"strings"
"time"
"unicode"
@@ -29,6 +30,12 @@ const (
postgresPoolWarmLimit = 64
)
type PostgresPoolOptions struct {
MaxConns int
IdleInTransactionTimeout time.Duration
LockTimeout time.Duration
}
func defaultAPIKeyScopes() []string {
return []string{"chat", "embedding", "rerank", "image", "image_vectorize", "video", "video_enhance", "music", "audio", "voice_clone"}
}
@@ -81,7 +88,11 @@ func Connect(ctx context.Context, databaseURL string) (*Store, error) {
}
func ConnectWithMaxConns(ctx context.Context, databaseURL string, maxConns int) (*Store, error) {
config, err := postgresPoolConfig(databaseURL, maxConns)
return ConnectWithPoolOptions(ctx, databaseURL, PostgresPoolOptions{MaxConns: maxConns})
}
func ConnectWithPoolOptions(ctx context.Context, databaseURL string, options PostgresPoolOptions) (*Store, error) {
config, err := postgresPoolConfigWithOptions(databaseURL, options)
if err != nil {
return nil, err
}
@@ -97,19 +108,37 @@ func ConnectWithMaxConns(ctx context.Context, databaseURL string, maxConns int)
}
func postgresPoolConfig(databaseURL string, maxConns ...int) (*pgxpool.Config, error) {
options := PostgresPoolOptions{}
if len(maxConns) > 0 {
options.MaxConns = maxConns[0]
}
return postgresPoolConfigWithOptions(databaseURL, options)
}
func postgresPoolConfigWithOptions(databaseURL string, options PostgresPoolOptions) (*pgxpool.Config, error) {
config, err := pgxpool.ParseConfig(databaseURL)
if err != nil {
return nil, err
}
config.ConnConfig.ConnectTimeout = postgresConnectTimeout
config.ConnConfig.RuntimeParams["application_name"] = postgresApplicationName
if len(maxConns) > 0 && maxConns[0] > 0 {
config.MaxConns = int32(maxConns[0])
config.MinIdleConns = int32(min(maxConns[0], postgresPoolWarmLimit))
if options.IdleInTransactionTimeout > 0 {
config.ConnConfig.RuntimeParams["idle_in_transaction_session_timeout"] = postgresDuration(options.IdleInTransactionTimeout)
}
if options.LockTimeout > 0 {
config.ConnConfig.RuntimeParams["lock_timeout"] = postgresDuration(options.LockTimeout)
}
if options.MaxConns > 0 {
config.MaxConns = int32(options.MaxConns)
config.MinIdleConns = int32(min(options.MaxConns, postgresPoolWarmLimit))
}
return config, nil
}
func postgresDuration(value time.Duration) string {
return strconv.FormatInt(max(value.Milliseconds(), 1), 10) + "ms"
}
func warmPostgresPool(ctx context.Context, pool *pgxpool.Pool, count int) error {
if count <= 0 {
return pool.Ping(ctx)
@@ -152,6 +181,11 @@ func IsPostgresUnavailable(err error) bool {
return pgconn.SafeToRetry(err)
}
func IsPostgresLockTimeout(err error) bool {
var postgresError *pgconn.PgError
return errors.As(err, &postgresError) && postgresError.Code == "55P03"
}
func (s *Store) Close() {
s.pool.Close()
}
@@ -938,7 +972,7 @@ func (s *Store) DeletePlatform(ctx context.Context, id string) error {
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
rows, err := tx.Query(ctx, `SELECT id::text FROM platform_models WHERE platform_id = $1::uuid`, id)
if err != nil {
@@ -1638,7 +1672,7 @@ func (s *Store) DeleteAPIKey(ctx context.Context, apiKeyID string, user *auth.Us
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
result, err := tx.Exec(ctx, `
UPDATE gateway_api_keys
@@ -1792,7 +1826,7 @@ func (s *Store) RegisterLocalUser(ctx context.Context, input LocalRegisterInput)
if err != nil {
return GatewayUser{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var tenantID string
userGroupID := ""
@@ -2065,7 +2099,7 @@ func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput,
if err != nil {
return CreateTaskResult{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
task, err := scanGatewayTask(tx.QueryRow(ctx, `
INSERT INTO gateway_tasks (
@@ -55,6 +55,29 @@ func TestPostgresPoolConfigWarmsBoundedIdleConnections(t *testing.T) {
}
}
func TestPostgresPoolConfigSetsBoundedTransactionTimeouts(t *testing.T) {
config, err := postgresPoolConfigWithOptions(
"postgresql://gateway:password@localhost:5432/gateway?sslmode=disable",
PostgresPoolOptions{
MaxConns: 32,
IdleInTransactionTimeout: 60 * time.Second,
LockTimeout: 30 * time.Second,
},
)
if err != nil {
t.Fatalf("parse PostgreSQL pool config: %v", err)
}
if got := config.ConnConfig.RuntimeParams["idle_in_transaction_session_timeout"]; got != "60000ms" {
t.Fatalf("idle transaction timeout = %q, want 60000ms", got)
}
if got := config.ConnConfig.RuntimeParams["lock_timeout"]; got != "30000ms" {
t.Fatalf("lock timeout = %q, want 30000ms", got)
}
if config.MaxConns != 32 || config.MinIdleConns != 32 {
t.Fatalf("pool bounds max=%d minIdle=%d, want 32/32", config.MaxConns, config.MinIdleConns)
}
}
func TestIsPostgresUnavailableClassifiesConnectivityFailures(t *testing.T) {
for _, testCase := range []struct {
name string
@@ -75,3 +98,12 @@ func TestIsPostgresUnavailableClassifiesConnectivityFailures(t *testing.T) {
t.Fatal("SQL syntax error was incorrectly classified as PostgreSQL unavailable")
}
}
func TestIsPostgresLockTimeout(t *testing.T) {
if !IsPostgresLockTimeout(&pgconn.PgError{Code: "55P03"}) {
t.Fatal("lock timeout was not classified")
}
if IsPostgresLockTimeout(&pgconn.PgError{Code: "57014"}) {
t.Fatal("query cancellation was classified as lock timeout")
}
}
+2 -2
View File
@@ -113,7 +113,7 @@ func (s *Store) CreatePricingRuleSet(ctx context.Context, input PricingRuleSetIn
if err != nil {
return PricingRuleSet{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
item, err := scanPricingRuleSet(tx.QueryRow(ctx, `
@@ -141,7 +141,7 @@ func (s *Store) UpdatePricingRuleSet(ctx context.Context, id string, input Prici
if err != nil {
return PricingRuleSet{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
item, err := scanPricingRuleSet(tx.QueryRow(ctx, `
+1 -1
View File
@@ -98,7 +98,7 @@ func (s *Store) RestorePlatformModelRuntimeStatus(ctx context.Context, platformM
if err != nil {
return ModelRateLimitStatus{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var restoredModelID string
if err := tx.QueryRow(ctx, `
+21 -10
View File
@@ -19,6 +19,8 @@ type RuntimeRecoveryResult struct {
CleanedTaskAdmissions int64 `json:"cleanedTaskAdmissions"`
}
const runtimeRecoveryBatchSize = 100
var ErrConcurrencyLeaseLost = errors.New("concurrency lease lost")
// CheckRateLimits performs a non-consuming admission preflight for fixed-window
@@ -104,7 +106,7 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
if err != nil {
return RateLimitResult{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
lockKeys := make([]string, 0)
lockKeySet := make(map[string]struct{})
@@ -414,7 +416,7 @@ func (s *Store) AttachRateLimitResultToAttempt(ctx context.Context, attemptID st
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
for _, reservation := range result.Reservations {
if reservation.ReservationID == "" {
@@ -458,7 +460,7 @@ func (s *Store) RecoverInterruptedRuntimeState(ctx context.Context) (RuntimeReco
if err != nil {
return RuntimeRecoveryResult{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-runtime-recovery', 0))`); err != nil {
return RuntimeRecoveryResult{}, err
}
@@ -542,7 +544,18 @@ WHERE status = 'running'
result.FailedAttempts = tag.RowsAffected()
asyncTaskRows, err := tx.Query(ctx, `
UPDATE gateway_tasks
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
SET status = 'queued',
error = NULL,
error_code = NULL,
@@ -555,11 +568,9 @@ SET status = 'queued',
next_run_at = now(),
finished_at = NULL,
updated_at = now()
WHERE async_mode = true
AND status = 'running'
AND execution_lease_expires_at IS NOT NULL
AND execution_lease_expires_at <= now()
RETURNING id::text`)
FROM recoverable_async_tasks recoverable
WHERE task.id = recoverable.id
RETURNING task.id::text`, runtimeRecoveryBatchSize)
if err != nil {
return RuntimeRecoveryResult{}, err
}
@@ -718,7 +729,7 @@ func (s *Store) finishRateLimitReservations(ctx context.Context, reservations []
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
for _, reservation := range reservations {
if reservation.ReservationID == "" {
+1 -3
View File
@@ -248,9 +248,7 @@ func (s *Store) ApplyCandidateFailureEffect(ctx context.Context, input Candidate
if err != nil {
return CandidateFailureEffectResult{}, err
}
defer func() {
_ = tx.Rollback(ctx)
}()
defer rollbackTransaction(tx)
lockKey := strings.TrimSpace(input.RequestedModel) + "\x1f" + strings.TrimSpace(input.ModelType)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1::text, 0))`, lockKey); err != nil {
@@ -185,7 +185,7 @@ func (s *Store) CreateSecurityEventConnection(ctx context.Context, input CreateS
if err != nil {
return SecurityEventConnection{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
_, err = tx.Exec(ctx, `
INSERT INTO gateway_security_event_connections(
connection_id,transmitter_issuer,endpoint_url,credential_ref,management_client_id,management_credential_ref,lifecycle_status,idempotency_key
@@ -196,7 +196,7 @@ INSERT INTO gateway_security_event_connections(
)
if err != nil {
if isUniqueViolation(err) {
_ = tx.Rollback(ctx)
rollbackTransaction(tx)
existing, getErr := s.SecurityEventConnection(ctx)
if getErr == nil && existing.IdempotencyKey == input.IdempotencyKey && existing.TransmitterIssuer == input.TransmitterIssuer {
return existing, nil
@@ -256,7 +256,7 @@ func (s *Store) SetSecurityEventManagementCredential(ctx context.Context, connec
if err != nil {
return SecurityEventConnection{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var oldReference sql.NullString
var lifecycle string
var version int64
@@ -315,7 +315,7 @@ func (s *Store) SetSecurityEventNextCredential(ctx context.Context, connectionID
if err != nil {
return SecurityEventConnection{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
tag, err := tx.Exec(ctx, `UPDATE gateway_security_event_connections
SET next_credential_ref=$2,lifecycle_status='rotating',last_error_category=NULL,version=version+1,updated_at=now()
WHERE connection_id=$1::uuid AND version=$3 AND next_credential_ref IS NULL
@@ -340,7 +340,7 @@ func (s *Store) PromoteSecurityEventCredential(ctx context.Context, connectionID
if err != nil {
return SecurityEventConnection{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var oldReference string
if err := tx.QueryRow(ctx, `SELECT credential_ref FROM gateway_security_event_connections
WHERE connection_id=$1::uuid AND next_credential_ref=$2 AND lifecycle_status='rotating' FOR UPDATE`, connectionID, nextReference).Scan(&oldReference); errors.Is(err, pgx.ErrNoRows) {
@@ -381,7 +381,7 @@ func (s *Store) DiscardPreparedSecurityEventConnection(ctx context.Context, conn
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var credentialReference string
var nextReference, managementReference sql.NullString
@@ -435,7 +435,7 @@ func (s *Store) FinalizeRetiringSecurityEventConnection(ctx context.Context, con
if err != nil {
return err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var credentialReference string
var nextReference, managementReference sql.NullString
+1 -7
View File
@@ -13,7 +13,6 @@ import (
const (
securityEventTransactionLockTimeout = 5 * time.Second
securityEventTransactionIdleTimeout = 15 * time.Second
transactionRollbackTimeout = 5 * time.Second
)
func (s *Store) beginSecurityEventTransaction(ctx context.Context) (pgx.Tx, error) {
@@ -38,12 +37,7 @@ func (s *Store) beginSecurityEventTransaction(ctx context.Context) (pgx.Tx, erro
}
func rollbackSecurityEventTransaction(tx pgx.Tx) {
if tx == nil {
return
}
ctx, cancel := context.WithTimeout(context.Background(), transactionRollbackTimeout)
defer cancel()
_ = tx.Rollback(ctx)
rollbackTransaction(tx)
}
type ApplySessionRevokedInput struct {
+1 -1
View File
@@ -174,7 +174,7 @@ func (s *Store) CleanupTaskHistory(ctx context.Context, analysisCutoff time.Time
batchSize = 1000
}
var result TaskHistoryCleanupResult
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil {
return err
}
+106 -17
View File
@@ -286,7 +286,7 @@ func (s *Store) ClaimTaskExecution(ctx context.Context, taskID string, execution
}
var task GatewayTask
manualReview := false
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
var queuedReady bool
var runningExpired bool
var production bool
@@ -573,6 +573,86 @@ WHERE id = $1::uuid
RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt, strings.TrimSpace(queueKey)))
}
func (s *Store) RequeueTaskBeforeUpstreamSubmission(
ctx context.Context,
taskID string,
executionToken string,
delay time.Duration,
) (GatewayTask, bool, error) {
if delay < time.Second {
delay = time.Second
}
if delay > 10*time.Minute {
delay = 10 * time.Minute
}
nextRunAt := time.Now().Add(delay)
var task GatewayTask
changed := false
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
var err error
task, err = scanGatewayTask(tx.QueryRow(ctx, `
UPDATE gateway_tasks task
SET status = 'queued',
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
next_run_at = $3::timestamptz,
error = NULL,
error_code = NULL,
error_message = NULL,
updated_at = now()
WHERE task.id = $1::uuid
AND task.status IN ('queued', 'running')
AND task.execution_token = $2::uuid
AND COALESCE(task.remote_task_id, '') = ''
AND NOT EXISTS (
SELECT 1
FROM gateway_task_attempts attempt
WHERE attempt.task_id = task.id
AND COALESCE(attempt.upstream_submission_status, 'not_submitted') <> 'not_submitted'
)
RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt))
if errors.Is(err, pgx.ErrNoRows) {
return nil
}
if err != nil {
return err
}
changed = true
if _, err := tx.Exec(ctx, `
DELETE FROM gateway_task_param_preprocessing_logs log
USING gateway_task_attempts attempt
WHERE log.attempt_id = attempt.id
AND attempt.task_id = $1::uuid
AND COALESCE(attempt.upstream_submission_status, 'not_submitted') = 'not_submitted'`, taskID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `
DELETE FROM gateway_task_attempts
WHERE task_id = $1::uuid
AND COALESCE(upstream_submission_status, 'not_submitted') = 'not_submitted'`, taskID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE task_id = $1::uuid
AND released_at IS NULL`, taskID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
return err
}
return notifyTaskAdmissionTx(ctx, tx, "*")
})
if err != nil {
return GatewayTask{}, false, err
}
return task, changed, nil
}
func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error {
if riverJobID <= 0 {
return nil
@@ -587,7 +667,7 @@ WHERE id = $1::uuid`, taskID, riverJobID)
func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, executionToken string, attemptID string, remoteTaskID string, payload map[string]any) error {
payloadJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(payload)))
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
return s.beginTransaction(ctx, func(tx pgx.Tx) error {
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET remote_task_id = NULLIF($3::text, ''),
@@ -644,7 +724,7 @@ func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message str
message = truncateUTF8Bytes(message, 2048)
var task GatewayTask
changed := false
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil {
return err
}
@@ -705,7 +785,7 @@ func (s *Store) CancelTaskBeforeUpstreamSubmission(
message = truncateUTF8Bytes(message, 2048)
var task GatewayTask
changed := false
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil {
return err
}
@@ -789,7 +869,7 @@ func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executio
message = truncateUTF8Bytes(message, 2048)
var task GatewayTask
changed := false
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
var err error
task, err = scanGatewayTask(tx.QueryRow(ctx, `
UPDATE gateway_tasks
@@ -851,14 +931,23 @@ func (s *Store) ListRecoverableAsyncTasks(ctx context.Context, limit int) ([]Asy
limit = 500
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, priority, next_run_at
FROM gateway_tasks
WHERE async_mode = true
SELECT task.id::text, task.priority, task.next_run_at
FROM gateway_tasks task
LEFT JOIN river_job job ON job.id = task.river_job_id
WHERE task.async_mode = true
AND (
status = 'queued'
OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()))
task.status = 'queued'
OR (
task.status = 'running'
AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
)
)
ORDER BY priority ASC, created_at ASC
AND (
task.river_job_id IS NULL
OR job.id IS NULL
OR job.state NOT IN ('available', 'pending', 'retryable', 'running', 'scheduled')
)
ORDER BY task.priority ASC, task.created_at ASC
LIMIT $1`, limit)
if err != nil {
return nil, err
@@ -883,7 +972,7 @@ func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptIn
if err != nil {
return "", err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
var attemptID string
err = tx.QueryRow(ctx, `
@@ -1301,7 +1390,7 @@ func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessIn
finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64)
}
currency := normalizeWalletCurrency(input.BillingCurrency)
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
if strings.TrimSpace(input.AttemptID) != "" {
if _, err := tx.Exec(ctx, `
UPDATE gateway_task_attempts
@@ -1430,7 +1519,7 @@ func (s *Store) FinishTaskManualReview(ctx context.Context, input FinishTaskManu
}
resultJSON, _ := json.Marshal(resultReport.Value)
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot)))
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
if strings.TrimSpace(input.AttemptID) != "" {
attemptStatus := "failed"
if status == "succeeded" {
@@ -1532,7 +1621,7 @@ func (s *Store) SettleTaskBilling(ctx context.Context, task GatewayTask) error {
"billingSummary": task.BillingSummary,
}
metadata, _ := json.Marshal(sanitizeJSONForStorage(metadataMap))
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
return s.beginTransaction(ctx, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_wallet_accounts (
gateway_tenant_id, gateway_user_id, tenant_id, tenant_key, user_id, currency
@@ -1662,7 +1751,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)
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
@@ -1817,7 +1906,7 @@ func (s *Store) AddTaskEvent(ctx context.Context, taskID string, eventType strin
if err != nil {
return TaskEvent{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT 1 FROM gateway_tasks WHERE id = $1::uuid FOR UPDATE`, taskID); err != nil {
return TaskEvent{}, err
}
+35 -1
View File
@@ -2,18 +2,52 @@ package store
import (
"context"
"errors"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
)
const transactionRollbackTimeout = 5 * time.Second
type Tx interface {
Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}
func (s *Store) InTx(ctx context.Context, fn func(Tx) error) error {
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
return s.beginTransaction(ctx, func(tx pgx.Tx) error {
return fn(tx)
})
}
func (s *Store) beginTransaction(ctx context.Context, fn func(pgx.Tx) error) (err error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return err
}
// A Worker shutdown cancels ctx before in-flight callbacks return. Use an
// independent bounded context so rollback still reaches PostgreSQL.
defer func() {
rollbackCtx, cancel := context.WithTimeout(context.Background(), transactionRollbackTimeout)
defer cancel()
rollbackErr := tx.Rollback(rollbackCtx)
if err == nil && rollbackErr != nil && !errors.Is(rollbackErr, pgx.ErrTxClosed) {
err = rollbackErr
}
}()
if err = fn(tx); err != nil {
return err
}
return tx.Commit(ctx)
}
func rollbackTransaction(tx pgx.Tx) {
if tx == nil {
return
}
ctx, cancel := context.WithTimeout(context.Background(), transactionRollbackTimeout)
defer cancel()
_ = tx.Rollback(ctx)
}
+3 -3
View File
@@ -163,7 +163,7 @@ func (s *Store) ReserveTaskBilling(ctx context.Context, task GatewayTask, user *
reservations := make([]WalletBillingReservation, 0, len(amounts))
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(pricingSnapshot))
requestFingerprint := walletString(pricingSnapshot["requestFingerprint"])
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
for currency, rawAmount := range amounts {
amount := roundMoney(rawAmount)
if amount <= 0 {
@@ -309,7 +309,7 @@ func (s *Store) reserveTaskBillingExact(ctx context.Context, task GatewayTask, g
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(pricingSnapshot))
requestFingerprint := walletString(pricingSnapshot["requestFingerprint"])
var reservations []WalletBillingReservation
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
var positive bool
if err := tx.QueryRow(ctx, `SELECT $1::numeric(38, 9) > 0`, amount).Scan(&positive); err != nil {
return fmt.Errorf("invalid exact reservation amount: %w", err)
@@ -408,7 +408,7 @@ func (s *Store) ReleaseTaskBillingReservations(ctx context.Context, reservations
if reason == "" {
reason = "task_not_settled"
}
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
return s.beginTransaction(ctx, func(tx pgx.Tx) error {
taskIDs := map[string]struct{}{}
for _, reservation := range reservations {
if strings.TrimSpace(reservation.AccountID) == "" {
+2 -2
View File
@@ -70,7 +70,7 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
if err != nil {
return WorkerAllocation{}, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
// Worker heartbeats and allocations are ephemeral coordination state that is
// refreshed every few seconds. Do not hold the global allocation lock while
// waiting for a synchronous replica to acknowledge the commit.
@@ -245,7 +245,7 @@ func (s *Store) RecoverOrphanedAsyncRiverJobs(
if err != nil {
return 0, err
}
defer tx.Rollback(ctx)
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-river-orphan-recovery', 0))`); err != nil {
return 0, err
}
@@ -24,6 +24,8 @@ data:
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS: "7"
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS: "300"
AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE: "1000"
AI_GATEWAY_DATABASE_IDLE_IN_TRANSACTION_TIMEOUT_SECONDS: "60"
AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS: "30"
CORS_ALLOWED_ORIGIN: https://ai.51easyai.com
AI_GATEWAY_PUBLIC_BASE_URL: https://ai.51easyai.com
AI_GATEWAY_WEB_BASE_URL: https://ai.51easyai.com