wangbo
e05922b0f4
实现本地三节点 K3s 同构环境、脱敏生产快照、Gemini 图片和多参考图视频协议模拟、统一验收报告及故障注入。\n\n新增 Worker 容量控制器、资源与连接预算、任务恢复保护，并将生产验收拆分为 validation 执行和人工 CAS 放量。\n\n验证包括 Go 全量测试、PostgreSQL HTTP 集成测试、go vet、OpenAPI、ShellCheck、前端检查、迁移及发布脚本测试。
528 lines
15 KiB
Go
528 lines
15 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
|
|
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.CapacityLimit == 0 {
|
|
input.CapacityLimit = input.DesiredCapacity
|
|
}
|
|
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, allocated_capacity, started_at, heartbeat_at, updated_at
|
|
)
|
|
VALUES ($1, $2, $3, $4, $5, 'active', $6, $7, 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,
|
|
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,
|
|
); 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 {
|
|
limit = desired
|
|
}
|
|
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
|
|
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
|
|
// 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 = now()
|
|
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
|
|
}
|