BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add river-backed async task queue

Browse Source
This commit is contained in:
wangbo
2026-05-12 10:11:54 +08:00
parent d69aaed444
commit 7e220b7477
30 changed files with 1342 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/tasks_runtime.go
+138 -14
View File
@@ -157,7 +157,7 @@ func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType st
_, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
model_type = NULLIF($2, ''),
model_type = NULLIF($2::text, ''),
normalized_request = $3::jsonb,
locked_at = now(),
heartbeat_at = now(),
@@ -166,6 +166,124 @@ WHERE id = $1::uuid`, taskID, modelType, string(normalizedJSON))
return err
}
func (s *Store) ClaimAsyncQueuedTask(ctx context.Context, workerID string) (GatewayTask, error) {
return scanGatewayTask(s.pool.QueryRow(ctx, `
WITH picked AS (
SELECT id AS task_id
FROM gateway_tasks
WHERE async_mode = true
AND status = 'queued'
AND next_run_at <= now()
ORDER BY priority ASC, created_at ASC
LIMIT 1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_tasks t
SET status = 'running',
locked_by = NULLIF($1::text, ''),
locked_at = now(),
heartbeat_at = now(),
updated_at = now()
FROM picked
WHERE t.id = picked.task_id
RETURNING `+gatewayTaskColumns, workerID))
}
func (s *Store) RequeueTask(ctx context.Context, taskID string, delay time.Duration) (GatewayTask, error) {
if delay < time.Second {
delay = time.Second
}
if delay > 10*time.Minute {
delay = 10 * time.Minute
}
nextRunAt := time.Now().Add(delay)
return scanGatewayTask(s.pool.QueryRow(ctx, `
UPDATE gateway_tasks
SET status = 'queued',
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
next_run_at = $2::timestamptz,
error = NULL,
error_code = NULL,
error_message = NULL,
updated_at = now()
WHERE id = $1::uuid
RETURNING `+gatewayTaskColumns, taskID, nextRunAt))
}
func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error {
if riverJobID <= 0 {
return nil
}
_, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2,
updated_at = now()
WHERE id = $1::uuid`, taskID, riverJobID)
return err
}
func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, attemptID string, remoteTaskID string, payload map[string]any) error {
payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload))
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET remote_task_id = NULLIF($2::text, ''),
remote_task_payload = $3::jsonb,
updated_at = now()
WHERE id = $1::uuid`,
taskID,
remoteTaskID,
string(payloadJSON),
); err != nil {
return err
}
if strings.TrimSpace(attemptID) == "" {
return nil
}
_, err := tx.Exec(ctx, `
UPDATE gateway_task_attempts
SET remote_task_id = NULLIF($2::text, ''),
response_snapshot = COALESCE(response_snapshot, '{}'::jsonb) || jsonb_build_object('remote_task_payload', $3::jsonb)
WHERE id = $1::uuid`,
attemptID,
remoteTaskID,
string(payloadJSON),
)
return err
})
}
func (s *Store) ListRecoverableAsyncTasks(ctx context.Context, limit int) ([]AsyncTaskQueueItem, error) {
if limit <= 0 {
limit = 500
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, priority
FROM gateway_tasks
WHERE async_mode = true
AND status IN ('queued', 'running')
ORDER BY priority ASC, created_at ASC
LIMIT $1`, limit)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]AsyncTaskQueueItem, 0)
for rows.Next() {
var item AsyncTaskQueueItem
if err := rows.Scan(&item.ID, &item.Priority); err != nil {
return nil, err
}
items = append(items, item)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) {
requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot))
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
@@ -182,7 +300,7 @@ INSERT INTO gateway_task_attempts (
status, simulated, request_snapshot, metrics
)
VALUES (
$1::uuid, $2, NULLIF($3, '')::uuid, NULLIF($4, '')::uuid, NULLIF($5, ''), $6,
$1::uuid, $2::int, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, NULLIF($5::text, ''), $6,
$7, $8, $9::jsonb, $10::jsonb
)
RETURNING id::text`,
@@ -202,7 +320,7 @@ RETURNING id::text`,
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET attempt_count = GREATEST(attempt_count, $2), updated_at = now()
SET attempt_count = GREATEST(attempt_count, $2::int), updated_at = now()
WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil {
return "", err
}
@@ -252,7 +370,7 @@ SET metrics = jsonb_set(
true
)
WHERE task_id = $1::uuid
AND attempt_no = $2`, taskID, attemptNo, string(entryJSON))
AND attempt_no = $2::int`, taskID, attemptNo, string(entryJSON))
return err
}
@@ -386,17 +504,17 @@ func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptIn
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_attempts
SET status = $2,
SET status = $2::text,
retryable = $3,
request_id = NULLIF($4, ''),
request_id = NULLIF($4::text, ''),
usage = $5::jsonb,
metrics = $6::jsonb,
response_snapshot = $7::jsonb,
response_started_at = $8::timestamptz,
response_finished_at = $9::timestamptz,
response_duration_ms = $10,
error_code = NULLIF($11, ''),
error_message = NULLIF($12, ''),
error_code = NULLIF($11::text, ''),
error_message = NULLIF($12::text, ''),
finished_at = now()
WHERE id = $1::uuid`,
input.AttemptID,
@@ -438,6 +556,9 @@ SET status = 'succeeded',
error = NULL,
error_code = NULL,
error_message = NULL,
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`,
@@ -561,14 +682,17 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
if _, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
error = NULLIF($2, ''),
error_code = NULLIF($3, ''),
error_message = NULLIF($2, ''),
request_id = NULLIF($4, ''),
error = NULLIF($2::text, ''),
error_code = NULLIF($3::text, ''),
error_message = NULLIF($2::text, ''),
request_id = NULLIF($4::text, ''),
metrics = $5::jsonb,
response_started_at = $6::timestamptz,
response_finished_at = $7::timestamptz,
response_duration_ms = $8,
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`,
@@ -604,7 +728,7 @@ WITH next_seq AS (
WHERE task_id = $1::uuid
)
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
SELECT $1::uuid, next_seq.seq, $2, NULLIF($3, ''), NULLIF($4, ''), $5, NULLIF($6, ''), $7::jsonb, $8
SELECT $1::uuid, next_seq.seq, $2::text, NULLIF($3::text, ''), NULLIF($4::text, ''), $5, NULLIF($6::text, ''), $7::jsonb, $8
FROM next_seq
RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''),
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at`,
@@ -688,7 +812,7 @@ func (s *Store) RecordClientRelease(ctx context.Context, clientID string, lastEr
_, err := s.pool.Exec(ctx, `
UPDATE runtime_client_states
SET running_count = GREATEST(running_count - 1, 0),
last_error = NULLIF($2, ''),
last_error = NULLIF($2::text, ''),
updated_at = now()
WHERE client_id = $1`, clientID, lastError)
return err