feat(routing): 引入多执行池智能调度
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。 实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。 新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
This commit is contained in:
@@ -0,0 +1,537 @@
|
||||
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()
|
||||
}
|
||||
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 - $2::interval
|
||||
ORDER BY pool_id, instance_id`, now, workerHeartbeatStaleAfter.String())
|
||||
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, ¤tSince,
|
||||
&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()
|
||||
}
|
||||
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 > $1 - $2::interval
|
||||
WHERE pool.state = 'active'
|
||||
GROUP BY pool.pool_id
|
||||
ORDER BY pool.pool_id`, now, workerHeartbeatStaleAfter.String())
|
||||
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)
|
||||
Reference in New Issue
Block a user