feat(queue): 增加非文本模型分布式准入队列

使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配，并将生产 API 与独立 Worker 角色拆分。补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本；未默认启用任何生产 queue_size 策略。已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收；合入最新主干后将重新执行发布门禁。
2026-07-29 16:15:43 +08:00
parent 2e5a90731b
commit 9e4fc7362d
55 changed files with 4565 additions and 237 deletions
@@ -0,0 +1,162 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+)
+
+const workerHeartbeatStaleAfter = 15 * time.Second
+
+type WorkerRegistrationInput struct {
+	InstanceID      string
+	PodUID          string
+	PodName         string
+	Site            string
+	Revision        string
+	DesiredCapacity int
+}
+
+type WorkerAllocation struct {
+	InstanceID      string
+	DesiredCapacity int
+	Allocated       int
+	ActiveInstances int
+	HeartbeatAt     time.Time
+}
+
+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")
+	}
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return WorkerAllocation{}, err
+	}
+	defer tx.Rollback(ctx)
+	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, allocated_capacity, started_at, heartbeat_at, updated_at
+)
+VALUES ($1, $2, $3, $4, $5, 'active', $6, 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 = 'active',
+    desired_capacity = EXCLUDED.desired_capacity,
+    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,
+	); err != nil {
+		return WorkerAllocation{}, err
+	}
+	if _, err := tx.Exec(ctx, `
+UPDATE gateway_worker_instances
+SET allocated_capacity = 0,
+    updated_at = now()
+WHERE status <> 'active'
+   OR heartbeat_at <= now() - $1::interval`, workerHeartbeatStaleAfter.String()); err != nil {
+		return WorkerAllocation{}, err
+	}
+
+	rows, err := tx.Query(ctx, `
+SELECT instance_id
+FROM gateway_worker_instances
+WHERE status = 'active'
+  AND heartbeat_at > now() - $1::interval
+ORDER BY instance_id ASC
+FOR UPDATE`, workerHeartbeatStaleAfter.String())
+	if err != nil {
+		return WorkerAllocation{}, err
+	}
+	activeIDs := make([]string, 0)
+	for rows.Next() {
+		var instanceID string
+		if err := rows.Scan(&instanceID); err != nil {
+			rows.Close()
+			return WorkerAllocation{}, err
+		}
+		activeIDs = append(activeIDs, instanceID)
+	}
+	if err := rows.Err(); err != nil {
+		rows.Close()
+		return WorkerAllocation{}, err
+	}
+	rows.Close()
+	if len(activeIDs) == 0 {
+		return WorkerAllocation{}, errors.New("worker registration was not active after heartbeat")
+	}
+
+	base := input.DesiredCapacity / len(activeIDs)
+	remainder := input.DesiredCapacity % len(activeIDs)
+	allocated := 0
+	for index, instanceID := range activeIDs {
+		capacity := base
+		if index < remainder {
+			capacity++
+		}
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_worker_instances
+SET desired_capacity = $2,
+    allocated_capacity = $3,
+    updated_at = now()
+WHERE instance_id = $1`, instanceID, input.DesiredCapacity, capacity); err != nil {
+			return WorkerAllocation{}, err
+		}
+		if 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,
+		ActiveInstances: len(activeIDs),
+		HeartbeatAt:     heartbeatAt,
+	}, nil
+}
+
+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,
+    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
+}