From 9ea60ddd9dcd667611974fa55568458aeb18e11a Mon Sep 17 00:00:00 2001 From: supreme0597 Date: Sat, 13 Jun 2026 17:02:06 +0800 Subject: [PATCH 01/18] docs: add channel flow control design notes --- docs/channel-flow-control-queue-design-v2.md | 1201 +++++++++++++ docs/channel-flow-control-queue-design-v3.md | 1699 ++++++++++++++++++ docs/channel-flow-control-queue-design.md | 1476 +++++++++++++++ tools/channel-flow-spike/README.md | 32 + tools/channel-flow-spike/main.go | 522 ++++++ 5 files changed, 4930 insertions(+) create mode 100644 docs/channel-flow-control-queue-design-v2.md create mode 100644 docs/channel-flow-control-queue-design-v3.md create mode 100644 docs/channel-flow-control-queue-design.md create mode 100644 tools/channel-flow-spike/README.md create mode 100644 tools/channel-flow-spike/main.go diff --git a/docs/channel-flow-control-queue-design-v2.md b/docs/channel-flow-control-queue-design-v2.md new file mode 100644 index 000000000000..d7e891dc2370 --- /dev/null +++ b/docs/channel-flow-control-queue-design-v2.md @@ -0,0 +1,1201 @@ +# Channel Flow Control and Queue Design Report v2 + +Date: 2026-06-13 + +Status: v2 draft after design audit + +Related documents: + +- `docs/channel-flow-control-queue-design.md` +- `/Users/laiyouxu/.gemini/antigravity-cli/brain/fdf12fcb-bc0f-48af-9f54-2dc1902d1eb9/flow-control-design-audit.md` + +## 1. v2 Executive Summary + +This v2 report answers the design audit questions and refines the implementation plan for channel-level flow control and queueing in new-api. + +The main corrections from v1 are: + +1. Flow control must run after channel setup and upstream model mapping, not as generic middleware. +2. Each retry attempt must acquire and release its own flow-control guard. +3. Billing should not pre-consume quota while a request is waiting in queue. Add a read-only billing precheck before queueing, then pre-consume only after a slot is acquired. +4. Raw `pool_id` should not be user-provided. The web UI exposes "Flow Pool"; the backend generates `pool_key`. +5. Queue length must be bounded. +6. Redis production backend is needed for multi-instance deployments. +7. Redis Lua is not strictly required for v2. The recommended v2 Redis implementation is `WATCH/MULTI` optimistic transactions plus short polling. Lua can be introduced later as a performance optimization. +8. Queue wakeup must not rely only on Redis Pub/Sub. Use poll-first logic; Pub/Sub is optional acceleration. +9. Add graceful shutdown, runtime config versioning, request body memory caps, metrics retention, and clear backend status warnings. + +Recommended initial settings for a 96-GPU upstream that supports 60 concurrent requests: + +```text +max_inflight: 60 +max_queue_size: 240 +queue_timeout_ms: 120000 +queue_policy: fifo +on_limit: queue +backend: redis in production, memory only for single instance/dev +``` + +## 2. Final Position on Lua + +### 2.1 Is Lua Required? + +No. Lua is not strictly required. + +The system needs atomic "check capacity then add running/waiting entry" semantics. Redis Lua is one way to do this, but not the only way. + +Available options: + +| Option | Atomic | Complexity | Performance | Maintainability | Recommendation | +|---|---:|---:|---:|---:|---| +| Go memory lock | Yes, single process only | Low | High | High | Use for dev/single instance | +| Redis `WATCH/MULTI` | Yes, with retries | Medium | Medium | Medium-high | Recommended v2 Redis backend | +| Redis Lua | Yes | High | High | Medium-low | Optional later optimization | +| Redis Streams only | Not enough by itself | High | Medium | Medium | Not v2 | +| DB row locks | Yes | High | Low/medium | Medium | Not recommended for hot path | + +### 2.2 Why Not Make Lua Mandatory in v2? + +new-api already has one Redis Lua token-bucket helper under `common/limiter`, but most Redis usage in the project is simple wrapper calls. Making a complex queue/semaphore system depend on Lua in the first implementation raises several risks: + +- Harder debugging and testing. +- Redis Cluster key-slot requirements. +- More operational knowledge required. +- Risk of long-running Lua scripts blocking Redis if cleanup scans too much. +- Harder to iterate while the product behavior is still being validated. + +### 2.3 Recommended v2 Redis Strategy + +Use Redis optimistic transactions: + +```text +WATCH running, waiting, config +read current state +MULTI + mutate running/waiting/request metadata +EXEC +if conflict -> retry with jitter +``` + +This gives atomicity without Lua. Under contention, `EXEC` may fail and retry. That is acceptable for v2 because: + +- The target scenario is queueing around an upstream bottleneck, not millions of requests per second. +- A little retry overhead is easier to operate than a complex Lua scheduler. +- We can cap transaction retries and fall back to short poll. + +### 2.4 When Should Lua Be Introduced? + +Lua should be considered in a later phase if metrics show: + +- Too many Redis transaction conflicts. +- Acquire latency from `WATCH/MULTI` becomes significant. +- Redis round trips become the bottleneck. +- Queue promotion needs to batch many waiters efficiently. + +If Lua is introduced later, it must follow these rules: + +- Use Redis hash tags so all keys for a pool share one slot: + +```text +flow:{pool_key}:running +flow:{pool_key}:waiting +flow:{pool_key}:config +``` + +Here `{pool_key}` is the Redis hash tag. The literal braces matter for Redis Cluster compatibility. + +- Limit cleanup work per script execution. +- Put script loading/execution behind a common helper, not scattered through business code. +- Add focused tests for each script. + +## 3. v2 Architecture Overview + +```text +Request + -> auth and request validation + -> token estimate and price estimate + -> billing precheck, read-only + -> channel selection + -> SetupContextForSelectedChannel + -> upstream model resolved + -> resolve Flow Pool binding + -> acquire Flow Guard + -> billing pre-consume, actual deduction + -> call upstream + -> release Flow Guard when attempt/stream/task completes + -> settle/refund billing as today +``` + +Important separation: + +```text +User rate limit: who may send how many requests +Billing: whether user/token/subscription can pay +Flow control: whether upstream resource pool has capacity +``` + +These should be separate services. + +## 4. Flow Pool Product Model + +### 4.1 User-visible Concept + +Users should not type `pool_id`. + +The admin UI exposes: + +```text +Flow Pool + name: "96-card DeepSeek-R1 production pool" + description + max_inflight + max_queue_size + queue_timeout_ms + queue_policy + bindings +``` + +The backend generates: + +```text +pool_key: flow_pool_8f3a2c... +``` + +Runtime Redis keys, logs, and metrics use `pool_key`. + +### 4.2 Binding to Channels and Upstream Models + +Binding must be explicit. + +Runtime resolution priority: + +```text +1. channel_id + upstream_model exact binding +2. channel_id binding +3. no binding -> no flow control +``` + +URL/base URL is only used to suggest possible bindings. It must not silently merge pools. + +Reason: + +- Same base URL can serve different physical GPU pools. +- Same base URL plus different key can map to different tenants. +- Same physical pool can have multiple URLs. +- Model mapping can change the actual upstream model. + +### 4.3 Web UI Placement + +Default frontend integration points: + +```text +web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx +web/default/src/features/channels/lib/channel-form.ts +web/default/src/features/channels/types.ts +``` + +Add a section in the channel create/update drawer: + +```text +Advanced Settings + -> Flow Control & Queue +``` + +Controls: + +```text +[Switch] Enable flow control and queue + +Resource pool + ( ) Create an independent pool for this channel + ( ) Bind to existing Flow Pool + ( ) Create new Flow Pool + +Binding scope + ( ) All upstream models on this channel + ( ) Selected upstream models + +Capacity + Max in-flight requests + Max queue size + Queue timeout + Queue policy + +Upstream identity preview + Channel type + Base URL + Published models + Model mapping + Suggested similar channels +``` + +Add a management tab: + +```text +Channels | Flow Pools +``` + +List columns: + +```text +Name +Bound channels +Running / max_inflight +Queued / max_queue_size +Wait P95 +Rejected / timeout +Backend +Health +``` + +## 5. Data Model + +Use DB tables, not only per-channel JSON. Shared pool configuration cannot be safely represented in multiple channel JSON blobs. + +### 5.1 `channel_flow_pools` + +```text +id int primary key +pool_key varchar unique, generated +name varchar +description text +enabled bool/int +max_inflight int +max_queue_size int +queue_timeout_ms int +queue_policy varchar, default "fifo" +on_limit varchar, default "queue" +max_context_tokens int, optional +max_context_chars int, optional +max_processing_ms int, optional +task_release_policy varchar, default "on_submit" +config_version bigint +created_time bigint +updated_time bigint +``` + +### 5.2 `channel_flow_pool_bindings` + +```text +id int primary key +pool_id int +channel_id int +upstream_model varchar, optional +match_mode varchar, "channel" | "channel_model" +enabled bool/int +created_time bigint +updated_time bigint +``` + +### 5.3 `channel_flow_metrics_minute` + +```text +id int primary key +bucket_ts bigint +pool_key varchar +channel_id int +model varchar +running_avg double or integer approximation +running_max int +queued_avg double or integer approximation +queued_max int +acquired_count int +queued_count int +released_count int +rejected_count int +timeout_count int +cancelled_count int +lease_renew_fail int +wait_ms_avg int +wait_ms_max int +wait_ms_p95 int, optional v2.1 +process_ms_avg int +process_ms_max int +process_ms_p95 int, optional v2.1 +created_time bigint +updated_time bigint +``` + +v2 should start with avg/max and counts. Percentiles can be added with an approximate histogram in v2.1. + +### 5.4 `channel_flow_events` + +Store only important events by default: + +```text +id +request_id +pool_key +channel_id +model +event_type queue_full | timeout | context_exceeded | lease_renew_failed | forced_release +reason +running +queued +wait_ms +process_ms +created_time +``` + +Retention: + +```text +FlowEventRetentionDays, default 7 +Per-pool daily cap, default 10000 events +created_time index for cleanup +``` + +## 6. Backend Interface + +Define an explicit backend interface. + +```go +type FlowBackend interface { + Acquire(ctx context.Context, req AcquireRequest) (FlowGuard, *AcquireDecision, error) + Status(ctx context.Context, poolKey string) (PoolStatus, error) + Close(ctx context.Context) error +} + +type FlowGuard interface { + Release(ctx context.Context) error + RenewLease(ctx context.Context) error + PoolKey() string + RequestID() string +} +``` + +Service layer: + +```go +type FlowController struct { + backend FlowBackend + poolStore PoolStore + metrics MetricsRecorder +} +``` + +Controller code should depend on `FlowController`, not directly on Redis or memory backend. + +## 7. Retry Loop Interaction + +### 7.1 Per-attempt Acquire/Release + +Each retry attempt may select a different channel and therefore a different Flow Pool. + +Rule: + +```text +Each attempt acquires exactly one guard. +That guard is released before the next retry attempt. +Successful streaming attempts hold guard until stream ends. +``` + +Pseudo-code: + +```go +var billingStarted bool + +for retry := 0; retry <= common.RetryTimes; retry++ { + channel, err := getChannel(...) + if err != nil { break } + + // SetupContextForSelectedChannel has already selected key and model mapping. + pool, ok := flow.ResolvePool(ctx, channel, upstreamModel) + guard, decision, err := flow.Acquire(ctx, pool, requestMeta) + if err != nil { + if decision.Temporary && pool.OnLimitAllowsFallback() { + markChannelTempUnavailableForThisRequest(channel.Id) + continue + } + return flowError(err) + } + + if !billingStarted { + err := billing.PreConsume(...) + if err != nil { + guard.Release(ctx) + return err + } + billingStarted = true + } + + err = callUpstream(...) + if isStreamSuccess { + wrapStreamWithGuard(guard) + return + } + + guard.Release(ctx) + + if err == nil { return } + if !shouldRetry(err) { break } +} +``` + +### 7.2 Temporary Unavailable vs Channel Failure + +Pool full is not a channel failure. + +Do not: + +```text +auto-ban channel +record as permanent failed channel +disable channel +``` + +Do: + +```text +mark channel/pool as temporarily unavailable only for this request attempt +``` + +### 7.3 `fallback_then_queue` + +`fallback_then_queue` is not recommended in MVP because current channel selection is iterative, not candidate-set based. + +MVP policies: + +```text +queue +reject +fallback +``` + +Add `fallback_then_queue` later after channel selector supports capacity-aware candidate enumeration. + +## 8. Upstream Model Resolution + +Flow Pool resolution must happen after: + +```text +middleware.Distribute() +SetupContextForSelectedChannel() +model mapping +upstream model name is available +``` + +Therefore flow control should not be implemented as generic Gin middleware. + +Recommended: + +```text +service/channel_flow.ResolvePool(c, channelID, upstreamModel) +``` + +Cache the resolved pool in request context for logs/metrics. + +## 9. Billing Lifecycle + +### 9.1 Problem + +Current relay flow pre-consumes quota before the retry loop. If flow control is added after channel selection, requests may wait in queue after quota has already been deducted. + +That is undesirable: + +- Queue timeout would require refund. +- Long waiting time holds user quota. +- Billing sessions remain open before upstream capacity is available. + +### 9.2 v2 Solution: Two-stage Billing + +Add a read-only billing precheck before queueing: + +```text +BillingPrecheck: + estimate quota + verify user/token/subscription likely has enough quota + no deduction +``` + +Then after Flow Guard is acquired: + +```text +PreConsumeBilling: + actual deduction/reservation + existing refund/settlement lifecycle +``` + +If pre-consume fails after acquire: + +```text +release guard immediately +return insufficient quota +``` + +### 9.3 Placement + +```text +Estimate tokens and price +BillingPrecheck +FlowControl Acquire +PreConsumeBilling +Call upstream +Settle/refund +Release guard +``` + +For stream, release guard on stream completion. + +## 10. Memory Backend v2 + +Use memory backend only for dev and single-instance deployments. + +Data structure: + +```text +map[poolKey]*slot + +slot: + mutex + config + normal queue + next sequence + +request: + request_id + state: waiting | dispatched + context_cost + enqueue_time + dispatch_time + notify channel + cancelled flag +``` + +Rules: + +- Queue itself is the source of truth. +- `state=dispatched` means running. +- `state=waiting` means queued. +- No independent running counter unless derived. +- Queue supports lazy cleanup of cancelled requests. +- Add `max_processing_ms` scanner to force-release leaked dispatched requests. + +Memory backend warning: + +```text +If Redis is disabled, show admin warning: +"Current Flow Control backend is local memory. Multi-instance deployments cannot guarantee global upstream concurrency limits." +``` + +## 11. Redis Backend v2 Without Lua + +### 11.1 Key Design + +Use hash tags for future Redis Cluster compatibility: + +```text +flow:{pool_key}:config +flow:{pool_key}:running +flow:{pool_key}:waiting +flow:{pool_key}:seq +flow:{pool_key}:req:{request_id} +flow:{pool_key}:events +``` + +All keys for one pool share the `{pool_key}` hash tag. + +### 11.2 Runtime Config in Redis + +On pool create/update, write config to Redis: + +```text +HSET flow:{pool_key}:config + enabled + max_inflight + max_queue_size + queue_timeout_ms + max_context_tokens + max_context_chars + max_processing_ms + config_version +``` + +Acquire reads config from Redis inside the `WATCH` transaction. This reduces inconsistent config across instances. + +### 11.3 Acquire Immediate or Enqueue + +Algorithm with optimistic transaction: + +```text +1. Generate request_id. +2. seq = INCR flow:{pool_key}:seq. +3. Cleanup a limited number of expired running leases. +4. WATCH running, waiting, config. +5. Read config, running count, waiting count. +6. If context exceeds limit -> UNWATCH, reject. +7. If running < max_inflight: + MULTI + ZADD running lease_expire_ms request_id + HSET req metadata state=running + EXPIRE req + EXEC + return guard +8. Else if waiting >= max_queue_size: + UNWATCH + reject queue_full +9. Else: + MULTI + ZADD waiting seq request_id + HSET req metadata state=waiting + EXPIRE req + EXEC + wait loop +10. If EXEC conflict, retry with jitter. +``` + +Bound transaction retries: + +```text +max_tx_retries = 8 +retry jitter = 5-30ms +``` + +If repeated conflicts occur: + +```text +return temporary busy, allow retry/fallback or short wait +``` + +### 11.4 Waiting Loop + +Do not rely only on Pub/Sub. + +Recommended v2 loop: + +```text +until deadline: + 1. Check whether request_id is already in running. + If yes -> return guard. + 2. TryPromoteSelf with WATCH/MULTI: + cleanup limited expired running leases + if capacity available and this request is at queue head: + move self from waiting to running + return guard + 3. Sleep poll interval with jitter. +``` + +Default poll: + +```text +initial: 100ms +normal: 250-500ms +max: 1000ms +jitter: +/- 20% +``` + +Optional optimization: + +```text +Release publishes a wakeup signal. +Waiter wakes early but still checks Redis state first. +Poll remains the correctness mechanism. +``` + +### 11.5 Release and Promotion + +Release: + +```text +1. WATCH running, waiting, config. +2. Read config and running count. +3. Read queue head candidates. +4. MULTI: + ZREM running request_id + for available capacity: + ZREM waiting candidate + ZADD running lease_expire candidate + HSET candidate state=running dispatch_time=now + PUBLISH wakeup, optional + EXEC +5. On conflict, retry with small cap. +``` + +Promotion must tolerate cancelled/stale waiters: + +- If candidate metadata missing, remove it. +- If candidate exceeded timeout, remove it. +- If candidate belongs to another instance, moving it to running is okay; that instance will discover it on poll. + +### 11.6 Why Poll-first Is Acceptable + +For a queue size of 240 and poll interval around 500ms: + +```text +approx additional Redis reads: 480/s in worst steady queue +``` + +This is acceptable for v2 and much easier to reason about than message-only wakeups. + +Pub/Sub can reduce latency but must not be required for correctness. + +## 12. Lease and Heartbeat + +### 12.1 Defaults + +```text +lease_ms: 60000 +renew_interval_ms: 20000 +renew_max_failures: 3 +``` + +### 12.2 Renew Failure Policy + +If lease renewal fails: + +```text +record warning metric +retry up to 3 times +do not terminate the user request +``` + +Reason: + +Killing an in-progress upstream request may be worse than temporarily allowing a slight overrun if Redis is unstable. + +Track: + +```text +flow_lease_renew_fail_total +flow_lease_expired_running_total +``` + +### 12.3 Memory Backend Leak Protection + +Memory backend has no Redis lease recovery. Add: + +```text +max_processing_ms +background scanner +forced release with warning event +``` + +If `max_processing_ms` is not configured: + +```text +default = max(queue_timeout_ms * 4, 30 minutes) +``` + +For stream/WebSocket, allow a larger configured value. + +## 13. Graceful Shutdown + +On gateway shutdown: + +```text +1. Mark local FlowController as draining. +2. New acquire calls return 503 service_draining. +3. Waiting local handlers are cancelled with 503. +4. Running requests are allowed to finish until shutdown timeout. +5. Redis backend releases or lets leases expire for local running requests. +6. Metrics record cancelled/drained counts. +``` + +Memory backend: + +- Waiting queue is local. Return 503 to waiting requests during shutdown. + +Redis backend: + +- Waiting handlers cancel and remove their request IDs from waiting. +- Running requests release if possible. +- If process exits abruptly, leases recover. + +## 14. Redis Unavailable Strategy + +Redis failure policy should be configurable. + +Options: + +| Policy | Behavior | Pros | Cons | +|---|---|---|---| +| fail_open | Disable flow control temporarily | Best availability | May overload upstream | +| fail_closed | Reject affected pool requests | Protects upstream | User-visible outage | +| local_memory | Use local fallback | Partial protection | Multi-instance overrun | + +Recommended default: + +```text +flow_control_redis_failure_policy = fail_open +``` + +But for private upstream pools that must never exceed capacity, admin can choose: + +```text +fail_closed +``` + +During failure: + +- Show warning in admin UI. +- Send admin notification. +- Record events and metrics. + +## 15. Request Body and Memory Management + +Queueing a request means the HTTP handler, parsed request, body storage, and context may remain in memory or temporary storage while waiting. + +Add safeguards: + +```text +max_queued_body_bytes_per_request +max_queued_body_bytes_per_pool +max_queued_context_tokens_per_pool +``` + +MVP practical default: + +```text +Do not add separate body-byte accounting in first code patch. +Do expose warning: + large request bodies + large queue size can increase memory pressure. +Use existing request body size limits. +Track queued context chars/tokens. +``` + +v2.1: + +- Add per-pool queued body byte estimate. +- Reject queue admission if pool queued memory is above threshold. + +## 16. Metrics and Trend Charts v2 + +### 16.1 Phase 1 Metrics + +Phase 1 must include realtime metrics. Without them, admins cannot validate whether flow control is working. + +Realtime: + +```text +running +max_inflight +queued +max_queue_size +oldest_wait_ms +backend +health +``` + +Minute aggregate v2: + +```text +running_avg +running_max +queued_avg +queued_max +acquired_count +queued_count +released_count +rejected_count +timeout_count +cancelled_count +wait_ms_avg +wait_ms_max +process_ms_avg +process_ms_max +``` + +Percentiles: + +```text +v2.1 use approximate histogram, not exact in-memory sorting +``` + +Candidate Go library: + +```text +github.com/HdrHistogram/hdrhistogram-go +``` + +### 16.2 Flow Pool Health State + +```text +Healthy: + running/max_inflight < 70%, queued == 0 + +Busy: + running/max_inflight >= 70%, queued == 0 + +Congested: + queued > 0 + +Critical: + queued/max_queue_size >= 80% + +Overloaded: + queue_full or timeout occurring + +Degraded: + Redis backend unavailable or lease renewal failures high +``` + +### 16.3 Error Response Metadata + +For queue full: + +```json +{ + "error": { + "message": "The upstream resource pool is busy. The waiting queue is full. Please retry later.", + "type": "rate_limit_error", + "code": "channel_flow_queue_full", + "metadata": { + "pool_running": 60, + "pool_max_inflight": 60, + "pool_queued": 240, + "pool_max_queue_size": 240, + "retry_after_seconds": 30 + } + } +} +``` + +Set HTTP header: + +```text +Retry-After: 30 +``` + +Do not expose sensitive pool names to normal users unless admin config allows it. + +## 17. Config Hot Update + +On pool config update: + +```text +1. DB transaction updates channel_flow_pools and increments config_version. +2. Update Redis config hash for that pool. +3. Invalidate in-memory cache on local instance. +4. Broadcast cache refresh if existing project mechanism supports it. +``` + +Runtime rules: + +- Reducing `max_inflight` does not cancel running requests. +- New dispatch stops until running drops below new max. +- Reducing `max_queue_size` does not kill already queued requests by default. +- New enqueue rejects if queue already exceeds new max. +- Disabling pool causes new acquire to reject or pass through based on policy; running requests drain. + +## 18. Task Relay + +Add: + +```text +task_release_policy: + on_submit + on_task_finish +``` + +`on_submit`: + +- Guard is released when upstream submit returns. +- Good for upstreams that have their own async queue. + +`on_task_finish`: + +- Guard remains associated with the local task record. +- Released when task reaches terminal state: success, failed, cancelled. +- Requires timeout/lease renewal for long tasks. + +v2 MVP: + +```text +Support on_submit. +Design data model for on_task_finish. +Implement on_task_finish in a later task-specific iteration. +``` + +Reason: + +The existing task system has multiple providers and polling paths. Holding capacity until task finish changes semantics and needs more focused testing. + +## 19. API Surface v2 + +### 19.1 Pool CRUD + +```text +GET /api/channel_flow/pools +POST /api/channel_flow/pools +GET /api/channel_flow/pools/:id +PUT /api/channel_flow/pools/:id +DELETE /api/channel_flow/pools/:id +``` + +### 19.2 Bindings + +```text +GET /api/channel_flow/pools/:id/bindings +POST /api/channel_flow/pools/:id/bindings +DELETE /api/channel_flow/bindings/:id +``` + +### 19.3 Status and Metrics + +```text +GET /api/channel_flow/pools/:id/status +GET /api/channel_flow/pools/:id/metrics?from=&to=&bucket=minute +GET /api/channel_flow/pools/:id/events +``` + +### 19.4 Suggestions + +```text +GET /api/channel_flow/suggestions?channel_id=123 +``` + +Suggestions are not automatic binding. + +## 20. i18n and Errors + +Backend error codes: + +```text +channel_flow_queue_full +channel_flow_queue_timeout +channel_flow_context_exceeded +channel_flow_draining +channel_flow_backend_unavailable +channel_flow_config_invalid +``` + +Messages must use existing backend i18n style where applicable. + +Frontend strings should be added for all supported default frontend locales: + +```text +en, zh, fr, ja, ru, vi +``` + +## 21. v2 Roadmap + +### Phase 1: Correct Single-instance MVP + +Required: + +- Flow Pool DB tables. +- Binding DB table. +- Memory backend. +- Explicit backend interface. +- Channel edit drawer Flow Control section. +- Flow Pools list. +- Realtime status. +- Bounded queue. +- Configurable queue timeout. +- Per-attempt acquire/release in normal relay. +- Stream release wrapper. +- Billing precheck before queue and preconsume after acquire. +- Admin warning for memory backend. + +Not included: + +- Redis backend. +- Percentile charts. +- `fallback_then_queue`. +- `on_task_finish`. + +### Phase 2: Trends and Operational Visibility + +- Minute aggregates. +- In-flight trend chart. +- Queue trend chart. +- Reject/timeout chart. +- Flow Pool health state. +- Event retention config. +- Error response metadata and Retry-After. + +### Phase 3: Redis Backend Without Lua + +- Redis `WATCH/MULTI` backend. +- Runtime config hash. +- Lease and renewal. +- Poll-first waiting loop. +- Optional Pub/Sub acceleration. +- Redis failure policy. +- Multi-instance tests. + +### Phase 4: Capacity-aware Routing + +- `fallback` support without permanent channel failure. +- Candidate-set selection. +- `fallback_then_queue`. +- Pool load-aware routing. + +### Phase 5: Advanced Controls + +- Approximate histogram percentiles. +- `on_task_finish`. +- VIP priority. +- Context/token in-flight limits. +- Optional Lua optimization if WATCH conflicts are high. + +## 22. Audit Issue Resolution Matrix + +| Audit issue | v2 resolution | +|---|---| +| Retry loop guard lifecycle unclear | Per-attempt acquire/release; pool full is temporary, not channel failure | +| Upstream model only known after setup | Flow resolution after `SetupContextForSelectedChannel`, not middleware | +| Memory backend unsafe in multi-instance | Admin warning; Redis required for production global capacity | +| Missing graceful shutdown | Add draining mode and queue cancellation | +| Lua complexity | Lua not mandatory; v2 Redis uses WATCH/MULTI | +| Pub/Sub message loss | Poll-first loop; Pub/Sub optional acceleration only | +| Lease renewal undefined | 60s lease, 20s renewal, warning after failures | +| Guard leak | `max_processing_ms` and scanner for memory backend | +| Cancellation removal performance | Lazy cleanup or linked queue in memory backend | +| Multi-instance notification starvation | Waiter always checks running state first | +| Config hot update | DB version + Redis config hash + cache invalidation | +| Redis unavailable | Configurable fail_open/fail_closed/local_memory | +| Backend interface unclear | Define `FlowBackend` and `FlowGuard` | +| Metrics percentile memory pressure | v2 avg/max; v2.1 approximate histogram | +| Event table growth | Retention days, created_time index, per-pool daily cap | +| Task relay support partial | v2 `on_submit`, later `on_task_finish` | +| User-facing queue info | Error metadata and Retry-After | +| Prometheus/OpenTelemetry | Not MVP; keep metrics API compatible for later exporter | +| Billing while queued | Billing precheck before queue, preconsume after acquire | +| Queued request body memory | Add warnings and future body-byte caps | +| Idempotency | request_id uniqueness and Redis request metadata | + +## 23. Key Implementation Decisions for Review + +Reviewers should explicitly approve or reject these choices: + +1. v2 Redis backend does not require Lua; use `WATCH/MULTI` first. +2. Flow Pool is a first-class DB entity; raw `pool_id` is not user-entered. +3. Flow resolution happens after channel setup and upstream model mapping. +4. Billing changes to two-stage precheck/preconsume. +5. Queue must have a hard `max_queue_size`. +6. Memory backend is allowed only with clear warning. +7. Redis failure default is `fail_open`, configurable to `fail_closed`. +8. `fallback_then_queue` is deferred until routing can inspect candidate pools. +9. Phase 1 includes realtime status; trend charts start in Phase 2. +10. Percentiles use approximate histograms later, not exact per-request arrays in v2. + +## 24. Recommended Next Step + +Before coding, produce a small technical spike for the Redis `WATCH/MULTI` backend: + +```text +Goal: + prove no more than max_inflight requests enter running across concurrent goroutines + +Scope: + Redis keys with hash tags + acquire immediate + enqueue + try promote self + release + queue timeout + +Load test: + 1000 concurrent acquire attempts + max_inflight = 60 + max_queue_size = 240 + +Pass condition: + running count never exceeds 60 except for documented lease-expiry edge cases + queue full and timeout behavior deterministic + transaction conflict rate measured +``` + +If conflict rate or latency is unacceptable, then design a Lua backend as Phase 3.5. + diff --git a/docs/channel-flow-control-queue-design-v3.md b/docs/channel-flow-control-queue-design-v3.md new file mode 100644 index 000000000000..2654e679de48 --- /dev/null +++ b/docs/channel-flow-control-queue-design-v3.md @@ -0,0 +1,1699 @@ +# Channel Flow Control and Queue Design Report v3 + +Date: 2026-06-13 + +Status: v3 design draft for external AI review + +Review target: + +- `docs/channel-flow-control-queue-design-v2.md` +- `/Users/laiyouxu/.gemini/antigravity-cli/brain/fdf12fcb-bc0f-48af-9f54-2dc1902d1eb9/flow-control-v2-review.md` + +Local reference code: + +- `controller/relay.go` +- `middleware/distributor.go` +- `service/billing.go` +- `service/billing_session.go` +- `middleware/rate-limit.go` +- `middleware/model-rate-limit.go` +- `/Users/laiyouxu/IdeaProjects/gateway/boom-gateway/boom-flowcontrol/src/lib.rs` +- `/Users/laiyouxu/IdeaProjects/gateway/boom-gateway/boom-routing/src/policy/load_helpers.rs` + +Official product references used for the market survey: + +- LiteLLM Proxy: https://docs.litellm.ai/docs/proxy/users +- Kong AI Rate Limiting Advanced: https://docs.konghq.com/hub/kong-inc/ai-rate-limiting-advanced/ +- Apache APISIX AI Rate Limiting: https://apisix.apache.org/docs/apisix/plugins/ai-rate-limiting/ +- Envoy AI Gateway usage-based rate limiting: https://aigateway.envoyproxy.io/docs/capabilities/traffic/usage-based-ratelimiting/ +- Portkey AI Gateway rate limits: https://portkey.ai/docs/product/ai-gateway/virtual-keys/rate-limits +- Cloudflare AI Gateway rate limiting: https://developers.cloudflare.com/ai-gateway/configuration/rate-limiting/ + +## 1. Executive Summary + +The target requirement is upstream resource-pool admission control, not ordinary user rate limiting. + +Example production scenario: + +```text +One upstream model pool has 96 GPUs. +The upstream can safely process 60 concurrent requests. +The 61st request must not enter the upstream. +The gateway should hold excess requests in a bounded queue, release them when capacity is available, and expose real-time plus historical in-flight/queued trends. +``` + +v2 had the right product direction, but the v2 review found several implementation risks. v3 changes the design in these important ways: + +1. Flow control still happens after channel selection and `SetupContextForSelectedChannel`, because only then do we know the actual channel, key, base URL, model mapping, and group context. +2. Billing is changed from "preconsume once before retry" to "precheck before queue, preconsume/reserve after acquire per selected attempt". This avoids charging while queued and reduces mismatch when retry selects a different channel/group. +3. Queue length must be hard bounded. Recommended initial config for the 60-concurrency pool is `max_inflight=60`, `max_queue_size=240`, `queue_timeout_ms=120000`. +4. `pool_id` is not user input. Admins create/select a Flow Pool by name; backend generates immutable `pool_key`. Bindings are explicit by `channel_id` and, later, optional upstream model. +5. Phase 1 should support only channel-level binding. `channel + upstream_model` binding is kept in the schema but enabled in Phase 2 to reduce the first test matrix. +6. Redis Lua is not mandatory. v3 uses a Redis `WATCH/MULTI` design first, but only after a Phase 0 spike proves conflict rate is acceptable. Lua remains the fallback if the spike fails. +7. Redis transactions must not watch config, and release must not promote a batch of waiters. Release only removes the running request and optionally publishes a wakeup signal. Waiting requests self-promote. +8. Waiting poll is adaptive by queue position. Only near-head requests run full promotion logic; tail requests poll slowly and do cheap checks. +9. Client disconnect, queue timeout, graceful shutdown, and stream completion must all release queue/running state idempotently. +10. Trend charts are part of the product requirement, not optional polish. Phase 1 includes realtime status and minimum minute-level trend data for running and queued counts; percentiles can be approximate in a later iteration. + +My recommendation: + +```text +Approve v3 as the implementation direction, but require Phase 0 Redis spike before coding the Redis backend. +Start with Memory backend plus channel-level Flow Pool binding and full lifecycle correctness. +Add Redis after the transaction conflict behavior is measured. +``` + +## 2. What Changed from v2 After Review + +| v2 review issue | v3 decision | +|---|---| +| Preconsume happens once, but retry may select a different channel/group | Precheck before queue only. After each selected attempt acquires a guard, recompute price context and create or extend the billing session with `Reserve`. | +| Polling estimate undercounted Redis ops | Adaptive poll by queue position. Only near-head waiters run `TryPromoteSelf`. Tail waiters avoid transactions. | +| Lease expiry can cause actual concurrency overrun | Explicitly documented as a bounded failure mode. Track `lease_expired_running_total` and accept temporary overrun rather than killing user requests. | +| Memory backend lazy cleanup can accumulate cancelled entries | Add compaction threshold: compact if cancelled/stale entries exceed 30 percent or 64 entries. | +| Event daily cap undefined | Use write-time per-pool daily counter and sampling after cap, not cleanup after uncontrolled writes. | +| `WATCH/MULTI` conflict can be high | Do Phase 0 spike. Minimize conflict by not watching config and by removing promotion from release. | +| Stale queue head can block later waiters | `TryPromoteSelf` scans a head window and removes stale/missing/timed-out candidates before deciding. | +| Billing TOCTOU during queue wait | Accept for MVP. If balance is consumed while waiting, release guard and return a clear insufficient quota message. Soft reservation is future work. | +| Config watched inside transaction | Config is read from cache/Redis outside transaction. It is never in `WATCH`. | +| Release promotes many waiters | Release performs only `ZREM running` plus optional wakeup publish. Waiters self-promote. | +| Graceful shutdown and LB timeout unclear | Acquire uses request context. Queue timeout should be less than upstream load balancer idle timeout. Client disconnect removes waiting entry. | +| `fail_open` can cause storm | Track bypass count and use recovery cooldown. Optionally use local memory safety valve during Redis outage. | +| `FlowGuard` not stream-aware | Add idempotent guard and stream/read-closer binding. Guard releases on non-stream return, stream end, stream drop, timeout, or cancel. | +| `AcquireDecision` lacks fields | Define complete decision shape for errors, metrics, logs, and UI. | +| Multi-tenant queue fairness missing | Add `max_queue_per_user` field, default off for compatibility, recommended on shared pools. | +| Roadmap too optimistic | Add Phase 0 spike, Phase 1 channel-only binding, Phase 3 split into Redis 3a/3b. | + +## 3. Existing new-api Limit Controls + +new-api already has several limit mechanisms, but none protects a shared upstream GPU pool with global in-flight capacity and queueing. + +| Existing mechanism | Scope | Implementation | What it protects | Gap for this feature | +|---|---|---|---|---| +| Global web API rate limit | IP | `middleware/rate-limit.go` with Redis list or memory limiter | Web/dashboard abuse | Not channel/pool aware | +| Global relay API rate limit | IP | `GlobalAPIRateLimit` | API abuse by client IP | Does not count in-flight upstream requests | +| Critical endpoint limit | IP | `CriticalRateLimit` | Login, reset, payment, token key endpoints | Not relay capacity | +| Search/email verification limits | User/IP | `SearchRateLimit`, `EmailVerificationRateLimit` | Expensive dashboard actions | Not upstream capacity | +| Model request rate limit | User/group | `middleware/model-rate-limit.go` with success count and Redis token bucket Lua | Request frequency per user/group | No queue, no channel pool, no total upstream concurrency | +| Billing quota | User/token/subscription | `service/billing.go`, `service/billing_session.go`, `service/quota.go` | Whether user can pay | Not an admission control semaphore | +| Channel retry/distribution | Channel/model/group | `middleware/distributor.go`, `controller/relay.go`, service selector | Selects a usable channel | Does not know pool capacity | + +Important separation: + +```text +User rate limit: who may send how many requests in a time window. +Billing: whether the request can be paid for. +Flow control: whether the selected upstream resource pool has capacity now. +``` + +The proposed feature should be a new service, not an extension of `ModelRequestRateLimit`. + +## 4. Market Survey Summary + +Mainstream AI gateways usually support one or more of these: + +- RPM/RPS request-window limits. +- TPM or token-aware limits. +- Budgets and spend caps. +- Per-key, per-user, team, model, or provider limits. +- Provider fallback and load-aware routing. +- Sometimes max parallel requests or scheduler queueing. + +The specific requirement here is narrower and stricter: + +```text +Protect a physical/logical upstream pool shared by multiple new-api channels. +Bound total in-flight requests. +Queue overflow requests. +Keep queue length finite. +Make in-flight and queued trends auditable. +``` + +| Gateway/product | Similar capability | Difference from new-api requirement | +|---|---|---| +| LiteLLM Proxy | User/team/key/model budgets and rate limits; max parallel style controls in proxy settings | Strong tenant-facing control, but new-api still needs explicit Flow Pool binding to its channel model | +| Kong AI Rate Limiting Advanced | AI-aware/token-aware rate limiting with gateway plugin model | Primarily request/token limiting, not necessarily a shared GPU-pool queue in the new-api channel selector | +| APISIX AI Rate Limiting | LLM token dimensions for prompt/completion/total tokens | Useful for token quota, not enough for GPU occupancy | +| Envoy AI Gateway | Provider traffic policies, fallback, usage-based rate limiting | Kubernetes/Gateway API architecture differs; concepts useful for future capacity-aware routing | +| Portkey AI Gateway | Virtual key rate limits and gateway policy controls | More tenant/key policy oriented | +| Cloudflare AI Gateway | Gateway-level rate limiting with fixed/sliding windows | Good edge policy, but not enough for per-upstream in-flight queueing | + +Conclusion: + +```text +This feature is justified as a first-class new-api capability. +It complements, rather than duplicates, existing user and token rate limits. +``` + +## 5. Gateway Project Reference + +The local `gateway` project is a useful reference, but it is not directly the same solution new-api needs. + +Observed design in `boom-flowcontrol`: + +- It models flow control per `deployment_id`. +- It keeps `vip_queue` and `normal_queue`. +- The queue itself is the source of truth; dispatched entries are in-flight. +- There is no separate counter that can leak. +- `FlowControlGuard` releases in `Drop`. +- `FlowControlledStream` holds the guard until stream end or stream drop. +- It exposes user request status such as waiting position and processing time. +- Routing can use in-flight plus queued load when choosing a deployment. + +What new-api should borrow: + +1. Queue-as-source-of-truth for memory backend. +2. Idempotent guard lifecycle. +3. Stream wrapper/guard binding. +4. Realtime waiting and processing status. +5. Load-aware routing as a later phase. + +What new-api cannot copy as-is: + +1. The gateway implementation is memory-local; new-api production deployments can be multi-instance. +2. new-api needs Redis backend for global capacity. +3. new-api has existing billing and retry semantics that must be integrated. +4. new-api channels are configured in DB and can share one physical upstream pool. +5. new-api needs a web admin CRUD model for Flow Pools and bindings. +6. v3 requires a hard `max_queue_size`; the gateway reference mainly uses timeout/context and in-flight limits. + +Therefore: + +```text +gateway is the right lifecycle model, not the final distributed backend. +``` + +## 6. Product Model + +### 6.1 Flow Pool + +A Flow Pool represents one upstream capacity domain. + +Examples: + +```text +96-card DeepSeek-R1 production pool +Azure East US GPT-4.1 shared deployment +Internal Qwen 72B cluster +``` + +Admins should not manually type raw `pool_id`. + +User-facing flow: + +```text +Admin opens channel edit drawer. +Admin enables Flow Control. +Admin selects an existing Flow Pool by name or creates a new pool. +Backend generates pool_key. +Backend stores explicit binding between channel and pool. +Runtime Redis keys, logs, and metrics use pool_key. +``` + +Identifiers: + +| Field | Purpose | +|---|---| +| `id` | DB primary key, internal only | +| `pool_key` | Backend-generated stable runtime key, unique, immutable | +| `name` | Admin-visible name | +| `description` | Admin-readable explanation | + +Example: + +```text +name: "DeepSeek R1 96-card pool" +pool_key: "flow_pool_8f3a2c7e" +``` + +### 6.2 Binding to Channel and Upstream URL + +Binding must be explicit. + +v3 resolution source of truth: + +```text +channel_flow_pool_bindings.channel_id -> channel_flow_pools.id +``` + +Phase 1: + +```text +match_mode = "channel" +All upstream models on this channel share the same pool. +``` + +Phase 2: + +```text +match_mode = "channel_model" +Binding key = channel_id + resolved upstream_model. +``` + +Base URL is not the binding source of truth. It is only used for UI suggestions and warnings. + +Reasons: + +- One base URL may serve multiple physical pools by key, tenant, or deployment name. +- One physical pool may have multiple base URLs. +- A channel may map public model names to different upstream model names. +- Base URL edits should not silently merge or split resource pools. + +UI may show: + +```text +This channel has the same base URL as channels #12 and #18. +Suggested existing Flow Pools: "DeepSeek R1 96-card pool". +No binding is changed until the admin explicitly selects one. +``` + +### 6.3 Queue Must Have a Hard Upper Bound + +Yes, queue length needs an upper bound. + +Without a hard cap: + +- Client HTTP connections can pile up. +- Gateway memory can grow without bound. +- Upstream recovery can be followed by a long stale backlog. +- User experience becomes unpredictable. +- One user can occupy all waiting capacity. + +Recommended first production config: + +```text +max_inflight: 60 +max_queue_size: 240 +queue_timeout_ms: 120000 +queue_policy: fifo +on_limit: queue +max_queue_per_user: 0 by default, recommended 20 for shared public pools +``` + +Default formula when admin creates a new pool: + +```text +max_queue_size = max_inflight * 4 +queue_timeout_ms = 120000 +``` + +UI should warn when: + +```text +max_queue_size > max_inflight * 10 +queue_timeout_ms > known/provided load balancer idle timeout +memory backend is used in multi-instance deployment +``` + +## 7. Web Admin Design + +### 7.1 Channel Drawer + +Add a section in the channel create/update drawer: + +```text +Advanced Settings + Flow Control & Queue +``` + +Controls: + +```text +[Switch] Enable flow control + +Flow Pool + [Select] Existing pool + [Button] Create new pool + +Binding scope + [Segmented] Entire channel + [Segmented disabled in Phase 1] Specific upstream models + +Capacity + Max in-flight requests + Max queue size + Queue timeout + Max queue per user (optional) + +Behavior + On limit: queue | reject | fallback + Redis failure: fail_open | fail_closed | local_memory + +Preview + Current channel ID + Base URL + Model mapping summary + Resolved binding + Similar channels by base URL +``` + +Phase 1 should disable or hide upstream-model binding. The schema can support it, but the UI should make it clear that the first release binds the entire channel. + +### 7.2 Flow Pools Page + +Add a management tab: + +```text +Channels | Flow Pools +``` + +List columns: + +```text +Name +Bound channels +Running / max_inflight +Queued / max_queue_size +Oldest wait +Wait avg/max +Rejected / timeout +Backend +Health +Updated +``` + +Detail page sections: + +```text +Overview + realtime running and queued status + pool config + backend health + +Bindings + bound channels + channel base URL + channel type + model mapping + +Trends + in-flight trend + queued trend + wait time trend + process time trend + reject/timeout/cancel trend + +Events + queue full + timeout + cancelled + lease expired + backend unavailable +``` + +### 7.3 Trend Chart Requirement + +The user explicitly needs in-flight and queued trend charts for traceability. + +Minimum v3 release must provide: + +```text +running_avg +running_max +queued_avg +queued_max +acquired_count +queued_count +released_count +rejected_count +timeout_count +cancelled_count +wait_ms_avg +wait_ms_max +process_ms_avg +process_ms_max +``` + +Chart views: + +```text +Last 15 minutes, bucket 10 seconds or 1 minute +Last 1 hour, bucket 1 minute +Last 24 hours, bucket 5 minutes or 1 hour +Custom time range, bucket selected by backend +``` + +Percentiles: + +```text +Phase 1: avg/max only. +Phase 2: approximate histogram for p50/p95/p99. +``` + +## 8. Runtime Placement in new-api + +Flow control should not be generic Gin middleware. + +It must run after: + +```text +channel selection +SetupContextForSelectedChannel +channel key selection +channel model mapping +upstream model resolution +``` + +Reason: + +- The selected channel can change during retry. +- The same client model may map to a different upstream model per channel. +- The selected channel key and group context can affect billing/logging. +- Pool binding is based on channel and later upstream model. + +Current relevant flow in `controller/relay.go`: + +```text +token estimate +ModelPriceHelper +PreConsumeBilling +retry loop: + getChannel + SetupContextForSelectedChannel + relayHandler +``` + +v3 target flow: + +```text +parse and validate request +estimate prompt tokens +billing precheck only, no deduction +retry loop: + getChannel + SetupContextForSelectedChannel + resolve upstream model + resolve Flow Pool binding + acquire Flow Guard with request context + recompute attempt price context if needed + preconsume or reserve billing for this selected attempt + call upstream + release guard on attempt failure, non-stream finish, stream finish/drop, timeout, or cancellation +settle/refund billing as today +``` + +Mermaid sequence: + +```mermaid +sequenceDiagram + participant Client + participant Relay as new-api relay + participant Selector as Channel selector + participant Flow as FlowController + participant Billing + participant Upstream + + Client->>Relay: request + Relay->>Billing: BillingPrecheck(no mutation) + loop retry attempts + Relay->>Selector: select channel + Selector-->>Relay: channel + Relay->>Relay: SetupContextForSelectedChannel + Relay->>Flow: ResolvePool(channel, upstream_model) + Relay->>Flow: Acquire(ctx) + Flow-->>Relay: guard or queue/reject + Relay->>Billing: PreConsume or Reserve(attempt quota) + Relay->>Upstream: call + alt success stream + Relay-->>Client: stream with guard wrapper + else success non-stream + Upstream-->>Relay: response + Relay->>Flow: Release + Relay-->>Client: response + else retryable failure + Upstream-->>Relay: error + Relay->>Flow: Release + end + end +``` + +## 9. Billing Lifecycle v3 + +### 9.1 Problem in Current Code + +`controller/relay.go` currently calculates price and calls `PreConsumeBilling` before the retry loop. But `getChannel` and `SetupContextForSelectedChannel` happen inside the retry loop. + +This means: + +- Flow control inserted after channel selection would happen after billing has already deducted quota. +- A queued request could hold user quota while waiting. +- Retry may switch channel/group context after the first billing estimate. +- v2's "preconsume once after first acquire" still leaves ambiguity if later retry uses a different selected channel/group. + +### 9.2 v3 Billing Rule + +Billing must not preconsume while queued. + +v3 splits billing into: + +```text +BillingPrecheck: + read-only, before queue + rejects obvious insufficient quota/subscription/token cases + does not mutate user quota, token quota, or subscription amount + +AttemptPreConsumeOrReserve: + after Flow Guard is acquired + uses the selected attempt's current RelayInfo and price context + creates BillingSession if this is the first billable attempt + calls BillingSession.Reserve(targetQuota) if a later attempt needs more quota +``` + +If the later attempt needs less quota: + +```text +Do not refund immediately. +Final SettleBilling handles actual quota and refund. +``` + +If billing fails after acquire: + +```text +release guard immediately +return insufficient quota +record queue_wait_then_billing_failed event +``` + +User-facing message: + +```text +排队期间余额或订阅额度已被其他请求消耗，请充值或稍后重试。 +``` + +This TOCTOU is acceptable for MVP because a soft reservation system would add significant complexity. It should be revisited after the first release. + +### 9.3 Pseudocode + +```go +priceEstimate, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta) +if err != nil { return err } + +if !priceEstimate.FreeModel { + if err := billing.Precheck(c, priceEstimate.QuotaToPreConsume, relayInfo); err != nil { + return err + } +} + +var billingStarted bool + +for retry := 0; retry <= common.RetryTimes; retry++ { + channel, err := getChannel(c, relayInfo, retryParam) + if err != nil { break } + + pool, ok := flow.ResolvePool(c, channel.Id, resolvedUpstreamModel(c, relayInfo)) + guard, decision, err := flow.Acquire(c.Request.Context(), acquireReq) + if err != nil { + return flow.ToAPIError(decision, err) + } + + attemptPrice, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta) + if err != nil { + guard.Release(context.Background()) + return err + } + + if !attemptPrice.FreeModel { + if !billingStarted { + err = service.PreConsumeBilling(c, attemptPrice.QuotaToPreConsume, relayInfo) + if err != nil { + guard.Release(context.Background()) + return err + } + billingStarted = true + } else if relayInfo.Billing != nil { + if err := relayInfo.Billing.Reserve(attemptPrice.QuotaToPreConsume); err != nil { + guard.Release(context.Background()) + return billingReserveError(err) + } + } + } + + err = callUpstreamWithGuard(c, relayInfo, guard) + if streamSuccess { + bindGuardToStream(guard) + return nil + } + + guard.Release(context.Background()) + if err == nil { return nil } + if !shouldRetry(c, err, remaining) { break } +} +``` + +### 9.4 What This Does Not Solve + +This does not solve provider-side duplicate billing when a retry happens after an upstream already consumed tokens but returned an error. That is an existing retry risk and should remain handled by current high-risk retry settings and logging. + +The flow-control feature should not expand scope into provider billing reconciliation. + +## 10. Backend Interface + +Controller code should depend on a service-level `FlowController`, not directly on Redis or memory structures. + +```go +type FlowBackend interface { + Acquire(ctx context.Context, req AcquireRequest) (FlowGuard, *AcquireDecision, error) + Status(ctx context.Context, poolKey string) (PoolStatus, error) + Close(ctx context.Context) error +} +``` + +`AcquireRequest`: + +```go +type AcquireRequest struct { + RequestID string + PoolKey string + ChannelID int + UpstreamModel string + UserID int + TokenID int + QueueTimeoutMs int64 + ContextTokens int + ContextChars int + CreatedAtMs int64 +} +``` + +`AcquireDecision`: + +```go +type AcquireDecision struct { + Admitted bool + Queued bool + QueuePos int + WaitedMs int64 + Temporary bool + RejectCode string + RunningNow int + QueuedNow int + RetryAfterS int + Backend string + PoolKey string + ConfigVersion int64 +} +``` + +`FlowGuard`: + +```go +type FlowGuard interface { + Release(ctx context.Context) error + RenewLease(ctx context.Context) error + PoolKey() string + RequestID() string + IsReleased() bool + + // For non-stream handlers, defer Release. + // For stream handlers, bind release to stream/read closer completion or drop. + BindRelease(release func()) + WrapReadCloser(rc io.ReadCloser) io.ReadCloser +} +``` + +Implementation notes: + +- `Release` must be idempotent. +- `Release` should be safe after queue timeout, client disconnect, or lease expiry. +- Stream wrapper is required for SSE, chunked streaming, WebSocket-like flows, and any handler that returns before upstream processing ends. + +## 11. Data Model + +Use DB tables, not per-channel JSON blobs. Shared pool config cannot be safely represented in multiple channel settings. + +All migrations must support SQLite, MySQL, and PostgreSQL. Prefer GORM models and avoid DB-specific JSONB or partial indexes in the initial implementation. + +### 11.1 `channel_flow_pools` + +```text +id int primary key +pool_key varchar unique, generated by backend +name varchar +description text +enabled bool/int +backend varchar, "memory" | "redis" +max_inflight int +max_queue_size int +max_queue_per_user int, 0 means disabled +queue_timeout_ms int +queue_policy varchar, default "fifo" +on_limit varchar, "queue" | "reject" | "fallback" +redis_failure_policy varchar, "fail_open" | "fail_closed" | "local_memory" +max_context_tokens int, optional +max_context_chars int, optional +max_processing_ms int, optional +lease_ms int, default 60000 +renew_interval_ms int, default 20000 +config_version bigint +created_time bigint +updated_time bigint +``` + +### 11.2 `channel_flow_pool_bindings` + +```text +id int primary key +pool_id int +channel_id int +upstream_model varchar, optional, Phase 2 +match_mode varchar, "channel" | "channel_model" +enabled bool/int +created_time bigint +updated_time bigint +``` + +Phase 1 runtime only uses: + +```text +channel_id + match_mode="channel" +``` + +### 11.3 `channel_flow_metrics_minute` + +```text +id int primary key +bucket_ts bigint +pool_key varchar +channel_id int +model varchar +running_avg double +running_max int +queued_avg double +queued_max int +acquired_count int +queued_count int +released_count int +rejected_count int +timeout_count int +cancelled_count int +billing_failed_count int +lease_renew_fail int +lease_expired_count int +wait_ms_avg int +wait_ms_max int +process_ms_avg int +process_ms_max int +created_time bigint +updated_time bigint +``` + +Phase 2 optional: + +```text +wait_ms_p50 +wait_ms_p95 +wait_ms_p99 +process_ms_p50 +process_ms_p95 +process_ms_p99 +``` + +### 11.4 `channel_flow_events` + +```text +id +request_id +pool_key +channel_id +model +user_id +token_id +event_type +reason +running +queued +queue_pos +wait_ms +process_ms +backend +created_time +``` + +Event types: + +```text +queue_full +queue_timeout +client_cancelled +service_draining +context_exceeded +billing_failed_after_wait +lease_renew_failed +lease_expired_running +backend_unavailable +forced_release +config_invalid +``` + +Retention: + +```text +FlowEventRetentionDays = 7 by default +Per-pool daily write cap = 10000 by default +After cap, sample writes at 1/N while keeping aggregate counters +``` + +## 12. Memory Backend + +Memory backend is for dev and single-instance deployments. + +Data structure: + +```text +map[poolKey]*slot + +slot: + mutex + config + queue []request + next sequence + +request: + request_id + state: waiting | running + user_id + context_tokens + context_chars + enqueue_time + dispatch_time + notify channel + cancelled flag +``` + +Rules: + +- Queue is the source of truth. +- Running entries are queue entries with `state=running`. +- No independent running counter unless derived under lock. +- `max_queue_size` counts waiting entries only. +- `max_queue_per_user` counts waiting entries by user when enabled. +- On cancellation, mark cancelled and notify dispatcher. +- Compact when stale/cancelled entries exceed 30 percent or 64 entries. +- Background scanner force-releases entries that exceed `max_processing_ms`. + +Admin warning: + +```text +Current Flow Control backend is local memory. Multi-instance deployments cannot guarantee global upstream concurrency limits. Use Redis for production pool-level capacity. +``` + +## 13. Redis Backend v3 Without Mandatory Lua + +### 13.1 Is Lua Required? + +No. Lua is not required as the first implementation. + +However, Redis concurrency must be validated before coding the production backend. The v3 position is: + +```text +Do a Phase 0 WATCH/MULTI spike. +If conflict rate and p99 acquire latency are acceptable, implement WATCH/MULTI. +If not, implement small Lua scripts for acquire/enqueue and self-promotion. +``` + +Lua should be treated as an optimization/atomicity packaging choice, not as a product requirement. + +### 13.2 Redis Keys + +Use Redis hash tags so keys for one pool share a slot in Redis Cluster: + +```text +flow:{pool_key}:config +flow:{pool_key}:running +flow:{pool_key}:waiting +flow:{pool_key}:waiting_deadline +flow:{pool_key}:seq +flow:{pool_key}:req:{request_id} +flow:{pool_key}:user_waiting +flow:{pool_key}:events:{yyyymmdd}:count +flow:{pool_key}:wakeup +``` + +Meaning: + +| Key | Type | Meaning | +|---|---|---| +| `config` | Hash | Runtime config snapshot | +| `running` | ZSET | request_id scored by lease expiration ms | +| `waiting` | ZSET | request_id scored by FIFO sequence | +| `waiting_deadline` | ZSET | request_id scored by queue deadline ms | +| `seq` | String counter | FIFO sequence | +| `req:{request_id}` | Hash | request metadata | +| `user_waiting` | Hash | user_id to waiting count for optional per-user cap | +| `events:*:count` | Counter | write-time event cap | +| `wakeup` | Pub/Sub channel | optional latency optimization | + +### 13.3 Config Is Not Watched + +Do not `WATCH` config. + +Reason: + +- Admin config changes are rare. +- Slightly stale config for one poll cycle is acceptable. +- Watching config causes all transactions to conflict on every config update. + +Runtime rule: + +```text +Read config from local cache or Redis before transaction. +WATCH only keys that must be protected for state transition. +``` + +If config changes: + +```text +DB config_version increments. +Redis config hash updates. +Local cache invalidates or refreshes. +Running requests are not killed. +New acquire/promotion sees new config after refresh. +``` + +### 13.4 Immediate Acquire or Enqueue + +High-level algorithm: + +```text +1. Read config outside transaction. +2. Validate context limits. +3. Cleanup a small batch of expired running leases and expired waiting entries. +4. If valid waiting queue exists, enqueue to preserve FIFO. +5. If no valid waiting and running < max_inflight, try immediate running admission. +6. Else enqueue if waiting < max_queue_size and per-user cap allows. +7. Else reject queue_full. +``` + +Transaction for immediate acquire: + +```text +WATCH running, waiting +read running count and waiting count +if running < max_inflight and waiting == 0: + MULTI + ZADD running lease_expire_ms request_id + HSET req metadata state=running dispatch_time=now + EXPIRE req + EXEC +else: + UNWATCH + enqueue or wait +``` + +Transaction for enqueue: + +```text +seq = INCR flow:{pool_key}:seq +deadline = now + queue_timeout_ms + +WATCH waiting, user_waiting +read waiting count and user waiting count +if waiting < max_queue_size and user cap ok: + MULTI + ZADD waiting seq request_id + ZADD waiting_deadline deadline request_id + HINCRBY user_waiting user_id 1 + HSET req metadata state=waiting enqueue_time=now deadline=deadline + EXPIRE req + EXEC +else: + UNWATCH + reject +``` + +Notes: + +- Cleanup is bounded per operation, for example 16 running and 64 waiting entries. +- If cleanup cannot remove enough stale entries before queue-full decision, a false queue-full can happen under extreme stale buildup. This is acceptable only if metrics make it visible; the cleanup budget can be increased. + +### 13.5 Waiting Loop + +Correctness does not rely on Pub/Sub. + +Every waiter uses request context and deadline: + +```text +until queue deadline or request context done: + check if request_id is already in running + check whether request_id still exists in waiting + calculate approximate queue position + if near head: + TryPromoteSelf + sleep adaptive interval with jitter or wake early on Pub/Sub +``` + +Adaptive poll: + +| Position | Behavior | +|---|---| +| already running | return guard | +| position <= 3 | poll 100-250 ms; run full `TryPromoteSelf` | +| position <= max_inflight | poll 300-700 ms; run promotion every few polls or on wakeup | +| tail | poll 1000-2000 ms; cheap state checks only | + +This reduces Redis ops. A queue of 240 waiters should not produce 240 concurrent `WATCH/MULTI` attempts every 500 ms. + +### 13.6 TryPromoteSelf + +`TryPromoteSelf` must handle stale head entries. + +Algorithm: + +```text +1. Read config outside transaction. +2. Cleanup small batch of expired running leases. +3. Fetch head window from waiting: ZRANGE waiting 0 9 WITHSCORES. +4. For each candidate before self: + - if metadata missing, ZREM waiting and waiting_deadline + - if deadline expired, remove candidate and decrement user_waiting + - if candidate is valid and not self, exit: not my turn +5. If self is first valid candidate and running < max_inflight: + WATCH running, waiting + re-check running count and that self is still in waiting head window + MULTI + ZREM waiting self + ZREM waiting_deadline self + HINCRBY user_waiting user_id -1 + ZADD running lease_expire_ms self + HSET req state=running dispatch_time=now + EXEC +6. On conflict, retry with small jitter and bounded attempts. +``` + +Bounded transaction retry: + +```text +max_tx_retries = 8 +retry jitter = 5-30 ms +``` + +If retries fail: + +```text +return temporary busy to the wait loop, not to the user immediately +``` + +### 13.7 Release + +Release must be simple. + +v3 release: + +```text +ZREM running request_id +HSET req state=released release_time=now +PUBLISH wakeup optional +``` + +Do not: + +```text +WATCH config +promote a batch of waiters +loop over available capacity +``` + +Why: + +- Release storms are common when many upstream requests finish together. +- Batch promotion in release causes large transactions and conflicts. +- Self-promotion by waiters keeps release cheap and predictable. +- Latency cost is at most one adaptive poll interval, usually below 250 ms for head waiters. + +### 13.8 Lease and Renewal + +Defaults: + +```text +lease_ms = 60000 +renew_interval_ms = 20000 +renew_max_failures = 3 +``` + +Renewal: + +```text +ZADD running new_lease_expire_ms request_id +HSET req last_renew_time=now +``` + +Known boundary behavior: + +```text +If lease expires while the upstream request is still running, another waiter may be promoted. +Actual upstream concurrency can temporarily exceed max_inflight by the number of expired-but-still-live requests. +This is accepted because killing an in-progress user request is worse than a temporary overrun during Redis/network instability. +``` + +Metrics: + +```text +flow_lease_renew_fail_total +flow_lease_expired_running_total +flow_actual_overrun_observed_total +``` + +### 13.9 Redis Failure Policy + +Configurable: + +| Policy | Behavior | Use case | +|---|---|---| +| `fail_open` | bypass flow control | public availability first | +| `fail_closed` | reject affected pool | strict private upstream protection | +| `local_memory` | local per-instance fallback | partial protection when Redis is unstable | + +v3 additions: + +```text +During fail_open, maintain local bypass counter per pool. +When Redis recovers and bypass_count > max_inflight * 2, enter 10s recovery cooldown. +During cooldown, new requests use normal flow control and are not bypassed. +``` + +For strict 96-GPU/60-concurrency private pools, recommended policy: + +```text +fail_closed +``` + +For public gateway availability: + +```text +fail_open with admin warning and recovery cooldown +``` + +### 13.10 When to Use Lua + +Use Lua if Phase 0 spike shows: + +```text +transaction conflict rate > 30 percent under target load +or p99 acquire/promotion latency is unacceptable +or Redis round trips become a bottleneck +``` + +If Lua is introduced, keep scripts small: + +```text +try_acquire_or_enqueue.lua +try_promote_self.lua +release.lua may stay plain Redis command +``` + +Do not implement a large scheduler script that scans unbounded queues. + +## 14. Client Disconnect, Timeout, and Shutdown + +Acquire must use: + +```go +c.Request.Context() +``` + +Rules: + +- If client disconnects while waiting, remove request from waiting and decrement per-user waiting count. +- If queue timeout fires, remove request from waiting and return queue timeout. +- If shutdown starts, reject new acquire with `service_draining`. +- Waiting local handlers should return 503 when the process is draining. +- Running requests should be allowed to finish until server shutdown timeout. +- If process dies abruptly, Redis leases recover running state. + +Load balancer guidance: + +```text +queue_timeout_ms should be lower than LB/proxy idle timeout. +If admin configures queue_timeout_ms higher than a known LB timeout, UI should warn. +``` + +## 15. Retry and Channel Failure Semantics + +Pool full is not a channel failure. + +Do not: + +```text +auto-ban channel +disable channel +record as permanent channel error +``` + +Do: + +```text +mark channel/pool temporarily unavailable for this request attempt +allow retry/fallback if policy says so +record flow-control-specific metrics +``` + +MVP policies: + +```text +queue +reject +fallback +``` + +Defer: + +```text +fallback_then_queue +``` + +Reason: + +Current channel selection is iterative. A clean `fallback_then_queue` needs candidate-set routing so the gateway can inspect all possible channels/pools before deciding whether to queue. + +## 16. Metrics, Trends, and Events + +### 16.1 Realtime Status + +Endpoint: + +```text +GET /api/channel_flow/pools/:id/status +``` + +Response: + +```json +{ + "pool_key": "flow_pool_8f3a2c7e", + "name": "DeepSeek R1 96-card pool", + "backend": "redis", + "health": "congested", + "running": 60, + "max_inflight": 60, + "queued": 137, + "max_queue_size": 240, + "oldest_wait_ms": 42100, + "lease_renew_failures": 0, + "config_version": 12 +} +``` + +Health states: + +```text +healthy: running < 70 percent and queued = 0 +busy: running >= 70 percent and queued = 0 +congested: queued > 0 +critical: queued / max_queue_size >= 80 percent +overloaded: queue_full or queue_timeout happening +degraded: Redis unavailable or lease renewal failures high +``` + +### 16.2 Trend APIs + +```text +GET /api/channel_flow/pools/:id/metrics?from=&to=&bucket= +GET /api/channel_flow/pools/:id/events?from=&to=&event_type= +``` + +Minimum charts: + +1. In-flight trend: `running_avg`, `running_max`. +2. Queue trend: `queued_avg`, `queued_max`. +3. Wait time trend: `wait_ms_avg`, `wait_ms_max`. +4. Process time trend: `process_ms_avg`, `process_ms_max`. +5. Outcome trend: acquired, released, rejected, timeout, cancelled, billing failed. + +### 16.3 Event Write Cap + +Do not write unlimited event rows. + +Write-time cap: + +```text +INCR event counter for pool/day +if counter <= daily cap: + write event +else: + sample at configured rate, e.g. 1/10 or 1/100 +always increment aggregate counters +``` + +This prevents uncontrolled table growth and avoids depending on delayed cleanup. + +### 16.4 Prometheus Naming Reserved + +Not required in MVP, but reserve names: + +```text +newapi_channel_flow_running +newapi_channel_flow_queued +newapi_channel_flow_acquired_total +newapi_channel_flow_rejected_total +newapi_channel_flow_timeout_total +newapi_channel_flow_cancelled_total +newapi_channel_flow_wait_ms +newapi_channel_flow_process_ms +newapi_channel_flow_lease_renew_fail_total +``` + +## 17. Error Codes and i18n + +Backend error codes: + +```text +channel_flow_queue_full +channel_flow_queue_timeout +channel_flow_context_exceeded +channel_flow_draining +channel_flow_backend_unavailable +channel_flow_config_invalid +channel_flow_billing_failed_after_wait +channel_flow_per_user_queue_full +``` + +Frontend/backend i18n keys: + +```text +channel_flow.queue_full +channel_flow.queue_timeout +channel_flow.context_exceeded +channel_flow.service_draining +channel_flow.backend_unavailable +channel_flow.config_invalid +channel_flow.billing_failed_after_wait +channel_flow.per_user_queue_full +channel_flow.memory_backend_warning +channel_flow.redis_degraded_warning +``` + +Queue full response: + +```json +{ + "error": { + "message": "The upstream resource pool is busy and the waiting queue is full. Please retry later.", + "type": "rate_limit_error", + "code": "channel_flow_queue_full", + "metadata": { + "pool_running": 60, + "pool_max_inflight": 60, + "pool_queued": 240, + "pool_max_queue_size": 240, + "retry_after_seconds": 30 + } + } +} +``` + +Set: + +```text +Retry-After: 30 +``` + +Normal users should not see sensitive pool names unless admin explicitly enables it. + +## 18. Config Hot Update + +On pool config update: + +```text +1. DB transaction updates channel_flow_pools and increments config_version. +2. Redis config hash updates. +3. Local cache invalidates or refreshes. +4. Status endpoint returns new config_version. +``` + +Runtime behavior: + +- Reducing `max_inflight` does not cancel running requests. +- New dispatch pauses until running drops below the new max. +- Reducing `max_queue_size` does not kill already queued requests by default. +- New enqueue rejects if valid queue length is already above the new max. +- Disabling a pool stops new acquire according to policy; running requests drain. +- Changing Redis failure policy takes effect on next acquire. + +## 19. API Surface + +Pool CRUD: + +```text +GET /api/channel_flow/pools +POST /api/channel_flow/pools +GET /api/channel_flow/pools/:id +PUT /api/channel_flow/pools/:id +DELETE /api/channel_flow/pools/:id +``` + +Bindings: + +```text +GET /api/channel_flow/pools/:id/bindings +POST /api/channel_flow/pools/:id/bindings +DELETE /api/channel_flow/bindings/:id +``` + +Status, metrics, events: + +```text +GET /api/channel_flow/pools/:id/status +GET /api/channel_flow/pools/:id/metrics?from=&to=&bucket= +GET /api/channel_flow/pools/:id/events?from=&to=&event_type= +``` + +Suggestions: + +```text +GET /api/channel_flow/suggestions?channel_id=123 +``` + +Suggestions are advisory only. They never create or modify bindings automatically. + +## 20. Roadmap v3 + +### Phase 0: Redis Transaction Spike + +Must happen before Redis backend implementation. + +Scope: + +```text +Redis keys with hash tags +immediate acquire +enqueue +adaptive wait loop +TryPromoteSelf +release as simple ZREM +queue timeout +client cancellation cleanup +conflict metrics +``` + +Load test: + +```text +1000 concurrent acquire attempts +max_inflight = 60 +max_queue_size = 240 +multiple gateway-like goroutines +release storm simulation +stale head simulation +``` + +Pass conditions: + +```text +running never exceeds 60 except documented lease-expiry boundary +queue full and timeout deterministic +transaction conflict rate measured and acceptable +p99 acquire/promotion latency acceptable +Redis ops/s acceptable for target deployment +``` + +If fail: + +```text +design small Lua scripts for acquire/enqueue and self-promotion +``` + +### Phase 1: Correct Single-instance Product Release + +Required: + +- Flow Pool DB tables. +- Binding table. +- Channel-level binding only. +- Memory backend. +- Backend interface and FlowController. +- Channel drawer Flow Control section. +- Flow Pools list/detail page. +- Realtime status. +- Minimum trend charts for running and queued. +- Bounded queue. +- Queue timeout. +- Optional `max_queue_per_user` field, default off. +- Per-attempt acquire/release in normal relay. +- Client disconnect detection. +- Idempotent guard. +- Stream release wrapper. +- Billing precheck before queue. +- Billing preconsume/reserve after acquire. +- Admin warning for memory backend. + +Not included: + +- Redis production backend. +- Upstream-model binding. +- `fallback_then_queue`. +- Exact percentile charts. +- `on_task_finish` async task holding. + +### Phase 2: Operational Visibility and Model Binding + +- Upstream model binding in UI and runtime. +- Event retention config. +- Event sampling after daily cap. +- Wait/process percentile approximation. +- Per-user queue cap UI defaults for shared pools. +- More detailed event filtering. +- Admin notifications for degraded backend. + +### Phase 3a: Redis Basic Backend + +- Redis acquire/enqueue. +- Redis waiting loop. +- Simple release. +- Queue timeout/cancel cleanup. +- Redis status endpoint. +- Multi-instance tests. + +### Phase 3b: Redis Lease and Recovery + +- Lease renewal. +- Expired lease cleanup. +- Recovery cooldown after fail_open. +- Optional Pub/Sub wakeup. +- Degraded backend metrics. + +### Phase 4: Capacity-aware Routing + +- Candidate-set channel selection. +- Pool load-aware routing. +- `fallback` without marking channel failed. +- `fallback_then_queue`. + +### Phase 5: Advanced Controls + +- VIP priority. +- Weighted fair queueing if needed. +- Context/token in-flight budget. +- `on_task_finish` for async tasks. +- Lua optimization if WATCH/MULTI metrics require it. +- Prometheus/OpenTelemetry exporter. + +## 21. Acceptance Criteria + +Functional: + +```text +With max_inflight=60, no more than 60 non-expired running requests are admitted. +With max_queue_size=240, the 301st simultaneous request gets queue_full or fallback behavior. +Queued requests are dispatched when running slots release. +Queue timeout removes waiting request. +Client disconnect removes waiting request. +Stream requests hold guard until stream ends or is dropped. +Retry attempt releases old guard before acquiring a new one. +Pool full is not recorded as channel failure or auto-ban. +Billing is not preconsumed while request is waiting. +Billing session is created/reserved only after guard acquire. +``` + +Observability: + +```text +Realtime status shows running and queued. +Trend chart shows in-flight and queued history. +Events show queue_full, timeout, cancellation, and backend failures. +Event volume is capped or sampled. +``` + +Multi-instance Redis: + +```text +Two or more gateway instances share the same pool capacity. +Release storm does not create high transaction conflict loops. +Stale queue head does not block valid later waiters forever. +Redis config changes do not invalidate all transactions. +``` + +Compatibility: + +```text +Migrations work on SQLite, MySQL, and PostgreSQL. +Memory backend remains usable for dev/single instance. +Existing user rate limit and billing behavior remain separate. +``` + +## 22. Open Decisions for Reviewer + +These should be explicitly approved or rejected before implementation: + +1. Phase 1 only supports channel-level binding; upstream-model binding moves to Phase 2. +2. Queue length is mandatory and defaults to `max_inflight * 4`. +3. `pool_key` is backend-generated and never typed manually by admin. +4. Base URL only provides suggestions; explicit binding is the source of truth. +5. Billing precheck is read-only and not a quota hold. +6. If quota is consumed while waiting, post-acquire billing failure is accepted for MVP with clear error. +7. Redis config is not watched in transactions. +8. Redis release does not promote waiters. +9. Waiters self-promote with adaptive polling. +10. Redis Lua is optional and depends on Phase 0 spike results. +11. Memory backend can ship first but must show a multi-instance warning. +12. Trend charts are included in the first usable product release. + +## 23. Final Recommendation + +v3 is implementable and cleaner than v2. + +The most important design simplification is: + +```text +Release only releases. +Waiters promote themselves. +Config is not part of Redis optimistic locking. +Billing starts only after capacity is actually acquired. +``` + +This keeps the first implementation understandable while preserving a clear upgrade path to Lua, capacity-aware routing, and advanced fairness. + +For the user's 96-GPU/60-concurrency upstream, the recommended initial production policy is: + +```text +Flow Pool: one explicit pool bound to all channels that share that physical upstream. +max_inflight: 60 +max_queue_size: 240 +queue_timeout_ms: 120000 +backend: redis +redis_failure_policy: fail_closed if the upstream must never exceed capacity, otherwise fail_open with cooldown +trend charts: running and queued avg/max from minute aggregates +``` + +## 24. Final Review Resolutions for Implementation + +The v3 final review approved implementation and raised several implementation-level details. These are the decisions to carry into code: + +1. `BillingSession.Reserve(targetQuota)` means "ensure total pre-reserved quota reaches targetQuota", not "add targetQuota again". This matches the existing `BillingSession.Reserve` implementation. +2. Phase 1 client disconnect handling uses `c.Request.Context()` plus `queue_timeout_ms` as the fallback. Active response flushing/probing is deferred. +3. Redis backend should avoid the `waiting` + `waiting_deadline` dual-ZSET design. Phase 0/Phase 3 should use a single `waiting` ZSET scored by enqueue timestamp/sequence-compatible ordering, with timeout derived from enqueue time. +4. Redis request metadata TTL must be `max(queue_timeout_ms, max_processing_ms) + 300s` so cleanup can still read user/channel metadata. +5. Redis per-user queue accounting must be updated in the same atomic transition as waiting enqueue/remove/promote. If metadata cannot be guaranteed, encode enough user identity into the waiting member format. +6. Memory backend event caps can be implemented with a per-pool/day in-memory counter when event persistence lands. +7. `fail_open` should use a local memory safety valve by default once Redis backend is implemented. Phase 1 does not ship Redis backend. +8. Strict FIFO may leave capacity idle for one poll interval in Redis mode. This is accepted for v3 because request processing time is much larger than the expected poll delay. + +Phase 1 implementation scope is therefore: + +```text +DB tables +Pool CRUD and channel-level binding APIs +Memory backend with bounded queue +Relay per-attempt acquire/release +Billing precheck before queue and preconsume/reserve after acquire +Docker build/start validation +``` diff --git a/docs/channel-flow-control-queue-design.md b/docs/channel-flow-control-queue-design.md new file mode 100644 index 000000000000..b50a93d2bf9a --- /dev/null +++ b/docs/channel-flow-control-queue-design.md @@ -0,0 +1,1476 @@ +# Channel Flow Control and Queue Design Report + +Date: 2026-06-13 + +Status: draft for architecture and product review + +## 1. Executive Summary + +This document describes a proposed channel-level flow control and queueing feature for new-api. + +The core requirement is not ordinary user rate limiting. The target scenario is an upstream model resource pool, for example a 96-GPU cluster, that can only safely process 60 concurrent requests. The gateway must prevent the 61st request from entering the upstream. It should hold excess requests in a bounded queue, release them when capacity is available, expose real-time and historical traffic trends, and make the configuration understandable in the web admin UI. + +The recommended model is: + +```text +Flow Pool + -> generated stable pool key, hidden from normal admin workflow + -> human-readable name configured by admin + -> bindings to channels and optional upstream models + -> max_inflight, max_queue_size, queue_timeout_ms + -> real-time and historical metrics +``` + +The previous term `pool_id` should not be exposed as a raw field for users to type. In the product UI, the user should create or select a "Flow Pool" and bind channels/upstream models to it. The backend generates a stable key for runtime use. + +Recommended first production configuration for a 60-concurrency upstream: + +```text +max_inflight: 60 +max_queue_size: 240 +queue_timeout_ms: 120000 +queue_policy: fifo +``` + +Queue length must have an upper bound. An unbounded queue only moves overload from the upstream to the gateway and eventually causes memory pressure, connection exhaustion, poor user experience, and retry storms. + +## 2. Requirement Background + +### 2.1 Problem Statement + +new-api currently supports several forms of rate limiting and protection: + +- IP-based web/API/critical endpoint throttling. +- Per-user search throttling. +- Per-user model request throttling. +- User/token/subscription quota pre-consumption and settlement. +- System CPU/memory/disk protection. + +These controls do not solve upstream capacity protection. + +The target business case: + +```text +An upstream model is served by a 96-GPU cluster. +The upstream cluster supports at most 60 concurrent requests. +new-api may receive more than 60 simultaneous requests for that model/channel. +The gateway must cap upstream concurrency at 60. +Excess requests should wait in a queue, not hit upstream. +The admin must be able to see in-flight and queued traffic trends. +``` + +### 2.2 Why Existing Rate Limits Are Not Enough + +Request-per-minute limits and user-level limits answer questions like: + +```text +How many requests can this user send in a window? +How many requests can this IP send in a window? +How many successful requests can this group make? +``` + +They do not answer: + +```text +How many requests are currently occupying the same upstream GPU pool? +How many requests are waiting before this request can enter upstream? +Which channel/model is building up queue pressure? +Did traffic spike because of one channel, one model, or one group? +``` + +This feature should therefore be treated as admission control: + +```text +Before entering upstream: + if running < max_inflight -> dispatch + else if queue has room -> wait + else -> reject +``` + +## 3. Terminology + +| Term | Meaning | +|---|---| +| Channel | Existing new-api channel record. It stores type, base URL, key, models, mappings, settings, etc. | +| Upstream URL | The base URL configured on a channel or default channel base URL. It is an input signal, not the source of truth for pooling. | +| Upstream model | The model actually sent to the provider after model mapping. | +| Flow Pool | A logical/physical upstream capacity pool. Example: "96-card DeepSeek-R1 production pool". | +| Pool key | Backend-generated stable key for runtime counters, for example `flow_pool_8f3a2c`. Users should not type it manually. | +| Binding | A relation between a flow pool and one or more channels, optionally narrowed to upstream models. | +| Fingerprint | A derived hint from channel type, normalized base URL, upstream model, and other provider-specific fields. It is used for recommendations only. | +| In-flight/running | Requests that already passed admission control and are currently occupying upstream capacity. | +| Queued | Requests waiting in the gateway before entering upstream. | +| Queue timeout | Maximum time a request may wait before the gateway returns an error. | + +## 4. External Research + +The market trend is that LLM gateways increasingly support more than request-per-minute throttling. Mature systems often combine request windows, token windows, concurrency protection, budgets, and provider fallback. + +### 4.1 LiteLLM Proxy + +LiteLLM Proxy supports user/team/key/model budget and rate limit concepts such as RPM, TPM, and max parallel requests. It also has queueing/prioritization capabilities in its scheduler. + +References: + +- https://docs.litellm.ai/docs/proxy/users +- https://docs.litellm.ai/docs/routing-load-balancing +- https://docs.litellm.ai/docs/scheduler + +Takeaway for new-api: + +- Keep tenant/user limits separate from upstream capacity limits. +- Queueing should be explicit and observable. +- Redis or another shared backend is needed for multi-instance deployments. + +### 4.2 Kong AI Gateway + +Kong AI Rate Limiting Advanced focuses on AI-aware rate limiting, including token-aware cost calculation and different counter strategies such as local, cluster, and Redis. + +Reference: + +- https://docs.konghq.com/hub/kong-inc/ai-rate-limiting-advanced/ + +Takeaway for new-api: + +- Production-grade gateway limits must define the storage consistency model. +- Local counters are not enough when multiple gateway instances serve the same upstream pool. + +### 4.3 APISIX AI Gateway + +APISIX has AI rate limiting support around LLM token dimensions such as total, prompt, and completion tokens. It can also work with upstream instance/fallback behavior. + +Reference: + +- https://apisix.apache.org/docs/apisix/plugins/ai-rate-limiting/ + +Takeaway for new-api: + +- Rate limiting and routing/fallback need to interact. +- If one upstream instance is saturated, the gateway can either queue on it, choose another instance, or reject. + +### 4.4 Envoy AI Gateway + +Envoy AI Gateway uses Envoy/Gateway API concepts and supports provider fallback and usage-based rate limiting in a Kubernetes-oriented architecture. + +References: + +- https://aigateway.envoyproxy.io/docs/capabilities/ +- https://aigateway.envoyproxy.io/docs/capabilities/traffic/provider-fallback +- https://aigateway.envoyproxy.io/docs/capabilities/traffic/usage-based-ratelimiting + +Takeaway for new-api: + +- Capacity control is naturally tied to backend/provider identity. +- Explicit backend identity is better than inferring everything from URL strings. + +### 4.5 Azure API Management GenAI Gateway + +Azure API Management provides GenAI gateway policies such as token limits and token metrics for Azure OpenAI and related LLM traffic. + +References: + +- https://learn.microsoft.com/en-us/azure/api-management/genai-gateway-capabilities +- https://learn.microsoft.com/en-us/azure/api-management/azure-openai-token-limit-policy + +Takeaway for new-api: + +- Token limits are useful but separate from concurrent GPU occupancy. +- The gateway should expose metrics for admin troubleshooting. + +### 4.6 Portkey AI Gateway + +Portkey supports virtual keys, provider/integration limits, load balancing, fallback, and retry behaviors. + +References: + +- https://portkey.ai/docs/product/ai-gateway/virtual-keys/rate-limits +- https://portkey.ai/docs/product/ai-gateway/load-balancing +- https://portkey.ai/docs/product/ai-gateway/fallbacks + +Takeaway for new-api: + +- Provider-level controls and fallback policies are part of the admin product surface. +- The UI should make limit ownership clear. + +### 4.7 Cloudflare AI Gateway + +Cloudflare AI Gateway supports gateway-level rate limiting with fixed/sliding window policies. + +Reference: + +- https://developers.cloudflare.com/ai-gateway/configuration/rate-limiting/ + +Takeaway for new-api: + +- Simple request-window throttling is useful, but insufficient for upstream GPU pool capacity. + +## 5. Local Reference: gateway Project + +The local `/Users/laiyouxu/IdeaProjects/gateway` project already implements a closely related flow control pattern. + +Relevant files: + +- `/Users/laiyouxu/IdeaProjects/gateway/config.example.yaml` +- `/Users/laiyouxu/IdeaProjects/gateway/boom-gateway/boom-config/src/lib.rs` +- `/Users/laiyouxu/IdeaProjects/gateway/boom-gateway/boom-flowcontrol/src/lib.rs` +- `/Users/laiyouxu/IdeaProjects/gateway/boom-gateway/boom-main/src/routes.rs` +- `/Users/laiyouxu/IdeaProjects/gateway/boom-gateway/boom-main/src/state.rs` +- `/Users/laiyouxu/IdeaProjects/gateway/boom-gateway/boom-routing/src/policy/load_helpers.rs` +- `/Users/laiyouxu/IdeaProjects/gateway/boom-gateway/boom-dashboard/src/handlers_admin.rs` + +### 5.1 What gateway Does Well + +The gateway project has a `FlowController` and per-deployment slots. Its flow control configuration uses a deployment identity: + +```yaml +model_info: + id: gpt4o-node-1 +flow_control: + model_queue_limit: 50 + model_context_limit: 5000000 +``` + +Important design points: + +- `model_queue_limit` is actually max in-flight concurrency, not max queue length. +- A deployment slot maintains two queues, VIP and normal. +- The queue itself is the source of truth. +- `dispatched = true` means in-flight. +- `dispatched = false` means waiting. +- This avoids maintaining separate counters that can leak. +- `FlowControlGuard` releases capacity on drop. +- `FlowControlledStream` releases capacity when stream ends. +- The dashboard exposes in-flight and queued status. +- Routing can consider total load: in-flight plus queued. + +### 5.2 Gaps in gateway Relevant to new-api + +The gateway implementation is a strong reference but not a complete target for new-api: + +| Area | gateway behavior | Recommended new-api behavior | +|---|---|---| +| Queue length | No explicit max queue size found | Must support `max_queue_size` | +| Queue timeout | Fixed 1200 seconds in route code | Per-pool configurable `queue_timeout_ms` | +| Storage | In-process memory | Memory backend for single instance, Redis backend for production | +| Pool identity | `deployment_id` | Flow Pool with generated `pool_key` and admin-visible name | +| URL binding | Deployment config based | Explicit binding table plus URL/model fingerprint suggestions | +| Multi-instance | Local process only | Redis lease/semaphore for global capacity | + +new-api should borrow the "queue as source of truth" idea, but add bounded queue, configurable timeout, explicit resource-pool management, and metrics storage. + +## 6. Current new-api Capability Review + +### 6.1 Existing Rate Limits and Protections + +| Feature | Granularity | Implementation | +|---|---|---| +| Global web limit | IP | `middleware/rate-limit.go` | +| Global API limit | IP | `middleware/rate-limit.go` | +| Critical endpoint limit | IP | `middleware/rate-limit.go` | +| Upload/download limit | IP | `middleware/rate-limit.go` | +| Search limit | authenticated user ID | `middleware/rate-limit.go` | +| Email verification limit | IP | `middleware/email-verification-rate-limit.go` | +| Model request limit | user ID, group override | `middleware/model-rate-limit.go`, `setting/rate_limit.go` | +| User/token/subscription quota | user/token/subscription | `service/billing.go`, `service/billing_session.go`, `service/quota.go` | +| System overload protection | process/system | `middleware/performance.go` | +| Notification send limit | user and notification type | `service/notify-limit.go` | + +### 6.2 Existing Routing and Channel Points + +Important existing files: + +- `router/relay-router.go` +- `middleware/distributor.go` +- `controller/relay.go` +- `model/channel.go` +- `dto/channel_settings.go` + +Current relevant behavior: + +- `/v1` and `/v1beta` use `ModelRequestRateLimit`. +- `/mj` and `/suno` do not currently use `ModelRequestRateLimit`. +- `Distribute()` selects a channel. +- `SetupContextForSelectedChannel()` stores channel metadata in context and selects multi-key key/index. +- `controller/relay.go` has retry loops for normal relay and task relay. +- `Channel` already has JSON fields `Setting` and `OtherSettings`. +- `ChannelInfo` supports multi-key status and random/polling key selection. + +### 6.3 Current Gap + +new-api currently does not have: + +- Channel-level max in-flight control. +- Shared resource-pool control across multiple channels. +- Per-channel or per-pool waiting queue. +- Configurable queue timeout. +- Configurable queue length. +- Stream/WebSocket-aware flow-control guard. +- Redis-based distributed semaphore/queue for upstream capacity. +- In-flight and queue trend charts. +- Flow-control event tracing. + +## 7. Design Goals + +### 7.1 Goals + +1. Cap upstream concurrency for a logical upstream resource pool. +2. Queue excess requests in a bounded FIFO queue. +3. Avoid sending more requests to upstream than the configured capacity. +4. Support multiple channels sharing the same upstream resource pool. +5. Support optional upstream-model-specific bindings. +6. Support normal HTTP, streaming, realtime/WebSocket, and task relay with correct release timing. +7. Provide real-time status and historical trend charts. +8. Provide admin-friendly configuration in the web UI. +9. Support single-instance memory mode and production Redis mode. +10. Keep existing user rate limits and billing quota behavior separate. + +### 7.2 Non-goals for the First Version + +The first version should not try to solve every traffic-shaping problem: + +- No complex weighted fair queueing. +- No full TPM token bucket implementation. +- No automatic GPU utilization integration. +- No automatic pool discovery from URLs without admin confirmation. +- No cross-region active-active queueing. +- No model-specific dynamic autoscaling. + +These can be later phases. + +## 8. Product Model: Flow Pool + +### 8.1 Why Users Should Not Type `pool_id` + +A raw `pool_id` is an implementation detail. Asking users to type it leads to confusion: + +```text +Where does this ID come from? +Is it the channel ID? +Is it the upstream URL? +Is it the model name? +Is it provided by the upstream? +``` + +The product should expose: + +```text +Flow Pool name: 96-card DeepSeek-R1 production pool +Flow Pool bindings: channels and upstream models +Flow Pool capacity: 60 concurrent requests +Queue: 240 requests, 120 seconds timeout +``` + +The backend should generate: + +```text +pool_key: flow_pool_8f3a2c... +``` + +This key is used in Redis/runtime storage, logs, metrics, and internal APIs. + +### 8.2 How Flow Pool Relates to Channel and URL + +The binding must be explicit. URL matching can only be a recommendation. + +Why URL alone is unsafe: + +- Same base URL may serve multiple independent model pools. +- Same base URL plus different API keys may map to different upstream tenants. +- Same physical pool may be available under multiple URLs. +- Azure-like providers need deployment names and API versions. +- Model mapping can change the actual upstream model. +- A private OpenAI-compatible gateway may multiplex different GPU pools behind one URL. + +Recommended binding truth: + +```text +Flow Pool -> channel_id +Flow Pool -> optional upstream_model +``` + +Runtime resolution priority: + +```text +1. Exact binding: channel_id + upstream_model +2. Channel binding: channel_id +3. No binding: no flow control, unless admin explicitly selected "independent channel pool" +``` + +URL/model fingerprint is only used to recommend a binding: + +```text +fingerprint = hash(channel_type + normalized_base_url + upstream_model + provider_specific_identity) +``` + +Admin UI may show: + +```text +Detected 3 channels with similar upstream identity. Bind them to the same Flow Pool? +``` + +But it should not silently merge them. + +## 9. Web Admin UX Design + +new-api default frontend already has a channels module under: + +```text +web/default/src/features/channels +``` + +Channel create/update is handled by: + +```text +web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx +``` + +The current drawer has sections such as Basic, API Access, Models, and Advanced Settings. The flow-control UI should fit into this existing pattern. + +### 9.1 Channel Edit Drawer + +Add a new section under Advanced Settings: + +```text +Advanced Settings + - Routing & Overrides + - Flow Control & Queue + - Request Overrides + - Upstream Model Automation +``` + +Suggested section content: + +```text +[Switch] Enable flow control and queue + +Resource pool + [Radio] This channel uses an independent pool + Runtime key is generated from channel ID after save. + + [Radio] Bind to an existing Flow Pool + [Select] 96-card DeepSeek-R1 production pool + Summary: 60 in-flight / queue 240 / timeout 120s / 4 bound channels + + [Radio] Create a new Flow Pool + Name: 96-card DeepSeek-R1 production pool + Max in-flight requests: 60 + Max queue size: 240 + Queue timeout: 120 seconds + Queue policy: FIFO + +Binding scope + [Radio] All models on this channel + [Radio] Only selected upstream models + [Multi-select] deepseek-r1, deepseek-v3 + +Upstream identity preview + Channel type: OpenAI Compatible + Base URL: https://example.com/v1 + Published models: deepseek-r1, deepseek-v3 + Upstream model mapping: deepseek-r1 -> deepseek-r1-prod + Suggested fingerprint: openai-compatible / example.com / deepseek-r1-prod +``` + +UX rule: + +- The raw `pool_key` should not be the main user-facing field. +- Advanced users may see it in a read-only details drawer for debugging. + +### 9.2 Flow Pools Management Page + +Add a tab or subpage under Channels: + +```text +Channels | Flow Pools +``` + +List view: + +```text +Name Bound channels In-flight Queue Wait P95 Rejected/Timeout Status +96-card R1 production pool 4 48 / 60 132 / 240 18.4s 2 / 7 Congested +Qwen-VL backup pool 1 6 / 20 0 / 80 0.2s 0 / 0 Healthy +``` + +Detail view: + +```text +Basic + Name + Description + Enabled + Generated pool key, read-only + +Capacity + Max in-flight requests + Max queue size + Queue timeout + Queue policy + Optional max in-flight context tokens/chars + +Bindings + Channel + Channel type + Base URL + Upstream model + Binding mode + +Realtime Status + Running + Queued + Oldest waiting seconds + Utilization + +Trends + In-flight requests + Queued requests + Wait duration P50/P95/P99 + Processing duration P50/P95/P99 + Rejections and timeouts +``` + +### 9.3 In-channel Status Entry + +In the channel table, add a compact indicator: + +```text +Flow: 48/60 running, 132 queued +``` + +Clicking it opens the Flow Pool detail. + +### 9.4 User-facing Error Messages + +When queue is full: + +```json +{ + "error": { + "message": "The upstream resource pool is busy. The waiting queue is full. Please retry later.", + "type": "rate_limit_error", + "code": "channel_flow_queue_full" + } +} +``` + +When queue times out: + +```json +{ + "error": { + "message": "The upstream resource pool is busy. The request waited too long in queue.", + "type": "rate_limit_error", + "code": "channel_flow_queue_timeout" + } +} +``` + +Recommended HTTP status: + +```text +429 for queue full and queue timeout +503 only for system overload or disabled upstream capacity +``` + +## 10. Backend Data Model + +Use GORM-compatible models and migrations. Keep SQLite, MySQL, and PostgreSQL compatibility. + +### 10.1 Flow Pool Table + +Suggested table: `channel_flow_pools` + +```text +id int primary key +pool_key varchar unique, generated by backend +name varchar +description text +enabled bool/int +max_inflight int +max_queue_size int +queue_timeout_ms int +queue_policy varchar, default "fifo" +max_context_tokens int, optional +max_context_chars int, optional +on_limit varchar, default "queue" +created_time bigint +updated_time bigint +``` + +Notes: + +- `pool_key` is not user-provided. +- `max_queue_size` should be required when enabled. +- `max_inflight` must be greater than 0 when enabled. +- `queue_timeout_ms` must be bounded by system max to avoid extremely long connection retention. +- Use plain text/varchar fields and GORM abstractions for cross-database compatibility. + +### 10.2 Binding Table + +Suggested table: `channel_flow_pool_bindings` + +```text +id int primary key +pool_id int +channel_id int +upstream_model varchar, optional +match_mode varchar, "channel" | "channel_model" +enabled bool/int +created_time bigint +updated_time bigint +``` + +Resolution: + +```text +if binding exists for channel_id + upstream_model: + use that pool +else if binding exists for channel_id: + use that pool +else: + pass through without flow control +``` + +If the admin selects "independent pool for this channel", the backend creates a pool and a binding for that channel. + +### 10.3 Metrics Aggregate Table + +Suggested table: `channel_flow_metrics_minute` + +```text +id int primary key +bucket_ts bigint +pool_key varchar +channel_id int +model varchar +running_avg double/integer approximation +running_max int +queued_avg double/integer approximation +queued_max int +acquired_count int +queued_count int +released_count int +rejected_count int +timeout_count int +cancelled_count int +wait_ms_p50 int +wait_ms_p95 int +wait_ms_p99 int +process_ms_p50 int +process_ms_p95 int +process_ms_p99 int +created_time bigint +updated_time bigint +``` + +For cross-database simplicity: + +- Avoid JSONB. +- Avoid database-specific percentile functions. +- Compute percentiles in memory before writing aggregate rows. + +### 10.4 Optional Event Trace Store + +For traffic tracing, minute-level metrics are not enough. Add a bounded event log: + +```text +channel_flow_events + id + request_id + pool_key + channel_id + model + event_type enter_queue | dispatch | release | reject | timeout | cancel + reason + queue_pos + running + queued + wait_ms + process_ms + created_time +``` + +This table can become large. Options: + +- Keep only error/timeout/reject events in DB. +- Keep full recent events in Redis with TTL. +- Add a system option to enable full event tracing temporarily. + +Recommended first version: + +```text +Always aggregate metrics. +Always store reject and timeout events. +Store dispatch/release events only when debug tracing is enabled. +``` + +## 11. Runtime Flow + +### 11.1 Insertion Point + +The flow controller should run after channel selection and before upstream call. + +Current normal relay loop is in: + +```text +controller/relay.go +``` + +Proposed position: + +```text +for retry: + channel = getChannel(...) + acquire flow control guard + call upstream + release guard when done +``` + +The flow controller should not live inside `ModelRequestRateLimit`. User rate limits and upstream capacity control are different concerns. + +### 11.2 Request Metadata Needed + +Acquire needs: + +```text +request_id +user_id +group +token_id +channel_id +channel_name +channel_type +is_multi_key +multi_key_index +origin_model +upstream_model +is_stream +estimated_prompt_tokens +estimated_context_chars +``` + +### 11.3 Pool Resolution + +Pseudo-code: + +```go +func ResolveFlowPool(channelID int, upstreamModel string) (*FlowPool, bool) { + if binding := findBinding(channelID, upstreamModel); binding != nil { + return binding.Pool, true + } + if binding := findChannelBinding(channelID); binding != nil { + return binding.Pool, true + } + return nil, false +} +``` + +If no pool is resolved, the request passes through without channel flow control. + +### 11.4 Acquire Algorithm + +```text +Input: + pool_key + request_id + context cost + timeout + max_inflight + max_queue_size + +Algorithm: + 1. If pool disabled -> pass through or reject based on config. + 2. If request context exceeds max_context -> reject immediately. + 3. If running < max_inflight -> mark dispatched and return guard. + 4. If waiting >= max_queue_size -> reject with 429 queue_full. + 5. Enqueue request. + 6. Wait until dispatched, client cancels, or timeout occurs. + 7. On dispatch -> return guard. + 8. On cancellation -> remove from queue. + 9. On timeout -> remove from queue, return 429 queue_timeout. +``` + +### 11.5 Release Algorithm + +```text +On request completion: + 1. Remove dispatched request from running state. + 2. Record processing duration. + 3. Dispatch next fitting request from queue. + 4. Record metrics. +``` + +### 11.6 Stream and WebSocket Release + +For streaming: + +```text +Acquire before upstream stream starts. +Hold guard while stream is open. +Release when stream ends or client disconnects. +``` + +For realtime/WebSocket: + +```text +Acquire before upstream realtime connection. +Hold guard while WebSocket is active. +Release on close/error/cancel. +``` + +This mirrors the good part of the gateway project's `FlowControlledStream`. + +### 11.7 Task Relay Release Policy + +Async task routes need special handling. + +There are two possible upstream semantics: + +```text +submit_only: + Upstream only accepts the task and queues/processes it internally. + Gateway slot can release after submit response returns. + +occupies_until_finished: + Upstream task occupies GPU capacity until task finishes. + Gateway slot must remain held until task reaches terminal status. +``` + +Add per-pool or per-channel option: + +```text +task_release_policy: "on_submit" | "on_task_finish" +``` + +Default should be `on_submit` for compatibility, but for a private 96-GPU pool the admin may need `on_task_finish`. + +## 12. Backend Implementation Options + +### 12.1 Memory Backend + +Memory backend is useful for: + +- Development. +- Single-instance deployments. +- Redis-disabled installations. + +Design: + +```text +map[pool_key]*Slot +Slot: + mutex + max_inflight + max_queue_size + max_context + queue []RequestState + +RequestState: + request_id + dispatched bool + context cost + enqueue time + dispatch time + notify channel +``` + +The queue should be the source of truth: + +```text +dispatched == true -> in-flight +dispatched == false -> waiting +``` + +Do not maintain an independent `running` counter if it can be derived from the queue. This avoids counter leaks. + +### 12.2 Redis Backend + +Redis backend is required for multi-instance production. + +Reason: + +```text +If 3 gateway instances each enforce max_inflight=60 locally, +the upstream may receive 180 concurrent requests. +``` + +Suggested Redis keys: + +```text +flow:{pool_key}:running ZSET request_id -> lease_expire_ms +flow:{pool_key}:waiting ZSET request_id -> sequence or enqueue time +flow:{pool_key}:request:{id} HASH request metadata, TTL +flow:{pool_key}:seq INCR sequence +flow:{pool_key}:notify Pub/Sub or stream for wakeups +``` + +Acquire should be Lua-backed: + +```text +1. Remove expired running leases. +2. If running count < max_inflight: + add to running with lease + return acquired +3. If waiting count >= max_queue_size: + return queue_full +4. Add to waiting queue. +5. Return queued with sequence. +``` + +Wait loop: + +```text +The waiter polls or waits for Pub/Sub notification. +Only queue head can move to running. +If timeout/cancel: + remove from waiting. +``` + +Release script: + +```text +1. Remove request from running. +2. Move as many waiting head items as fit into running. +3. Publish wakeup events. +``` + +Lease handling: + +- Non-streaming requests can use a lease slightly longer than request timeout. +- Streaming/WebSocket requests need heartbeat renewal. +- If an instance crashes, leases expire and slots recover. + +### 12.3 Redis vs Memory Behavior + +| Area | Memory | Redis | +|---|---|---| +| Single instance | Good | Good | +| Multiple instances | Incorrect global capacity | Correct global capacity | +| Crash recovery | Lost state | Lease recovery | +| Implementation complexity | Lower | Higher | +| Recommended production default | No | Yes | + +## 13. Queue Length: Why It Must Have an Upper Bound + +Queue length should never be infinite. + +Risks of unbounded queues: + +- Gateway memory grows with request bodies and waiting contexts. +- HTTP connections remain open for a long time. +- Client timeouts cause cancellation churn. +- Retries amplify pressure. +- Waiting time becomes unbounded and user experience degrades. +- Admin cannot reason about worst-case capacity. + +Recommended default: + +```text +max_queue_size = max_inflight * 4 +queue_timeout_ms = 120000 +``` + +For the 60-concurrency scenario: + +```text +max_inflight = 60 +max_queue_size = 240 +queue_timeout_ms = 120000 +``` + +Sizing formula: + +```text +upstream throughput ~= max_inflight / average_processing_seconds +reasonable queue size ~= upstream throughput * max_acceptable_wait_seconds +``` + +Example: + +```text +max_inflight = 60 +average processing time = 30s +throughput ~= 2 requests/s +acceptable wait = 120s +queue size ~= 240 +``` + +UI should not allow `max_queue_size = unlimited`. If administrators need larger queues, they should explicitly raise the number. + +## 14. Retry and Fallback Interaction + +Flow control must define how it interacts with retry and channel fallback. + +Recommended `on_limit` policies: + +| Policy | Behavior | Use case | +|---|---|---| +| queue | Wait in the selected pool queue. | Single upstream pool, capacity must be preserved. | +| reject | Return 429 immediately when full. | Low-latency APIs. | +| fallback | Treat full pool as unavailable and try another channel. | Multiple equivalent upstream pools. | +| fallback_then_queue | Try other pools first, queue only if all candidates are full. | Multiple pools with shared SLA. | + +For the 96-GPU/60-concurrency scenario, recommended default: + +```text +on_limit = queue +``` + +If there are several equivalent GPU pools, use: + +```text +on_limit = fallback_then_queue +``` + +## 15. Metrics, Trend Charts, and Traceability + +### 15.1 Real-time Metrics + +Expose real-time status per pool: + +```text +running +max_inflight +queued +max_queue_size +oldest_wait_ms +utilization = running / max_inflight +queue_utilization = queued / max_queue_size +``` + +API example: + +```text +GET /api/channel_flow/pools/:id/status +``` + +Response: + +```json +{ + "pool_key": "flow_pool_8f3a2c", + "name": "96-card DeepSeek-R1 production pool", + "running": 48, + "max_inflight": 60, + "queued": 132, + "max_queue_size": 240, + "oldest_wait_ms": 18400, + "utilization": 0.8, + "queue_utilization": 0.55 +} +``` + +### 15.2 Historical Metrics + +Minute-level aggregation should support: + +- In-flight trend. +- Queue depth trend. +- Wait duration percentiles. +- Processing duration percentiles. +- Rejected and timeout trend. +- Per-channel contribution. +- Per-model contribution. +- Per-group contribution if available. + +API examples: + +```text +GET /api/channel_flow/pools/:id/metrics?from=...&to=...&bucket=minute +GET /api/channel_flow/pools/:id/events?limit=200&type=timeout,reject +``` + +### 15.3 Dashboard Charts + +Recommended charts: + +1. In-flight requests: + +```text +line: running +horizontal line: max_inflight +``` + +2. Queue depth: + +```text +line: queued +horizontal line: max_queue_size +``` + +3. Wait duration: + +```text +lines: p50, p95, p99 +``` + +4. Processing duration: + +```text +lines: p50, p95, p99 +``` + +5. Rejections and timeouts: + +```text +stacked bars: + queue_full + queue_timeout + context_exceeded + cancelled +``` + +6. Top contributors: + +```text +by channel +by upstream model +by user group +``` + +### 15.4 Traceability + +Every request that enters flow control should have a `request_id`. + +For review/debugging, trace events should show: + +```text +request_id +pool +channel +model +event timeline: + enter_queue at T1 + dispatch at T2 + release at T3 +wait_ms = T2 - T1 +process_ms = T3 - T2 +``` + +Recommended default storage: + +- Always store aggregate metrics. +- Store reject and timeout events in DB. +- Store recent detailed events in Redis with TTL. +- Add admin switch for temporary full tracing. + +## 16. API Design + +### 16.1 Flow Pool CRUD + +```text +GET /api/channel_flow/pools +POST /api/channel_flow/pools +GET /api/channel_flow/pools/:id +PUT /api/channel_flow/pools/:id +DELETE /api/channel_flow/pools/:id +``` + +Create request: + +```json +{ + "name": "96-card DeepSeek-R1 production pool", + "description": "Private upstream cluster A", + "enabled": true, + "max_inflight": 60, + "max_queue_size": 240, + "queue_timeout_ms": 120000, + "queue_policy": "fifo", + "max_context_tokens": 0, + "on_limit": "queue" +} +``` + +Response includes generated `pool_key`: + +```json +{ + "id": 1, + "pool_key": "flow_pool_8f3a2c", + "name": "96-card DeepSeek-R1 production pool" +} +``` + +### 16.2 Bindings + +```text +GET /api/channel_flow/pools/:id/bindings +POST /api/channel_flow/pools/:id/bindings +DELETE /api/channel_flow/bindings/:id +``` + +Create binding: + +```json +{ + "channel_id": 123, + "match_mode": "channel_model", + "upstream_model": "deepseek-r1-prod" +} +``` + +### 16.3 Suggestions + +```text +GET /api/channel_flow/suggestions?channel_id=123 +``` + +Response: + +```json +{ + "channel_id": 123, + "base_url": "https://example.com/v1", + "suggested_fingerprint": "openai-compatible/example.com/deepseek-r1-prod", + "similar_channels": [ + { + "channel_id": 124, + "name": "R1 backup key", + "base_url": "https://example.com/v1", + "models": ["deepseek-r1"] + } + ] +} +``` + +This API should not auto-bind. It only helps administrators avoid misconfiguration. + +### 16.4 Status and Metrics + +```text +GET /api/channel_flow/pools/:id/status +GET /api/channel_flow/pools/:id/metrics +GET /api/channel_flow/pools/:id/events +``` + +## 17. Integration With new-api Files + +Suggested backend additions: + +```text +dto/channel_flow.go +model/channel_flow_pool.go +model/channel_flow_binding.go +model/channel_flow_metric.go +service/channel_flow/ + controller.go + memory_backend.go + redis_backend.go + metrics.go +controller/channel_flow.go +router/api-router.go +``` + +Suggested frontend additions: + +```text +web/default/src/features/channels/components/drawers/sections/channel-flow-control-section.tsx +web/default/src/features/channels/components/dialogs/flow-pool-detail-dialog.tsx +web/default/src/features/channels/components/flow-pools-table.tsx +web/default/src/features/channels/hooks/use-channel-flow-pools.ts +web/default/src/features/channels/lib/channel-flow.ts +``` + +Existing form integration points: + +```text +web/default/src/features/channels/lib/channel-form.ts +web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx +web/default/src/features/channels/types.ts +``` + +Do not force users to edit raw JSON settings for flow control. + +## 18. Validation Rules + +Backend validation: + +```text +name: required +max_inflight: required and > 0 when enabled +max_queue_size: required and >= 0 +queue_timeout_ms: required and between 1000 and configured system max +queue_policy: fifo for v1 +on_limit: queue | reject | fallback | fallback_then_queue +``` + +Recommended hard caps: + +```text +max_inflight <= 100000 +max_queue_size <= 100000 +queue_timeout_ms <= 3600000 +``` + +UI validation: + +```text +If max_inflight = 60, suggest max_queue_size = 240. +Warn if queue size > max_inflight * 10. +Warn if queue timeout > client timeout. +Warn if multiple pools appear to bind the same channel/model. +Warn if similar URLs are not bound together. +``` + +## 19. Failure Modes and Safeguards + +| Failure | Safeguard | +|---|---| +| Request cancelled while waiting | Remove from queue in cleanup. | +| Request cancelled while streaming | Drop guard and release slot. | +| Gateway instance crashes | Redis lease expires running entry. | +| Queue grows too large | `max_queue_size` hard cap. | +| Wait too long | `queue_timeout_ms`. | +| Pool mis-bound by URL | Explicit bindings, URL only suggests. | +| Multiple gateway instances | Redis backend. | +| Metrics table grows too large | Retention policy and aggregation. | +| Admin reduces max_inflight below current running | Do not kill running requests; only stop new dispatch until running drops. | +| Admin disables pool | Existing running requests continue; new acquire rejects or passes according to policy. | + +## 20. Implementation Roadmap + +### Phase 1: Single-instance MVP + +- Add Flow Pool and Binding models. +- Add admin APIs. +- Add memory backend. +- Add acquire/release around normal relay. +- Add stream-safe release. +- Add queue length and timeout. +- Add real-time status API. +- Add basic UI in channel edit drawer. + +### Phase 2: Metrics and Dashboard + +- Add minute-level metrics aggregation. +- Add flow pool list and detail page. +- Add in-flight and queue trend charts. +- Add reject/timeout event list. +- Add channel table flow-status indicator. + +### Phase 3: Redis Production Backend + +- Add Redis Lua scripts. +- Add distributed running lease. +- Add waiting queue and wakeup mechanism. +- Add stream heartbeat lease renewal. +- Add crash recovery tests. + +### Phase 4: Advanced Routing + +- Add `fallback` and `fallback_then_queue`. +- Make channel selection aware of pool load. +- Add pool utilization to routing decision. +- Add optional VIP priority. + +### Phase 5: Token/Context Enhancements + +- Add max in-flight context tokens/chars. +- Add TPM-like token window if needed. +- Add per-model or per-group overrides inside a pool. + +## 21. Test Plan + +### 21.1 Unit Tests + +- Acquire dispatches immediately when capacity exists. +- Acquire queues when capacity full. +- Queue full rejects. +- Queue timeout removes waiting request. +- Cancellation removes waiting request. +- Release dispatches next request. +- Admin lowering max_inflight does not corrupt state. +- Context-exceeded request rejects immediately. + +### 21.2 Integration Tests + +- 100 concurrent requests with `max_inflight = 60` never dispatch more than 60 upstream calls. +- Stream request holds slot until stream ends. +- Client disconnect releases slot. +- Queue order is FIFO. +- Metrics record running max and queued max correctly. +- Retry/fallback policies behave as configured. + +### 21.3 Redis Tests + +- Multiple processes share the same max_inflight. +- Running lease expires after simulated crash. +- Heartbeat keeps long stream alive. +- Release wakes queued requests. +- Timeout removes waiting request atomically. + +### 21.4 UI Tests + +- Create Flow Pool from channel drawer. +- Bind channel to existing Flow Pool. +- Bind channel + upstream model to Flow Pool. +- Suggested similar channels are shown but not auto-bound. +- Trend chart renders when metrics exist. +- Validation warnings appear for risky queue values. + +## 22. Open Questions for Review + +1. Should the default when no binding exists be "no flow control" or "auto independent channel pool"? + - Recommendation: no flow control unless explicitly enabled. + +2. Should queue timeout return 429 or 503? + - Recommendation: 429 for flow-control pressure, 503 for system overload. + +3. Should async task slots be released on submit or task finish? + - Recommendation: make it configurable by pool/channel. + +4. Should `max_context_tokens` use estimated prompt tokens or raw input chars in v1? + - Recommendation: start with input chars or estimated prompt tokens already available in relay; refine later. + +5. Should Redis backend be required when Redis is enabled globally? + - Recommendation: yes, if Redis is enabled use Redis backend for flow control. + +6. Should VIP priority be included in v1? + - Recommendation: not in MVP unless there is an immediate product requirement. + +7. Should flow pool config live in DB tables or channel JSON settings? + - Recommendation: DB tables, because shared pools cannot be safely represented by per-channel JSON. + +## 23. Recommended Decision + +Implement channel flow control as a first-class Flow Pool feature. + +Do not expose raw `pool_id` as a user-filled field. The admin creates/selects a Flow Pool by name, binds channels and optional upstream models, and the backend generates a stable runtime `pool_key`. + +Use explicit bindings as the source of truth. Use URL/upstream model fingerprints only for suggestions and warnings. + +For the first usable release, implement: + +```text +Flow Pool CRUD +Channel/model bindings +max_inflight +max_queue_size +queue_timeout_ms +FIFO queue +normal and stream release +real-time status +minute trend metrics +basic dashboard charts +``` + +For production safety, add Redis backend before recommending this for multi-instance deployments. + +## 24. Reviewer Checklist + +Use this checklist when reviewing the design: + +- Does the design protect a 96-GPU upstream with a strict 60-concurrency cap? +- Does it avoid relying on URL-only inference? +- Is the origin of pool identity clear? +- Can multiple channels share the same upstream capacity pool? +- Is the queue bounded? +- Is queue timeout configurable? +- Does stream/WebSocket release happen at the correct time? +- Does the design work in multi-instance deployments? +- Can administrators configure it from the web UI without editing raw JSON? +- Can administrators see in-flight and queued trends? +- Can operators trace queue-full and queue-timeout events? +- Are DB changes compatible with SQLite, MySQL, and PostgreSQL? + diff --git a/tools/channel-flow-spike/README.md b/tools/channel-flow-spike/README.md new file mode 100644 index 000000000000..7d8cb3eb21c1 --- /dev/null +++ b/tools/channel-flow-spike/README.md @@ -0,0 +1,32 @@ +# Channel Flow Redis Phase 0 Spike + +This spike validates the Redis backend shape for channel-level flow control +before productionizing it in `service/channel_flow.go`. + +It intentionally does not use Lua. The experiment uses: + +- one `running` ZSET, scored by lease expiry timestamp; +- one `waiting` ZSET, scored by Redis `INCR` sequence; +- `WATCH` / `MULTI` on `running` and `waiting`; +- release as `ZREM running ` only; +- waiter self-promotion by polling and promoting itself when it is queue head. + +Run example: + +```bash +go run ./tools/channel-flow-spike \ + -redis redis://localhost:6379/0 \ + -concurrency 1000 \ + -max-inflight 60 \ + -max-queue 240 \ + -queue-timeout 10s +``` + +The output is a JSON summary with conflict rate, p50/p95/p99 acquire latency, +peak running/queued counts, and the `max_inflight` invariant result. + +Production Redis backend work should only proceed if this spike stays within the +target SLO for the expected deployment concurrency. If `tx_conflicts` or p99 are +too high, benchmark a Lua version or redesign the queue before wiring Redis into +the live relay path. + diff --git a/tools/channel-flow-spike/main.go b/tools/channel-flow-spike/main.go new file mode 100644 index 000000000000..bf0a0e841699 --- /dev/null +++ b/tools/channel-flow-spike/main.go @@ -0,0 +1,522 @@ +package main + +import ( + "context" + "errors" + "flag" + "fmt" + "math/rand" + "os" + "sort" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/go-redis/redis/v8" +) + +type spikeConfig struct { + RedisURL string + PoolKey string + Concurrency int + MaxInflight int + MaxQueueSize int + QueueTimeout time.Duration + LeaseTTL time.Duration + HoldTime time.Duration + HoldJitter time.Duration + PollMin time.Duration + PollMax time.Duration + SampleInterval time.Duration + Cleanup bool +} + +type redisKeys struct { + Running string `json:"running"` + Waiting string `json:"waiting"` + Seq string `json:"seq"` +} + +type acquireDecision struct { + Admitted bool + Queued bool + Rejected string + Waited time.Duration + QueueScore float64 +} + +type redisFlowProbe struct { + rdb *redis.Client + keys redisKeys + cfg spikeConfig +} + +type spikeMetrics struct { + WatchAttempts int64 + TxConflicts int64 + Admitted int64 + Immediate int64 + Queued int64 + QueueFull int64 + QueueTimeout int64 + Errors int64 + PeakRunning int64 + PeakQueued int64 + + mu sync.Mutex + latencies []int64 +} + +type spikeSummary struct { + Config map[string]any `json:"config"` + Keys redisKeys `json:"keys"` + Totals map[string]any `json:"totals"` + LatencyMs map[string]int64 `json:"latency_ms"` + InvariantOK bool `json:"invariant_ok"` + DurationMs int64 `json:"duration_ms"` + Recommendation string `json:"recommendation"` +} + +const ( + rejectQueueFull = "queue_full" + rejectQueueTimeout = "queue_timeout" +) + +func main() { + cfg := parseFlags() + if cfg.RedisURL == "" { + fmt.Fprintln(os.Stderr, "missing -redis or REDIS_CONN_STRING") + os.Exit(2) + } + if cfg.Concurrency <= 0 || cfg.MaxInflight <= 0 { + fmt.Fprintln(os.Stderr, "-concurrency and -max-inflight must be positive") + os.Exit(2) + } + + opt, err := redis.ParseURL(cfg.RedisURL) + if err != nil { + fmt.Fprintf(os.Stderr, "parse redis url: %v\n", err) + os.Exit(2) + } + opt.PoolSize = max(cfg.Concurrency/4, 10) + rdb := redis.NewClient(opt) + defer rdb.Close() + + ctx, cancel := context.WithTimeout(context.Background(), cfg.QueueTimeout+cfg.HoldTime+30*time.Second) + defer cancel() + if err := rdb.Ping(ctx).Err(); err != nil { + fmt.Fprintf(os.Stderr, "redis ping failed: %v\n", err) + os.Exit(2) + } + + keys := redisKeys{ + Running: fmt.Sprintf("new-api:flow-spike:%s:running", cfg.PoolKey), + Waiting: fmt.Sprintf("new-api:flow-spike:%s:waiting", cfg.PoolKey), + Seq: fmt.Sprintf("new-api:flow-spike:%s:seq", cfg.PoolKey), + } + probe := &redisFlowProbe{rdb: rdb, keys: keys, cfg: cfg} + _ = probe.cleanup(ctx) + defer func() { + if cfg.Cleanup { + _ = probe.cleanup(context.Background()) + } + }() + + metrics := &spikeMetrics{ + latencies: make([]int64, 0, cfg.Concurrency), + } + startedAt := time.Now() + stopSampling := make(chan struct{}) + var samplerDone sync.WaitGroup + samplerDone.Add(1) + go samplePeaks(ctx, &samplerDone, probe, metrics, stopSampling) + + var workers sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < cfg.Concurrency; i++ { + workers.Add(1) + go func(index int) { + defer workers.Done() + <-start + runWorker(ctx, probe, metrics, index) + }(i) + } + close(start) + workers.Wait() + close(stopSampling) + samplerDone.Wait() + + summary := buildSummary(cfg, keys, metrics, time.Since(startedAt)) + data, err := common.Marshal(summary) + if err != nil { + fmt.Fprintf(os.Stderr, "marshal summary: %v\n", err) + os.Exit(1) + } + fmt.Println(string(data)) + if !summary.InvariantOK { + os.Exit(1) + } +} + +func parseFlags() spikeConfig { + defaultRedisURL := os.Getenv("REDIS_CONN_STRING") + defaultPoolKey := fmt.Sprintf("pool-%d", time.Now().Unix()) + cfg := spikeConfig{} + flag.StringVar(&cfg.RedisURL, "redis", defaultRedisURL, "Redis URL, defaults to REDIS_CONN_STRING") + flag.StringVar(&cfg.PoolKey, "pool-key", defaultPoolKey, "temporary Redis key suffix for this run") + flag.IntVar(&cfg.Concurrency, "concurrency", 1000, "number of simultaneous acquire attempts") + flag.IntVar(&cfg.MaxInflight, "max-inflight", 60, "running lease cap") + flag.IntVar(&cfg.MaxQueueSize, "max-queue", 240, "waiting queue cap") + flag.DurationVar(&cfg.QueueTimeout, "queue-timeout", 10*time.Second, "per-request queue timeout") + flag.DurationVar(&cfg.LeaseTTL, "lease", 30*time.Second, "running lease TTL") + flag.DurationVar(&cfg.HoldTime, "hold", 250*time.Millisecond, "simulated upstream processing time") + flag.DurationVar(&cfg.HoldJitter, "hold-jitter", 150*time.Millisecond, "additional random processing time") + flag.DurationVar(&cfg.PollMin, "poll-min", 5*time.Millisecond, "minimum waiter self-promote poll interval") + flag.DurationVar(&cfg.PollMax, "poll-max", 25*time.Millisecond, "maximum waiter self-promote poll interval") + flag.DurationVar(&cfg.SampleInterval, "sample-interval", 10*time.Millisecond, "peak sampler interval") + flag.BoolVar(&cfg.Cleanup, "cleanup", true, "delete temporary Redis keys after the run") + flag.Parse() + if cfg.QueueTimeout <= 0 { + cfg.QueueTimeout = 10 * time.Second + } + if cfg.LeaseTTL <= 0 { + cfg.LeaseTTL = 30 * time.Second + } + if cfg.PollMin <= 0 { + cfg.PollMin = 5 * time.Millisecond + } + if cfg.PollMax < cfg.PollMin { + cfg.PollMax = cfg.PollMin + } + if cfg.SampleInterval <= 0 { + cfg.SampleInterval = 10 * time.Millisecond + } + return cfg +} + +func runWorker(ctx context.Context, probe *redisFlowProbe, metrics *spikeMetrics, index int) { + requestID := fmt.Sprintf("req-%d-%d", time.Now().UnixNano(), index) + acquireCtx, cancel := context.WithTimeout(ctx, probe.cfg.QueueTimeout) + defer cancel() + startedAt := time.Now() + decision, err := probe.acquire(acquireCtx, requestID, metrics) + latencyMs := time.Since(startedAt).Milliseconds() + metrics.recordLatency(latencyMs) + if err != nil { + switch decision.Rejected { + case rejectQueueFull: + atomic.AddInt64(&metrics.QueueFull, 1) + case rejectQueueTimeout: + atomic.AddInt64(&metrics.QueueTimeout, 1) + default: + atomic.AddInt64(&metrics.Errors, 1) + } + return + } + if !decision.Admitted { + atomic.AddInt64(&metrics.Errors, 1) + return + } + atomic.AddInt64(&metrics.Admitted, 1) + if decision.Queued { + atomic.AddInt64(&metrics.Queued, 1) + } else { + atomic.AddInt64(&metrics.Immediate, 1) + } + hold := probe.cfg.HoldTime + randomDuration(probe.cfg.HoldJitter) + time.Sleep(hold) + if err := probe.release(context.Background(), requestID); err != nil { + atomic.AddInt64(&metrics.Errors, 1) + } +} + +func (p *redisFlowProbe) acquire(ctx context.Context, requestID string, metrics *spikeMetrics) (acquireDecision, error) { + enqueued := false + queuedAt := time.Time{} + sequenceScore := float64(0) + + for { + if err := ctx.Err(); err != nil { + if enqueued { + _ = p.rdb.ZRem(context.Background(), p.keys.Waiting, requestID).Err() + } + return acquireDecision{Rejected: rejectQueueTimeout}, err + } + _ = p.cleanupExpiredRunning(ctx) + + decision, done, err := p.tryAcquireOnce(ctx, requestID, enqueued, sequenceScore, queuedAt, metrics) + if err == nil && done { + if decision.Rejected != "" { + return decision, errors.New(decision.Rejected) + } + return decision, nil + } + if err != nil && !errors.Is(err, redis.TxFailedErr) { + return decision, err + } + if errors.Is(err, redis.TxFailedErr) { + atomic.AddInt64(&metrics.TxConflicts, 1) + } + if !enqueued && decision.Queued { + enqueued = true + queuedAt = time.Now() + sequenceScore = decision.QueueScore + } + if decision.Rejected == rejectQueueFull { + return decision, fmt.Errorf("queue full") + } + time.Sleep(p.pollDelay()) + } +} + +func (p *redisFlowProbe) tryAcquireOnce( + ctx context.Context, + requestID string, + enqueued bool, + sequenceScore float64, + queuedAt time.Time, + metrics *spikeMetrics, +) (acquireDecision, bool, error) { + atomic.AddInt64(&metrics.WatchAttempts, 1) + decision := acquireDecision{} + err := p.rdb.Watch(ctx, func(tx *redis.Tx) error { + running, err := tx.ZCard(ctx, p.keys.Running).Result() + if err != nil { + return err + } + waiting, err := tx.ZCard(ctx, p.keys.Waiting).Result() + if err != nil { + return err + } + if !enqueued { + if running < int64(p.cfg.MaxInflight) && waiting == 0 { + _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { + pipe.ZAdd(ctx, p.keys.Running, &redis.Z{ + Score: float64(time.Now().Add(p.cfg.LeaseTTL).UnixMilli()), + Member: requestID, + }) + return nil + }) + if err == nil { + decision = acquireDecision{Admitted: true} + } + return err + } + if p.cfg.MaxQueueSize > 0 && waiting >= int64(p.cfg.MaxQueueSize) { + decision = acquireDecision{Rejected: rejectQueueFull} + return nil + } + score, scoreErr := p.nextSequence(ctx, sequenceScore) + if scoreErr != nil { + return scoreErr + } + _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { + pipe.ZAdd(ctx, p.keys.Waiting, &redis.Z{ + Score: score, + Member: requestID, + }) + return nil + }) + if err == nil { + decision = acquireDecision{Queued: true, QueueScore: score} + } + return err + } + + rank, err := tx.ZRank(ctx, p.keys.Waiting, requestID).Result() + if errors.Is(err, redis.Nil) { + return nil + } + if err != nil { + return err + } + if rank != 0 || running >= int64(p.cfg.MaxInflight) { + return nil + } + _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { + pipe.ZRem(ctx, p.keys.Waiting, requestID) + pipe.ZAdd(ctx, p.keys.Running, &redis.Z{ + Score: float64(time.Now().Add(p.cfg.LeaseTTL).UnixMilli()), + Member: requestID, + }) + return nil + }) + if err == nil { + decision = acquireDecision{ + Admitted: true, + Queued: true, + Waited: time.Since(queuedAt), + } + } + return err + }, p.keys.Running, p.keys.Waiting) + if err != nil { + return decision, false, err + } + if decision.Admitted || decision.Rejected != "" { + return decision, true, nil + } + return decision, false, nil +} + +func (p *redisFlowProbe) nextSequence(ctx context.Context, existing float64) (float64, error) { + if existing > 0 { + return existing, nil + } + seq, err := p.rdb.Incr(ctx, p.keys.Seq).Result() + return float64(seq), err +} + +func (p *redisFlowProbe) release(ctx context.Context, requestID string) error { + return p.rdb.ZRem(ctx, p.keys.Running, requestID).Err() +} + +func (p *redisFlowProbe) cleanupExpiredRunning(ctx context.Context) error { + return p.rdb.ZRemRangeByScore(ctx, p.keys.Running, "-inf", fmt.Sprintf("%d", time.Now().UnixMilli())).Err() +} + +func (p *redisFlowProbe) cleanup(ctx context.Context) error { + return p.rdb.Del(ctx, p.keys.Running, p.keys.Waiting, p.keys.Seq).Err() +} + +func (p *redisFlowProbe) pollDelay() time.Duration { + window := p.cfg.PollMax - p.cfg.PollMin + if window <= 0 { + return p.cfg.PollMin + } + return p.cfg.PollMin + time.Duration(rand.Int63n(int64(window))) +} + +func randomDuration(maxDuration time.Duration) time.Duration { + if maxDuration <= 0 { + return 0 + } + return time.Duration(rand.Int63n(int64(maxDuration))) +} + +func samplePeaks(ctx context.Context, wg *sync.WaitGroup, probe *redisFlowProbe, metrics *spikeMetrics, stop <-chan struct{}) { + defer wg.Done() + ticker := time.NewTicker(probe.cfg.SampleInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-stop: + return + case <-ticker.C: + running, err := probe.rdb.ZCard(ctx, probe.keys.Running).Result() + if err == nil { + updatePeak(&metrics.PeakRunning, running) + } + queued, err := probe.rdb.ZCard(ctx, probe.keys.Waiting).Result() + if err == nil { + updatePeak(&metrics.PeakQueued, queued) + } + } + } +} + +func updatePeak(target *int64, value int64) { + for { + current := atomic.LoadInt64(target) + if value <= current { + return + } + if atomic.CompareAndSwapInt64(target, current, value) { + return + } + } +} + +func (m *spikeMetrics) recordLatency(latencyMs int64) { + m.mu.Lock() + defer m.mu.Unlock() + m.latencies = append(m.latencies, latencyMs) +} + +func buildSummary(cfg spikeConfig, keys redisKeys, metrics *spikeMetrics, duration time.Duration) spikeSummary { + latency := latencySummary(metrics.latencies) + watchAttempts := atomic.LoadInt64(&metrics.WatchAttempts) + conflicts := atomic.LoadInt64(&metrics.TxConflicts) + peakRunning := atomic.LoadInt64(&metrics.PeakRunning) + invariantOK := cfg.MaxInflight <= 0 || peakRunning <= int64(cfg.MaxInflight) + recommendation := "WATCH/MULTI shape is acceptable for the next backend prototype if conflict rate and p99 are within SLO." + if watchAttempts > 0 && float64(conflicts)/float64(watchAttempts) > 0.2 { + recommendation = "Conflict rate is high; benchmark Lua or redesign before productionizing Redis backend." + } + if !invariantOK { + recommendation = "Invariant failed; do not use this Redis algorithm without fixing over-admission." + } + return spikeSummary{ + Config: map[string]any{ + "concurrency": cfg.Concurrency, + "max_inflight": cfg.MaxInflight, + "max_queue_size": cfg.MaxQueueSize, + "queue_timeout": cfg.QueueTimeout.String(), + "lease": cfg.LeaseTTL.String(), + "hold": cfg.HoldTime.String(), + "hold_jitter": cfg.HoldJitter.String(), + "poll_min": cfg.PollMin.String(), + "poll_max": cfg.PollMax.String(), + "sample_interval": cfg.SampleInterval.String(), + }, + Keys: keys, + Totals: map[string]any{ + "watch_attempts": watchAttempts, + "tx_conflicts": conflicts, + "conflict_rate": ratio(conflicts, watchAttempts), + "admitted": atomic.LoadInt64(&metrics.Admitted), + "immediate": atomic.LoadInt64(&metrics.Immediate), + "queued": atomic.LoadInt64(&metrics.Queued), + "queue_full": atomic.LoadInt64(&metrics.QueueFull), + "queue_timeout": atomic.LoadInt64(&metrics.QueueTimeout), + "errors": atomic.LoadInt64(&metrics.Errors), + "peak_running": peakRunning, + "peak_queued": atomic.LoadInt64(&metrics.PeakQueued), + }, + LatencyMs: latency, + InvariantOK: invariantOK, + DurationMs: duration.Milliseconds(), + Recommendation: recommendation, + } +} + +func latencySummary(values []int64) map[string]int64 { + if len(values) == 0 { + return map[string]int64{"p50": 0, "p95": 0, "p99": 0, "max": 0} + } + sortedValues := append([]int64(nil), values...) + sort.Slice(sortedValues, func(i int, j int) bool { + return sortedValues[i] < sortedValues[j] + }) + return map[string]int64{ + "p50": percentile(sortedValues, 0.50), + "p95": percentile(sortedValues, 0.95), + "p99": percentile(sortedValues, 0.99), + "max": sortedValues[len(sortedValues)-1], + } +} + +func percentile(sortedValues []int64, percentileValue float64) int64 { + if len(sortedValues) == 0 { + return 0 + } + index := int(float64(len(sortedValues)-1) * percentileValue) + return sortedValues[index] +} + +func ratio(numerator int64, denominator int64) float64 { + if denominator == 0 { + return 0 + } + return float64(numerator) / float64(denominator) +} + +func max(left int, right int) int { + if left > right { + return left + } + return right +} From c7bf8321cbfabaaee5859e9433351dcb402b8431 Mon Sep 17 00:00:00 2001 From: supreme0597 Date: Sat, 13 Jun 2026 17:03:15 +0800 Subject: [PATCH 02/18] feat: add channel flow control backend --- controller/channel_flow.go | 290 ++++++++++++ controller/relay.go | 40 +- model/channel_flow.go | 281 ++++++++++++ model/main.go | 8 + router/api-router.go | 15 + service/billing.go | 61 +++ service/channel_flow.go | 801 ++++++++++++++++++++++++++++++++++ service/channel_flow_redis.go | 650 +++++++++++++++++++++++++++ service/channel_flow_test.go | 500 +++++++++++++++++++++ types/error.go | 10 + 10 files changed, 2654 insertions(+), 2 deletions(-) create mode 100644 controller/channel_flow.go create mode 100644 model/channel_flow.go create mode 100644 service/channel_flow.go create mode 100644 service/channel_flow_redis.go create mode 100644 service/channel_flow_test.go diff --git a/controller/channel_flow.go b/controller/channel_flow.go new file mode 100644 index 000000000000..26497d0a96c6 --- /dev/null +++ b/controller/channel_flow.go @@ -0,0 +1,290 @@ +package controller + +import ( + "errors" + "fmt" + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +type channelFlowPoolRequest struct { + Id int `json:"id"` + Name string `json:"name"` + Description string `json:"description"` + Enabled *bool `json:"enabled"` + Backend string `json:"backend"` + MaxInflight int `json:"max_inflight"` + MaxQueueSize int `json:"max_queue_size"` + MaxQueuePerUser int `json:"max_queue_per_user"` + QueueTimeoutMs int64 `json:"queue_timeout_ms"` + QueuePolicy string `json:"queue_policy"` + OnLimit string `json:"on_limit"` + RedisFailurePolicy string `json:"redis_failure_policy"` + MaxContextTokens int `json:"max_context_tokens"` + MaxContextChars int `json:"max_context_chars"` + MaxProcessingMs int64 `json:"max_processing_ms"` + LeaseMs int64 `json:"lease_ms"` + RenewIntervalMs int64 `json:"renew_interval_ms"` +} + +type channelFlowBindingRequest struct { + ChannelId int `json:"channel_id"` + UpstreamModel string `json:"upstream_model"` + MatchMode string `json:"match_mode"` + Enabled *bool `json:"enabled"` +} + +func ListChannelFlowPools(c *gin.Context) { + pageInfo := common.GetPageQuery(c) + var pools []*model.ChannelFlowPool + query := model.DB.Model(&model.ChannelFlowPool{}) + if keyword := c.Query("keyword"); keyword != "" { + query = query.Where("name LIKE ? OR pool_key LIKE ?", "%"+keyword+"%", "%"+keyword+"%") + } + var total int64 + if err := query.Count(&total).Error; err != nil { + common.ApiError(c, err) + return + } + if err := query.Order("id DESC").Limit(pageInfo.GetPageSize()).Offset(pageInfo.GetStartIdx()).Find(&pools).Error; err != nil { + common.ApiError(c, err) + return + } + pageInfo.SetTotal(int(total)) + pageInfo.SetItems(pools) + common.ApiSuccess(c, pageInfo) +} + +func GetChannelFlowPool(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + pool, err := model.GetChannelFlowPoolByID(id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, pool) +} + +func CreateChannelFlowPool(c *gin.Context) { + var req channelFlowPoolRequest + if err := c.ShouldBindJSON(&req); err != nil { + common.ApiError(c, err) + return + } + pool := channelFlowPoolFromRequest(req, nil) + if err := pool.Validate(); err != nil { + common.ApiError(c, err) + return + } + if err := model.DB.Create(pool).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, pool) +} + +func UpdateChannelFlowPool(c *gin.Context) { + var req channelFlowPoolRequest + if err := c.ShouldBindJSON(&req); err != nil { + common.ApiError(c, err) + return + } + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + if req.Id <= 0 { + req.Id = id + } + if req.Id != id { + common.ApiErrorMsg(c, "Flow Pool ID 与 URL 不一致") + return + } + if id <= 0 { + common.ApiErrorMsg(c, "缺少 Flow Pool ID") + return + } + pool, err := model.GetChannelFlowPoolByID(id) + if err != nil { + common.ApiError(c, err) + return + } + updated := channelFlowPoolFromRequest(req, pool) + if err := updated.Validate(); err != nil { + common.ApiError(c, err) + return + } + if err := model.DB.Save(updated).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, updated) +} + +func DeleteChannelFlowPool(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + count, err := model.CountChannelFlowPoolBindings(id) + if err != nil { + common.ApiError(c, err) + return + } + if count > 0 { + common.ApiError(c, fmt.Errorf("Flow Pool 仍有绑定渠道，请先删除绑定")) + return + } + if err := model.DB.Delete(&model.ChannelFlowPool{}, id).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, nil) +} + +func GetChannelFlowPoolStatus(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + pool, err := model.GetChannelFlowPoolByID(id) + if err != nil { + common.ApiError(c, err) + return + } + status, err := service.GetChannelFlowPoolStatus(c.Request.Context(), *pool) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, status) +} + +func ListChannelFlowPoolBindings(c *gin.Context) { + poolID, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + var bindings []*model.ChannelFlowPoolBinding + if err := model.DB.Where("pool_id = ?", poolID).Order("id DESC").Find(&bindings).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, bindings) +} + +func CreateChannelFlowPoolBinding(c *gin.Context) { + poolID, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + if _, err := model.GetChannelFlowPoolByID(poolID); err != nil { + common.ApiError(c, err) + return + } + var req channelFlowBindingRequest + if err := c.ShouldBindJSON(&req); err != nil { + common.ApiError(c, err) + return + } + enabled := true + if req.Enabled != nil { + enabled = *req.Enabled + } + binding := &model.ChannelFlowPoolBinding{ + PoolId: poolID, + ChannelId: req.ChannelId, + UpstreamModel: req.UpstreamModel, + MatchMode: req.MatchMode, + Enabled: enabled, + } + if err := binding.Validate(); err != nil { + common.ApiError(c, err) + return + } + if _, err := model.GetChannelById(binding.ChannelId, false); err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + common.ApiErrorMsg(c, "渠道不存在，无法绑定 Flow Pool") + return + } + common.ApiError(c, err) + return + } + if binding.MatchMode != model.ChannelFlowMatchModeChannel { + common.ApiErrorMsg(c, "Phase 1 仅支持按渠道绑定，upstream_model 绑定将在后续阶段开放") + return + } + var existing int64 + if err := model.DB.Model(&model.ChannelFlowPoolBinding{}). + Where("channel_id = ? AND match_mode = ? AND enabled = ?", binding.ChannelId, model.ChannelFlowMatchModeChannel, true). + Count(&existing).Error; err != nil { + common.ApiError(c, err) + return + } + if existing > 0 { + common.ApiErrorMsg(c, "该渠道已绑定 Flow Pool，请先删除原绑定") + return + } + if err := model.DB.Create(binding).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, binding) +} + +func DeleteChannelFlowPoolBinding(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + if err := model.DB.Delete(&model.ChannelFlowPoolBinding{}, id).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, nil) +} + +func channelFlowPoolFromRequest(req channelFlowPoolRequest, existing *model.ChannelFlowPool) *model.ChannelFlowPool { + pool := &model.ChannelFlowPool{} + if existing != nil { + *pool = *existing + } + pool.Name = req.Name + pool.Description = req.Description + if req.Enabled != nil { + pool.Enabled = *req.Enabled + } else if existing == nil { + pool.Enabled = true + } + pool.Backend = req.Backend + pool.MaxInflight = req.MaxInflight + pool.MaxQueueSize = req.MaxQueueSize + pool.MaxQueuePerUser = req.MaxQueuePerUser + pool.QueueTimeoutMs = req.QueueTimeoutMs + pool.QueuePolicy = req.QueuePolicy + pool.OnLimit = req.OnLimit + pool.RedisFailurePolicy = req.RedisFailurePolicy + pool.MaxContextTokens = req.MaxContextTokens + pool.MaxContextChars = req.MaxContextChars + pool.MaxProcessingMs = req.MaxProcessingMs + pool.LeaseMs = req.LeaseMs + pool.RenewIntervalMs = req.RenewIntervalMs + pool.Normalize() + return pool +} diff --git a/controller/relay.go b/controller/relay.go index 1d14dcc6f880..3fa786044ca7 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -1,6 +1,7 @@ package controller import ( + "context" "errors" "fmt" "io" @@ -159,9 +160,9 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { // common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta) if priceData.FreeModel { - logger.LogInfo(c, fmt.Sprintf("模型 %s 免费，跳过预扣费", relayInfo.OriginModelName)) + logger.LogInfo(c, fmt.Sprintf("模型 %s 免费，跳过计费预检查", relayInfo.OriginModelName)) } else { - newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo) + newAPIError = service.PrecheckBilling(c, priceData.QuotaToPreConsume, relayInfo) if newAPIError != nil { return } @@ -197,8 +198,39 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } addUsedChannel(c, channel.Id) + flowGuard, _, flowErr := service.AcquireChannelFlowGuard(c, channel.Id, relayInfo) + if flowErr != nil { + newAPIError = flowErr + break + } + + attemptPriceData, priceErr := helper.ModelPriceHelper(c, relayInfo, tokens, meta) + if priceErr != nil { + if flowGuard != nil { + _ = flowGuard.Release(context.Background()) + } + newAPIError = types.NewError(priceErr, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest)) + break + } + if !attemptPriceData.FreeModel { + if relayInfo.Billing == nil { + newAPIError = service.PreConsumeBilling(c, attemptPriceData.QuotaToPreConsume, relayInfo) + } else if reserveErr := relayInfo.Billing.Reserve(attemptPriceData.QuotaToPreConsume); reserveErr != nil { + newAPIError = types.NewErrorWithStatusCode(reserveErr, types.ErrorCodeChannelFlowBillingFailedAfterWait, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) + } + if newAPIError != nil { + if flowGuard != nil { + _ = flowGuard.Release(context.Background()) + } + break + } + } + bodyStorage, bodyErr := common.GetBodyStorage(c) if bodyErr != nil { + if flowGuard != nil { + _ = flowGuard.Release(context.Background()) + } // Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path) if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) { newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry()) @@ -220,6 +252,10 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = relayHandler(c, relayInfo) } + if flowGuard != nil { + _ = flowGuard.Release(context.Background()) + } + if newAPIError == nil { relayInfo.LastError = nil return diff --git a/model/channel_flow.go b/model/channel_flow.go new file mode 100644 index 000000000000..0d1a68a5662f --- /dev/null +++ b/model/channel_flow.go @@ -0,0 +1,281 @@ +package model + +import ( + "fmt" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" +) + +const ( + ChannelFlowBackendMemory = "memory" + ChannelFlowBackendRedis = "redis" + + ChannelFlowQueuePolicyFIFO = "fifo" + + ChannelFlowOnLimitQueue = "queue" + ChannelFlowOnLimitReject = "reject" + ChannelFlowOnLimitFallback = "fallback" + + ChannelFlowRedisFailureFailOpen = "fail_open" + ChannelFlowRedisFailureFailClosed = "fail_closed" + ChannelFlowRedisFailureLocalMemory = "local_memory" + + ChannelFlowMatchModeChannel = "channel" + ChannelFlowMatchModeChannelModel = "channel_model" +) + +type ChannelFlowPool struct { + Id int `json:"id"` + PoolKey string `json:"pool_key" gorm:"type:varchar(64);uniqueIndex"` + Name string `json:"name" gorm:"type:varchar(128);index"` + Description string `json:"description" gorm:"type:text"` + Enabled bool `json:"enabled" gorm:"default:true"` + Backend string `json:"backend" gorm:"type:varchar(32);default:'memory'"` + MaxInflight int `json:"max_inflight" gorm:"default:0"` + MaxQueueSize int `json:"max_queue_size" gorm:"default:0"` + MaxQueuePerUser int `json:"max_queue_per_user" gorm:"default:0"` + QueueTimeoutMs int64 `json:"queue_timeout_ms" gorm:"bigint;default:120000"` + QueuePolicy string `json:"queue_policy" gorm:"type:varchar(32);default:'fifo'"` + OnLimit string `json:"on_limit" gorm:"type:varchar(32);default:'queue'"` + RedisFailurePolicy string `json:"redis_failure_policy" gorm:"type:varchar(32);default:'fail_open'"` + MaxContextTokens int `json:"max_context_tokens" gorm:"default:0"` + MaxContextChars int `json:"max_context_chars" gorm:"default:0"` + MaxProcessingMs int64 `json:"max_processing_ms" gorm:"bigint;default:0"` + LeaseMs int64 `json:"lease_ms" gorm:"bigint;default:60000"` + RenewIntervalMs int64 `json:"renew_interval_ms" gorm:"bigint;default:20000"` + ConfigVersion int64 `json:"config_version" gorm:"bigint;default:1"` + CreatedTime int64 `json:"created_time" gorm:"bigint"` + UpdatedTime int64 `json:"updated_time" gorm:"bigint"` +} + +type ChannelFlowPoolBinding struct { + Id int `json:"id"` + PoolId int `json:"pool_id" gorm:"index"` + ChannelId int `json:"channel_id" gorm:"index"` + UpstreamModel string `json:"upstream_model" gorm:"type:varchar(191);default:''"` + MatchMode string `json:"match_mode" gorm:"type:varchar(32);default:'channel'"` + Enabled bool `json:"enabled" gorm:"default:true"` + CreatedTime int64 `json:"created_time" gorm:"bigint"` + UpdatedTime int64 `json:"updated_time" gorm:"bigint"` +} + +type ChannelFlowMetricMinute struct { + Id int `json:"id"` + BucketTs int64 `json:"bucket_ts" gorm:"bigint;index"` + PoolKey string `json:"pool_key" gorm:"type:varchar(64);index"` + ChannelId int `json:"channel_id" gorm:"index"` + Model string `json:"model" gorm:"type:varchar(191);index"` + RunningAvg float64 `json:"running_avg"` + RunningMax int `json:"running_max"` + QueuedAvg float64 `json:"queued_avg"` + QueuedMax int `json:"queued_max"` + AcquiredCount int `json:"acquired_count"` + QueuedCount int `json:"queued_count"` + ReleasedCount int `json:"released_count"` + RejectedCount int `json:"rejected_count"` + TimeoutCount int `json:"timeout_count"` + CancelledCount int `json:"cancelled_count"` + BillingFailedCount int `json:"billing_failed_count"` + LeaseRenewFail int `json:"lease_renew_fail"` + LeaseExpiredCount int `json:"lease_expired_count"` + WaitMsAvg int64 `json:"wait_ms_avg" gorm:"bigint"` + WaitMsMax int64 `json:"wait_ms_max" gorm:"bigint"` + ProcessMsAvg int64 `json:"process_ms_avg" gorm:"bigint"` + ProcessMsMax int64 `json:"process_ms_max" gorm:"bigint"` + CreatedTime int64 `json:"created_time" gorm:"bigint"` + UpdatedTime int64 `json:"updated_time" gorm:"bigint"` +} + +type ChannelFlowEvent struct { + Id int `json:"id"` + RequestId string `json:"request_id" gorm:"type:varchar(64);index"` + PoolKey string `json:"pool_key" gorm:"type:varchar(64);index"` + ChannelId int `json:"channel_id" gorm:"index"` + Model string `json:"model" gorm:"type:varchar(191);index"` + UserId int `json:"user_id" gorm:"index"` + TokenId int `json:"token_id" gorm:"index"` + EventType string `json:"event_type" gorm:"type:varchar(64);index"` + Reason string `json:"reason" gorm:"type:text"` + Running int `json:"running"` + Queued int `json:"queued"` + QueuePos int `json:"queue_pos"` + WaitMs int64 `json:"wait_ms" gorm:"bigint"` + ProcessMs int64 `json:"process_ms" gorm:"bigint"` + Backend string `json:"backend" gorm:"type:varchar(32)"` + CreatedTime int64 `json:"created_time" gorm:"bigint;index"` +} + +func (p *ChannelFlowPool) Normalize() { + p.Name = strings.TrimSpace(p.Name) + p.Description = strings.TrimSpace(p.Description) + if p.Backend == "" { + p.Backend = ChannelFlowBackendMemory + } + if p.QueuePolicy == "" { + p.QueuePolicy = ChannelFlowQueuePolicyFIFO + } + if p.OnLimit == "" { + p.OnLimit = ChannelFlowOnLimitQueue + } + if p.RedisFailurePolicy == "" { + p.RedisFailurePolicy = ChannelFlowRedisFailureFailOpen + } + if p.QueueTimeoutMs <= 0 { + p.QueueTimeoutMs = 120000 + } + if p.LeaseMs <= 0 { + p.LeaseMs = 60000 + } + if p.RenewIntervalMs <= 0 { + p.RenewIntervalMs = 20000 + } + if p.MaxQueueSize <= 0 && p.MaxInflight > 0 { + p.MaxQueueSize = p.MaxInflight * 4 + } +} + +func (p *ChannelFlowPool) Validate() error { + p.Normalize() + if p.Name == "" { + return fmt.Errorf("flow pool name cannot be empty") + } + if p.MaxInflight < 0 || p.MaxQueueSize < 0 || p.MaxQueuePerUser < 0 { + return fmt.Errorf("flow pool limits cannot be negative") + } + if p.MaxInflight == 0 && p.MaxContextTokens == 0 && p.MaxContextChars == 0 { + return fmt.Errorf("max_inflight or context limit must be configured") + } + switch p.Backend { + case ChannelFlowBackendMemory, ChannelFlowBackendRedis: + default: + return fmt.Errorf("invalid flow pool backend: %s", p.Backend) + } + switch p.QueuePolicy { + case ChannelFlowQueuePolicyFIFO: + default: + return fmt.Errorf("invalid flow pool queue_policy: %s", p.QueuePolicy) + } + switch p.OnLimit { + case ChannelFlowOnLimitQueue, ChannelFlowOnLimitReject, ChannelFlowOnLimitFallback: + default: + return fmt.Errorf("invalid flow pool on_limit: %s", p.OnLimit) + } + switch p.RedisFailurePolicy { + case ChannelFlowRedisFailureFailOpen, ChannelFlowRedisFailureFailClosed, ChannelFlowRedisFailureLocalMemory: + default: + return fmt.Errorf("invalid flow pool redis_failure_policy: %s", p.RedisFailurePolicy) + } + return nil +} + +func (p *ChannelFlowPool) BeforeCreate(_ *gorm.DB) error { + now := time.Now().Unix() + if p.CreatedTime == 0 { + p.CreatedTime = now + } + if p.UpdatedTime == 0 { + p.UpdatedTime = now + } + if p.ConfigVersion == 0 { + p.ConfigVersion = 1 + } + if p.PoolKey == "" { + p.PoolKey = GenerateChannelFlowPoolKey() + } + return p.Validate() +} + +func (p *ChannelFlowPool) BeforeUpdate(_ *gorm.DB) error { + p.UpdatedTime = time.Now().Unix() + p.ConfigVersion++ + return p.Validate() +} + +func (b *ChannelFlowPoolBinding) Normalize() { + b.UpstreamModel = strings.TrimSpace(b.UpstreamModel) + if b.MatchMode == "" { + b.MatchMode = ChannelFlowMatchModeChannel + } + if b.MatchMode == ChannelFlowMatchModeChannel { + b.UpstreamModel = "" + } +} + +func (b *ChannelFlowPoolBinding) Validate() error { + b.Normalize() + if b.PoolId <= 0 { + return fmt.Errorf("pool_id is required") + } + if b.ChannelId <= 0 { + return fmt.Errorf("channel_id is required") + } + switch b.MatchMode { + case ChannelFlowMatchModeChannel: + case ChannelFlowMatchModeChannelModel: + if b.UpstreamModel == "" { + return fmt.Errorf("upstream_model is required for channel_model binding") + } + default: + return fmt.Errorf("invalid flow pool binding match_mode: %s", b.MatchMode) + } + return nil +} + +func (b *ChannelFlowPoolBinding) BeforeCreate(_ *gorm.DB) error { + now := time.Now().Unix() + if b.CreatedTime == 0 { + b.CreatedTime = now + } + if b.UpdatedTime == 0 { + b.UpdatedTime = now + } + return b.Validate() +} + +func (b *ChannelFlowPoolBinding) BeforeUpdate(_ *gorm.DB) error { + b.UpdatedTime = time.Now().Unix() + return b.Validate() +} + +func GenerateChannelFlowPoolKey() string { + return "flow_pool_" + strings.ToLower(common.GetRandomString(12)) +} + +func GetChannelFlowPoolByID(id int) (*ChannelFlowPool, error) { + var pool ChannelFlowPool + if err := DB.First(&pool, id).Error; err != nil { + return nil, err + } + return &pool, nil +} + +func GetChannelFlowPoolByKey(poolKey string) (*ChannelFlowPool, error) { + var pool ChannelFlowPool + if err := DB.Where("pool_key = ?", poolKey).First(&pool).Error; err != nil { + return nil, err + } + return &pool, nil +} + +func GetChannelFlowPoolBindingForChannel(channelID int) (*ChannelFlowPoolBinding, *ChannelFlowPool, error) { + var binding ChannelFlowPoolBinding + if err := DB.Where("channel_id = ? AND match_mode = ? AND enabled = ?", channelID, ChannelFlowMatchModeChannel, true). + Order("id ASC"). + First(&binding).Error; err != nil { + return nil, nil, err + } + pool, err := GetChannelFlowPoolByID(binding.PoolId) + if err != nil { + return nil, nil, err + } + return &binding, pool, nil +} + +func CountChannelFlowPoolBindings(poolID int) (int64, error) { + var count int64 + err := DB.Model(&ChannelFlowPoolBinding{}).Where("pool_id = ?", poolID).Count(&count).Error + return count, err +} diff --git a/model/main.go b/model/main.go index 6d9002462873..87e6ef0b9377 100644 --- a/model/main.go +++ b/model/main.go @@ -281,6 +281,10 @@ func migrateDB() error { &CustomOAuthProvider{}, &UserOAuthBinding{}, &PerfMetric{}, + &ChannelFlowPool{}, + &ChannelFlowPoolBinding{}, + &ChannelFlowMetricMinute{}, + &ChannelFlowEvent{}, ) if err != nil { return err @@ -330,6 +334,10 @@ func migrateDBFast() error { {&CustomOAuthProvider{}, "CustomOAuthProvider"}, {&UserOAuthBinding{}, "UserOAuthBinding"}, {&PerfMetric{}, "PerfMetric"}, + {&ChannelFlowPool{}, "ChannelFlowPool"}, + {&ChannelFlowPoolBinding{}, "ChannelFlowPoolBinding"}, + {&ChannelFlowMetricMinute{}, "ChannelFlowMetricMinute"}, + {&ChannelFlowEvent{}, "ChannelFlowEvent"}, } // 动态计算migration数量，确保errChan缓冲区足够大 errChan := make(chan error, len(migrations)) diff --git a/router/api-router.go b/router/api-router.go index baf7cda20152..d1c5dacf787b 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -228,6 +228,7 @@ func SetApiRouter(router *gin.Engine) { channelRoute := apiRouter.Group("/channel") channelRoute.Use(middleware.AdminAuth()) { + channelRoute.GET("", controller.GetAllChannels) channelRoute.GET("/", controller.GetAllChannels) channelRoute.GET("/search", controller.SearchChannels) channelRoute.GET("/models", controller.ChannelListModels) @@ -238,6 +239,7 @@ func SetApiRouter(router *gin.Engine) { channelRoute.GET("/test/:id", controller.TestChannel) channelRoute.GET("/update_balance", controller.UpdateAllChannelsBalance) channelRoute.GET("/update_balance/:id", controller.UpdateChannelBalance) + channelRoute.POST("", controller.AddChannel) channelRoute.POST("/", controller.AddChannel) channelRoute.PUT("/", controller.UpdateChannel) channelRoute.DELETE("/disabled", controller.DeleteDisabledChannel) @@ -264,6 +266,19 @@ func SetApiRouter(router *gin.Engine) { channelRoute.POST("/upstream_updates/detect", controller.DetectChannelUpstreamModelUpdates) channelRoute.POST("/upstream_updates/detect_all", controller.DetectAllChannelUpstreamModelUpdates) } + channelFlowRoute := apiRouter.Group("/channel_flow") + channelFlowRoute.Use(middleware.AdminAuth()) + { + channelFlowRoute.GET("/pools", controller.ListChannelFlowPools) + channelFlowRoute.POST("/pools", controller.CreateChannelFlowPool) + channelFlowRoute.GET("/pools/:id", controller.GetChannelFlowPool) + channelFlowRoute.PUT("/pools/:id", controller.UpdateChannelFlowPool) + channelFlowRoute.DELETE("/pools/:id", controller.DeleteChannelFlowPool) + channelFlowRoute.GET("/pools/:id/status", controller.GetChannelFlowPoolStatus) + channelFlowRoute.GET("/pools/:id/bindings", controller.ListChannelFlowPoolBindings) + channelFlowRoute.POST("/pools/:id/bindings", controller.CreateChannelFlowPoolBinding) + channelFlowRoute.DELETE("/bindings/:id", controller.DeleteChannelFlowPoolBinding) + } tokenRoute := apiRouter.Group("/token") tokenRoute.Use(middleware.UserAuth()) { diff --git a/service/billing.go b/service/billing.go index 81daeed82c29..1168a26a78b7 100644 --- a/service/billing.go +++ b/service/billing.go @@ -2,8 +2,11 @@ package service import ( "fmt" + "net/http" + "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" @@ -25,6 +28,64 @@ func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycom return nil } +// PrecheckBilling performs a read-only quota sanity check before a request enters +// a flow-control queue. It intentionally does not reserve or deduct quota. +func PrecheckBilling(c *gin.Context, estimatedQuota int, relayInfo *relaycommon.RelayInfo) *types.NewAPIError { + if relayInfo == nil || estimatedQuota <= 0 { + return nil + } + trustQuota := common.GetTrustQuota() + if !relayInfo.TokenUnlimited { + tokenQuota := c.GetInt("token_quota") + if tokenQuota <= trustQuota && tokenQuota < estimatedQuota { + return types.NewErrorWithStatusCode( + fmt.Errorf("令牌额度不足, 剩余额度: %s, 预计需要额度: %s", logger.FormatQuota(tokenQuota), logger.FormatQuota(estimatedQuota)), + types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, + types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) + } + } + + userQuota, err := model.GetUserQuota(relayInfo.UserId, false) + if err != nil { + return types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry()) + } + relayInfo.UserQuota = userQuota + walletOK := userQuota > 0 && userQuota >= estimatedQuota + hasSub, subErr := model.HasActiveUserSubscription(relayInfo.UserId) + if subErr != nil { + return types.NewError(subErr, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry()) + } + + switch common.NormalizeBillingPreference(relayInfo.UserSetting.BillingPreference) { + case "wallet_only": + if !walletOK { + return insufficientPrecheckError(userQuota, estimatedQuota) + } + case "subscription_only": + if !hasSub { + return insufficientPrecheckError(userQuota, estimatedQuota) + } + case "wallet_first": + if !walletOK && !hasSub { + return insufficientPrecheckError(userQuota, estimatedQuota) + } + case "subscription_first": + fallthrough + default: + if !hasSub && !walletOK { + return insufficientPrecheckError(userQuota, estimatedQuota) + } + } + return nil +} + +func insufficientPrecheckError(userQuota int, estimatedQuota int) *types.NewAPIError { + return types.NewErrorWithStatusCode( + fmt.Errorf("额度不足, 剩余额度: %s, 预计需要额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(estimatedQuota)), + types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, + types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) +} + // --------------------------------------------------------------------------- // SettleBilling — 后结算辅助函数 // --------------------------------------------------------------------------- diff --git a/service/channel_flow.go b/service/channel_flow.go new file mode 100644 index 000000000000..31180b9d01b6 --- /dev/null +++ b/service/channel_flow.go @@ -0,0 +1,801 @@ +package service + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +const ( + FlowDecisionRejectQueueFull = "queue_full" + FlowDecisionRejectQueueTimeout = "queue_timeout" + FlowDecisionRejectContextExceeded = "context_exceeded" + FlowDecisionRejectPerUserQueueFull = "per_user_queue_full" + FlowDecisionRejectBackendDisabled = "backend_disabled" +) + +type AcquireRequest struct { + RequestID string + Pool model.ChannelFlowPool + ChannelID int + UpstreamModel string + UserID int + TokenID int + ContextTokens int + ContextChars int + CreatedAtMs int64 + QueueTimeoutMs int64 +} + +type AcquireDecision struct { + Admitted bool `json:"admitted"` + Queued bool `json:"queued"` + QueuePos int `json:"queue_pos"` + WaitedMs int64 `json:"waited_ms"` + Temporary bool `json:"temporary"` + RejectCode string `json:"reject_code"` + RunningNow int `json:"running_now"` + QueuedNow int `json:"queued_now"` + RetryAfterS int `json:"retry_after_seconds"` + Backend string `json:"backend"` + PoolKey string `json:"pool_key"` + ConfigVersion int64 `json:"config_version"` +} + +type PoolStatus struct { + PoolKey string `json:"pool_key"` + Name string `json:"name"` + Backend string `json:"backend"` + Health string `json:"health"` + Running int `json:"running"` + MaxInflight int `json:"max_inflight"` + Queued int `json:"queued"` + MaxQueueSize int `json:"max_queue_size"` + OldestWaitMs int64 `json:"oldest_wait_ms"` + ConfigVersion int64 `json:"config_version"` + LeaseRenewFailures int `json:"lease_renew_failures"` +} + +type FlowBackend interface { + Acquire(ctx context.Context, req AcquireRequest) (FlowGuard, *AcquireDecision, error) + Status(ctx context.Context, pool model.ChannelFlowPool) (PoolStatus, error) + Close(ctx context.Context) error +} + +type FlowGuard interface { + Release(ctx context.Context) error + RenewLease(ctx context.Context) error + PoolKey() string + RequestID() string + IsReleased() bool + BindRelease(release func()) + WrapReadCloser(rc io.ReadCloser) io.ReadCloser +} + +type FlowController struct { + memoryBackend FlowBackend + redisBackend FlowBackend +} + +var defaultChannelFlowController = NewFlowController(NewMemoryFlowBackend(), NewRedisFlowBackend()) + +func NewFlowController(backends ...FlowBackend) *FlowController { + controller := &FlowController{} + if len(backends) > 0 { + controller.memoryBackend = backends[0] + } + if len(backends) > 1 { + controller.redisBackend = backends[1] + } + if controller.memoryBackend == nil { + controller.memoryBackend = NewMemoryFlowBackend() + } + if controller.redisBackend == nil { + controller.redisBackend = NewRedisFlowBackend() + } + return controller +} + +func GetChannelFlowController() *FlowController { + return defaultChannelFlowController +} + +func (fc *FlowController) Acquire(ctx context.Context, req AcquireRequest) (FlowGuard, *AcquireDecision, error) { + backend := fc.backendForPool(req.Pool) + if backend == nil { + return nil, nil, fmt.Errorf("channel flow backend is not initialized") + } + return backend.Acquire(ctx, req) +} + +func (fc *FlowController) Status(ctx context.Context, pool model.ChannelFlowPool) (PoolStatus, error) { + backend := fc.backendForPool(pool) + if backend == nil { + return PoolStatus{}, fmt.Errorf("channel flow backend is not initialized") + } + return backend.Status(ctx, pool) +} + +func (fc *FlowController) Close(ctx context.Context) error { + if fc == nil { + return nil + } + if fc.memoryBackend != nil { + if err := fc.memoryBackend.Close(ctx); err != nil { + return err + } + } + if fc.redisBackend != nil && fc.redisBackend != fc.memoryBackend { + return fc.redisBackend.Close(ctx) + } + return nil +} + +func (fc *FlowController) backendForPool(pool model.ChannelFlowPool) FlowBackend { + if fc == nil { + return nil + } + if pool.Backend == model.ChannelFlowBackendRedis { + return fc.redisBackend + } + return fc.memoryBackend +} + +func ResolveChannelFlowPool(channelID int) (*model.ChannelFlowPoolBinding, *model.ChannelFlowPool, bool, error) { + binding, pool, err := model.GetChannelFlowPoolBindingForChannel(channelID) + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil, false, nil + } + if err != nil { + return nil, nil, false, err + } + if pool == nil || !pool.Enabled { + return binding, pool, false, nil + } + return binding, pool, true, nil +} + +func AcquireChannelFlowGuard(c *gin.Context, channelID int, info *relaycommon.RelayInfo) (FlowGuard, *AcquireDecision, *types.NewAPIError) { + if c == nil || info == nil { + return nil, nil, nil + } + _, pool, ok, err := ResolveChannelFlowPool(channelID) + if err != nil { + return nil, nil, types.NewError(err, types.ErrorCodeChannelFlowConfigInvalid, types.ErrOptionWithSkipRetry()) + } + if !ok || pool == nil { + return nil, nil, nil + } + if passThrough, fallbackPool, apiErr := resolveRedisFlowUnavailable(c.Request.Context(), pool); apiErr != nil || passThrough { + return nil, nil, apiErr + } else if fallbackPool != nil { + pool = fallbackPool + } + upstreamModel := info.OriginModelName + if info.ChannelMeta != nil && info.UpstreamModelName != "" { + upstreamModel = info.UpstreamModelName + } + req := AcquireRequest{ + RequestID: c.GetString(common.RequestIdKey), + Pool: *pool, + ChannelID: channelID, + UpstreamModel: upstreamModel, + UserID: info.UserId, + TokenID: info.TokenId, + ContextTokens: info.GetEstimatePromptTokens(), + CreatedAtMs: time.Now().UnixMilli(), + QueueTimeoutMs: pool.QueueTimeoutMs, + } + if req.RequestID == "" { + req.RequestID = common.GetUUID() + } + guard, decision, acquireErr := GetChannelFlowController().Acquire(c.Request.Context(), req) + if acquireErr != nil { + if passThrough, fallbackPool, apiErr := handleRedisFlowAcquireError(c.Request.Context(), *pool, decision, acquireErr); apiErr != nil || passThrough { + return nil, decision, apiErr + } else if fallbackPool != nil { + req.Pool = *fallbackPool + guard, decision, acquireErr = GetChannelFlowController().Acquire(c.Request.Context(), req) + if acquireErr == nil { + return guard, decision, nil + } + } + return nil, decision, flowDecisionToAPIError(decision, acquireErr) + } + return guard, decision, nil +} + +func GetChannelFlowPoolStatus(ctx context.Context, pool model.ChannelFlowPool) (PoolStatus, error) { + if passThrough, fallbackPool, _ := resolveRedisFlowUnavailable(ctx, &pool); passThrough { + return degradedRedisFlowStatus(pool), nil + } else if fallbackPool != nil { + return GetChannelFlowController().Status(ctx, *fallbackPool) + } + status, err := GetChannelFlowController().Status(ctx, pool) + if err == nil { + return status, nil + } + if errors.Is(err, ErrRedisFlowBackendUnavailable) { + switch pool.RedisFailurePolicy { + case model.ChannelFlowRedisFailureLocalMemory: + return GetChannelFlowController().Status(ctx, localMemoryFallbackFlowPool(pool)) + default: + return degradedRedisFlowStatus(pool), nil + } + } + return status, err +} + +func localMemoryFallbackFlowPool(pool model.ChannelFlowPool) model.ChannelFlowPool { + pool.Backend = model.ChannelFlowBackendMemory + return pool +} + +func resolveRedisFlowUnavailable(ctx context.Context, pool *model.ChannelFlowPool) (passThrough bool, fallbackPool *model.ChannelFlowPool, apiErr *types.NewAPIError) { + if pool == nil || pool.Backend != model.ChannelFlowBackendRedis || IsRedisFlowBackendAvailable(ctx) { + return false, nil, nil + } + switch pool.RedisFailurePolicy { + case model.ChannelFlowRedisFailureFailClosed: + decision := newFlowDecision(*pool, false, false) + decision.RejectCode = FlowDecisionRejectBackendDisabled + decision.Temporary = true + return false, nil, flowDecisionToAPIError(decision, fmt.Errorf("channel flow redis backend is unavailable")) + case model.ChannelFlowRedisFailureLocalMemory: + fallback := localMemoryFallbackFlowPool(*pool) + return false, &fallback, nil + default: + return true, nil, nil + } +} + +func handleRedisFlowAcquireError(ctx context.Context, pool model.ChannelFlowPool, decision *AcquireDecision, acquireErr error) (passThrough bool, fallbackPool *model.ChannelFlowPool, apiErr *types.NewAPIError) { + if pool.Backend != model.ChannelFlowBackendRedis || !errors.Is(acquireErr, ErrRedisFlowBackendUnavailable) { + return false, nil, nil + } + switch pool.RedisFailurePolicy { + case model.ChannelFlowRedisFailureFailClosed: + if decision == nil { + decision = newFlowDecision(pool, false, false) + } + decision.RejectCode = FlowDecisionRejectBackendDisabled + return false, nil, flowDecisionToAPIError(decision, acquireErr) + case model.ChannelFlowRedisFailureLocalMemory: + fallback := localMemoryFallbackFlowPool(pool) + return false, &fallback, nil + default: + return true, nil, nil + } +} + +func degradedRedisFlowStatus(pool model.ChannelFlowPool) PoolStatus { + pool.Normalize() + return PoolStatus{ + PoolKey: pool.PoolKey, + Name: pool.Name, + Backend: pool.Backend, + Health: "degraded", + MaxInflight: pool.MaxInflight, + MaxQueueSize: pool.MaxQueueSize, + ConfigVersion: pool.ConfigVersion, + } +} + +func newFlowDecision(pool model.ChannelFlowPool, admitted bool, queued bool) *AcquireDecision { + return &AcquireDecision{ + Admitted: admitted, + Queued: queued, + Temporary: true, + RetryAfterS: retryAfterSeconds(pool.QueueTimeoutMs), + Backend: pool.Backend, + PoolKey: pool.PoolKey, + ConfigVersion: pool.ConfigVersion, + } +} + +func retryAfterSeconds(timeoutMs int64) int { + if timeoutMs <= 0 { + return 30 + } + seconds := int((timeoutMs + 999) / 1000) + if seconds < 1 { + return 1 + } + if seconds > 30 { + return 30 + } + return seconds +} + +func flowDecisionToAPIError(decision *AcquireDecision, err error) *types.NewAPIError { + if err == nil { + err = fmt.Errorf("channel flow control rejected request") + } + errorCode := types.ErrorCodeChannelFlowQueueFull + statusCode := http.StatusTooManyRequests + if decision != nil { + switch decision.RejectCode { + case FlowDecisionRejectQueueTimeout: + errorCode = types.ErrorCodeChannelFlowQueueTimeout + case FlowDecisionRejectContextExceeded: + errorCode = types.ErrorCodeChannelFlowContextExceeded + statusCode = http.StatusBadRequest + case FlowDecisionRejectPerUserQueueFull: + errorCode = types.ErrorCodeChannelFlowPerUserQueueFull + case FlowDecisionRejectBackendDisabled: + errorCode = types.ErrorCodeChannelFlowBackendUnavailable + statusCode = http.StatusServiceUnavailable + } + } + openAIError := types.OpenAIError{ + Message: err.Error(), + Type: "rate_limit_error", + Code: errorCode, + } + if decision != nil { + metadata, marshalErr := common.Marshal(map[string]any{ + "pool_running": decision.RunningNow, + "pool_queued": decision.QueuedNow, + "queue_pos": decision.QueuePos, + "waited_ms": decision.WaitedMs, + "reject_code": decision.RejectCode, + "retry_after_seconds": decision.RetryAfterS, + "channel_flow_backend": decision.Backend, + "channel_flow_pool_key": decision.PoolKey, + }) + if marshalErr == nil { + openAIError.Metadata = metadata + } + } + return types.WithOpenAIError(openAIError, statusCode, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) +} + +type memoryFlowBackend struct { + mu sync.RWMutex + slots map[string]*memoryFlowSlot +} + +type memoryFlowSlot struct { + mu sync.Mutex + config model.ChannelFlowPool + queue []*memoryFlowRequest + nextSeq int64 + eventCounts map[string]int +} + +type memoryFlowRequestState string + +const ( + memoryFlowStateWaiting memoryFlowRequestState = "waiting" + memoryFlowStateRunning memoryFlowRequestState = "running" + memoryFlowStateReleased memoryFlowRequestState = "released" +) + +type memoryFlowRequest struct { + id string + seq int64 + userID int + channelID int + upstreamModel string + state memoryFlowRequestState + enqueuedAt time.Time + dispatchedAt time.Time + notify chan struct{} + cancelled bool +} + +type memoryFlowGuard struct { + backend *memoryFlowBackend + slot *memoryFlowSlot + poolKey string + requestID string + released atomic.Bool + releaseFunc atomic.Value +} + +type flowReadCloser struct { + io.ReadCloser + guard FlowGuard +} + +func NewMemoryFlowBackend() FlowBackend { + return &memoryFlowBackend{ + slots: make(map[string]*memoryFlowSlot), + } +} + +func (b *memoryFlowBackend) Acquire(ctx context.Context, req AcquireRequest) (FlowGuard, *AcquireDecision, error) { + if ctx == nil { + ctx = context.Background() + } + req.Pool.Normalize() + if req.QueueTimeoutMs <= 0 { + req.QueueTimeoutMs = req.Pool.QueueTimeoutMs + } + decision := newFlowDecision(req.Pool, false, false) + slot := b.getSlot(req.Pool) + now := time.Now() + + slot.mu.Lock() + slot.config = req.Pool + slot.cleanupLocked(now) + running, queued, _ := slot.statsLocked(now) + decision.RunningNow = running + decision.QueuedNow = queued + if req.Pool.MaxContextTokens > 0 && req.ContextTokens > req.Pool.MaxContextTokens { + decision.RejectCode = FlowDecisionRejectContextExceeded + slot.mu.Unlock() + return nil, decision, fmt.Errorf("request context tokens %d exceeds flow pool max_context_tokens %d", req.ContextTokens, req.Pool.MaxContextTokens) + } + if slot.hasCapacityLocked() && queued == 0 { + request := slot.newRequestLocked(req, memoryFlowStateRunning, now) + request.dispatchedAt = now + slot.queue = append(slot.queue, request) + running, queued, _ = slot.statsLocked(now) + decision.Admitted = true + decision.RunningNow = running + decision.QueuedNow = queued + decision.WaitedMs = 0 + guard := &memoryFlowGuard{backend: b, slot: slot, poolKey: req.Pool.PoolKey, requestID: request.id} + slot.mu.Unlock() + return guard, decision, nil + } + if req.Pool.OnLimit != model.ChannelFlowOnLimitQueue { + decision.RejectCode = FlowDecisionRejectQueueFull + slot.mu.Unlock() + return nil, decision, fmt.Errorf("channel flow pool is busy") + } + if req.Pool.MaxQueueSize > 0 && queued >= req.Pool.MaxQueueSize { + decision.RejectCode = FlowDecisionRejectQueueFull + slot.mu.Unlock() + return nil, decision, fmt.Errorf("channel flow queue is full") + } + if req.Pool.MaxQueuePerUser > 0 && slot.userWaitingLocked(req.UserID) >= req.Pool.MaxQueuePerUser { + decision.RejectCode = FlowDecisionRejectPerUserQueueFull + slot.mu.Unlock() + return nil, decision, fmt.Errorf("channel flow per-user queue is full") + } + + request := slot.newRequestLocked(req, memoryFlowStateWaiting, now) + slot.queue = append(slot.queue, request) + slot.dispatchLocked(now) + running, queued, _ = slot.statsLocked(now) + admittedAfterDispatch := request.state == memoryFlowStateRunning + decision.Queued = request.state == memoryFlowStateWaiting + decision.Admitted = admittedAfterDispatch + decision.QueuePos = slot.positionLocked(request.id) + decision.RunningNow = running + decision.QueuedNow = queued + slot.mu.Unlock() + + if admittedAfterDispatch { + decision.WaitedMs = 0 + return &memoryFlowGuard{backend: b, slot: slot, poolKey: req.Pool.PoolKey, requestID: request.id}, decision, nil + } + + timer := time.NewTimer(time.Duration(req.QueueTimeoutMs) * time.Millisecond) + defer timer.Stop() + select { + case <-ctx.Done(): + waitedMs, runningNow, queuedNow := slot.cancelWaiting(request.id) + decision.WaitedMs = waitedMs + decision.RunningNow = runningNow + decision.QueuedNow = queuedNow + decision.RejectCode = FlowDecisionRejectQueueTimeout + return nil, decision, ctx.Err() + case <-timer.C: + waitedMs, runningNow, queuedNow := slot.cancelWaiting(request.id) + decision.WaitedMs = waitedMs + decision.RunningNow = runningNow + decision.QueuedNow = queuedNow + decision.RejectCode = FlowDecisionRejectQueueTimeout + return nil, decision, fmt.Errorf("channel flow queue timeout") + case <-request.notify: + dispatchedAt := request.dispatchedAt + if dispatchedAt.IsZero() { + dispatchedAt = time.Now() + } + decision.Admitted = true + decision.Queued = true + decision.WaitedMs = dispatchedAt.Sub(request.enqueuedAt).Milliseconds() + slot.mu.Lock() + running, queued, _ = slot.statsLocked(time.Now()) + decision.RunningNow = running + decision.QueuedNow = queued + decision.QueuePos = 0 + slot.mu.Unlock() + return &memoryFlowGuard{backend: b, slot: slot, poolKey: req.Pool.PoolKey, requestID: request.id}, decision, nil + } +} + +func (b *memoryFlowBackend) Status(_ context.Context, pool model.ChannelFlowPool) (PoolStatus, error) { + pool.Normalize() + slot := b.getSlot(pool) + slot.mu.Lock() + defer slot.mu.Unlock() + slot.config = pool + now := time.Now() + slot.cleanupLocked(now) + running, queued, oldestWaitMs := slot.statsLocked(now) + return PoolStatus{ + PoolKey: pool.PoolKey, + Name: pool.Name, + Backend: pool.Backend, + Health: flowHealth(running, pool.MaxInflight, queued, pool.MaxQueueSize), + Running: running, + MaxInflight: pool.MaxInflight, + Queued: queued, + MaxQueueSize: pool.MaxQueueSize, + OldestWaitMs: oldestWaitMs, + ConfigVersion: pool.ConfigVersion, + }, nil +} + +func (b *memoryFlowBackend) Close(_ context.Context) error { + return nil +} + +func (b *memoryFlowBackend) getSlot(pool model.ChannelFlowPool) *memoryFlowSlot { + b.mu.RLock() + slot := b.slots[pool.PoolKey] + b.mu.RUnlock() + if slot != nil { + return slot + } + b.mu.Lock() + defer b.mu.Unlock() + if slot = b.slots[pool.PoolKey]; slot != nil { + return slot + } + slot = &memoryFlowSlot{ + config: pool, + eventCounts: make(map[string]int), + } + b.slots[pool.PoolKey] = slot + return slot +} + +func (s *memoryFlowSlot) newRequestLocked(req AcquireRequest, state memoryFlowRequestState, now time.Time) *memoryFlowRequest { + s.nextSeq++ + return &memoryFlowRequest{ + id: req.RequestID, + seq: s.nextSeq, + userID: req.UserID, + channelID: req.ChannelID, + upstreamModel: req.UpstreamModel, + state: state, + enqueuedAt: now, + notify: make(chan struct{}, 1), + } +} + +func (s *memoryFlowSlot) hasCapacityLocked() bool { + if s.config.MaxInflight <= 0 { + return true + } + running := 0 + for _, req := range s.queue { + if req.state == memoryFlowStateRunning && !req.cancelled { + running++ + } + } + return running < s.config.MaxInflight +} + +func (s *memoryFlowSlot) dispatchLocked(now time.Time) { + for s.hasCapacityLocked() { + dispatched := false + for _, req := range s.queue { + if req.state != memoryFlowStateWaiting || req.cancelled { + continue + } + req.state = memoryFlowStateRunning + req.dispatchedAt = now + select { + case req.notify <- struct{}{}: + default: + } + dispatched = true + break + } + if !dispatched { + return + } + } +} + +func (s *memoryFlowSlot) statsLocked(now time.Time) (running int, queued int, oldestWaitMs int64) { + for _, req := range s.queue { + if req.cancelled || req.state == memoryFlowStateReleased { + continue + } + switch req.state { + case memoryFlowStateRunning: + running++ + case memoryFlowStateWaiting: + queued++ + waitMs := now.Sub(req.enqueuedAt).Milliseconds() + if oldestWaitMs == 0 || waitMs > oldestWaitMs { + oldestWaitMs = waitMs + } + } + } + return running, queued, oldestWaitMs +} + +func (s *memoryFlowSlot) positionLocked(requestID string) int { + position := 0 + for _, req := range s.queue { + if req.cancelled || req.state != memoryFlowStateWaiting { + continue + } + position++ + if req.id == requestID { + return position + } + } + return 0 +} + +func (s *memoryFlowSlot) userWaitingLocked(userID int) int { + if userID <= 0 { + return 0 + } + count := 0 + for _, req := range s.queue { + if req.userID == userID && req.state == memoryFlowStateWaiting && !req.cancelled { + count++ + } + } + return count +} + +func (s *memoryFlowSlot) cancelWaiting(requestID string) (waitedMs int64, runningNow int, queuedNow int) { + s.mu.Lock() + defer s.mu.Unlock() + now := time.Now() + for _, req := range s.queue { + if req.id == requestID && req.state == memoryFlowStateWaiting && !req.cancelled { + req.cancelled = true + req.state = memoryFlowStateReleased + waitedMs = now.Sub(req.enqueuedAt).Milliseconds() + break + } + } + s.compactIfNeededLocked() + s.dispatchLocked(now) + runningNow, queuedNow, _ = s.statsLocked(now) + return waitedMs, runningNow, queuedNow +} + +func (s *memoryFlowSlot) release(requestID string) error { + s.mu.Lock() + defer s.mu.Unlock() + now := time.Now() + for _, req := range s.queue { + if req.id == requestID && req.state == memoryFlowStateRunning { + req.state = memoryFlowStateReleased + req.cancelled = true + break + } + } + s.compactIfNeededLocked() + s.dispatchLocked(now) + return nil +} + +func (s *memoryFlowSlot) cleanupLocked(now time.Time) { + if s.config.MaxProcessingMs > 0 { + for _, req := range s.queue { + if req.state == memoryFlowStateRunning && !req.dispatchedAt.IsZero() && now.Sub(req.dispatchedAt).Milliseconds() > s.config.MaxProcessingMs { + req.state = memoryFlowStateReleased + req.cancelled = true + } + } + } + s.compactIfNeededLocked() +} + +func (s *memoryFlowSlot) compactIfNeededLocked() { + if len(s.queue) == 0 { + return + } + stale := 0 + for _, req := range s.queue { + if req.cancelled || req.state == memoryFlowStateReleased { + stale++ + } + } + if stale < 64 && stale*100 < len(s.queue)*30 { + return + } + compact := s.queue[:0] + for _, req := range s.queue { + if req.cancelled || req.state == memoryFlowStateReleased { + continue + } + compact = append(compact, req) + } + s.queue = compact +} + +func (g *memoryFlowGuard) Release(ctx context.Context) error { + if g == nil || g.released.Swap(true) { + return nil + } + if release, ok := g.releaseFunc.Load().(func()); ok && release != nil { + release() + } + if g.slot == nil { + return nil + } + return g.slot.release(g.requestID) +} + +func (g *memoryFlowGuard) RenewLease(_ context.Context) error { + return nil +} + +func (g *memoryFlowGuard) PoolKey() string { + if g == nil { + return "" + } + return g.poolKey +} + +func (g *memoryFlowGuard) RequestID() string { + if g == nil { + return "" + } + return g.requestID +} + +func (g *memoryFlowGuard) IsReleased() bool { + return g == nil || g.released.Load() +} + +func (g *memoryFlowGuard) BindRelease(release func()) { + if g == nil || release == nil { + return + } + g.releaseFunc.Store(release) +} + +func (g *memoryFlowGuard) WrapReadCloser(rc io.ReadCloser) io.ReadCloser { + if rc == nil { + return nil + } + return &flowReadCloser{ReadCloser: rc, guard: g} +} + +func (rc *flowReadCloser) Close() error { + err := rc.ReadCloser.Close() + _ = rc.guard.Release(context.Background()) + return err +} + +func flowHealth(running int, maxInflight int, queued int, maxQueueSize int) string { + if queued > 0 { + if maxQueueSize > 0 && queued*100/maxQueueSize >= 80 { + return "critical" + } + return "congested" + } + if maxInflight > 0 && running*100/maxInflight >= 70 { + return "busy" + } + return "healthy" +} diff --git a/service/channel_flow_redis.go b/service/channel_flow_redis.go new file mode 100644 index 000000000000..6b85ec61dadb --- /dev/null +++ b/service/channel_flow_redis.go @@ -0,0 +1,650 @@ +package service + +import ( + "context" + "errors" + "fmt" + "io" + "strconv" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/go-redis/redis/v8" +) + +var ErrRedisFlowBackendUnavailable = errors.New("channel flow redis backend unavailable") + +const ( + redisFlowNamespace = "new-api:channel-flow:v1" + redisFlowPollMin = 10 * time.Millisecond + redisFlowPollMax = 50 * time.Millisecond + redisFlowCleanupBatch = 128 + redisFlowRequestTTLExtra = time.Hour +) + +type redisFlowBackend struct { + pollMin time.Duration + pollMax time.Duration +} + +type redisFlowKeys struct { + Running string + Waiting string + Deadline string + Seq string + Base string +} + +type redisFlowGuard struct { + backend *redisFlowBackend + pool model.ChannelFlowPool + poolKey string + requestID string + released atomic.Bool + releaseFunc atomic.Value +} + +type redisAcquireAttempt struct { + decision redisAcquireDecision + done bool +} + +type redisAcquireDecision struct { + admitted bool + queued bool + rejectCode string + queuePos int + waitedMs int64 + score float64 + runningNow int + queuedNow int +} + +func NewRedisFlowBackend() FlowBackend { + return &redisFlowBackend{ + pollMin: redisFlowPollMin, + pollMax: redisFlowPollMax, + } +} + +func IsRedisFlowBackendAvailable(ctx context.Context) bool { + return common.RedisEnabled && common.RDB != nil +} + +func (b *redisFlowBackend) Acquire(ctx context.Context, req AcquireRequest) (FlowGuard, *AcquireDecision, error) { + if ctx == nil { + ctx = context.Background() + } + req.Pool.Normalize() + if req.QueueTimeoutMs <= 0 { + req.QueueTimeoutMs = req.Pool.QueueTimeoutMs + } + if req.RequestID == "" { + req.RequestID = common.GetUUID() + } + decision := newFlowDecision(req.Pool, false, false) + if req.Pool.MaxContextTokens > 0 && req.ContextTokens > req.Pool.MaxContextTokens { + decision.RejectCode = FlowDecisionRejectContextExceeded + return nil, decision, fmt.Errorf("request context tokens %d exceeds flow pool max_context_tokens %d", req.ContextTokens, req.Pool.MaxContextTokens) + } + + rdb, err := b.client() + if err != nil { + decision.RejectCode = FlowDecisionRejectBackendDisabled + return nil, decision, err + } + + timeout := time.Duration(req.QueueTimeoutMs) * time.Millisecond + acquireCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + keys := redisKeysForPool(req.Pool) + enqueued := false + queuedAt := time.Time{} + sequenceScore := float64(0) + + for { + if err := acquireCtx.Err(); err != nil { + if enqueued { + _ = b.removeWaiting(context.Background(), rdb, keys, req.RequestID, req.UserID) + } + decision.RejectCode = FlowDecisionRejectQueueTimeout + if !queuedAt.IsZero() { + decision.WaitedMs = time.Since(queuedAt).Milliseconds() + } + status, statusErr := b.Status(context.Background(), req.Pool) + if statusErr == nil { + decision.RunningNow = status.Running + decision.QueuedNow = status.Queued + } + return nil, decision, fmt.Errorf("channel flow queue timeout") + } + + _ = b.cleanupExpired(acquireCtx, rdb, keys) + attempt, err := b.tryAcquireOnce(acquireCtx, rdb, keys, req, enqueued, sequenceScore, queuedAt) + if err != nil { + if errors.Is(err, redis.TxFailedErr) { + continue + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + continue + } + if isRedisFlowUnavailableErr(err) { + decision.RejectCode = FlowDecisionRejectBackendDisabled + return nil, decision, fmt.Errorf("%w: %v", ErrRedisFlowBackendUnavailable, err) + } + return nil, decision, err + } + decision.RunningNow = attempt.decision.runningNow + decision.QueuedNow = attempt.decision.queuedNow + decision.QueuePos = attempt.decision.queuePos + + if attempt.done { + if attempt.decision.rejectCode != "" { + decision.RejectCode = attempt.decision.rejectCode + return nil, decision, redisRejectError(attempt.decision.rejectCode) + } + if attempt.decision.admitted { + decision.Admitted = true + decision.Queued = attempt.decision.queued + decision.WaitedMs = attempt.decision.waitedMs + return &redisFlowGuard{ + backend: b, + pool: req.Pool, + poolKey: req.Pool.PoolKey, + requestID: req.RequestID, + }, decision, nil + } + } + + if !enqueued && attempt.decision.queued { + enqueued = true + queuedAt = time.Now() + sequenceScore = attempt.decision.score + } + + if err := sleepRedisFlowPoll(acquireCtx, b.pollDelay()); err != nil { + continue + } + } +} + +func (b *redisFlowBackend) Status(ctx context.Context, pool model.ChannelFlowPool) (PoolStatus, error) { + if ctx == nil { + ctx = context.Background() + } + pool.Normalize() + rdb, err := b.client() + if err != nil { + return PoolStatus{}, err + } + keys := redisKeysForPool(pool) + _ = b.cleanupExpired(ctx, rdb, keys) + + running, err := rdb.ZCard(ctx, keys.Running).Result() + if err != nil { + return PoolStatus{}, redisFlowUnavailable(err) + } + queued, err := rdb.ZCard(ctx, keys.Waiting).Result() + if err != nil { + return PoolStatus{}, redisFlowUnavailable(err) + } + oldestWaitMs := int64(0) + oldest, err := rdb.ZRange(ctx, keys.Waiting, 0, 0).Result() + if err != nil && !errors.Is(err, redis.Nil) { + return PoolStatus{}, redisFlowUnavailable(err) + } + if len(oldest) > 0 { + enqueuedAt, _ := b.requestInt64(ctx, rdb, keys, oldest[0], "enqueued_at_ms") + if enqueuedAt > 0 { + oldestWaitMs = time.Now().UnixMilli() - enqueuedAt + } + } + return PoolStatus{ + PoolKey: pool.PoolKey, + Name: pool.Name, + Backend: pool.Backend, + Health: flowHealth(int(running), pool.MaxInflight, int(queued), pool.MaxQueueSize), + Running: int(running), + MaxInflight: pool.MaxInflight, + Queued: int(queued), + MaxQueueSize: pool.MaxQueueSize, + OldestWaitMs: oldestWaitMs, + ConfigVersion: pool.ConfigVersion, + }, nil +} + +func (b *redisFlowBackend) Close(_ context.Context) error { + return nil +} + +func (b *redisFlowBackend) tryAcquireOnce( + ctx context.Context, + rdb *redis.Client, + keys redisFlowKeys, + req AcquireRequest, + enqueued bool, + sequenceScore float64, + queuedAt time.Time, +) (redisAcquireAttempt, error) { + attempt := redisAcquireAttempt{} + watchKeys := []string{keys.Running, keys.Waiting} + if req.UserID > 0 { + watchKeys = append(watchKeys, keys.userWaiting(req.UserID)) + } + err := rdb.Watch(ctx, func(tx *redis.Tx) error { + running, err := tx.ZCard(ctx, keys.Running).Result() + if err != nil { + return err + } + waiting, err := tx.ZCard(ctx, keys.Waiting).Result() + if err != nil { + return err + } + attempt.decision.runningNow = int(running) + attempt.decision.queuedNow = int(waiting) + + if !enqueued { + if redisFlowHasCapacity(running, req.Pool.MaxInflight) && waiting == 0 { + expiresAtMs := time.Now().Add(redisLeaseDuration(req.Pool)).UnixMilli() + _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { + pipe.ZAdd(ctx, keys.Running, &redis.Z{ + Score: float64(expiresAtMs), + Member: req.RequestID, + }) + b.writeRequestMeta(ctx, pipe, keys, req, "running", 0, expiresAtMs) + return nil + }) + if err == nil { + attempt.done = true + attempt.decision.admitted = true + attempt.decision.runningNow = int(running) + 1 + attempt.decision.queuedNow = int(waiting) + } + return err + } + if req.Pool.OnLimit != model.ChannelFlowOnLimitQueue { + attempt.done = true + attempt.decision.rejectCode = FlowDecisionRejectQueueFull + return nil + } + if req.Pool.MaxQueueSize > 0 && waiting >= int64(req.Pool.MaxQueueSize) { + attempt.done = true + attempt.decision.rejectCode = FlowDecisionRejectQueueFull + return nil + } + if req.Pool.MaxQueuePerUser > 0 && req.UserID > 0 { + userWaiting, err := tx.ZCard(ctx, keys.userWaiting(req.UserID)).Result() + if err != nil { + return err + } + if userWaiting >= int64(req.Pool.MaxQueuePerUser) { + attempt.done = true + attempt.decision.rejectCode = FlowDecisionRejectPerUserQueueFull + return nil + } + } + score, err := b.nextSequence(ctx, rdb, keys, sequenceScore) + if err != nil { + return err + } + enqueuedAtMs := time.Now().UnixMilli() + deadlineMs := enqueuedAtMs + req.QueueTimeoutMs + _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { + pipe.ZAdd(ctx, keys.Waiting, &redis.Z{ + Score: score, + Member: req.RequestID, + }) + pipe.ZAdd(ctx, keys.Deadline, &redis.Z{ + Score: float64(deadlineMs), + Member: req.RequestID, + }) + if req.UserID > 0 { + pipe.ZAdd(ctx, keys.userWaiting(req.UserID), &redis.Z{ + Score: score, + Member: req.RequestID, + }) + } + b.writeRequestMeta(ctx, pipe, keys, req, "waiting", enqueuedAtMs, 0) + return nil + }) + if err == nil { + attempt.decision.queued = true + attempt.decision.score = score + attempt.decision.queuePos = int(waiting) + 1 + attempt.decision.runningNow = int(running) + attempt.decision.queuedNow = int(waiting) + 1 + } + return err + } + + rank, err := tx.ZRank(ctx, keys.Waiting, req.RequestID).Result() + if errors.Is(err, redis.Nil) { + attempt.done = true + attempt.decision.rejectCode = FlowDecisionRejectQueueTimeout + return nil + } + if err != nil { + return err + } + attempt.decision.queuePos = int(rank) + 1 + if rank != 0 || !redisFlowHasCapacity(running, req.Pool.MaxInflight) { + return nil + } + expiresAtMs := time.Now().Add(redisLeaseDuration(req.Pool)).UnixMilli() + _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { + pipe.ZRem(ctx, keys.Waiting, req.RequestID) + pipe.ZRem(ctx, keys.Deadline, req.RequestID) + if req.UserID > 0 { + pipe.ZRem(ctx, keys.userWaiting(req.UserID), req.RequestID) + } + pipe.ZAdd(ctx, keys.Running, &redis.Z{ + Score: float64(expiresAtMs), + Member: req.RequestID, + }) + b.writeRequestMeta(ctx, pipe, keys, req, "running", 0, expiresAtMs) + return nil + }) + if err == nil { + attempt.done = true + attempt.decision.admitted = true + attempt.decision.queued = true + attempt.decision.waitedMs = time.Since(queuedAt).Milliseconds() + attempt.decision.queuePos = 0 + attempt.decision.runningNow = int(running) + 1 + attempt.decision.queuedNow = maxInt(0, int(waiting)-1) + } + return err + }, watchKeys...) + return attempt, err +} + +func (b *redisFlowBackend) removeWaiting(ctx context.Context, rdb *redis.Client, keys redisFlowKeys, requestID string, userID int) error { + pipe := rdb.TxPipeline() + pipe.ZRem(ctx, keys.Waiting, requestID) + pipe.ZRem(ctx, keys.Deadline, requestID) + if userID > 0 { + pipe.ZRem(ctx, keys.userWaiting(userID), requestID) + } + pipe.Del(ctx, keys.request(requestID)) + _, err := pipe.Exec(ctx) + return redisFlowUnavailable(err) +} + +func (b *redisFlowBackend) release(ctx context.Context, pool model.ChannelFlowPool, requestID string) error { + rdb, err := b.client() + if err != nil { + return err + } + keys := redisKeysForPool(pool) + pipe := rdb.TxPipeline() + pipe.ZRem(ctx, keys.Running, requestID) + pipe.Del(ctx, keys.request(requestID)) + _, err = pipe.Exec(ctx) + return redisFlowUnavailable(err) +} + +func (b *redisFlowBackend) renew(ctx context.Context, pool model.ChannelFlowPool, requestID string) error { + rdb, err := b.client() + if err != nil { + return err + } + keys := redisKeysForPool(pool) + exists, err := rdb.ZScore(ctx, keys.Running, requestID).Result() + if errors.Is(err, redis.Nil) { + return nil + } + if err != nil { + return redisFlowUnavailable(err) + } + if exists <= 0 { + return nil + } + expiresAtMs := time.Now().Add(redisLeaseDuration(pool)).UnixMilli() + pipe := rdb.TxPipeline() + pipe.ZAdd(ctx, keys.Running, &redis.Z{ + Score: float64(expiresAtMs), + Member: requestID, + }) + pipe.HSet(ctx, keys.request(requestID), "expires_at_ms", strconv.FormatInt(expiresAtMs, 10)) + pipe.Expire(ctx, keys.request(requestID), redisRequestTTL(pool)) + _, err = pipe.Exec(ctx) + return redisFlowUnavailable(err) +} + +func (b *redisFlowBackend) cleanupExpired(ctx context.Context, rdb *redis.Client, keys redisFlowKeys) error { + nowMs := time.Now().UnixMilli() + if err := rdb.ZRemRangeByScore(ctx, keys.Running, "-inf", strconv.FormatInt(nowMs, 10)).Err(); err != nil { + return redisFlowUnavailable(err) + } + expired, err := rdb.ZRangeByScore(ctx, keys.Deadline, &redis.ZRangeBy{ + Min: "-inf", + Max: strconv.FormatInt(nowMs, 10), + Offset: 0, + Count: redisFlowCleanupBatch, + }).Result() + if err != nil { + return redisFlowUnavailable(err) + } + if len(expired) == 0 { + return nil + } + pipe := rdb.TxPipeline() + for _, requestID := range expired { + userID, _ := b.requestInt(ctx, rdb, keys, requestID, "user_id") + pipe.ZRem(ctx, keys.Waiting, requestID) + pipe.ZRem(ctx, keys.Deadline, requestID) + if userID > 0 { + pipe.ZRem(ctx, keys.userWaiting(userID), requestID) + } + pipe.Del(ctx, keys.request(requestID)) + } + _, err = pipe.Exec(ctx) + return redisFlowUnavailable(err) +} + +func (b *redisFlowBackend) nextSequence(ctx context.Context, rdb *redis.Client, keys redisFlowKeys, existing float64) (float64, error) { + if existing > 0 { + return existing, nil + } + seq, err := rdb.Incr(ctx, keys.Seq).Result() + if err != nil { + return 0, redisFlowUnavailable(err) + } + return float64(seq), nil +} + +func (b *redisFlowBackend) writeRequestMeta(ctx context.Context, pipe redis.Pipeliner, keys redisFlowKeys, req AcquireRequest, state string, enqueuedAtMs int64, expiresAtMs int64) { + data := map[string]interface{}{ + "state": state, + "user_id": strconv.Itoa(req.UserID), + "channel_id": strconv.Itoa(req.ChannelID), + "upstream_model": req.UpstreamModel, + } + if enqueuedAtMs > 0 { + data["enqueued_at_ms"] = strconv.FormatInt(enqueuedAtMs, 10) + } + if expiresAtMs > 0 { + data["expires_at_ms"] = strconv.FormatInt(expiresAtMs, 10) + } + pipe.HSet(ctx, keys.request(req.RequestID), data) + pipe.Expire(ctx, keys.request(req.RequestID), redisRequestTTL(req.Pool)) +} + +func (b *redisFlowBackend) requestInt(ctx context.Context, rdb *redis.Client, keys redisFlowKeys, requestID string, field string) (int, error) { + value, err := b.requestInt64(ctx, rdb, keys, requestID, field) + return int(value), err +} + +func (b *redisFlowBackend) requestInt64(ctx context.Context, rdb *redis.Client, keys redisFlowKeys, requestID string, field string) (int64, error) { + value, err := rdb.HGet(ctx, keys.request(requestID), field).Result() + if errors.Is(err, redis.Nil) { + return 0, nil + } + if err != nil { + return 0, redisFlowUnavailable(err) + } + parsed, err := strconv.ParseInt(value, 10, 64) + if err != nil { + return 0, nil + } + return parsed, nil +} + +func (b *redisFlowBackend) client() (*redis.Client, error) { + if !common.RedisEnabled || common.RDB == nil { + return nil, ErrRedisFlowBackendUnavailable + } + return common.RDB, nil +} + +func (b *redisFlowBackend) pollDelay() time.Duration { + window := b.pollMax - b.pollMin + if window <= 0 { + return b.pollMin + } + return b.pollMin + time.Duration(time.Now().UnixNano()%int64(window)) +} + +func (g *redisFlowGuard) Release(ctx context.Context) error { + if g == nil || g.released.Swap(true) { + return nil + } + if release, ok := g.releaseFunc.Load().(func()); ok && release != nil { + release() + } + if ctx == nil { + ctx = context.Background() + } + return g.backend.release(ctx, g.pool, g.requestID) +} + +func (g *redisFlowGuard) RenewLease(ctx context.Context) error { + if g == nil || g.released.Load() { + return nil + } + if ctx == nil { + ctx = context.Background() + } + return g.backend.renew(ctx, g.pool, g.requestID) +} + +func (g *redisFlowGuard) PoolKey() string { + if g == nil { + return "" + } + return g.poolKey +} + +func (g *redisFlowGuard) RequestID() string { + if g == nil { + return "" + } + return g.requestID +} + +func (g *redisFlowGuard) IsReleased() bool { + return g == nil || g.released.Load() +} + +func (g *redisFlowGuard) BindRelease(release func()) { + if g == nil || release == nil { + return + } + g.releaseFunc.Store(release) +} + +func (g *redisFlowGuard) WrapReadCloser(rc io.ReadCloser) io.ReadCloser { + if rc == nil { + return nil + } + return &flowReadCloser{ReadCloser: rc, guard: g} +} + +func redisKeysForPool(pool model.ChannelFlowPool) redisFlowKeys { + base := fmt.Sprintf("%s:%s", redisFlowNamespace, pool.PoolKey) + return redisFlowKeys{ + Base: base, + Running: base + ":running", + Waiting: base + ":waiting", + Deadline: base + ":deadline", + Seq: base + ":seq", + } +} + +func (k redisFlowKeys) request(requestID string) string { + return k.Base + ":request:" + requestID +} + +func (k redisFlowKeys) userWaiting(userID int) string { + return fmt.Sprintf("%s:user:%d:waiting", k.Base, userID) +} + +func redisFlowHasCapacity(running int64, maxInflight int) bool { + return maxInflight <= 0 || running < int64(maxInflight) +} + +func redisLeaseDuration(pool model.ChannelFlowPool) time.Duration { + pool.Normalize() + return time.Duration(pool.LeaseMs) * time.Millisecond +} + +func redisRequestTTL(pool model.ChannelFlowPool) time.Duration { + pool.Normalize() + ttl := time.Duration(pool.QueueTimeoutMs)*time.Millisecond + redisLeaseDuration(pool) + redisFlowRequestTTLExtra + if pool.MaxProcessingMs > 0 { + ttl += time.Duration(pool.MaxProcessingMs) * time.Millisecond + } + if ttl < 5*time.Minute { + return 5 * time.Minute + } + return ttl +} + +func sleepRedisFlowPoll(ctx context.Context, delay time.Duration) error { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func redisRejectError(code string) error { + switch code { + case FlowDecisionRejectPerUserQueueFull: + return fmt.Errorf("channel flow per-user queue is full") + case FlowDecisionRejectQueueTimeout: + return fmt.Errorf("channel flow queue timeout") + default: + return fmt.Errorf("channel flow queue is full") + } +} + +func redisFlowUnavailable(err error) error { + if err == nil { + return nil + } + if errors.Is(err, redis.Nil) { + return nil + } + if errors.Is(err, ErrRedisFlowBackendUnavailable) { + return err + } + return fmt.Errorf("%w: %v", ErrRedisFlowBackendUnavailable, err) +} + +func isRedisFlowUnavailableErr(err error) bool { + return errors.Is(redisFlowUnavailable(err), ErrRedisFlowBackendUnavailable) +} + +func maxInt(left int, right int) int { + if left > right { + return left + } + return right +} diff --git a/service/channel_flow_test.go b/service/channel_flow_test.go new file mode 100644 index 000000000000..f2bf6a46aa43 --- /dev/null +++ b/service/channel_flow_test.go @@ -0,0 +1,500 @@ +package service + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/go-redis/redis/v8" +) + +func testFlowPool() model.ChannelFlowPool { + return model.ChannelFlowPool{ + PoolKey: "flow_pool_test", + Name: "test pool", + Enabled: true, + Backend: model.ChannelFlowBackendMemory, + MaxInflight: 1, + MaxQueueSize: 1, + QueueTimeoutMs: 500, + QueuePolicy: model.ChannelFlowQueuePolicyFIFO, + OnLimit: model.ChannelFlowOnLimitQueue, + ConfigVersion: 1, + } +} + +func TestMemoryFlowBackendReleaseDispatchesWaitingRequest(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + + guard1, decision1, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + t.Fatalf("first acquire failed: %v", err) + } + if guard1 == nil || decision1 == nil || !decision1.Admitted { + t.Fatalf("first acquire should be admitted immediately, decision=%+v guard=%v", decision1, guard1) + } + + resultCh := make(chan error, 1) + go func() { + guard2, decision2, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + resultCh <- err + return + } + if guard2 == nil || decision2 == nil || !decision2.Admitted || !decision2.Queued { + resultCh <- context.Canceled + return + } + _ = guard2.Release(context.Background()) + resultCh <- nil + }() + + time.Sleep(50 * time.Millisecond) + if err := guard1.Release(context.Background()); err != nil { + t.Fatalf("release failed: %v", err) + } + + select { + case err := <-resultCh: + if err != nil { + t.Fatalf("waiting acquire failed: %v", err) + } + case <-time.After(time.Second): + t.Fatal("waiting acquire was not dispatched after release") + } +} + +func TestMemoryFlowBackendRejectsWhenQueueFull(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + t.Fatalf("first acquire failed: %v", err) + } + defer guard1.Release(context.Background()) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + waitingStarted := make(chan struct{}) + go func() { + close(waitingStarted) + _, _, _ = backend.Acquire(ctx, AcquireRequest{ + RequestID: "req-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + }() + <-waitingStarted + time.Sleep(50 * time.Millisecond) + + _, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-3", + Pool: pool, + UserID: 3, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err == nil { + t.Fatal("third acquire should fail when queue is full") + } + if decision == nil || decision.RejectCode != FlowDecisionRejectQueueFull { + t.Fatalf("unexpected decision: %+v", decision) + } + cancel() +} + +func TestMemoryFlowBackendAllowsQueueUpToMaxQueueSize(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + pool.MaxQueueSize = 2 + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + t.Fatalf("first acquire failed: %v", err) + } + defer guard1.Release(context.Background()) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + for i := 2; i <= 3; i++ { + requestID := i + waitingStarted := make(chan struct{}) + go func() { + close(waitingStarted) + _, _, _ = backend.Acquire(ctx, AcquireRequest{ + RequestID: "req-" + string(rune('0'+requestID)), + Pool: pool, + UserID: requestID, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + }() + <-waitingStarted + time.Sleep(50 * time.Millisecond) + } + + _, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-4", + Pool: pool, + UserID: 4, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err == nil { + t.Fatal("fourth acquire should fail when total queue is full") + } + if decision == nil || decision.RejectCode != FlowDecisionRejectQueueFull { + t.Fatalf("unexpected decision: %+v", decision) + } + if decision.QueuedNow != 2 { + t.Fatalf("queued count should be 2, got decision=%+v", decision) + } +} + +func TestMemoryFlowBackendRejectsWhenPerUserQueueFull(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + pool.MaxQueueSize = 2 + pool.MaxQueuePerUser = 1 + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + t.Fatalf("first acquire failed: %v", err) + } + defer guard1.Release(context.Background()) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + waitingStarted := make(chan struct{}) + go func() { + close(waitingStarted) + _, _, _ = backend.Acquire(ctx, AcquireRequest{ + RequestID: "req-2", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + }() + <-waitingStarted + time.Sleep(50 * time.Millisecond) + + _, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-3", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err == nil { + t.Fatal("third acquire should fail when per-user queue is full") + } + if decision == nil || decision.RejectCode != FlowDecisionRejectPerUserQueueFull { + t.Fatalf("unexpected decision: %+v", decision) + } +} + +func TestRedisLocalMemoryFallbackStatusUsesMemoryBackend(t *testing.T) { + pool := testFlowPool() + pool.PoolKey = "flow_pool_redis_local_memory_status" + pool.Backend = model.ChannelFlowBackendRedis + pool.RedisFailurePolicy = model.ChannelFlowRedisFailureLocalMemory + fallbackPool := localMemoryFallbackFlowPool(pool) + + guard, _, err := GetChannelFlowController().Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-local-memory-status-1", + Pool: fallbackPool, + UserID: 1, + QueueTimeoutMs: fallbackPool.QueueTimeoutMs, + }) + if err != nil { + t.Fatalf("fallback acquire failed: %v", err) + } + defer guard.Release(context.Background()) + + status, err := GetChannelFlowPoolStatus(context.Background(), pool) + if err != nil { + t.Fatalf("status failed: %v", err) + } + if status.Backend != model.ChannelFlowBackendMemory { + t.Fatalf("status should report effective memory backend, got %+v", status) + } + if status.Running != 1 || status.MaxInflight != pool.MaxInflight { + t.Fatalf("status should read memory fallback counters, got %+v", status) + } +} + +func TestRedisFlowBackendReleaseDispatchesWaitingRequest(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxQueueSize = 2 + + guard1, decision1, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + t.Fatalf("first redis acquire failed: %v", err) + } + if guard1 == nil || decision1 == nil || !decision1.Admitted { + t.Fatalf("first redis acquire should be admitted, decision=%+v guard=%v", decision1, guard1) + } + + resultCh := make(chan error, 1) + go func() { + guard2, decision2, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-req-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + resultCh <- err + return + } + if guard2 == nil || decision2 == nil || !decision2.Admitted || !decision2.Queued { + resultCh <- fmt.Errorf("waiting redis acquire was not queued then admitted: decision=%+v guard=%v", decision2, guard2) + return + } + _ = guard2.Release(context.Background()) + resultCh <- nil + }() + + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 1 + }) + + if err := guard1.Release(context.Background()); err != nil { + t.Fatalf("redis release failed: %v", err) + } + + select { + case err := <-resultCh: + if err != nil { + t.Fatalf("waiting redis acquire failed: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("waiting redis acquire was not dispatched after release") + } +} + +func TestRedisFlowBackendAllowsQueueUpToMaxQueueSize(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxQueueSize = 2 + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-queue-limit-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + t.Fatalf("first redis acquire failed: %v", err) + } + defer guard1.Release(context.Background()) + + waitCtx, cancel := context.WithCancel(context.Background()) + resultCh := make(chan error, 2) + for i := 2; i <= 3; i++ { + requestID := i + go func() { + guard, decision, err := backend.Acquire(waitCtx, AcquireRequest{ + RequestID: fmt.Sprintf("redis-queue-limit-%d", requestID), + Pool: pool, + UserID: requestID, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if guard != nil { + _ = guard.Release(context.Background()) + } + if err == nil { + resultCh <- fmt.Errorf("queued request %d was admitted before release: decision=%+v", requestID, decision) + return + } + resultCh <- nil + }() + } + + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 2 + }) + + _, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-queue-limit-4", + Pool: pool, + UserID: 4, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err == nil { + t.Fatal("fourth redis acquire should fail when total queue is full") + } + if decision == nil || decision.RejectCode != FlowDecisionRejectQueueFull { + t.Fatalf("unexpected redis decision: %+v", decision) + } + if decision.QueuedNow != 2 { + t.Fatalf("redis queued count should be 2, got decision=%+v", decision) + } + + cancel() + for i := 0; i < 2; i++ { + select { + case err := <-resultCh: + if err != nil { + t.Fatal(err) + } + case <-time.After(2 * time.Second): + t.Fatal("queued redis acquire did not exit after cancellation") + } + } +} + +func TestRedisFlowBackendRejectsWhenPerUserQueueFull(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxQueueSize = 2 + pool.MaxQueuePerUser = 1 + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-user-req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil { + t.Fatalf("first redis acquire failed: %v", err) + } + defer guard1.Release(context.Background()) + + waitCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { + _, _, _ = backend.Acquire(waitCtx, AcquireRequest{ + RequestID: "redis-user-req-2", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + }() + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 1 + }) + + _, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-user-req-3", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err == nil { + t.Fatal("third redis acquire should fail when per-user queue is full") + } + if decision == nil || decision.RejectCode != FlowDecisionRejectPerUserQueueFull { + t.Fatalf("unexpected redis decision: %+v", decision) + } +} + +func newRedisFlowBackendForTest(t *testing.T) (*redisFlowBackend, model.ChannelFlowPool, func()) { + t.Helper() + redisURL := os.Getenv("REDIS_CONN_STRING") + if redisURL == "" { + t.Skip("REDIS_CONN_STRING is not set") + } + opt, err := redis.ParseURL(redisURL) + if err != nil { + t.Fatalf("parse redis url: %v", err) + } + client := redis.NewClient(opt) + if err := client.Ping(context.Background()).Err(); err != nil { + _ = client.Close() + t.Skipf("redis is not available: %v", err) + } + + oldRedisEnabled := common.RedisEnabled + oldRDB := common.RDB + common.RedisEnabled = true + common.RDB = client + + pool := testFlowPool() + pool.PoolKey = fmt.Sprintf("flow_pool_redis_test_%d", time.Now().UnixNano()) + pool.Backend = model.ChannelFlowBackendRedis + pool.RedisFailurePolicy = model.ChannelFlowRedisFailureFailClosed + pool.QueueTimeoutMs = 1500 + pool.LeaseMs = 2000 + backend := NewRedisFlowBackend().(*redisFlowBackend) + cleanupRedisFlowKeys(t, client, pool) + + return backend, pool, func() { + cleanupRedisFlowKeys(t, client, pool) + common.RedisEnabled = oldRedisEnabled + common.RDB = oldRDB + _ = client.Close() + } +} + +func cleanupRedisFlowKeys(t *testing.T, client *redis.Client, pool model.ChannelFlowPool) { + t.Helper() + keys := redisKeysForPool(pool) + pattern := keys.Base + ":*" + ctx := context.Background() + var cursor uint64 + for { + found, nextCursor, err := client.Scan(ctx, cursor, pattern, 100).Result() + if err != nil { + t.Fatalf("scan redis flow keys: %v", err) + } + cursor = nextCursor + if len(found) > 0 { + if err := client.Del(ctx, found...).Err(); err != nil { + t.Fatalf("delete redis flow keys: %v", err) + } + } + if cursor == 0 { + return + } + } +} + +func eventuallyFlowStatus(t *testing.T, backend FlowBackend, pool model.ChannelFlowPool, predicate func(PoolStatus) bool) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + var last PoolStatus + var lastErr error + for time.Now().Before(deadline) { + last, lastErr = backend.Status(context.Background(), pool) + if lastErr == nil && predicate(last) { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatalf("status predicate not met, last=%+v err=%v", last, lastErr) +} diff --git a/types/error.go b/types/error.go index 9717401ae7b2..910198ace4e9 100644 --- a/types/error.go +++ b/types/error.go @@ -85,6 +85,16 @@ const ( // quota error ErrorCodeInsufficientUserQuota ErrorCode = "insufficient_user_quota" ErrorCodePreConsumeTokenQuotaFailed ErrorCode = "pre_consume_token_quota_failed" + + // channel flow control errors + ErrorCodeChannelFlowQueueFull ErrorCode = "channel_flow_queue_full" + ErrorCodeChannelFlowQueueTimeout ErrorCode = "channel_flow_queue_timeout" + ErrorCodeChannelFlowContextExceeded ErrorCode = "channel_flow_context_exceeded" + ErrorCodeChannelFlowDraining ErrorCode = "channel_flow_draining" + ErrorCodeChannelFlowBackendUnavailable ErrorCode = "channel_flow_backend_unavailable" + ErrorCodeChannelFlowConfigInvalid ErrorCode = "channel_flow_config_invalid" + ErrorCodeChannelFlowBillingFailedAfterWait ErrorCode = "channel_flow_billing_failed_after_wait" + ErrorCodeChannelFlowPerUserQueueFull ErrorCode = "channel_flow_per_user_queue_full" ) type NewAPIError struct { From 4053eadabef631b9f4ddddcb40adfd657960c241 Mon Sep 17 00:00:00 2001 From: supreme0597 Date: Sat, 13 Jun 2026 17:04:05 +0800 Subject: [PATCH 03/18] feat(web): add flow pool management page --- web/default/src/features/channel-flow/api.ts | 119 +++++ .../components/binding-form-sheet.tsx | 360 ++++++++++++++ .../components/pool-bindings-panel.tsx | 135 ++++++ .../components/pool-form-sheet.tsx | 449 ++++++++++++++++++ .../channel-flow/components/pool-list.tsx | 175 +++++++ .../components/pool-status-panel.tsx | 247 ++++++++++ .../src/features/channel-flow/index.tsx | 372 +++++++++++++++ .../src/features/channel-flow/lib/form.ts | 152 ++++++ .../src/features/channel-flow/lib/index.ts | 22 + .../features/channel-flow/lib/query-keys.ts | 30 ++ .../src/features/channel-flow/types.ts | 108 +++++ web/default/src/hooks/use-sidebar-data.ts | 6 + web/default/src/i18n/locales/en.json | 89 +++- web/default/src/i18n/locales/fr.json | 89 +++- web/default/src/i18n/locales/ja.json | 89 +++- web/default/src/i18n/locales/ru.json | 89 +++- web/default/src/i18n/locales/vi.json | 89 +++- web/default/src/i18n/locales/zh.json | 89 +++- web/default/src/routeTree.gen.ts | 22 + .../_authenticated/flow-pools/index.tsx | 37 ++ 20 files changed, 2750 insertions(+), 18 deletions(-) create mode 100644 web/default/src/features/channel-flow/api.ts create mode 100644 web/default/src/features/channel-flow/components/binding-form-sheet.tsx create mode 100644 web/default/src/features/channel-flow/components/pool-bindings-panel.tsx create mode 100644 web/default/src/features/channel-flow/components/pool-form-sheet.tsx create mode 100644 web/default/src/features/channel-flow/components/pool-list.tsx create mode 100644 web/default/src/features/channel-flow/components/pool-status-panel.tsx create mode 100644 web/default/src/features/channel-flow/index.tsx create mode 100644 web/default/src/features/channel-flow/lib/form.ts create mode 100644 web/default/src/features/channel-flow/lib/index.ts create mode 100644 web/default/src/features/channel-flow/lib/query-keys.ts create mode 100644 web/default/src/features/channel-flow/types.ts create mode 100644 web/default/src/routes/_authenticated/flow-pools/index.tsx diff --git a/web/default/src/features/channel-flow/api.ts b/web/default/src/features/channel-flow/api.ts new file mode 100644 index 000000000000..78e5946dd41e --- /dev/null +++ b/web/default/src/features/channel-flow/api.ts @@ -0,0 +1,119 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +import { api } from '@/lib/api' +import type { + ApiResponse, + ChannelFlowBindingPayload, + ChannelFlowPool, + ChannelFlowPoolBinding, + ChannelFlowPoolPayload, + ChannelFlowPoolStatus, + PageResponse, +} from './types' + +const channelFlowActionConfig = { + skipBusinessError: true, + skipErrorHandler: true, +} + +export type ListChannelFlowPoolsParams = { + p?: number + page_size?: number + keyword?: string +} + +export async function listChannelFlowPools( + params: ListChannelFlowPoolsParams = {} +): Promise>> { + const res = await api.get('/api/channel_flow/pools', { params }) + return res.data +} + +export async function createChannelFlowPool( + payload: ChannelFlowPoolPayload +): Promise> { + const res = await api.post( + '/api/channel_flow/pools', + payload, + channelFlowActionConfig + ) + return res.data +} + +export async function updateChannelFlowPool( + poolId: number, + payload: ChannelFlowPoolPayload +): Promise> { + const res = await api.put( + `/api/channel_flow/pools/${poolId}`, + payload, + channelFlowActionConfig + ) + return res.data +} + +export async function deleteChannelFlowPool( + poolId: number +): Promise { + const res = await api.delete( + `/api/channel_flow/pools/${poolId}`, + channelFlowActionConfig + ) + return res.data +} + +export async function getChannelFlowPoolStatus( + poolId: number +): Promise> { + const res = await api.get(`/api/channel_flow/pools/${poolId}/status`, { + disableDuplicate: true, + }) + return res.data +} + +export async function listChannelFlowPoolBindings( + poolId: number +): Promise> { + const res = await api.get(`/api/channel_flow/pools/${poolId}/bindings`) + return res.data +} + +export async function createChannelFlowPoolBinding( + poolId: number, + payload: ChannelFlowBindingPayload +): Promise> { + const res = await api.post( + `/api/channel_flow/pools/${poolId}/bindings`, + payload, + channelFlowActionConfig + ) + return res.data +} + +export async function deleteChannelFlowPoolBinding( + bindingId: number +): Promise { + const res = await api.delete( + `/api/channel_flow/bindings/${bindingId}`, + channelFlowActionConfig + ) + return res.data +} + diff --git a/web/default/src/features/channel-flow/components/binding-form-sheet.tsx b/web/default/src/features/channel-flow/components/binding-form-sheet.tsx new file mode 100644 index 000000000000..f7ddd59ac1aa --- /dev/null +++ b/web/default/src/features/channel-flow/components/binding-form-sheet.tsx @@ -0,0 +1,360 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +import { useEffect, useMemo, useState } from 'react' +import { type Resolver, useForm } from 'react-hook-form' +import { zodResolver } from '@hookform/resolvers/zod' +import { useQuery } from '@tanstack/react-query' +import { useTranslation } from 'react-i18next' +import { Check, ChevronsUpDown, Loader2 } from 'lucide-react' +import { getChannels } from '@/features/channels/api' +import { + getChannelStatusBadge, + getChannelTypeLabel, +} from '@/features/channels/lib/channel-utils' +import type { Channel } from '@/features/channels/types' +import { cn } from '@/lib/utils' +import { Button } from '@/components/ui/button' +import { Badge } from '@/components/ui/badge' +import { + Command, + CommandEmpty, + CommandGroup, + CommandInput, + CommandItem, + CommandList, +} from '@/components/ui/command' +import { + Form, + FormControl, + FormDescription, + FormField, + FormItem, + FormLabel, + FormMessage, +} from '@/components/ui/form' +import { + Popover, + PopoverContent, + PopoverTrigger, +} from '@/components/ui/popover' +import { + Sheet, + SheetContent, + SheetDescription, + SheetFooter, + SheetHeader, + SheetTitle, +} from '@/components/ui/sheet' +import { Switch } from '@/components/ui/switch' +import { + channelFlowBindingFormSchema, + defaultBindingFormValues, + type ChannelFlowBindingFormValues, +} from '../lib' +import type { ChannelFlowPool, ChannelFlowPoolBinding } from '../types' + +type BindingFormSheetProps = { + open: boolean + onOpenChange: (open: boolean) => void + pool?: ChannelFlowPool | null + bindings: ChannelFlowPoolBinding[] + submitting: boolean + onSubmit: (values: ChannelFlowBindingFormValues) => void +} + +const CHANNEL_SELECTOR_PAGE_SIZE = 200 + +export function BindingFormSheet(props: BindingFormSheetProps) { + const { t } = useTranslation() + const form = useForm({ + resolver: zodResolver( + channelFlowBindingFormSchema + ) as unknown as Resolver, + defaultValues: defaultBindingFormValues, + }) + const selectedChannelId = form.watch('channel_id') + const channelsQuery = useQuery({ + queryKey: [ + 'channel-flow', + 'binding-channel-options', + CHANNEL_SELECTOR_PAGE_SIZE, + ], + queryFn: () => + getChannels({ + p: 1, + page_size: CHANNEL_SELECTOR_PAGE_SIZE, + id_sort: true, + }), + enabled: props.open, + }) + + const channels = channelsQuery.data?.data?.items ?? [] + const availableChannels = useMemo(() => { + const boundChannelIds = new Set( + props.bindings + .filter((binding) => binding.enabled) + .map((binding) => binding.channel_id) + ) + return channels.filter((channel) => !boundChannelIds.has(channel.id)) + }, [channels, props.bindings]) + + useEffect(() => { + if (!props.open) return + form.reset(defaultBindingFormValues) + }, [form, props.open]) + + return ( + + + + {t('Bind channel')} + + {props.pool + ? t('Pool: {{name}}', { name: props.pool.name }) + : t('Select a Flow Pool first')} + + + +
+ + ( + +
+ {t('Enabled')} + + {t('Disabled bindings are retained but ignored by routing.')} + +
+ + + + + )} + /> + + ( + + {t('Channel')} + + + + + {t( + 'The channel keeps its own upstream Base URL and model mapping; this binding only attaches pool capacity to that channel.' + )} + + + + )} + /> + + ( + + {t('Binding mode')} + +
+ {t('Channel')} + Phase 1 + +
+ + + {t('Phase 1 supports channel-level binding only.')} + + + + )} + /> + + + + + + +
+ + ) +} + +type ChannelPickerProps = { + channels: Channel[] + loading: boolean + value: number + onValueChange: (value: number) => void +} + +function ChannelPicker(props: ChannelPickerProps) { + const { t } = useTranslation() + const [open, setOpen] = useState(false) + const [searchValue, setSearchValue] = useState('') + const selectedChannel = props.channels.find( + (channel) => channel.id === props.value + ) + + const filteredChannels = useMemo(() => { + const search = searchValue.trim().toLowerCase() + if (!search) return props.channels + + return props.channels.filter((channel) => { + const typeLabel = t(getChannelTypeLabel(channel.type)).toLowerCase() + return [ + String(channel.id), + channel.name, + channel.base_url || '', + channel.models || '', + typeLabel, + ].some((value) => value.toLowerCase().includes(search)) + }) + }, [props.channels, searchValue, t]) + + const handleSelect = (channelId: number) => { + props.onValueChange(channelId) + setOpen(false) + setSearchValue('') + } + + return ( + + + } + > + {selectedChannel ? ( + + ) : ( + {t('Channel')} + )} + + + event.stopPropagation()} + onTouchMove={(event) => event.stopPropagation()} + onPointerDown={(event) => event.stopPropagation()} + > + + + + + {props.loading ? t('Loading') : t('No Channels Found')} + + + {filteredChannels.map((channel) => ( + handleSelect(channel.id)} + className='data-[selected=true]:bg-muted items-start gap-3 rounded-lg px-3 py-3 transition-colors' + > + + + + ))} + + + + + + ) +} + +function ChannelOptionContent({ + channel, + compact = false, +}: { + channel: Channel + compact?: boolean +}) { + const { t } = useTranslation() + const status = getChannelStatusBadge(channel.status) + const typeLabel = t(getChannelTypeLabel(channel.type)) + + return ( + + + {channel.name} + + #{channel.id} + + + + {typeLabel} + {!compact && {t(status.label)}} + {channel.base_url && ( + + {t('Base URL')}: {channel.base_url} + + )} + + + ) +} diff --git a/web/default/src/features/channel-flow/components/pool-bindings-panel.tsx b/web/default/src/features/channel-flow/components/pool-bindings-panel.tsx new file mode 100644 index 000000000000..edd24cd152ac --- /dev/null +++ b/web/default/src/features/channel-flow/components/pool-bindings-panel.tsx @@ -0,0 +1,135 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +import { useTranslation } from 'react-i18next' +import { Plus, Trash2 } from 'lucide-react' +import { Badge } from '@/components/ui/badge' +import { Button } from '@/components/ui/button' +import { Skeleton } from '@/components/ui/skeleton' +import { StaticDataTable } from '@/components/data-table/static/static-data-table' +import type { ChannelFlowPool, ChannelFlowPoolBinding } from '../types' + +type PoolBindingsPanelProps = { + pool?: ChannelFlowPool | null + bindings: ChannelFlowPoolBinding[] + loading: boolean + deletingBindingId?: number | null + onAddBinding: () => void + onDeleteBinding: (binding: ChannelFlowPoolBinding) => void +} + +export function PoolBindingsPanel(props: PoolBindingsPanelProps) { + const { t } = useTranslation() + + return ( +
+
+
+

{t('Channel bindings')}

+

+ {t('Bindings attach pool capacity to channels; upstream URLs remain configured on each channel.')} +

+
+ +
+ + {props.loading ? ( +
+ {Array.from({ length: 3 }).map((_, index) => ( + + ))} +
+ ) : ( + binding.id} + columns={[ + { + id: 'channel', + header: t('Channel'), + cell: (binding) => ( +
+
+ #{binding.channel_id} +
+
+ {t('Pool ID')} #{binding.pool_id} +
+
+ ), + }, + { + id: 'mode', + header: t('Mode'), + className: 'hidden sm:table-cell', + cellClassName: 'hidden sm:table-cell', + cell: (binding) => ( + + {binding.match_mode === 'channel_model' + ? t('Channel and model') + : t('Channel')} + + ), + }, + { + id: 'enabled', + header: t('Status'), + cell: (binding) => ( + + {binding.enabled ? t('Enabled') : t('Disabled')} + + ), + }, + { + id: 'actions', + header: '', + className: 'w-16 text-right', + cellClassName: 'text-right', + cell: (binding) => ( + + ), + }, + ]} + emptyContent={ + + {props.pool + ? t('No channels bound to this Flow Pool') + : t('Select a Flow Pool to view bindings')} + + } + /> + )} +
+ ) +} diff --git a/web/default/src/features/channel-flow/components/pool-form-sheet.tsx b/web/default/src/features/channel-flow/components/pool-form-sheet.tsx new file mode 100644 index 000000000000..db912ca58068 --- /dev/null +++ b/web/default/src/features/channel-flow/components/pool-form-sheet.tsx @@ -0,0 +1,449 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +import { useEffect } from 'react' +import { type Resolver, type UseFormReturn, useForm } from 'react-hook-form' +import { zodResolver } from '@hookform/resolvers/zod' +import { useTranslation } from 'react-i18next' +import { Loader2 } from 'lucide-react' +import { Button } from '@/components/ui/button' +import { + Form, + FormControl, + FormDescription, + FormField, + FormItem, + FormLabel, + FormMessage, +} from '@/components/ui/form' +import { Input } from '@/components/ui/input' +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from '@/components/ui/select' +import { + Sheet, + SheetContent, + SheetDescription, + SheetFooter, + SheetHeader, + SheetTitle, +} from '@/components/ui/sheet' +import { Switch } from '@/components/ui/switch' +import { Textarea } from '@/components/ui/textarea' +import { + channelFlowPoolFormSchema, + defaultPoolFormValues, + poolToFormValues, + type ChannelFlowPoolFormValues, +} from '../lib' +import type { ChannelFlowPool } from '../types' + +type PoolFormSheetProps = { + open: boolean + onOpenChange: (open: boolean) => void + pool?: ChannelFlowPool | null + submitting: boolean + onSubmit: (values: ChannelFlowPoolFormValues) => void +} + +const numberFields = [ + 'max_inflight', + 'max_queue_size', + 'max_queue_per_user', + 'queue_timeout_ms', + 'max_context_tokens', + 'max_context_chars', + 'max_processing_ms', + 'lease_ms', + 'renew_interval_ms', +] as const + +type SelectOption = { + value: T + label: string +} + +function getOptionLabel( + options: SelectOption[], + value: T +) { + return options.find((option) => option.value === value)?.label ?? value +} + +export function PoolFormSheet(props: PoolFormSheetProps) { + const { t } = useTranslation() + const form = useForm({ + resolver: zodResolver( + channelFlowPoolFormSchema + ) as unknown as Resolver, + defaultValues: defaultPoolFormValues, + }) + const backend = form.watch('backend') + const isEditMode = Boolean(props.pool?.id) + const backendOptions: SelectOption[] = [ + { value: 'memory', label: t('Memory') }, + { value: 'redis', label: t('Redis (experimental)') }, + ] + const onLimitOptions: SelectOption[] = [ + { value: 'queue', label: t('Queue') }, + { value: 'reject', label: t('Reject') }, + { value: 'fallback', label: t('Fallback') }, + ] + const queuePolicyOptions: SelectOption< + ChannelFlowPoolFormValues['queue_policy'] + >[] = [{ value: 'fifo', label: t('FIFO') }] + const redisFailurePolicyOptions: SelectOption< + ChannelFlowPoolFormValues['redis_failure_policy'] + >[] = [ + { value: 'fail_open', label: t('Fail open') }, + { value: 'fail_closed', label: t('Fail closed') }, + { value: 'local_memory', label: t('Local memory fallback') }, + ] + + useEffect(() => { + if (!props.open) return + form.reset(poolToFormValues(props.pool)) + }, [form, props.open, props.pool]) + + return ( + + + + + {isEditMode ? t('Edit Flow Pool') : t('Create Flow Pool')} + + + {t( + 'Flow Pools cap total upstream concurrency and keep excess requests in a bounded queue.' + )} + + + +
+ + ( + +
+ {t('Enabled')} + + {t('Disabled pools keep their bindings but do not gate traffic.')} + +
+ + + + + )} + /> + +
+ ( + + {t('Pool name')} + + + + + + )} + /> + + ( + + {t('Backend')} + + + + )} + /> +
+ + ( + + {t('Description')} + +