feat(routing): 引入多执行池智能调度
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。 实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。 新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
This commit is contained in:
@@ -151,7 +151,7 @@ func TestExpiredExecutionLeaseDoesNotReplayAmbiguousUpstreamSubmission(t *testin
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_submission_unknown" {
|
||||
if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_timeout" {
|
||||
t.Fatalf("manual review task=%+v", review)
|
||||
}
|
||||
var attemptStatus string
|
||||
@@ -162,7 +162,7 @@ FROM gateway_task_attempts
|
||||
WHERE id=$1::uuid`, attemptID).Scan(&attemptStatus, &attemptErrorCode); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if attemptStatus != "failed" || attemptErrorCode != "upstream_submission_unknown" {
|
||||
if attemptStatus != "failed" || attemptErrorCode != "upstream_timeout" {
|
||||
t.Fatalf("manual review attempt status=%s error=%s", attemptStatus, attemptErrorCode)
|
||||
}
|
||||
var outboxStatus string
|
||||
@@ -174,11 +174,84 @@ FROM settlement_outbox
|
||||
WHERE task_id=$1::uuid AND event_type='task.billing.review'`, created.ID).Scan(&outboxStatus, &outboxAction, &reviewReason); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if outboxStatus != "manual_review" || outboxAction != "release" || reviewReason != "upstream_submission_unknown" {
|
||||
if outboxStatus != "manual_review" || outboxAction != "release" || reviewReason != "upstream_timeout" {
|
||||
t.Fatalf("review outbox status=%s action=%s reason=%s", outboxStatus, outboxAction, reviewReason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveInterruptedTaskExecutionPreservesSubmissionFence(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
user := &auth.User{ID: "interrupted-worker-" + uuid.NewString(), Source: "gateway"}
|
||||
|
||||
type testCase struct {
|
||||
name string
|
||||
submissionStatus string
|
||||
remoteTaskID string
|
||||
wantManualReview bool
|
||||
}
|
||||
for _, tc := range []testCase{
|
||||
{name: "not submitted requeues", submissionStatus: "not_submitted"},
|
||||
{name: "ambiguous submission stops", submissionStatus: "submitting", wantManualReview: true},
|
||||
{name: "known remote task requeues for polling", submissionStatus: "response_received", remoteTaskID: "remote-" + uuid.NewString()},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "interrupted-worker-model", RunMode: "production", Async: true,
|
||||
Request: map[string]any{"model": "interrupted-worker-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
})
|
||||
|
||||
executionToken := uuid.NewString()
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, executionToken, 5*time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
|
||||
TaskID: created.ID, ExecutionToken: executionToken, AttemptNo: 1, Status: "running",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SetAttemptUpstreamSubmissionStatusForExecution(
|
||||
ctx, attemptID, created.ID, executionToken, tc.submissionStatus,
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tc.remoteTaskID != "" {
|
||||
if err := db.SetTaskRemoteTask(ctx, created.ID, executionToken, attemptID, tc.remoteTaskID, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
resolved, err := db.ResolveInterruptedTaskExecution(ctx, created.ID, executionToken)
|
||||
if tc.wantManualReview {
|
||||
if !errors.Is(err, ErrTaskExecutionManualReview) {
|
||||
t.Fatalf("resolve error=%v, want ErrTaskExecutionManualReview", err)
|
||||
}
|
||||
failed, getErr := db.GetTask(ctx, created.ID)
|
||||
if getErr != nil {
|
||||
t.Fatal(getErr)
|
||||
}
|
||||
if failed.Status != "failed" || failed.ErrorCode != "upstream_timeout" || failed.SubmissionState != "submission_confirmation_pending" {
|
||||
t.Fatalf("ambiguous task=%+v", failed)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resolved.Status != "queued" || resolved.ExecutionToken != "" || resolved.RemoteTaskID != tc.remoteTaskID {
|
||||
t.Fatalf("resolved task=%+v", resolved)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueuedPreparationDoesNotReplayInterruptedUpstreamSubmission(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
@@ -244,7 +317,7 @@ WHERE id=$1::uuid`, created.ID); err != nil {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_submission_unknown" {
|
||||
if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_timeout" {
|
||||
t.Fatalf("manual review task=%+v", review)
|
||||
}
|
||||
})
|
||||
@@ -365,7 +438,7 @@ func TestFinishTaskManualReviewCreatesVisibleBillingRecord(t *testing.T) {
|
||||
}
|
||||
if _, err := db.FinishTaskManualReview(ctx, FinishTaskManualReviewInput{
|
||||
TaskID: created.ID, ExecutionToken: token, AttemptID: attemptID, TaskStatus: "failed",
|
||||
Code: "upstream_submission_unknown", Message: "upstream submission result is unknown",
|
||||
Code: "upstream_timeout", Message: "upstream submission confirmation timed out",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -374,7 +447,7 @@ func TestFinishTaskManualReviewCreatesVisibleBillingRecord(t *testing.T) {
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM settlement_outbox
|
||||
WHERE task_id=$1::uuid AND status='manual_review'
|
||||
AND action='release' AND manual_review_reason='upstream_submission_unknown'
|
||||
AND action='release' AND manual_review_reason='upstream_timeout'
|
||||
)`, created.ID).Scan(&visible); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
@@ -42,6 +42,10 @@ func (s *Store) TryAcquireCapacityControllerLeadership(ctx context.Context) (Lea
|
||||
return s.tryAcquireLeadership(ctx, capacityControllerLeadershipKey)
|
||||
}
|
||||
|
||||
func (s *Store) TryAcquireRouteProbeLeadership(ctx context.Context, poolID, routeProfileKey string) (Leadership, bool, error) {
|
||||
return s.tryAcquireLeadership(ctx, "easyai-gateway-route-probe:"+poolID+":"+routeProfileKey)
|
||||
}
|
||||
|
||||
func (s *Store) tryAcquireLeadership(ctx context.Context, key string) (Leadership, bool, error) {
|
||||
conn, err := s.pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
|
||||
@@ -634,6 +634,15 @@ type GatewayTask struct {
|
||||
CompatibilitySubmitHTTPStatus int `json:"compatibilitySubmitHttpStatus,omitempty"`
|
||||
CompatibilitySubmitHeaders map[string]any `json:"compatibilitySubmitHeaders,omitempty"`
|
||||
CompatibilitySubmitBody map[string]any `json:"compatibilitySubmitBody,omitempty"`
|
||||
AssignedPoolID string `json:"-"`
|
||||
AssignedWorkerID string `json:"-"`
|
||||
RouteProfileKey string `json:"-"`
|
||||
RoutingPlatformID string `json:"-"`
|
||||
RoutingPlatformModelID string `json:"-"`
|
||||
RoutingVersion string `json:"-"`
|
||||
RoutingReason string `json:"-"`
|
||||
RoutingSnapshot map[string]any `json:"-"`
|
||||
SubmissionState string `json:"-"`
|
||||
Attempts []TaskAttempt `json:"attempts,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
@@ -662,6 +671,10 @@ COALESCE(error_code, ''), COALESCE(error_message, ''), COALESCE(public_error, '{
|
||||
COALESCE(compatibility_protocol, ''), COALESCE(compatibility_public_id, ''),
|
||||
COALESCE(compatibility_source_protocol, ''), COALESCE(compatibility_submit_http_status, 0),
|
||||
COALESCE(compatibility_submit_headers, '{}'::jsonb), COALESCE(compatibility_submit_body, '{}'::jsonb),
|
||||
COALESCE(assigned_pool_id, ''), COALESCE(assigned_worker_id, ''), COALESCE(route_profile_key, ''),
|
||||
COALESCE(routing_platform_id::text, ''), COALESCE(routing_platform_model_id::text, ''),
|
||||
COALESCE(routing_version, ''), COALESCE(routing_reason, ''), COALESCE(routing_snapshot, '{}'::jsonb),
|
||||
COALESCE(submission_state, 'not_started'),
|
||||
created_at, updated_at, COALESCE(finished_at::text, '')`
|
||||
|
||||
type TaskEvent struct {
|
||||
@@ -2243,6 +2256,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
var remoteTaskPayloadBytes []byte
|
||||
var compatibilitySubmitHeadersBytes []byte
|
||||
var compatibilitySubmitBodyBytes []byte
|
||||
var routingSnapshotBytes []byte
|
||||
var publicErrorBytes []byte
|
||||
if err := scanner.Scan(
|
||||
&task.ID,
|
||||
@@ -2306,6 +2320,15 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
&task.CompatibilitySubmitHTTPStatus,
|
||||
&compatibilitySubmitHeadersBytes,
|
||||
&compatibilitySubmitBodyBytes,
|
||||
&task.AssignedPoolID,
|
||||
&task.AssignedWorkerID,
|
||||
&task.RouteProfileKey,
|
||||
&task.RoutingPlatformID,
|
||||
&task.RoutingPlatformModelID,
|
||||
&task.RoutingVersion,
|
||||
&task.RoutingReason,
|
||||
&routingSnapshotBytes,
|
||||
&task.SubmissionState,
|
||||
&task.CreatedAt,
|
||||
&task.UpdatedAt,
|
||||
&task.FinishedAt,
|
||||
@@ -2322,6 +2345,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
task.PricingSnapshot = decodeObject(pricingSnapshotBytes)
|
||||
task.CompatibilitySubmitHeaders = decodeObject(compatibilitySubmitHeadersBytes)
|
||||
task.CompatibilitySubmitBody = decodeObject(compatibilitySubmitBodyBytes)
|
||||
task.RoutingSnapshot = decodeObject(routingSnapshotBytes)
|
||||
if len(publicErrorBytes) > 0 {
|
||||
var snapshot publicerror.Error
|
||||
if json.Unmarshal(publicErrorBytes, &snapshot) == nil && snapshot.Code != "" {
|
||||
|
||||
@@ -691,13 +691,13 @@ ON CONFLICT (task_id, event_type) DO NOTHING`, releaseBillingTaskIDs); err != ni
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts attempt
|
||||
SET status = 'failed',
|
||||
retryable = task.error_code IS DISTINCT FROM 'upstream_submission_unknown',
|
||||
retryable = task.submission_state IS DISTINCT FROM 'submission_confirmation_pending',
|
||||
error_code = CASE
|
||||
WHEN task.error_code = 'upstream_submission_unknown' THEN 'upstream_submission_unknown'
|
||||
WHEN task.submission_state = 'submission_confirmation_pending' THEN 'upstream_timeout'
|
||||
ELSE 'execution_lease_expired'
|
||||
END,
|
||||
error_message = CASE
|
||||
WHEN task.error_code = 'upstream_submission_unknown' THEN 'upstream submission result is unknown'
|
||||
WHEN task.submission_state = 'submission_confirmation_pending' THEN 'upstream submission confirmation timed out'
|
||||
ELSE 'attempt execution lease expired after worker heartbeat became stale'
|
||||
END,
|
||||
finished_at = now()
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
)
|
||||
|
||||
type RouteProbeTarget struct {
|
||||
Candidate RuntimeModelCandidate
|
||||
Hot bool
|
||||
}
|
||||
|
||||
func (s *Store) ListRouteProbeTargets(ctx context.Context) ([]RouteProbeTarget, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT platform.id::text,
|
||||
platform.provider,
|
||||
COALESCE(NULLIF(platform.config->>'specType', ''), NULLIF(catalog.provider_type, ''), NULLIF(platform.config->>'sourceSpecType', ''), platform.provider),
|
||||
COALESCE(
|
||||
NULLIF(platform.base_url, ''),
|
||||
NULLIF(platform.config->>'endpoint', ''),
|
||||
NULLIF(platform.config->>'baseURL', ''),
|
||||
NULLIF(platform.config->>'base_url', ''),
|
||||
NULLIF(catalog.default_base_url, ''),
|
||||
''
|
||||
),
|
||||
platform.config,
|
||||
model.id::text,
|
||||
COALESCE(model.model_type->>0, ''),
|
||||
EXISTS (
|
||||
SELECT 1
|
||||
FROM gateway_task_attempts attempt
|
||||
WHERE attempt.platform_model_id = model.id
|
||||
AND attempt.started_at > now() - interval '10 minutes'
|
||||
)
|
||||
FROM integration_platforms platform
|
||||
JOIN platform_models model ON model.platform_id = platform.id
|
||||
LEFT JOIN model_catalog_providers catalog
|
||||
ON catalog.provider_key = platform.provider OR catalog.provider_code = platform.provider
|
||||
WHERE platform.status = 'enabled'
|
||||
AND platform.deleted_at IS NULL
|
||||
AND model.enabled = true
|
||||
AND COALESCE(
|
||||
NULLIF(platform.base_url, ''),
|
||||
NULLIF(platform.config->>'endpoint', ''),
|
||||
NULLIF(platform.config->>'baseURL', ''),
|
||||
NULLIF(platform.config->>'base_url', ''),
|
||||
NULLIF(catalog.default_base_url, '')
|
||||
) IS NOT NULL
|
||||
ORDER BY platform.id, model.id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]RouteProbeTarget, 0)
|
||||
for rows.Next() {
|
||||
var item RouteProbeTarget
|
||||
var config []byte
|
||||
if err := rows.Scan(
|
||||
&item.Candidate.PlatformID, &item.Candidate.Provider, &item.Candidate.SpecType,
|
||||
&item.Candidate.BaseURL, &config, &item.Candidate.PlatformModelID,
|
||||
&item.Candidate.ModelType, &item.Hot,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.Candidate.PlatformConfig = decodeObject(config)
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
@@ -392,8 +392,8 @@ func markTaskExecutionManualReviewTx(
|
||||
UPDATE gateway_task_attempts
|
||||
SET status = 'failed',
|
||||
retryable = false,
|
||||
error_code = 'upstream_submission_unknown',
|
||||
error_message = 'upstream submission result is unknown',
|
||||
error_code = 'upstream_timeout',
|
||||
error_message = 'upstream submission confirmation timed out',
|
||||
finished_at = COALESCE(finished_at, now()),
|
||||
upstream_submission_updated_at = now()
|
||||
WHERE task_id = $1::uuid
|
||||
@@ -407,14 +407,15 @@ SET status = 'failed',
|
||||
billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END,
|
||||
billing_updated_at = now(),
|
||||
error = NULL,
|
||||
error_code = 'upstream_submission_unknown',
|
||||
error_message = 'upstream submission result is unknown',
|
||||
error_code = 'upstream_timeout',
|
||||
error_message = 'upstream submission confirmation timed out',
|
||||
locked_by = NULL,
|
||||
locked_at = NULL,
|
||||
heartbeat_at = NULL,
|
||||
execution_token = NULL,
|
||||
execution_lease_expires_at = NULL,
|
||||
remote_task_payload = '{}'::jsonb,
|
||||
submission_state = 'submission_confirmation_pending',
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil {
|
||||
@@ -424,7 +425,7 @@ WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil {
|
||||
return nil
|
||||
}
|
||||
payloadJSON, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "classification": "upstream_submission_unknown",
|
||||
"taskId": taskID, "classification": "upstream_timeout",
|
||||
})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO settlement_outbox (
|
||||
@@ -432,7 +433,7 @@ INSERT INTO settlement_outbox (
|
||||
status, next_attempt_at, manual_review_reason
|
||||
)
|
||||
SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currency,
|
||||
pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_submission_unknown'
|
||||
pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_timeout'
|
||||
FROM gateway_tasks
|
||||
WHERE id = $1::uuid
|
||||
ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil {
|
||||
@@ -866,6 +867,75 @@ WHERE task_id = $1::uuid
|
||||
return task, changed, nil
|
||||
}
|
||||
|
||||
// ResolveInterruptedTaskExecution atomically preserves the no-duplicate-submit
|
||||
// invariant when a Worker disappears or its execution context is cancelled.
|
||||
// A known remote task is safe to requeue for polling takeover. An ambiguous
|
||||
// submission without a remote ID is moved to manual review instead of replayed.
|
||||
func (s *Store) ResolveInterruptedTaskExecution(
|
||||
ctx context.Context,
|
||||
taskID string,
|
||||
executionToken string,
|
||||
) (GatewayTask, error) {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
defer rollbackTransaction(tx)
|
||||
var remoteTaskID string
|
||||
var hasGatewayUser bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(remote_task_id, ''), gateway_user_id IS NOT NULL
|
||||
FROM gateway_tasks
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'
|
||||
AND execution_token = $2::uuid
|
||||
FOR UPDATE`, taskID, executionToken).Scan(&remoteTaskID, &hasGatewayUser); err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return GatewayTask{}, ErrTaskExecutionLeaseLost
|
||||
}
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
ambiguous, err := taskExecutionRequiresManualReviewTx(ctx, tx, taskID)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
if remoteTaskID == "" && ambiguous {
|
||||
if err := markTaskExecutionManualReviewTx(ctx, tx, taskID, hasGatewayUser); err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
return GatewayTask{}, ErrTaskExecutionManualReview
|
||||
}
|
||||
nextRunAt := time.Now().Add(time.Second)
|
||||
queued, err := scanGatewayTask(tx.QueryRow(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'queued',
|
||||
locked_by = NULL,
|
||||
locked_at = NULL,
|
||||
heartbeat_at = NULL,
|
||||
execution_token = NULL,
|
||||
execution_lease_expires_at = NULL,
|
||||
next_run_at = $3,
|
||||
error = NULL,
|
||||
error_code = NULL,
|
||||
error_message = NULL,
|
||||
public_error = NULL,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'
|
||||
AND execution_token = $2::uuid
|
||||
RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt))
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
return queued, nil
|
||||
}
|
||||
|
||||
func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error {
|
||||
if riverJobID <= 0 {
|
||||
return nil
|
||||
|
||||
@@ -2,6 +2,7 @@ package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -9,31 +10,37 @@ import (
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const workerHeartbeatStaleAfter = 30 * time.Second
|
||||
const workerHeartbeatStaleAfter = 15 * time.Second
|
||||
|
||||
type WorkerRegistrationInput struct {
|
||||
InstanceID string
|
||||
PodUID string
|
||||
PodName string
|
||||
Site string
|
||||
Revision string
|
||||
DesiredCapacity int
|
||||
CapacityLimit int
|
||||
LoadMode string
|
||||
SafeCapacity int
|
||||
HeavyCapacity int
|
||||
ActiveTasks int
|
||||
PreparingTasks int
|
||||
WaitingUpstreamTasks int
|
||||
FinalizingTasks int
|
||||
PressureState string
|
||||
PressureReason string
|
||||
LoadSampledAt time.Time
|
||||
HeartbeatStaleAfter time.Duration
|
||||
InstanceID string
|
||||
WorkerID string
|
||||
PoolID string
|
||||
Endpoint string
|
||||
Labels map[string]string
|
||||
Capabilities map[string]any
|
||||
ProtocolVersion string
|
||||
OrchestratorInstanceRef string
|
||||
Revision string
|
||||
DesiredCapacity int
|
||||
CapacityLimit int
|
||||
LoadMode string
|
||||
SafeCapacity int
|
||||
HeavyCapacity int
|
||||
ActiveTasks int
|
||||
PreparingTasks int
|
||||
WaitingUpstreamTasks int
|
||||
FinalizingTasks int
|
||||
PressureState string
|
||||
PressureReason string
|
||||
LoadSampledAt time.Time
|
||||
HeartbeatStaleAfter time.Duration
|
||||
}
|
||||
|
||||
type WorkerAllocation struct {
|
||||
InstanceID string
|
||||
WorkerID string
|
||||
PoolID string
|
||||
DesiredCapacity int
|
||||
Allocated int
|
||||
GlobalAllocated int
|
||||
@@ -63,6 +70,31 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
|
||||
if input.InstanceID == "" {
|
||||
return WorkerAllocation{}, errors.New("worker instance ID is required")
|
||||
}
|
||||
input.WorkerID = strings.TrimSpace(input.WorkerID)
|
||||
if input.WorkerID == "" {
|
||||
input.WorkerID = input.InstanceID
|
||||
}
|
||||
input.PoolID = strings.TrimSpace(input.PoolID)
|
||||
if input.PoolID == "" {
|
||||
input.PoolID = "legacy-default"
|
||||
}
|
||||
if strings.TrimSpace(input.ProtocolVersion) == "" {
|
||||
input.ProtocolVersion = "v1"
|
||||
}
|
||||
if input.Labels == nil {
|
||||
input.Labels = map[string]string{}
|
||||
}
|
||||
if input.Capabilities == nil {
|
||||
input.Capabilities = map[string]any{}
|
||||
}
|
||||
labels, err := json.Marshal(input.Labels)
|
||||
if err != nil {
|
||||
return WorkerAllocation{}, err
|
||||
}
|
||||
capabilities, err := json.Marshal(input.Capabilities)
|
||||
if err != nil {
|
||||
return WorkerAllocation{}, err
|
||||
}
|
||||
if input.DesiredCapacity < 0 {
|
||||
return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative")
|
||||
}
|
||||
@@ -110,22 +142,38 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
|
||||
return WorkerAllocation{}, err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_execution_pools (pool_id, labels, capabilities, state, updated_at)
|
||||
VALUES ($1, $2::jsonb, $3::jsonb, 'active', now())
|
||||
ON CONFLICT (pool_id) DO UPDATE
|
||||
SET labels = EXCLUDED.labels,
|
||||
capabilities = EXCLUDED.capabilities,
|
||||
updated_at = now()
|
||||
WHERE gateway_execution_pools.state <> 'disabled'`, input.PoolID, labels, capabilities); err != nil {
|
||||
return WorkerAllocation{}, err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_worker_instances (
|
||||
instance_id, pod_uid, pod_name, site, revision, status,
|
||||
instance_id, worker_id, pool_id, endpoint, labels, capabilities, protocol_version,
|
||||
orchestrator_instance_ref, revision, status,
|
||||
desired_capacity, capacity_limit, hard_capacity_limit, safe_capacity, heavy_capacity,
|
||||
active_tasks, preparing_tasks, waiting_upstream_tasks, finalizing_tasks,
|
||||
pressure_state, pressure_reason, load_sampled_at,
|
||||
allocated_capacity, started_at, heartbeat_at, updated_at
|
||||
)
|
||||
VALUES (
|
||||
$1, $2, $3, $4, $5, 'active', $6, $7, $8, $9, $10,
|
||||
$11, $12, $13, $14, $15, $16, $17,
|
||||
$1, $2, $3, $4, $5::jsonb, $6::jsonb, $7,
|
||||
$8, $9, 'active', $10, $11, $12, $13, $14,
|
||||
$15, $16, $17, $18, $19, $20, $21,
|
||||
0, now(), now(), now()
|
||||
)
|
||||
ON CONFLICT (instance_id) DO UPDATE
|
||||
SET pod_uid = EXCLUDED.pod_uid,
|
||||
pod_name = EXCLUDED.pod_name,
|
||||
site = EXCLUDED.site,
|
||||
SET worker_id = EXCLUDED.worker_id,
|
||||
pool_id = EXCLUDED.pool_id,
|
||||
endpoint = EXCLUDED.endpoint,
|
||||
labels = EXCLUDED.labels,
|
||||
capabilities = EXCLUDED.capabilities,
|
||||
protocol_version = EXCLUDED.protocol_version,
|
||||
orchestrator_instance_ref = EXCLUDED.orchestrator_instance_ref,
|
||||
revision = EXCLUDED.revision,
|
||||
status = CASE
|
||||
WHEN gateway_worker_instances.status = 'draining' THEN 'draining'
|
||||
@@ -146,9 +194,13 @@ SET pod_uid = EXCLUDED.pod_uid,
|
||||
heartbeat_at = now(),
|
||||
updated_at = now()`,
|
||||
input.InstanceID,
|
||||
strings.TrimSpace(input.PodUID),
|
||||
strings.TrimSpace(input.PodName),
|
||||
strings.TrimSpace(input.Site),
|
||||
input.WorkerID,
|
||||
input.PoolID,
|
||||
strings.TrimRight(strings.TrimSpace(input.Endpoint), "/"),
|
||||
labels,
|
||||
capabilities,
|
||||
strings.TrimSpace(input.ProtocolVersion),
|
||||
strings.TrimSpace(input.OrchestratorInstanceRef),
|
||||
strings.TrimSpace(input.Revision),
|
||||
input.DesiredCapacity,
|
||||
input.CapacityLimit,
|
||||
@@ -235,6 +287,8 @@ WHERE instance_id = $1`, input.InstanceID).Scan(&heartbeatAt); err != nil {
|
||||
}
|
||||
return WorkerAllocation{
|
||||
InstanceID: input.InstanceID,
|
||||
WorkerID: input.WorkerID,
|
||||
PoolID: input.PoolID,
|
||||
DesiredCapacity: input.DesiredCapacity,
|
||||
Allocated: allocated,
|
||||
GlobalAllocated: globalAllocated,
|
||||
@@ -310,28 +364,30 @@ WHERE instance_id = $1
|
||||
}
|
||||
|
||||
type WorkerInstanceRuntime struct {
|
||||
InstanceID string `json:"instanceId"`
|
||||
PodUID string `json:"podUid,omitempty"`
|
||||
PodName string `json:"podName,omitempty"`
|
||||
Site string `json:"site,omitempty"`
|
||||
Revision string `json:"revision,omitempty"`
|
||||
Status string `json:"status"`
|
||||
Allocated int `json:"allocatedCapacity"`
|
||||
CapacityLimit int `json:"capacityLimit"`
|
||||
HardCapacityLimit int `json:"hardCapacityLimit"`
|
||||
SafeCapacity int `json:"safeCapacity"`
|
||||
HeavyCapacity int `json:"heavyCapacity"`
|
||||
ReportedActiveTasks int `json:"reportedActiveTasks"`
|
||||
PreparingTasks int `json:"preparingTasks"`
|
||||
WaitingUpstreamTasks int `json:"waitingUpstreamTasks"`
|
||||
FinalizingTasks int `json:"finalizingTasks"`
|
||||
PressureState string `json:"pressureState"`
|
||||
PressureReason string `json:"pressureReason,omitempty"`
|
||||
LoadSampledAt *time.Time `json:"loadSampledAt,omitempty"`
|
||||
RunningTasks int `json:"runningTasks"`
|
||||
ActiveLeases int `json:"activeLeases"`
|
||||
HeartbeatAt time.Time `json:"heartbeatAt"`
|
||||
DrainingAt *time.Time `json:"drainingAt,omitempty"`
|
||||
InstanceID string `json:"instanceId"`
|
||||
PodUID string `json:"podUid,omitempty"`
|
||||
PodName string `json:"podName,omitempty"`
|
||||
OrchestratorInstanceRef string `json:"-"`
|
||||
Site string `json:"site,omitempty"`
|
||||
PoolID string `json:"poolId,omitempty"`
|
||||
Revision string `json:"revision,omitempty"`
|
||||
Status string `json:"status"`
|
||||
Allocated int `json:"allocatedCapacity"`
|
||||
CapacityLimit int `json:"capacityLimit"`
|
||||
HardCapacityLimit int `json:"hardCapacityLimit"`
|
||||
SafeCapacity int `json:"safeCapacity"`
|
||||
HeavyCapacity int `json:"heavyCapacity"`
|
||||
ReportedActiveTasks int `json:"reportedActiveTasks"`
|
||||
PreparingTasks int `json:"preparingTasks"`
|
||||
WaitingUpstreamTasks int `json:"waitingUpstreamTasks"`
|
||||
FinalizingTasks int `json:"finalizingTasks"`
|
||||
PressureState string `json:"pressureState"`
|
||||
PressureReason string `json:"pressureReason,omitempty"`
|
||||
LoadSampledAt *time.Time `json:"loadSampledAt,omitempty"`
|
||||
RunningTasks int `json:"runningTasks"`
|
||||
ActiveLeases int `json:"activeLeases"`
|
||||
HeartbeatAt time.Time `json:"heartbeatAt"`
|
||||
DrainingAt *time.Time `json:"drainingAt,omitempty"`
|
||||
}
|
||||
|
||||
type WorkerQueueRuntime struct {
|
||||
@@ -340,6 +396,12 @@ type WorkerQueueRuntime struct {
|
||||
OldestWaitSeconds float64 `json:"oldestWaitSeconds"`
|
||||
}
|
||||
|
||||
type PoolQueueRuntime struct {
|
||||
PoolID string `json:"poolId"`
|
||||
Queued int `json:"queued"`
|
||||
Running int `json:"running"`
|
||||
}
|
||||
|
||||
type WorkerClusterRuntime struct {
|
||||
Workers []WorkerInstanceRuntime `json:"workers"`
|
||||
Queue WorkerQueueRuntime `json:"queue"`
|
||||
@@ -380,12 +442,40 @@ WHERE status IN ('queued', 'running')
|
||||
return snapshot, err
|
||||
}
|
||||
|
||||
func (s *Store) ListPoolQueueRuntime(ctx context.Context) ([]PoolQueueRuntime, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT assigned_pool_id,
|
||||
count(*) FILTER (WHERE status = 'queued')::int,
|
||||
count(*) FILTER (WHERE status = 'running')::int
|
||||
FROM gateway_tasks
|
||||
WHERE assigned_pool_id IS NOT NULL
|
||||
AND status IN ('queued', 'running')
|
||||
AND run_mode IN ('production', 'acceptance', 'acceptance_canary')
|
||||
GROUP BY assigned_pool_id
|
||||
ORDER BY assigned_pool_id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]PoolQueueRuntime, 0)
|
||||
for rows.Next() {
|
||||
var item PoolQueueRuntime
|
||||
if err := rows.Scan(&item.PoolID, &item.Queued, &item.Running); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.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.orchestrator_instance_ref,
|
||||
worker.site,
|
||||
worker.pool_id,
|
||||
worker.revision,
|
||||
worker.status,
|
||||
worker.allocated_capacity,
|
||||
@@ -414,7 +504,7 @@ LEFT JOIN gateway_concurrency_leases lease
|
||||
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`,
|
||||
ORDER BY worker.pool_id ASC, worker.status DESC, worker.instance_id ASC`,
|
||||
runtimeWorkerStaleAfter.String(),
|
||||
)
|
||||
if err != nil {
|
||||
@@ -428,7 +518,9 @@ ORDER BY worker.site ASC, worker.status DESC, worker.instance_id ASC`,
|
||||
&instance.InstanceID,
|
||||
&instance.PodUID,
|
||||
&instance.PodName,
|
||||
&instance.OrchestratorInstanceRef,
|
||||
&instance.Site,
|
||||
&instance.PoolID,
|
||||
&instance.Revision,
|
||||
&instance.Status,
|
||||
&instance.Allocated,
|
||||
|
||||
Reference in New Issue
Block a user