Files
easyai-ai-gateway/apps/api/internal/store/worker_registry.go
T
wangbo 484b05d005 fix(queue): 回收终态任务的孤儿 River Job
原因:Worker 滚动切换后,执行租约恢复会终态化提交结果不确定的任务,但失活 Worker 持有的 River Job 仍需等待一小时通用救援窗口。\n\n影响:仅在任务已终态、执行租约清空、River owner 无活跃 Worker 且超过 45 秒时,将孤儿 Job 幂等标记为 completed;活跃 owner 不受影响。\n\n验证:PostgreSQL 集成测试覆盖失活 owner、活跃 owner 与重复回收;API 全量 go test 和 go vet 通过。
2026-08-04 18:51:47 +08:00

711 lines
22 KiB
Go

package store
import (
"context"
"errors"
"strings"
"time"
"github.com/jackc/pgx/v5"
)
const workerHeartbeatStaleAfter = 30 * time.Second
type WorkerRegistrationInput struct {
InstanceID string
PodUID string
PodName string
Site string
Revision string
DesiredCapacity int
CapacityLimit int
LoadMode string
SafeCapacity int
HeavyCapacity int
ActiveTasks int
PreparingTasks int
WaitingUpstreamTasks int
FinalizingTasks int
PressureState string
PressureReason string
LoadSampledAt time.Time
HeartbeatStaleAfter time.Duration
}
type WorkerAllocation struct {
InstanceID string
DesiredCapacity int
Allocated int
GlobalAllocated int
ActiveInstances int
HeartbeatAt time.Time
}
type activeWorkerCapacity struct {
InstanceID string
CapacityLimit int
}
// ActiveWorkerCapacity returns the cluster-wide execution capacity currently
// owned by live Worker instances.
func (s *Store) ActiveWorkerCapacity(ctx context.Context) (int, error) {
var capacity int
err := s.pool.QueryRow(ctx, `
SELECT COALESCE(SUM(allocated_capacity), 0)::int
FROM gateway_worker_instances
WHERE status = 'active'
AND heartbeat_at > now() - $1::interval`, workerHeartbeatStaleAfter.String()).Scan(&capacity)
return capacity, err
}
func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistrationInput) (WorkerAllocation, error) {
input.InstanceID = strings.TrimSpace(input.InstanceID)
if input.InstanceID == "" {
return WorkerAllocation{}, errors.New("worker instance ID is required")
}
if input.DesiredCapacity < 0 {
return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative")
}
if input.CapacityLimit < 0 {
return WorkerAllocation{}, errors.New("worker capacity limit cannot be negative")
}
if input.SafeCapacity < 0 || input.HeavyCapacity < 0 || input.ActiveTasks < 0 || input.PreparingTasks < 0 || input.WaitingUpstreamTasks < 0 || input.FinalizingTasks < 0 {
return WorkerAllocation{}, errors.New("worker load values cannot be negative")
}
if input.ActiveTasks != input.PreparingTasks+input.WaitingUpstreamTasks+input.FinalizingTasks {
return WorkerAllocation{}, errors.New("worker active task count must equal phase task counts")
}
if input.CapacityLimit == 0 {
input.CapacityLimit = input.DesiredCapacity
}
hardCapacityLimit := input.CapacityLimit
if strings.EqualFold(strings.TrimSpace(input.LoadMode), "adaptive") {
input.CapacityLimit = min(input.CapacityLimit, input.SafeCapacity)
}
pressureState := strings.ToLower(strings.TrimSpace(input.PressureState))
switch pressureState {
case "normal", "busy", "critical":
default:
pressureState = "unknown"
}
if input.LoadSampledAt.IsZero() {
input.LoadSampledAt = time.Now()
}
staleAfter := input.HeartbeatStaleAfter
if staleAfter < workerHeartbeatStaleAfter {
staleAfter = workerHeartbeatStaleAfter
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return WorkerAllocation{}, err
}
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.
if _, err := tx.Exec(ctx, `SELECT set_config('synchronous_commit', 'off', true)`); err != nil {
return WorkerAllocation{}, err
}
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-worker-capacity', 0))`); err != nil {
return WorkerAllocation{}, err
}
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_worker_instances (
instance_id, pod_uid, pod_name, site, revision, status,
desired_capacity, capacity_limit, hard_capacity_limit, safe_capacity, heavy_capacity,
active_tasks, preparing_tasks, waiting_upstream_tasks, finalizing_tasks,
pressure_state, pressure_reason, load_sampled_at,
allocated_capacity, started_at, heartbeat_at, updated_at
)
VALUES (
$1, $2, $3, $4, $5, 'active', $6, $7, $8, $9, $10,
$11, $12, $13, $14, $15, $16, $17,
0, now(), now(), now()
)
ON CONFLICT (instance_id) DO UPDATE
SET pod_uid = EXCLUDED.pod_uid,
pod_name = EXCLUDED.pod_name,
site = EXCLUDED.site,
revision = EXCLUDED.revision,
status = CASE
WHEN gateway_worker_instances.status = 'draining' THEN 'draining'
ELSE 'active'
END,
desired_capacity = EXCLUDED.desired_capacity,
capacity_limit = EXCLUDED.capacity_limit,
hard_capacity_limit = EXCLUDED.hard_capacity_limit,
safe_capacity = EXCLUDED.safe_capacity,
heavy_capacity = EXCLUDED.heavy_capacity,
active_tasks = EXCLUDED.active_tasks,
preparing_tasks = EXCLUDED.preparing_tasks,
waiting_upstream_tasks = EXCLUDED.waiting_upstream_tasks,
finalizing_tasks = EXCLUDED.finalizing_tasks,
pressure_state = EXCLUDED.pressure_state,
pressure_reason = EXCLUDED.pressure_reason,
load_sampled_at = EXCLUDED.load_sampled_at,
heartbeat_at = now(),
updated_at = now()`,
input.InstanceID,
strings.TrimSpace(input.PodUID),
strings.TrimSpace(input.PodName),
strings.TrimSpace(input.Site),
strings.TrimSpace(input.Revision),
input.DesiredCapacity,
input.CapacityLimit,
hardCapacityLimit,
input.SafeCapacity,
input.HeavyCapacity,
input.ActiveTasks,
input.PreparingTasks,
input.WaitingUpstreamTasks,
input.FinalizingTasks,
pressureState,
strings.TrimSpace(input.PressureReason),
input.LoadSampledAt,
); err != nil {
return WorkerAllocation{}, err
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_worker_instances
SET status = 'stale',
allocated_capacity = 0,
updated_at = now()
WHERE status = 'active'
AND heartbeat_at <= now() - $1::interval`, staleAfter.String()); err != nil {
return WorkerAllocation{}, err
}
rows, err := tx.Query(ctx, `
SELECT instance_id, capacity_limit
FROM gateway_worker_instances
WHERE status = 'active'
AND heartbeat_at > now() - $1::interval
ORDER BY instance_id ASC
FOR UPDATE`, staleAfter.String())
if err != nil {
return WorkerAllocation{}, err
}
activeWorkers := make([]activeWorkerCapacity, 0)
for rows.Next() {
var worker activeWorkerCapacity
if err := rows.Scan(&worker.InstanceID, &worker.CapacityLimit); err != nil {
rows.Close()
return WorkerAllocation{}, err
}
activeWorkers = append(activeWorkers, worker)
}
if err := rows.Err(); err != nil {
rows.Close()
return WorkerAllocation{}, err
}
rows.Close()
if len(activeWorkers) == 0 {
return WorkerAllocation{}, errors.New("worker registration was not active after heartbeat")
}
allocations, globalAllocated := allocateWorkerCapacities(activeWorkers, input.DesiredCapacity)
allocated := 0
for _, worker := range activeWorkers {
capacity := allocations[worker.InstanceID]
if _, err := tx.Exec(ctx, `
UPDATE gateway_worker_instances
SET desired_capacity = $2,
allocated_capacity = $3,
updated_at = now()
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 {
allocated = capacity
}
}
var heartbeatAt time.Time
if err := tx.QueryRow(ctx, `
SELECT heartbeat_at
FROM gateway_worker_instances
WHERE instance_id = $1`, input.InstanceID).Scan(&heartbeatAt); err != nil {
return WorkerAllocation{}, err
}
if err := tx.Commit(ctx); err != nil {
return WorkerAllocation{}, err
}
return WorkerAllocation{
InstanceID: input.InstanceID,
DesiredCapacity: input.DesiredCapacity,
Allocated: allocated,
GlobalAllocated: globalAllocated,
ActiveInstances: len(activeWorkers),
HeartbeatAt: heartbeatAt,
}, nil
}
func allocateWorkerCapacities(workers []activeWorkerCapacity, desired int) (map[string]int, int) {
allocations := make(map[string]int, len(workers))
if desired <= 0 || len(workers) == 0 {
return allocations, 0
}
remaining := desired
for remaining > 0 {
progressed := false
for _, worker := range workers {
limit := worker.CapacityLimit
if limit <= 0 {
continue
}
if allocations[worker.InstanceID] >= limit {
continue
}
allocations[worker.InstanceID]++
remaining--
progressed = true
if remaining == 0 {
break
}
}
if !progressed {
break
}
}
return allocations, desired - remaining
}
func (s *Store) MarkWorkerDraining(ctx context.Context, instanceID string) error {
result, err := s.pool.Exec(ctx, `
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))
if err != nil {
return err
}
if result.RowsAffected() == 0 {
return pgx.ErrNoRows
}
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 `json:"instanceId"`
PodUID string `json:"podUid,omitempty"`
PodName string `json:"podName,omitempty"`
Site string `json:"site,omitempty"`
Revision string `json:"revision,omitempty"`
Status string `json:"status"`
Allocated int `json:"allocatedCapacity"`
CapacityLimit int `json:"capacityLimit"`
HardCapacityLimit int `json:"hardCapacityLimit"`
SafeCapacity int `json:"safeCapacity"`
HeavyCapacity int `json:"heavyCapacity"`
ReportedActiveTasks int `json:"reportedActiveTasks"`
PreparingTasks int `json:"preparingTasks"`
WaitingUpstreamTasks int `json:"waitingUpstreamTasks"`
FinalizingTasks int `json:"finalizingTasks"`
PressureState string `json:"pressureState"`
PressureReason string `json:"pressureReason,omitempty"`
LoadSampledAt *time.Time `json:"loadSampledAt,omitempty"`
RunningTasks int `json:"runningTasks"`
ActiveLeases int `json:"activeLeases"`
HeartbeatAt time.Time `json:"heartbeatAt"`
DrainingAt *time.Time `json:"drainingAt,omitempty"`
}
type WorkerQueueRuntime struct {
Queued int `json:"queued"`
Running int `json:"running"`
OldestWaitSeconds float64 `json:"oldestWaitSeconds"`
}
type WorkerClusterRuntime struct {
Workers []WorkerInstanceRuntime `json:"workers"`
Queue WorkerQueueRuntime `json:"queue"`
CapturedAt time.Time `json:"capturedAt"`
}
func (s *Store) GetWorkerClusterRuntime(ctx context.Context) (WorkerClusterRuntime, error) {
workers, err := s.ListWorkerInstanceRuntime(ctx)
if err != nil {
return WorkerClusterRuntime{}, err
}
queue, err := s.WorkerQueueRuntime(ctx)
if err != nil {
return WorkerClusterRuntime{}, err
}
return WorkerClusterRuntime{Workers: workers, Queue: queue, CapturedAt: time.Now()}, nil
}
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,
worker.hard_capacity_limit,
worker.safe_capacity,
worker.heavy_capacity,
worker.active_tasks,
worker.preparing_tasks,
worker.waiting_upstream_tasks,
worker.finalizing_tasks,
worker.pressure_state,
worker.pressure_reason,
worker.load_sampled_at,
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.HardCapacityLimit,
&instance.SafeCapacity,
&instance.HeavyCapacity,
&instance.ReportedActiveTasks,
&instance.PreparingTasks,
&instance.WaitingUpstreamTasks,
&instance.FinalizingTasks,
&instance.PressureState,
&instance.PressureReason,
&instance.LoadSampledAt,
&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
// the admission, while unrelated queue followers are no longer head-of-line
// blocked behind stale runtime ownership.
func (s *Store) YieldStaleAsyncTaskAdmissions(
ctx context.Context,
staleAfter time.Duration,
limit int,
) (int64, error) {
if staleAfter < workerHeartbeatStaleAfter {
staleAfter = workerHeartbeatStaleAfter
}
if limit <= 0 || limit > 1000 {
limit = 100
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return 0, err
}
defer rollbackTransaction(tx)
var yielded int64
if err := tx.QueryRow(ctx, `
WITH stale AS MATERIALIZED (
SELECT admission.task_id
FROM gateway_task_admissions admission
JOIN gateway_tasks task ON task.id = admission.task_id
JOIN river_job job ON job.id = task.river_job_id
WHERE admission.status = 'waiting'
AND admission.mode = 'async'
AND task.async_mode = true
AND task.status = 'queued'
AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
AND job.queue = 'gateway_tasks'
AND job.kind = 'gateway_task_run'
AND job.state = 'running'
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,
pg_advisory_xact_lock(hashtextextended('task-admission:' || task_id::text, 0))
FROM stale
),
deleted AS (
DELETE FROM gateway_task_admissions admission
USING locked
WHERE admission.task_id = locked.task_id
AND admission.status = 'waiting'
AND admission.mode = 'async'
RETURNING admission.task_id
)
SELECT count(*)::bigint
FROM deleted`, staleAfter.String(), limit).Scan(&yielded); err != nil {
return 0, err
}
if err := tx.Commit(ctx); err != nil {
return 0, err
}
if yielded > 0 {
s.notifyTaskAdmissionBestEffort(ctx, "*")
}
return yielded, nil
}
// RecoverOrphanedAsyncRiverJobs retries stale River jobs after the gateway
// execution lease has already returned the task to queued. A process heartbeat
// does not prove that an individual River job goroutine is still alive; the
// task execution lease is the ownership fence.
func (s *Store) RecoverOrphanedAsyncRiverJobs(
ctx context.Context,
workerStaleAfter time.Duration,
limit int,
) (int64, error) {
if workerStaleAfter < workerHeartbeatStaleAfter {
workerStaleAfter = workerHeartbeatStaleAfter
}
if limit <= 0 || limit > 1000 {
limit = 100
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return 0, err
}
defer rollbackTransaction(tx)
var recovered int64
if err := tx.QueryRow(ctx, `
WITH orphaned AS MATERIALIZED (
SELECT job.id AS job_id,
task.id AS task_id,
job.attempt
FROM river_job job
JOIN gateway_tasks task ON task.river_job_id = job.id
WHERE job.queue = 'gateway_tasks'
AND job.kind = 'gateway_task_run'
AND job.state = 'running'
AND job.attempted_at <= now() - $1::interval
AND task.async_mode = true
AND task.status = 'queued'
AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
ORDER BY job.attempted_at ASC, job.id ASC
LIMIT $2
FOR UPDATE OF job SKIP LOCKED
),
released_leases AS (
UPDATE gateway_concurrency_leases lease
SET released_at = statement_timestamp()
FROM orphaned
WHERE lease.task_id = orphaned.task_id
AND lease.released_at IS NULL
),
cleared_admissions AS (
DELETE FROM gateway_task_admissions admission
USING orphaned
WHERE admission.task_id = orphaned.task_id
),
reset_tasks AS (
UPDATE gateway_tasks task
SET locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
next_run_at = now(),
updated_at = now()
FROM orphaned
WHERE task.id = orphaned.task_id
),
recovered_jobs AS (
UPDATE river_job job
SET errors = array_append(
job.errors,
jsonb_build_object(
'at', now(),
'attempt', orphaned.attempt,
'error', 'Orphaned gateway job recovered after worker heartbeat expired',
'trace', ''
)
),
finalized_at = NULL,
scheduled_at = now(),
state = 'retryable'
FROM orphaned
WHERE job.id = orphaned.job_id
RETURNING job.id
)
SELECT count(*)::bigint
FROM recovered_jobs`, workerStaleAfter.String(), limit).Scan(&recovered); err != nil {
return 0, err
}
if err := tx.Commit(ctx); err != nil {
return 0, err
}
return recovered, nil
}
// FinalizeOrphanedTerminalAsyncRiverJobs closes River jobs whose gateway task
// is already terminal after the owning Worker disappeared. River deliberately
// uses a one-hour rescue horizon for long provider polls, so relying on its
// generic rescue loop would leave rollout-interrupted terminal jobs reported
// as running long after their execution lease and business resources were
// released.
func (s *Store) FinalizeOrphanedTerminalAsyncRiverJobs(
ctx context.Context,
workerStaleAfter time.Duration,
limit int,
) (int64, error) {
if workerStaleAfter < workerHeartbeatStaleAfter {
workerStaleAfter = workerHeartbeatStaleAfter
}
if limit <= 0 || limit > 1000 {
limit = 100
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return 0, err
}
defer rollbackTransaction(tx)
var finalized int64
if err := tx.QueryRow(ctx, `
WITH orphaned AS MATERIALIZED (
SELECT job.id AS job_id,
task.id AS task_id,
job.attempt
FROM river_job job
JOIN gateway_tasks task ON task.river_job_id = job.id
WHERE job.queue = 'gateway_tasks'
AND job.kind = 'gateway_task_run'
AND job.state = 'running'
AND job.attempted_at <= now() - $1::interval
AND task.async_mode = true
AND task.status IN ('succeeded', 'failed', 'cancelled')
AND task.execution_token IS NULL
AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
AND NOT EXISTS (
SELECT 1
FROM gateway_worker_instances worker
WHERE worker.status = 'active'
AND worker.heartbeat_at > now() - $2::interval
AND EXISTS (
SELECT 1
FROM unnest(job.attempted_by) attempted_owner
WHERE attempted_owner LIKE worker.instance_id || '-exec-%'
)
)
ORDER BY job.attempted_at ASC, job.id ASC
LIMIT $3
FOR UPDATE OF job SKIP LOCKED
), finalized_jobs AS (
UPDATE river_job job
SET errors = array_append(
COALESCE(job.errors, ARRAY[]::jsonb[]),
jsonb_build_object(
'at', now(),
'attempt', orphaned.attempt,
'error', 'Terminal gateway job finalized after Worker owner disappeared',
'trace', ''
)
),
finalized_at = now(),
state = 'completed'
FROM orphaned
WHERE job.id = orphaned.job_id
RETURNING job.id
)
SELECT count(*)::bigint
FROM finalized_jobs`, workerStaleAfter.String(), workerHeartbeatStaleAfter.String(), limit).Scan(&finalized); err != nil {
return 0, err
}
if err := tx.Commit(ctx); err != nil {
return 0, err
}
return finalized, nil
}