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

Fix realtime queued task counts

Browse Source
This commit is contained in:
wangbo
2026-05-12 10:41:07 +08:00
parent 2a91b31d12
commit 9ea83be718
6 changed files with 128 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/httpapi/core_flow_integration_test.go
+18
View File
@@ -783,6 +783,24 @@ WHERE reference_type = 'gateway_task'
if len(asyncRateLimitDetail.Attempts) != 0 {
t.Fatalf("async rate-limited task should wait in queue without recording a failed attempt: %+v", asyncRateLimitDetail)
}
var modelRateLimits struct {
Items []struct {
ModelName string `json:"modelName"`
ModelAlias string `json:"modelAlias"`
QueuedTasks float64 `json:"queuedTasks"`
} `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/runtime/model-rate-limits", loginResponse.AccessToken, nil, http.StatusOK, &modelRateLimits)
var queuedTasks float64
for _, item := range modelRateLimits.Items {
if item.ModelName == rateLimitedModel || item.ModelAlias == rateLimitedModel {
queuedTasks = item.QueuedTasks
break
}
}
if queuedTasks < 1 {
t.Fatalf("realtime load should count async rate-limited task as queued, got %v in %+v", queuedTasks, modelRateLimits.Items)
}
asyncRateLimitCompleted := waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, asyncRateLimitTask.TaskID, []string{"succeeded"}, time.Duration(rateLimitWindowSeconds+3)*time.Second)
if asyncRateLimitCompleted.Status != "succeeded" {
t.Fatalf("async rate-limited task should be pulled from queue after the limit window resets, got %+v", asyncRateLimitCompleted)
apps/api/internal/runner/service.go
+7 -5
View File
@@ -121,11 +121,13 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
maxFailoverDuration := maxFailoverDurationForCandidates(candidates, runnerPolicy)
attemptNo := task.AttemptCount
var lastErr error
var lastCandidate store.RuntimeModelCandidate
candidatesLoop:
for index, candidate := range candidates {
if index >= maxPlatforms {
break
}
lastCandidate = candidate
clientAttempts := clientAttemptsForCandidate(candidate)
var candidateErr error
for clientAttempt := 1; clientAttempt <= clientAttempts; clientAttempt++ {
@@ -180,7 +182,7 @@ candidatesLoop:
lastErr = err
candidateErr = err
if task.AsyncMode && store.RateLimitRetryable(err) {
queued, delay, queueErr := s.requeueRateLimitedTask(ctx, task, err)
queued, delay, queueErr := s.requeueRateLimitedTask(ctx, task, err, candidate)
if queueErr != nil {
return Result{}, queueErr
}
@@ -297,7 +299,7 @@ candidatesLoop:
return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: 0}
}
if task.AsyncMode && errors.Is(lastErr, store.ErrRateLimited) && store.RateLimitRetryable(lastErr) {
queued, delay, queueErr := s.requeueRateLimitedTask(ctx, task, lastErr)
queued, delay, queueErr := s.requeueRateLimitedTask(ctx, task, lastErr, lastCandidate)
if queueErr != nil {
return Result{}, queueErr
}
@@ -525,12 +527,12 @@ func (s *Service) failTask(ctx context.Context, taskID string, code string, mess
return failed, nil
}
func (s *Service) requeueRateLimitedTask(ctx context.Context, task store.GatewayTask, cause error) (store.GatewayTask, time.Duration, error) {
func (s *Service) requeueRateLimitedTask(ctx context.Context, task store.GatewayTask, cause error, candidate store.RuntimeModelCandidate) (store.GatewayTask, time.Duration, error) {
delay := localRateLimitRetryAfter(cause)
if delay <= 0 {
delay = 5 * time.Second
}
queued, err := s.store.RequeueTask(ctx, task.ID, delay)
queued, err := s.store.RequeueTask(ctx, task.ID, delay, candidate.QueueKey)
if err != nil {
return store.GatewayTask{}, 0, err
}
@@ -546,7 +548,7 @@ func (s *Service) requeueRateLimitedTask(ctx context.Context, task store.Gateway
}
func (s *Service) requeueInterruptedAsyncTask(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) {
queued, err := s.store.RequeueTask(ctx, task.ID, 0)
queued, err := s.store.RequeueTask(ctx, task.ID, 0, "")
if err != nil {
return store.GatewayTask{}, err
}
apps/api/internal/store/rate_limit_status.go
+21 -7
View File
@@ -62,17 +62,31 @@ LEFT JOIN (
GROUP BY scope_key
) con ON con.scope_key = m.id::text
LEFT JOIN (
SELECT latest.platform_model_id, COUNT(*) AS waiting
SELECT queued_sources.platform_model_id, COUNT(DISTINCT queued_sources.task_id) AS waiting
FROM (
SELECT DISTINCT ON (a.task_id) a.task_id, a.platform_model_id::text AS platform_model_id
SELECT t.id::text AS task_id, qm.id::text AS platform_model_id
FROM gateway_tasks t
JOIN gateway_task_attempts a ON a.task_id = t.id
JOIN integration_platforms qp ON TRUE
JOIN platform_models qm ON qm.platform_id = qp.id
WHERE t.async_mode = true
AND t.status = 'queued'
AND a.platform_model_id IS NOT NULL
ORDER BY a.task_id, a.attempt_no DESC, a.started_at DESC
) latest
GROUP BY latest.platform_model_id
AND NULLIF(t.model_type, '') IS NOT NULL
AND qm.model_type @> jsonb_build_array(t.model_type)
AND t.queue_key = qp.platform_key || ':' || t.model_type || ':' || COALESCE(NULLIF(qm.provider_model_name, ''), qm.model_name)
AND NOT EXISTS (SELECT 1 FROM gateway_task_attempts existing_attempt WHERE existing_attempt.task_id = t.id)
UNION ALL
SELECT latest.task_id, latest.platform_model_id
FROM (
SELECT DISTINCT ON (a.task_id) a.task_id, a.platform_model_id::text AS platform_model_id
FROM gateway_tasks t
JOIN gateway_task_attempts a ON a.task_id = t.id
WHERE t.async_mode = true
AND t.status = 'queued'
AND a.platform_model_id IS NOT NULL
ORDER BY a.task_id, a.attempt_no DESC, a.started_at DESC
) latest
) queued_sources
GROUP BY queued_sources.platform_model_id
) queued ON queued.platform_model_id = m.id::text
LEFT JOIN (
SELECT DISTINCT ON (scope_key) scope_key, used_value, reserved_value, reset_at
apps/api/internal/store/tasks_runtime.go
+3 -2
View File
@@ -189,7 +189,7 @@ WHERE t.id = picked.task_id
RETURNING `+gatewayTaskColumns, workerID))
}
func (s *Store) RequeueTask(ctx context.Context, taskID string, delay time.Duration) (GatewayTask, error) {
func (s *Store) RequeueTask(ctx context.Context, taskID string, delay time.Duration, queueKey string) (GatewayTask, error) {
if delay < time.Second {
delay = time.Second
}
@@ -204,12 +204,13 @@ SET status = 'queued',
locked_at = NULL,
heartbeat_at = NULL,
next_run_at = $2::timestamptz,
queue_key = COALESCE(NULLIF($3::text, ''), queue_key),
error = NULL,
error_code = NULL,
error_message = NULL,
updated_at = now()
WHERE id = $1::uuid
RETURNING `+gatewayTaskColumns, taskID, nextRunAt))
RETURNING `+gatewayTaskColumns, taskID, nextRunAt, strings.TrimSpace(queueKey)))
}
func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error {