easyai-ai-gateway/apps/api/internal/store/worker_registry.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
	HeartbeatStaleAfter time.Duration
}

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")
	}
	staleAfter := input.HeartbeatStaleAfter
	if staleAfter < workerHeartbeatStaleAfter {
		staleAfter = workerHeartbeatStaleAfter
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return WorkerAllocation{}, err
	}
	defer tx.Rollback(ctx)
	// 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, 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`, staleAfter.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`, staleAfter.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
}