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/rate_limits.go
+134 -19
View File
@@ -3,6 +3,7 @@ package store
import (
"context"
"errors"
"fmt"
"time"
"github.com/jackc/pgx/v5"
@@ -13,6 +14,7 @@ type RuntimeRecoveryResult struct {
ReleasedRateReservations int64 `json:"releasedRateReservations"`
FailedAttempts int64 `json:"failedAttempts"`
FailedTasks int64 `json:"failedTasks"`
RequeuedAsyncTasks int64 `json:"requeuedAsyncTasks"`
}
func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID string, reservations []RateLimitReservation) (RateLimitResult, error) {
@@ -28,7 +30,11 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
continue
}
if reservation.Metric == "" || reservation.Amount > reservation.Limit {
return RateLimitResult{}, ErrRateLimited
return RateLimitResult{}, &RateLimitExceededError{
Metric: reservation.Metric,
Message: fmt.Sprintf("rate limit exceeded: %s request amount %.0f is greater than limit %.0f", reservation.Metric, reservation.Amount, reservation.Limit),
Retryable: false,
}
}
if reservation.WindowSeconds <= 0 {
reservation.WindowSeconds = 60
@@ -55,8 +61,10 @@ func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, atte
reservation.LeaseTTLSeconds = 120
}
var active float64
var nextAvailableAt time.Time
if err := tx.QueryRow(ctx, `
SELECT COALESCE(SUM(lease_value), 0)::float8
SELECT COALESCE(SUM(lease_value), 0)::float8,
COALESCE(MIN(expires_at), now() + ($3::int * interval '1 second'))
FROM gateway_concurrency_leases
WHERE scope_type = $1
AND scope_key = $2
@@ -64,11 +72,17 @@ WHERE scope_type = $1
AND expires_at > now()`,
reservation.ScopeType,
reservation.ScopeKey,
).Scan(&active); err != nil {
reservation.LeaseTTLSeconds,
).Scan(&active, &nextAvailableAt); err != nil {
return "", err
}
if active+reservation.Amount > reservation.Limit {
return "", ErrRateLimited
return "", &RateLimitExceededError{
Metric: reservation.Metric,
Message: fmt.Sprintf("rate limit exceeded: concurrent active %.0f plus request %.0f is greater than limit %.0f", active, reservation.Amount, reservation.Limit),
RetryAfter: concurrencyRetryAfter(nextAvailableAt),
Retryable: true,
}
}
var leaseID string
if err := tx.QueryRow(ctx, `
@@ -92,13 +106,16 @@ func reserveCounterWindow(ctx context.Context, tx pgx.Tx, taskID string, attempt
reservedAmount := reservation.Amount
var windowStart time.Time
err := tx.QueryRow(ctx, `
WITH bounds AS (
SELECT
to_timestamp(floor(extract(epoch FROM now()) / $7::int) * $7::int) AS window_start,
to_timestamp(floor(extract(epoch FROM now()) / $7::int) * $7::int) + ($7::int * interval '1 second') AS reset_at
)
INSERT INTO gateway_rate_limit_counters (
scope_type, scope_key, metric, window_start, limit_value, used_value, reserved_value, reset_at
)
VALUES (
$1, $2, $3, date_trunc('minute', now()), $4, $5, $6,
date_trunc('minute', now()) + ($7::int * interval '1 second')
)
SELECT $1, $2, $3, bounds.window_start, $4, $5, $6, bounds.reset_at
FROM bounds
ON CONFLICT (scope_type, scope_key, metric, window_start) DO UPDATE
SET limit_value = EXCLUDED.limit_value,
used_value = gateway_rate_limit_counters.used_value + EXCLUDED.used_value,
@@ -117,7 +134,28 @@ RETURNING window_start`,
).Scan(&windowStart)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return RateLimitReservation{}, ErrRateLimited
resetAt := time.Now().Add(time.Duration(reservation.WindowSeconds) * time.Second)
_ = tx.QueryRow(ctx, `
WITH bounds AS (
SELECT to_timestamp(floor(extract(epoch FROM now()) / $4::int) * $4::int) AS window_start
)
SELECT counters.reset_at
FROM gateway_rate_limit_counters counters
JOIN bounds ON counters.window_start = bounds.window_start
WHERE scope_type = $1
AND scope_key = $2
AND metric = $3`,
reservation.ScopeType,
reservation.ScopeKey,
reservation.Metric,
reservation.WindowSeconds,
).Scan(&resetAt)
return RateLimitReservation{}, &RateLimitExceededError{
Metric: reservation.Metric,
Message: fmt.Sprintf("rate limit exceeded: %s window has no remaining capacity", reservation.Metric),
RetryAfter: retryAfterUntil(resetAt),
Retryable: true,
}
}
return RateLimitReservation{}, err
}
@@ -144,6 +182,28 @@ RETURNING id::text`,
return reservation, nil
}
func retryAfterUntil(when time.Time) time.Duration {
if when.IsZero() {
return 0
}
duration := time.Until(when)
if duration < time.Second {
return time.Second
}
return duration
}
func concurrencyRetryAfter(leaseExpiresAt time.Time) time.Duration {
if leaseExpiresAt.IsZero() {
return time.Second
}
duration := time.Until(leaseExpiresAt)
if duration <= time.Second {
return time.Second
}
return time.Second
}
func (s *Store) CommitRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64) error {
return s.finishRateLimitReservations(ctx, reservations, actualByMetric, "committed", "success")
}
@@ -189,23 +249,26 @@ RETURNING scope_type, scope_key, metric, window_start, reserved_amount::float8`)
if err != nil {
return RuntimeRecoveryResult{}, err
}
reservations := make([]RateLimitReservation, 0)
for rows.Next() {
var reservation RateLimitReservation
if err := rows.Scan(&reservation.ScopeType, &reservation.ScopeKey, &reservation.Metric, &reservation.WindowStart, &reservation.Amount); err != nil {
rows.Close()
return RuntimeRecoveryResult{}, err
}
if err := releaseCounterReservation(ctx, tx, reservation.ScopeType, reservation.ScopeKey, reservation.Metric, reservation.WindowStart, reservation.Amount); err != nil {
rows.Close()
return RuntimeRecoveryResult{}, err
}
result.ReleasedRateReservations++
reservations = append(reservations, reservation)
}
if err := rows.Err(); err != nil {
rows.Close()
return RuntimeRecoveryResult{}, err
}
rows.Close()
for _, reservation := range reservations {
if err := releaseCounterReservation(ctx, tx, reservation.ScopeType, reservation.ScopeKey, reservation.Metric, reservation.WindowStart, reservation.Amount); err != nil {
return RuntimeRecoveryResult{}, err
}
}
result.ReleasedRateReservations = int64(len(reservations))
tag, err := tx.Exec(ctx, `
UPDATE gateway_concurrency_leases
@@ -220,7 +283,7 @@ WHERE released_at IS NULL
tag, err = tx.Exec(ctx, `
UPDATE gateway_task_attempts
SET status = 'failed',
retryable = false,
retryable = true,
error_code = 'server_restarted',
error_message = 'attempt interrupted by service restart',
finished_at = now()
@@ -230,6 +293,57 @@ WHERE status = 'running'`)
}
result.FailedAttempts = tag.RowsAffected()
asyncTaskRows, err := tx.Query(ctx, `
UPDATE gateway_tasks
SET status = 'queued',
error = NULL,
error_code = NULL,
error_message = NULL,
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
next_run_at = now(),
finished_at = NULL,
updated_at = now()
WHERE async_mode = true
AND status = 'running'
RETURNING id::text`)
if err != nil {
return RuntimeRecoveryResult{}, err
}
asyncTaskIDs := make([]string, 0)
for asyncTaskRows.Next() {
var taskID string
if err := asyncTaskRows.Scan(&taskID); err != nil {
asyncTaskRows.Close()
return RuntimeRecoveryResult{}, err
}
asyncTaskIDs = append(asyncTaskIDs, taskID)
}
if err := asyncTaskRows.Err(); err != nil {
asyncTaskRows.Close()
return RuntimeRecoveryResult{}, err
}
asyncTaskRows.Close()
for _, taskID := range asyncTaskIDs {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
VALUES (
$1::uuid,
COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1),
'task.recovered',
'queued',
'recovered',
0.2,
'async task recovered after service restart',
'{"code":"server_restarted"}'::jsonb,
false
)`, taskID); err != nil {
return RuntimeRecoveryResult{}, err
}
}
result.RequeuedAsyncTasks = int64(len(asyncTaskIDs))
taskRows, err := tx.Query(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
@@ -238,7 +352,8 @@ SET status = 'failed',
error_message = 'task interrupted by service restart',
finished_at = now(),
updated_at = now()
WHERE status IN ('queued', 'running')
WHERE async_mode = false
AND status = 'running'
RETURNING id::text`)
if err != nil {
return RuntimeRecoveryResult{}, err
@@ -301,9 +416,9 @@ func (s *Store) finishRateLimitReservations(ctx context.Context, reservations []
var stored RateLimitReservation
err := tx.QueryRow(ctx, `
UPDATE gateway_rate_limit_reservations
SET status = $2,
reason = NULLIF($3, ''),
actual_amount = CASE WHEN $2 = 'committed' THEN $4 ELSE actual_amount END,
SET status = $2::text,
reason = NULLIF($3::text, ''),
actual_amount = CASE WHEN $2::text = 'committed' THEN $4 ELSE actual_amount END,
finalized_at = now(),
updated_at = now()
WHERE id = $1::uuid