easyai-ai-gateway/apps/api/internal/store/execution_pools.go

package store

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"strings"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

type TaskRoutingDecision struct {
	TaskID          string
	PoolID          string
	WorkerID        string
	RouteProfileKey string
	PlatformID      string
	PlatformModelID string
	RoutingVersion  string
	Reason          string
	Snapshot        map[string]any
}

type WorkerExecutionLease struct {
	LeaseID    string
	TaskID     string
	PoolID     string
	WorkerID   string
	InstanceID string
	Endpoint   string
	Nonce      string
	ExpiresAt  time.Time
}

func (s *Store) UpsertExecutionPool(ctx context.Context, pool executionpool.ExecutionPool) error {
	pool.ID = strings.TrimSpace(pool.ID)
	if pool.ID == "" {
		return errors.New("execution pool ID is required")
	}
	if pool.Labels == nil {
		pool.Labels = map[string]string{}
	}
	if pool.Capabilities == nil {
		pool.Capabilities = map[string]any{}
	}
	labels, err := json.Marshal(pool.Labels)
	if err != nil {
		return err
	}
	capabilities, err := json.Marshal(pool.Capabilities)
	if err != nil {
		return err
	}
	state := pool.State
	if state == "" {
		state = executionpool.PoolActive
	}
	_, err = s.pool.Exec(ctx, `
INSERT INTO gateway_execution_pools (pool_id, labels, capabilities, state, updated_at)
VALUES ($1, $2::jsonb, $3::jsonb, $4, now())
ON CONFLICT (pool_id) DO UPDATE
SET labels = EXCLUDED.labels,
    capabilities = EXCLUDED.capabilities,
    state = CASE
      WHEN gateway_execution_pools.state = 'disabled' THEN 'disabled'
      ELSE EXCLUDED.state
    END,
    updated_at = now()`, pool.ID, labels, capabilities, string(state))
	return err
}

func (s *Store) ListExecutionPools(ctx context.Context) ([]executionpool.ExecutionPool, error) {
	rows, err := s.pool.Query(ctx, `
SELECT pool_id, labels, capabilities, state, updated_at
FROM gateway_execution_pools
ORDER BY pool_id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]executionpool.ExecutionPool, 0)
	for rows.Next() {
		var item executionpool.ExecutionPool
		var labels, capabilities []byte
		if err := rows.Scan(&item.ID, &labels, &capabilities, &item.State, &item.UpdatedAt); err != nil {
			return nil, err
		}
		item.Labels = decodeStringMap(labels)
		item.Capabilities = decodeObject(capabilities)
		items = append(items, item)
	}
	return items, rows.Err()
}

func (s *Store) ListWorkers(ctx context.Context, now time.Time) ([]executionpool.WorkerDescriptor, error) {
	if now.IsZero() {
		now = time.Now()
	}
	cutoff := now.Add(-workerHeartbeatStaleAfter)
	rows, err := s.pool.Query(ctx, `
SELECT worker_id, instance_id, pool_id, endpoint, protocol_version, revision,
       capabilities, allocated_capacity, safe_capacity, heavy_capacity,
       active_tasks, pressure_state, heartbeat_at, load_sampled_at
FROM gateway_worker_instances
WHERE status = 'active'
  AND heartbeat_at > $1::timestamptz
ORDER BY pool_id, instance_id`, cutoff)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]executionpool.WorkerDescriptor, 0)
	for rows.Next() {
		var item executionpool.WorkerDescriptor
		var capabilities []byte
		var loadSampledAt *time.Time
		if err := rows.Scan(
			&item.WorkerID, &item.InstanceID, &item.PoolID, &item.Endpoint,
			&item.ProtocolVersion, &item.Revision, &capabilities, &item.Allocated,
			&item.SafeCapacity, &item.HeavyCapacity, &item.ActiveTasks,
			&item.PressureState, &item.HeartbeatAt, &loadSampledAt,
		); err != nil {
			return nil, err
		}
		item.Capabilities = decodeObject(capabilities)
		if loadSampledAt != nil {
			item.LoadSampledAt = *loadSampledAt
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

func (s *Store) RecordRouteHealth(ctx context.Context, health executionpool.RouteHealth) error {
	_, err := s.pool.Exec(ctx, `
INSERT INTO gateway_route_health (
  pool_id, route_profile_key, state, success_rate,
  connect_tls_p95_ms, first_byte_p95_ms, upload_bytes_per_second, jitter_p95_ms,
  consecutive_failures, consecutive_successes, sample_count, sampled_at, expires_at, updated_at
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, now())
ON CONFLICT (pool_id, route_profile_key) DO UPDATE
SET state = EXCLUDED.state,
    success_rate = EXCLUDED.success_rate,
    connect_tls_p95_ms = EXCLUDED.connect_tls_p95_ms,
    first_byte_p95_ms = EXCLUDED.first_byte_p95_ms,
    upload_bytes_per_second = EXCLUDED.upload_bytes_per_second,
    jitter_p95_ms = EXCLUDED.jitter_p95_ms,
    consecutive_failures = EXCLUDED.consecutive_failures,
    consecutive_successes = EXCLUDED.consecutive_successes,
    sample_count = EXCLUDED.sample_count,
    sampled_at = EXCLUDED.sampled_at,
    expires_at = EXCLUDED.expires_at,
    updated_at = now()
WHERE gateway_route_health.sampled_at <= EXCLUDED.sampled_at`,
		health.PoolID, health.RouteProfileKey, string(health.State), health.SuccessRate,
		health.ConnectTLSP95.Milliseconds(), health.FirstByteP95.Milliseconds(), health.UploadBytesPerSecond,
		health.JitterP95.Milliseconds(), health.ConsecutiveFailures, health.ConsecutiveSuccesses,
		health.SampleCount, health.SampledAt, health.ExpiresAt)
	return err
}

func (s *Store) RecordRouteObservation(ctx context.Context, observation executionpool.RouteObservation) error {
	if observation.SampleCount < 1 || observation.SuccessCount < 0 || observation.SuccessCount > observation.SampleCount {
		return errors.New("invalid passive route observation")
	}
	successRate := float64(observation.SuccessCount) / float64(observation.SampleCount)
	_, err := s.pool.Exec(ctx, `
UPDATE gateway_route_health
SET success_rate = 0.8 * success_rate + 0.2 * $3,
    connect_tls_p95_ms = CASE
      WHEN $4::bigint > 0 THEN (0.8 * connect_tls_p95_ms + 0.2 * $4::bigint)::bigint
      ELSE connect_tls_p95_ms
    END,
    upload_bytes_per_second = CASE
      WHEN $5::double precision > 0 THEN 0.8 * upload_bytes_per_second + 0.2 * $5::double precision
      ELSE upload_bytes_per_second
    END,
    sample_count = sample_count + $6,
    updated_at = now()
WHERE pool_id = $1
  AND route_profile_key = $2`,
		strings.TrimSpace(observation.PoolID), strings.TrimSpace(observation.RouteProfileKey),
		successRate, observation.ConnectTLSP95.Milliseconds(), observation.UploadBytesPerSecond,
		observation.SampleCount,
	)
	return err
}

func (s *Store) RequestRouteProbe(ctx context.Context, routeProfileKey string) error {
	_, err := s.pool.Exec(ctx, `
INSERT INTO gateway_route_probe_requests (route_profile_key, requested_at, expires_at)
VALUES ($1, now(), now() + interval '30 seconds')
ON CONFLICT (route_profile_key) DO UPDATE
SET requested_at = now(), expires_at = now() + interval '30 seconds'`, strings.TrimSpace(routeProfileKey))
	return err
}

func (s *Store) ListRequestedRouteProbes(ctx context.Context) (map[string]time.Time, error) {
	rows, err := s.pool.Query(ctx, `
SELECT route_profile_key, requested_at
FROM gateway_route_probe_requests
WHERE expires_at > now()`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make(map[string]time.Time)
	for rows.Next() {
		var key string
		var requestedAt time.Time
		if err := rows.Scan(&key, &requestedAt); err != nil {
			return nil, err
		}
		items[key] = requestedAt
	}
	return items, rows.Err()
}

func (s *Store) ListRouteHealth(ctx context.Context, routeProfileKey string, _ time.Time) ([]executionpool.RouteHealth, error) {
	return s.listRouteHealth(ctx, "WHERE route_profile_key = $1", strings.TrimSpace(routeProfileKey))
}

func (s *Store) ListAllRouteHealth(ctx context.Context) ([]executionpool.RouteHealth, error) {
	return s.listRouteHealth(ctx, "", nil)
}

func (s *Store) ResolveRoutePreference(
	ctx context.Context,
	routeProfileKey string,
	proposedPoolID string,
	scores map[string]float64,
	now time.Time,
) (executionpool.RoutePreference, error) {
	if now.IsZero() {
		now = time.Now()
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return executionpool.RoutePreference{}, err
	}
	defer rollbackTransaction(tx)
	if _, err := tx.Exec(ctx, `
INSERT INTO gateway_route_preferences (route_profile_key, updated_at)
VALUES ($1, $2)
ON CONFLICT (route_profile_key) DO NOTHING`, strings.TrimSpace(routeProfileKey), now); err != nil {
		return executionpool.RoutePreference{}, err
	}
	var previous executionpool.RoutePreference
	var currentSince *time.Time
	if err := tx.QueryRow(ctx, `
SELECT route_profile_key, current_pool_id, current_since,
       challenger_pool_id, challenger_wins, updated_at
FROM gateway_route_preferences
WHERE route_profile_key = $1
FOR UPDATE`, strings.TrimSpace(routeProfileKey)).Scan(
		&previous.RouteProfileKey, &previous.CurrentPoolID, &currentSince,
		&previous.ChallengerPoolID, &previous.ChallengerWins, &previous.UpdatedAt,
	); err != nil {
		return executionpool.RoutePreference{}, err
	}
	if currentSince != nil {
		previous.CurrentSince = *currentSince
	}
	next := executionpool.AdvanceRoutePreference(previous, proposedPoolID, scores, now)
	if _, err := tx.Exec(ctx, `
UPDATE gateway_route_preferences
SET current_pool_id = $2,
    current_since = $3,
    challenger_pool_id = $4,
    challenger_wins = $5,
    updated_at = $6
WHERE route_profile_key = $1`, next.RouteProfileKey, next.CurrentPoolID, nullablePreferenceTime(next.CurrentSince),
		next.ChallengerPoolID, next.ChallengerWins, next.UpdatedAt); err != nil {
		return executionpool.RoutePreference{}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return executionpool.RoutePreference{}, err
	}
	return next, nil
}

func nullablePreferenceTime(value time.Time) any {
	if value.IsZero() {
		return nil
	}
	return value
}

func (s *Store) listRouteHealth(ctx context.Context, where string, argument any) ([]executionpool.RouteHealth, error) {
	query := `
SELECT pool_id, route_profile_key, state, success_rate,
       connect_tls_p95_ms, first_byte_p95_ms, upload_bytes_per_second, jitter_p95_ms,
       consecutive_failures, consecutive_successes, sample_count, sampled_at, expires_at
FROM gateway_route_health
` + where + `
ORDER BY route_profile_key, pool_id`
	var rows pgx.Rows
	var err error
	if where == "" {
		rows, err = s.pool.Query(ctx, query)
	} else {
		rows, err = s.pool.Query(ctx, query, argument)
	}
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]executionpool.RouteHealth, 0)
	for rows.Next() {
		var item executionpool.RouteHealth
		var connectMS, firstByteMS, jitterMS int64
		if err := rows.Scan(
			&item.PoolID, &item.RouteProfileKey, &item.State, &item.SuccessRate,
			&connectMS, &firstByteMS, &item.UploadBytesPerSecond, &jitterMS,
			&item.ConsecutiveFailures, &item.ConsecutiveSuccesses, &item.SampleCount,
			&item.SampledAt, &item.ExpiresAt,
		); err != nil {
			return nil, err
		}
		item.ConnectTLSP95 = time.Duration(connectMS) * time.Millisecond
		item.FirstByteP95 = time.Duration(firstByteMS) * time.Millisecond
		item.JitterP95 = time.Duration(jitterMS) * time.Millisecond
		items = append(items, item)
	}
	return items, rows.Err()
}

func (s *Store) PublishDesiredCapacity(ctx context.Context, desired executionpool.DesiredCapacity) error {
	_, err := s.pool.Exec(ctx, `
INSERT INTO gateway_pool_capacity_desires (pool_id, desired, reason, valid_until, updated_at)
VALUES ($1, $2, $3, $4, now())
ON CONFLICT (pool_id) DO UPDATE
SET desired = EXCLUDED.desired,
    reason = EXCLUDED.reason,
    valid_until = EXCLUDED.valid_until,
    updated_at = now()`, desired.PoolID, desired.Desired, desired.Reason, desired.ValidUntil)
	return err
}

func (s *Store) ListCapacity(ctx context.Context, now time.Time) ([]executionpool.CapacitySnapshot, error) {
	if now.IsZero() {
		now = time.Now()
	}
	cutoff := now.Add(-workerHeartbeatStaleAfter)
	rows, err := s.pool.Query(ctx, `
SELECT pool.pool_id,
       COUNT(worker.instance_id)::int,
       COALESCE(SUM(worker.allocated_capacity), 0)::int,
       COALESCE(SUM(worker.safe_capacity), 0)::int,
       COALESCE(SUM(worker.active_tasks), 0)::int,
       COALESCE(SUM(worker.heavy_capacity), 0)::int,
       COALESCE(SUM(worker.preparing_tasks + worker.finalizing_tasks), 0)::int,
       COALESCE(BOOL_OR(worker.pressure_state = 'critical'), false),
       COALESCE(MAX(worker.load_sampled_at), $1)
FROM gateway_execution_pools pool
LEFT JOIN gateway_worker_instances worker
 ON worker.pool_id = pool.pool_id
 AND worker.status = 'active'
 AND worker.heartbeat_at > $2::timestamptz
WHERE pool.state = 'active'
GROUP BY pool.pool_id
ORDER BY pool.pool_id`, now, cutoff)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]executionpool.CapacitySnapshot, 0)
	for rows.Next() {
		var item executionpool.CapacitySnapshot
		if err := rows.Scan(
			&item.PoolID, &item.WorkerCount, &item.Allocated, &item.SafeCapacity,
			&item.ActiveTasks, &item.HeavyCapacity, &item.HeavyTasks, &item.Critical,
			&item.SampledAt,
		); err != nil {
			return nil, err
		}
		if item.SafeCapacity > 0 {
			item.ResourceHeadroom = float64(max(item.SafeCapacity-item.ActiveTasks, 0)) / float64(item.SafeCapacity)
		}
		item.Stability = 1
		items = append(items, item)
	}
	return items, rows.Err()
}

func (s *Store) AssignTaskRouting(ctx context.Context, decision TaskRoutingDecision) error {
	snapshot, err := json.Marshal(decision.Snapshot)
	if err != nil {
		return err
	}
	_, err = s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET assigned_pool_id = NULLIF($2, ''),
    assigned_worker_id = NULLIF($3, ''),
    route_profile_key = NULLIF($4, ''),
	    routing_platform_id = NULLIF($5, '')::uuid,
	    routing_platform_model_id = NULLIF($6, '')::uuid,
	    routing_version = NULLIF($7, ''),
	    routing_reason = NULLIF($8, ''),
	    routing_snapshot = $9::jsonb,
    updated_at = now()
WHERE id = $1::uuid`, decision.TaskID, decision.PoolID, decision.WorkerID,
		decision.RouteProfileKey, decision.PlatformID, decision.PlatformModelID,
		decision.RoutingVersion, decision.Reason, snapshot)
	return err
}

func (s *Store) SetTaskSubmissionState(ctx context.Context, taskID, executionToken, state string) error {
	command, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET submission_state = $3, updated_at = now()
WHERE id = $1::uuid
  AND execution_token = $2::uuid
  AND status = 'running'`, taskID, executionToken, state)
	if err != nil {
		return err
	}
	if command.RowsAffected() == 0 {
		return ErrTaskExecutionLeaseLost
	}
	return nil
}

func (s *Store) ReserveWorkerExecution(ctx context.Context, taskID, poolID, nonce string, ttl time.Duration) (WorkerExecutionLease, error) {
	if ttl <= 0 || ttl > time.Minute {
		ttl = 30 * time.Second
	}
	nonceDigest := sha256.Sum256([]byte(nonce))
	nonceHash := hex.EncodeToString(nonceDigest[:])
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return WorkerExecutionLease{}, err
	}
	defer rollbackTransaction(tx)
	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "gateway-worker-execution:"+strings.TrimSpace(poolID)); err != nil {
		return WorkerExecutionLease{}, err
	}
	var lease WorkerExecutionLease
	lease.LeaseID = uuid.NewString()
	lease.TaskID = taskID
	lease.PoolID = poolID
	lease.Nonce = nonce
	lease.ExpiresAt = time.Now().Add(ttl)
	err = tx.QueryRow(ctx, `
WITH active_leases AS (
  SELECT instance_id, COUNT(*)::int AS count
  FROM gateway_worker_execution_leases
  WHERE released_at IS NULL AND expires_at > now()
  GROUP BY instance_id
), candidate AS (
  SELECT worker.worker_id, worker.instance_id, worker.endpoint
  FROM gateway_worker_instances worker
  LEFT JOIN active_leases lease ON lease.instance_id = worker.instance_id
  WHERE worker.pool_id = $1
    AND worker.status = 'active'
    AND worker.heartbeat_at > now() - $2::interval
    AND worker.protocol_version = $3
    AND NULLIF(worker.endpoint, '') IS NOT NULL
    AND worker.pressure_state <> 'critical'
    AND worker.active_tasks + COALESCE(lease.count, 0) < LEAST(worker.allocated_capacity, worker.safe_capacity)
  ORDER BY worker.active_tasks + COALESCE(lease.count, 0), worker.instance_id
  LIMIT 1
  FOR UPDATE OF worker SKIP LOCKED
)
SELECT worker_id, instance_id, endpoint FROM candidate`, poolID, workerHeartbeatStaleAfter.String(), executionpool.ProtocolVersion).Scan(
		&lease.WorkerID, &lease.InstanceID, &lease.Endpoint,
	)
	if err != nil {
		return WorkerExecutionLease{}, err
	}
	if _, err := tx.Exec(ctx, `
INSERT INTO gateway_worker_execution_leases (
  lease_id, task_id, pool_id, worker_id, instance_id, nonce_hash, expires_at
)
VALUES ($1::uuid, $2::uuid, $3, $4, $5, $6, $7)`, lease.LeaseID, taskID, poolID,
		lease.WorkerID, lease.InstanceID, nonceHash, lease.ExpiresAt); err != nil {
		return WorkerExecutionLease{}, err
	}
	if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET assigned_pool_id = $2,
    assigned_worker_id = $3,
    updated_at = now()
WHERE id = $1::uuid`, taskID, poolID, lease.WorkerID); err != nil {
		return WorkerExecutionLease{}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return WorkerExecutionLease{}, err
	}
	return lease, nil
}

func (s *Store) ConsumeWorkerExecutionLease(ctx context.Context, leaseID, nonce string) (WorkerExecutionLease, error) {
	digest := sha256.Sum256([]byte(nonce))
	nonceHash := hex.EncodeToString(digest[:])
	var lease WorkerExecutionLease
	err := s.pool.QueryRow(ctx, `
UPDATE gateway_worker_execution_leases
SET state = 'running', updated_at = now()
WHERE lease_id = $1::uuid
  AND nonce_hash = $2
  AND state = 'reserved'
  AND released_at IS NULL
  AND expires_at > now()
RETURNING lease_id::text, task_id::text, pool_id, worker_id, instance_id, expires_at`, leaseID, nonceHash).Scan(
		&lease.LeaseID, &lease.TaskID, &lease.PoolID, &lease.WorkerID, &lease.InstanceID, &lease.ExpiresAt,
	)
	return lease, err
}

func (s *Store) ReleaseWorkerExecutionLease(ctx context.Context, leaseID string) error {
	command, err := s.pool.Exec(ctx, `
UPDATE gateway_worker_execution_leases
SET state = 'released', released_at = now(), updated_at = now()
WHERE lease_id = $1::uuid AND released_at IS NULL`, leaseID)
	if err != nil {
		return err
	}
	if command.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}

func decodeStringMap(payload []byte) map[string]string {
	result := map[string]string{}
	_ = json.Unmarshal(payload, &result)
	return result
}

var _ executionpool.WorkerDirectory = (*Store)(nil)
var _ executionpool.RouteHealthRepository = (*Store)(nil)
var _ executionpool.CapacityProvider = (*Store)(nil)