wangbo
c28bf74230
保留平台模型 RPM、TPM 和并发策略语义，增加 PostgreSQL 集群级租约、饱和候选重选和多平台自动负载，避免突发任务固定等待首个平台。\n\n新增 Worker 实时负载采样、自适应 active/heavy 容量、心跳与管理端指标，并扩展本地 acceptance runner，覆盖三 Worker、同模型三平台 2/4/6 并发和 48 个带图视频突发任务。\n\n验证：go test ./...、go vet ./...、PostgreSQL 跨 Store 集成测试、gofmt、bash -n、ShellCheck 及本地集群 provider-burst 验收通过；48/48 成功，无越限、重复提交、重复计费、重复回调或终态资源泄漏。
632 lines
19 KiB
Go
632 lines
19 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
|
|
}
|