fix(queue): 防止过期执行占位阻塞与重复提交
队列恢复后,仍在运行标记中的 River job 可能保留 waiting admission,形成全局 FIFO 队头阻塞;同时旧执行在租约被接管后仍可能创建 attempt 或切换到 submitting。 新增 stale admission 自动让出机制,并用当前 execution token 对 attempt 创建和上游提交状态切换做 fencing。任务、River job 和结算状态均不在让出流程中改写。 验证:go vet ./...;env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;真实 PostgreSQL 集成测试覆盖 stale admission 让出与旧 token 提交拦截。
This commit is contained in:
@@ -702,12 +702,19 @@ WHERE id = $1::uuid
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) SetAttemptUpstreamSubmissionStatus(ctx context.Context, attemptID string, status string) error {
|
||||
func validateAttemptUpstreamSubmissionStatus(status string) error {
|
||||
switch status {
|
||||
case "not_submitted", "submitting", "response_received":
|
||||
default:
|
||||
return fmt.Errorf("invalid upstream submission status %q", status)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) SetAttemptUpstreamSubmissionStatus(ctx context.Context, attemptID string, status string) error {
|
||||
if err := validateAttemptUpstreamSubmissionStatus(status); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET upstream_submission_status = $2,
|
||||
@@ -716,6 +723,42 @@ WHERE id = $1::uuid`, attemptID, status)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) SetAttemptUpstreamSubmissionStatusForExecution(
|
||||
ctx context.Context,
|
||||
attemptID string,
|
||||
taskID string,
|
||||
executionToken string,
|
||||
status string,
|
||||
) error {
|
||||
if err := validateAttemptUpstreamSubmissionStatus(status); err != nil {
|
||||
return err
|
||||
}
|
||||
tag, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts attempt
|
||||
SET upstream_submission_status = $4,
|
||||
upstream_submission_updated_at = now()
|
||||
FROM gateway_tasks task
|
||||
WHERE attempt.id = $1::uuid
|
||||
AND attempt.task_id = task.id
|
||||
AND attempt.task_id = $2::uuid
|
||||
AND attempt.status = 'running'
|
||||
AND task.status = 'running'
|
||||
AND task.execution_token = $3::uuid
|
||||
AND task.execution_lease_expires_at > now()`,
|
||||
attemptID,
|
||||
taskID,
|
||||
executionToken,
|
||||
status,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return ErrTaskExecutionLeaseLost
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message string) (GatewayTask, bool, error) {
|
||||
message = strings.TrimSpace(message)
|
||||
if message == "" {
|
||||
@@ -980,6 +1023,24 @@ func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptIn
|
||||
}
|
||||
defer rollbackTransaction(tx)
|
||||
|
||||
if strings.TrimSpace(input.ExecutionToken) != "" {
|
||||
var ownsExecution bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT status = 'running'
|
||||
AND execution_token = $2::uuid
|
||||
AND execution_lease_expires_at > now()
|
||||
FROM gateway_tasks
|
||||
WHERE id = $1::uuid
|
||||
FOR UPDATE`,
|
||||
input.TaskID,
|
||||
input.ExecutionToken,
|
||||
).Scan(&ownsExecution); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !ownsExecution {
|
||||
return "", ErrTaskExecutionLeaseLost
|
||||
}
|
||||
}
|
||||
var attemptID string
|
||||
err = tx.QueryRow(ctx, `
|
||||
INSERT INTO gateway_task_attempts (
|
||||
|
||||
Reference in New Issue
Block a user