diff --git a/.gitignore b/.gitignore index 75f5c463387..c8d1d20c85d 100644 --- a/.gitignore +++ b/.gitignore @@ -33,6 +33,7 @@ data/ .gocache-temp .gopath .test +*.test token_estimator_test.go skills-lock.json .playwright-mcp diff --git a/controller/channel_flow.go b/controller/channel_flow.go new file mode 100644 index 00000000000..2d13b0a5e9b --- /dev/null +++ b/controller/channel_flow.go @@ -0,0 +1,338 @@ +package controller + +import ( + "errors" + "fmt" + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + channelflowmetrics "github.com/QuantumNous/new-api/pkg/channel_flow_metrics" + "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"` + MaxInflightPerUser int `json:"max_inflight_per_user"` + 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"` + ScheduleMode string `json:"schedule_mode"` + ScheduleTimezone string `json:"schedule_timezone"` + EffectiveStartTime int64 `json:"effective_start_time"` + EffectiveEndTime int64 `json:"effective_end_time"` + ScheduleWindows string `json:"schedule_windows"` +} + +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 GetChannelFlowPoolTrend(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 + } + hours := 6 + if rawHours := c.Query("hours"); rawHours != "" { + if parsed, parseErr := strconv.Atoi(rawHours); parseErr == nil { + hours = parsed + } + } + minutes := 0 + if rawMinutes := c.Query("minutes"); rawMinutes != "" { + if parsed, parseErr := strconv.Atoi(rawMinutes); parseErr == nil { + minutes = parsed + } + } + trend, err := channelflowmetrics.Query(channelflowmetrics.QueryParams{ + PoolKey: pool.PoolKey, + Hours: hours, + Minutes: minutes, + }) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, trend) +} + +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.MaxInflightPerUser = req.MaxInflightPerUser + 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.ScheduleMode = req.ScheduleMode + pool.ScheduleTimezone = req.ScheduleTimezone + pool.EffectiveStartTime = req.EffectiveStartTime + pool.EffectiveEndTime = req.EffectiveEndTime + pool.ScheduleWindows = req.ScheduleWindows + pool.Normalize() + return pool +} diff --git a/controller/channel_flow_test.go b/controller/channel_flow_test.go new file mode 100644 index 00000000000..02eda2ed65a --- /dev/null +++ b/controller/channel_flow_test.go @@ -0,0 +1,21 @@ +package controller + +import ( + "testing" + + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/require" +) + +func TestChannelFlowPoolFromRequestIncludesMaxInflightPerUser(t *testing.T) { + pool := channelFlowPoolFromRequest(channelFlowPoolRequest{ + Name: "fair pool", + Backend: model.ChannelFlowBackendMemory, + MaxInflight: 4, + MaxInflightPerUser: 2, + QueuePolicy: model.ChannelFlowQueuePolicyFIFO, + OnLimit: model.ChannelFlowOnLimitQueue, + }, nil) + + require.Equal(t, 2, pool.MaxInflightPerUser, "max_inflight_per_user should be copied from request") +} diff --git a/controller/model_list_test.go b/controller/model_list_test.go index 97d27cae5c6..3a066784073 100644 --- a/controller/model_list_test.go +++ b/controller/model_list_test.go @@ -224,6 +224,7 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { recorder := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(recorder) ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true) common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{ "zz-token-tiered-visible-model": true, diff --git a/controller/relay.go b/controller/relay.go index 1d14dcc6f88..3325a0f452f 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,47 @@ 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 + } + defer func(guard service.FlowGuard) { + if r := recover(); r != nil { + if guard != nil { + _ = guard.Release(context.Background()) + } + panic(r) + } + }(flowGuard) + + 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 +260,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = relayHandler(c, relayInfo) } + if flowGuard != nil { + service.RecordChannelFlowOutcome(flowGuard, channel.Id, relayInfo, newAPIError == nil) + _ = flowGuard.Release(context.Background()) + } + if newAPIError == nil { relayInfo.LastError = nil return diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index e75befaeea8..fd345528830 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -54,6 +54,8 @@ services: POSTGRES_USER: root POSTGRES_PASSWORD: 123456 POSTGRES_DB: new-api + ports: + - "127.0.0.1:5432:5432" volumes: - dev_pg_data:/var/lib/postgresql/data networks: 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 00000000000..592d83c4e84 --- /dev/null +++ b/docs/channel-flow-control-queue-design-v2.md @@ -0,0 +1,1200 @@ +# 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` +- `docs/flow-control-design-audit.md` (internal audit document) + +## 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 00000000000..80385555d48 --- /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` +- `docs/flow-control-v2-review.md` (internal audit document) + +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` +- `../boom-gateway/boom-flowcontrol/src/lib.rs` (gateway reference) +- `../boom-gateway/boom-routing/src/policy/load_helpers.rs` (gateway reference) + +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 00000000000..1e88f770877 --- /dev/null +++ b/docs/channel-flow-control-queue-design.md @@ -0,0 +1,1477 @@ +# Channel Flow Control and Queue Design Report + +> **Superseded** — This document is historical context only. The implementation follows [channel-flow-control-queue-design-v3.md](./channel-flow-control-queue-design-v3.md). + +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 gateway project at `../boom-gateway/` implements a closely related flow control pattern. + +Relevant files: + +- `../boom-gateway/config.example.yaml` +- `../boom-gateway/boom-config/src/lib.rs` +- `../boom-gateway/boom-flowcontrol/src/lib.rs` +- `../boom-gateway/boom-main/src/routes.rs` +- `../boom-gateway/boom-main/src/state.rs` +- `../boom-gateway/boom-routing/src/policy/load_helpers.rs` +- `../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/docs/channel-flow-next-phase-plan.md b/docs/channel-flow-next-phase-plan.md new file mode 100644 index 00000000000..4ef787b85c1 --- /dev/null +++ b/docs/channel-flow-next-phase-plan.md @@ -0,0 +1,798 @@ +# 渠道流控与排队功能后续阶段交接计划 + +日期：2026-06-15 + +状态：Phase 1 已实现后的后续任务交接文档 + +当前分支： + +```text +codex/flow-pool-scheduling-controls +``` + +当前已推送基线： + +```text +57437580 feat: add flow pool scheduling controls +``` + +远程分支： + +```text +https://github.com/supreme0597/new-api/tree/codex/flow-pool-scheduling-controls +``` + +相关设计文档： + +- `docs/channel-flow-control-queue-design-v3.md` +- `docs/channel-flow-control-queue-design-v2.md` +- `docs/channel-flow-control-queue-design.md` + +## 1. 背景与目标 + +本功能解决的是“上游资源池容量保护”，不是普通用户 RPM/TPM 限速。 + +典型场景： + +```text +一个上游模型资源池背后有 96 张卡。 +上游最多只能稳定支持 60 个并发请求。 +第 61 个请求不应该继续打到上游。 +网关需要根据配置进行排队、拒绝或 fallback。 +管理员需要看到实时在途、排队、拒绝、超时、成功等指标，并能追溯趋势。 +``` + +核心目标： + +- 限制同一个上游资源池的总在途请求数。 +- 超过容量时进入有界队列。 +- 队列长度必须有上限。 +- 支持单用户最大排队数，避免单个用户占满队列。 +- 支持 Redis 后端，保证多实例部署下全局限流一致。 +- 支持流控池生效时间，方便按日期范围或周计划启停。 +- 支持实时状态和历史趋势图，便于排查流量问题。 + +## 2. 当前已完成内容 + +当前分支已经包含第一版可用实现。 + +### 2.1 后端 + +已实现文件： + +- `model/channel_flow.go` + - Flow Pool 数据模型。 + - 渠道绑定模型。 + - 分钟级指标模型。 + - 事件模型。 + - 生效时间字段、校验和判断逻辑。 +- `controller/channel_flow.go` + - Flow Pool 管理接口。 + - 渠道绑定管理接口。 + - 状态接口。 + - 趋势接口。 +- `service/channel_flow.go` + - Flow Controller 抽象。 + - Memory backend。 + - FlowGuard 生命周期。 + - Redis 不可用策略。 + - Lease 续租 hook。 + - 指标记录。 +- `service/channel_flow_redis.go` + - Redis backend 初版实现。 + - acquire、enqueue、release、status、renew lease。 +- `service/channel_flow_status_sampler.go` + - 定时采样流控池实时状态，用于趋势图。 +- `pkg/channel_flow_metrics` + - 指标聚合、分钟桶写入和查询。 +- `controller/relay.go` + - 渠道选择后获取 FlowGuard。 + - 排队前只做 billing precheck。 + - acquire 成功后再做 billing reserve/preconsume。 + - 单次请求结束后记录 outcome 并 release guard。 + +### 2.2 前端 + +已实现目录： + +- `web/default/src/features/channel-flow` + - Flow Pool 列表页。 + - 创建/编辑 Flow Pool 表单。 + - 渠道绑定面板。 + - 实时状态面板。 + - 趋势图面板。 + +当前支持的生效时间： + +- 始终生效。 +- 日期范围。 +- 每周重复，支持跨午夜窗口。 + +当前左侧卡片展示： + +- 当前生效 / 当前未生效。 +- 生效时间摘要。 +- 后端类型。 +- 满载策略。 +- 紧凑容量标签，例如 `容量 60+240`。 + +当前状态刷新选项： + +- 关闭。 +- 1 秒。 +- 2 秒。 +- 5 秒。 +- 10 秒。 +- 30 秒。 + +不提供 500 ms 轮询。后续如果需要更实时状态，建议用 SSE 或 WebSocket。 + +### 2.3 已验证内容 + +提交前已执行： + +```text +bun run build +bun run i18n:sync +GOCACHE=/private/tmp/new-api-go-build go test ./model ./service ./pkg/channel_flow_metrics -count=1 +git diff --check +``` + +已手工验证： + +- 使用本地管理员账号登录。 +- Flow Pool 页面可打开。 +- 始终生效池显示当前生效。 +- 未来日期范围显示当前未生效。 +- 周一到周五 09:00-18:00 在当前窗口内显示当前生效。 +- 左侧卡片时间摘要与右侧状态一致。 +- 保存后立刻编辑，生效时间字段能正确回填。 +- 临时测试流控池已删除。 + +## 3. 必须保留的产品决策 + +后续同事继续开发时，除非产品明确变更，否则请保持以下决策。 + +### 3.1 `pool_id` 不是用户输入 + +用户和管理员不应该手动输入 runtime `pool_id`。 + +正确方式： + +```text +管理员创建 Flow Pool + -> 后端生成 pool_key + -> 管理员把渠道绑定到 Flow Pool + -> 运行时按 channel_id 解析绑定关系 +``` + +### 3.2 绑定必须显式配置 + +不要用上游 URL 自动合并流控池。 + +原因： + +- 同一个 URL 可能服务多个物理资源池。 +- 同一个 URL + 不同 key 可能对应不同租户。 +- 同一个物理池也可能有多个 URL。 +- 模型映射后真实上游模型可能变化。 + +URL 可以作为 UI 上的辅助提示，但不能作为运行时绑定依据。 + +### 3.3 同一个渠道同一时间只能属于一个启用的 Flow Pool + +当前 controller 已阻止正常 UI 创建重复启用绑定。 + +后续还需要做： + +- 脏数据检测。 +- UI 冲突提示。 +- 是否增加跨数据库兼容的唯一约束或启动健康检查。 + +### 3.4 队列长度必须有上限 + +不能做无限排队。 + +必须保留： + +- `max_queue_size` +- `queue_timeout_ms` +- `max_queue_per_user` + +### 3.5 生效时间外是绕过，而不是随机尝试另一个池 + +当前语义： + +```text +渠道绑定了 Flow Pool + -> Flow Pool 未启用或当前不在生效时间 + -> 该请求绕过这个 Flow Pool +``` + +不要在运行时随机匹配另一个 Flow Pool。否则排查会非常困难。 + +### 3.6 多实例生产必须使用 Redis + +Memory backend 只适合： + +- 本地开发。 +- 单实例部署。 +- Redis 故障时的临时 local_memory fallback。 + +生产多实例要保证全局并发上限，必须使用 Redis backend。 + +## 4. 当前已知风险与缺口 + +### 4.1 流式请求生命周期还需要专项审计 + +当前代码已经有： + +- `FlowGuard.RenewLease` +- `startChannelFlowLeaseRenewer` +- `FlowGuard.WrapReadCloser` + +但下一阶段仍必须专项审计： + +- relay helper 是否真的阻塞到下游 stream 完成。 +- WebSocket realtime 是否持有 guard 到连接结束。 +- 客户端断开时是否及时 release。 +- 上游 stream 报错时是否 release。 +- retry 前一次 attempt 是否一定 release。 +- `WrapReadCloser` 是否需要接入某些 provider path。 + +这是下一阶段 P0。 + +### 4.2 Redis backend 需要生产化验证 + +Redis backend 已有初版，但还不能直接视为生产稳定。 + +需要验证： + +- 高并发 acquire/enqueue/release 原子性。 +- FIFO 公平性。 +- 队列超时清理。 +- 单用户队列上限。 +- lease 过期清理。 +- Redis 故障策略。 +- 多实例共享 Redis 时，总在途是否严格受控。 + +Lua 不是一开始必须做。 + +建议： + +```text +先测当前 Redis 实现的冲突率和正确性。 +如果 WATCH/MULTI 冲突高，或组合操作难以证明正确，再把核心 hot path 改成 Lua。 +``` + +### 4.3 绑定冲突治理不足 + +当前 controller 能阻止正常新增重复绑定，但还缺： + +- 已存在脏数据检测。 +- 管理页冲突告警。 +- 修复入口。 +- 跨 SQLite/MySQL/PostgreSQL 的 DB 约束方案评估。 + +### 4.4 可观测性还不够完整 + +当前已有状态和趋势图，但线上排查还需要： + +- 事件明细。 +- 配置变更审计。 +- 绑定变更审计。 +- Redis 降级和 fallback 事件。 +- lease renew 失败明细。 +- 排队超时、拒绝原因拆解。 + +### 4.5 缺少压测和多实例 E2E + +必须补： + +- 单实例压测。 +- Redis 多实例压测。 +- stream 长连接压测。 +- client abort 释放验证。 +- queue full / queue timeout / per-user queue full 验证。 + +## 5. 下一阶段推荐计划 + +### Phase A：Relay 生命周期与流式请求正确性 + +优先级：P0 + +目标： + +```text +保证每个 FlowGuard 都只在真实请求生命周期内持有，并且最终只释放一次。 +``` + +任务： + +1. 审计 relay helper。 + - `relayHandler` + - `geminiRelayHandler` + - `relay.ClaudeHelper` + - `relay.WssHelper` + - OpenAI-compatible streaming + - Responses API streaming + - Claude streaming + - Gemini streaming + +2. 审计 `controller/relay.go` 的 release 路径。 + - pricing error + - billing reserve error + - request body read error + - upstream helper error + - retry + - panic/recover + - client cancel + +3. 明确 stream 生命周期。 + - 如果 helper 会阻塞到 stream 完成，则写测试或注释固化这个约定。 + - 如果 helper 会提前返回，则必须用 `FlowGuard.WrapReadCloser` 或等价方式绑定 stream close。 + +4. 验证 Redis lease 续租。 + - acquire 后启动。 + - release 后停止。 + - 续租失败记录 metric。 + - 长 stream 超过 `lease_ms` 时，slot 不应被错误释放。 + +验收标准： + +- 没有已知 guard 泄漏路径。 +- 没有已知 stream 未结束就 release 的路径。 +- `Release` 幂等。 +- 长 stream 测试通过。 +- client abort 后容量可恢复。 + +建议测试： + +```text +go test ./service -run ChannelFlow +go test ./controller -run Relay +``` + +如果 controller 测试成本过高，可以先加 fake relay lifecycle test。 + +### Phase B：Redis backend 正确性与多实例验证 + +优先级：P0 + +目标： + +```text +证明 Redis backend 能在多实例部署下限制全局容量。 +``` + +任务： + +1. 增加 Redis 集成测试。 + - 建议通过环境变量开启，避免默认测试依赖 Redis。 + +```text +CHANNEL_FLOW_REDIS_TEST=1 +REDIS_CONN_STRING=redis://localhost:6379/... +``` + +2. 高并发 acquire/enqueue 测试。 + - `max_inflight=1` + - `max_queue_size=2` + - 只允许 1 个 running。 + - 最多 2 个 waiting。 + - 第 4 个请求应被拒绝。 + +3. release promotion 测试。 + - release running 后，队首 queued 能进入 running。 + - FIFO 顺序稳定。 + - 已取消或超时的队首不会阻塞后续请求。 + +4. 单用户队列上限测试。 + - 同一个用户不能超过 `max_queue_per_user`。 + - 不同用户仍可使用全局队列剩余空间。 + +5. lease expiry 测试。 + - running 过期后能被清理。 + - 如有指标设计，记录 lease expired。 + +6. Redis 故障策略测试。 + - `fail_open` + - `fail_closed` + - `local_memory` + +7. 冲突率评估。 + - 如果当前 Redis 实现使用 WATCH/MULTI，需要统计事务冲突和重试次数。 + - 冲突过高时再考虑 Lua。 + +验收标准： + +- 多实例总 running 不超过 `max_inflight`。 +- queue 不超过 `max_queue_size`。 +- per-user queue cap 生效。 +- Redis 故障表现符合配置。 + +### Phase C：绑定冲突治理 + +优先级：P1 + +目标： + +```text +让重复绑定和歧义绑定可见、可阻止、可修复。 +``` + +任务： + +1. 增加冲突检测服务。 + - 检测同一 `channel_id + match_mode` 下多个 enabled binding。 + +2. 前端展示冲突提示。 + - Flow Pool 页面展示。 + - 绑定弹窗展示。 + - 后续可以在渠道编辑页展示。 + +3. 评估 DB 约束。 + - 跨 SQLite/MySQL/PostgreSQL 的部分唯一索引并不简单。 + - 如果不加 DB 约束，需要启动检查和 admin health check。 + +4. 优化绑定体验。 + - 继续避免用户手输渠道 ID。 + - 显示渠道名、ID、类型、base URL。 + - 支持搜索。 + +验收标准： + +- UI 不能创建重复启用绑定。 +- 已存在脏数据能被管理员看到。 +- runtime 对脏数据的行为确定且有文档说明。 + +### Phase D：观测与审计增强 + +优先级：P1 + +目标： + +```text +管理员能从页面解释一次 429、排队、超时或 Redis 降级。 +``` + +任务： + +1. 增加事件明细 API。 + - pool + - channel + - model + - user + - event type + - time range + - request id + +2. 增加事件表格。 + - 最近排队。 + - 最近拒绝。 + - 最近超时。 + - 最近 release。 + - wait_ms / process_ms。 + +3. 增加配置审计。 + - Flow Pool 创建、修改、删除。 + - 绑定创建、删除。 + - 生效时间变化。 + - 容量变化。 + +4. 完善趋势图。 + - 拒绝原因拆分。 + - timeout/cancelled/billing_failed 展示。 + - 后续可增加 p95 wait time。 + +验收标准： + +管理员能回答： + +- 为什么这个请求 429？ +- 是不是池满了？ +- 是不是不在生效时间？ +- Redis 是否降级？ +- 这个请求命中了哪个 channel 和 Flow Pool？ + +### Phase E：渠道编辑页集成 + +优先级：P2 + +目标： + +```text +管理员可以在渠道创建/编辑流程里直接管理 Flow Pool 绑定。 +``` + +任务： + +1. 渠道编辑页增加 Flow Control 区块。 + - 显示当前绑定。 + - 选择已有 Flow Pool。 + - 跳转或弹窗创建 Flow Pool。 + +2. 增加上下文提示。 + - 可以基于 base URL 提示候选池。 + - 不能自动绑定。 + - 明确说明 URL 不是流控池归属依据。 + +3. 后续再支持 channel + upstream_model 绑定。 + - 当前 Phase 1 是 channel-level binding。 + - model-level binding 等 Redis 和生命周期稳定后再做。 + +验收标准： + +- 管理员不需要输入渠道 ID。 +- 渠道页面能看到当前绑定状态。 +- 不会因为相同 URL 自动合并池。 + +### Phase F：上线准备 + +优先级：P1，上生产前必须完成 + +目标： + +```text +让运维和管理员能安全启用、观察和回滚。 +``` + +任务： + +1. 编写运维文档。 + - 推荐配置。 + - Redis backend 要求。 + - queue timeout 建议。 + - failure policy 建议。 + +2. 编写部署检查清单。 + - SQLite/MySQL/PostgreSQL migration 验证。 + - Redis 连通性。 + - 初始 pool 先 disabled 或小流量 canary。 + - 指标保留策略。 + +3. 编写灰度方案。 + - 先绑定非核心渠道。 + - 小队列。 + - 观察 reject、timeout、lease renew failure、Redis health。 + - 再扩大覆盖范围。 + +4. 编写回滚方案。 + - 禁用 Flow Pool。 + - 删除绑定。 + - 必要时调整 Redis failure policy。 + +验收标准： + +- 管理员不用读代码也能启用功能。 +- 回滚不需要直接改数据库。 + +## 6. 推荐执行顺序 + +建议后续同事按这个顺序执行： + +```text +1. Phase A：relay 生命周期和 stream guard 审计。 +2. Phase B：Redis 多实例正确性测试。 +3. Phase C：绑定冲突治理。 +4. Phase D：事件明细和审计。 +5. Phase E：渠道编辑页集成。 +6. Phase F：上线文档和灰度方案。 +``` + +不要先继续做 UI 小优化。 + +当前最大风险不是页面样式，而是： + +- guard 提前释放。 +- guard 泄漏。 +- Redis 并发竞态。 +- 多实例超发。 +- 长流式请求 lease 过期。 + +## 7. 测试矩阵 + +### 7.1 单元测试 + +已有或应扩展： + +```text +model/channel_flow_schedule_test.go +service/channel_flow_test.go +pkg/channel_flow_metrics/metrics_test.go +``` + +继续补充： + +- 周计划跨午夜。 +- 日期范围边界。 +- disabled pool bypass。 +- memory FIFO。 +- memory queue timeout。 +- memory per-user queue cap。 +- guard idempotent release。 +- metrics 聚合。 + +### 7.2 Redis 集成测试 + +需要覆盖： + +- acquire capacity。 +- enqueue capacity。 +- queue full rejection。 +- per-user queue full rejection。 +- release promotion。 +- queue timeout cleanup。 +- lease renewal。 +- lease expiry cleanup。 +- Redis outage policies。 + +### 7.3 Relay 生命周期测试 + +需要覆盖： + +- 非流式成功。 +- 非流式上游错误。 +- 流式成功。 +- 流式客户端中断。 +- 流式上游中断。 +- 第一次渠道失败后 retry。 +- acquire 后 billing 失败。 + +### 7.4 手工 UI 测试 + +本地管理员账号： + +```text +admin_user / example_password +``` + +检查项： + +- 创建池。 +- 编辑池。 +- 删除未绑定池。 +- 搜索和绑定渠道。 +- 重复绑定被阻止。 +- 始终生效。 +- 未来日期范围未生效。 +- 当前日期范围生效。 +- 每周生效窗口生效。 +- 每周非窗口未生效。 +- 刷新频率显示文字，不显示 `1000/2000`。 +- 趋势范围切换正常。 + +### 7.5 压测场景 + +基础场景： + +```text +max_inflight=1 +max_queue_size=2 +max_queue_per_user=2 +queue_timeout_ms=30000 +backend=redis +``` + +预期： + +- 第 1 个请求进入上游。 +- 第 2、3 个请求排队。 +- 第 4 个请求返回 429。 +- 第 1 个请求释放后，第 2 个请求进入 running。 +- 客户端中断后 slot 可恢复。 + +多实例场景： + +```text +启动两个 new-api 实例 +连接同一个 Redis +同时向两个实例发请求 +全局 running 不超过 max_inflight +``` + +## 8. 运行时语义 + +后续实现应保持这个语义： + +```text +请求进入 relay + -> 选择渠道 + -> 按 channel_id 解析 Flow Pool + -> 如果无绑定、disabled、或不在生效时间：绕过 flow control + -> acquire FlowGuard + -> acquire 后做 billing reserve/preconsume + -> 调用上游 + -> 记录 outcome + -> release guard +``` + +retry 语义： + +```text +每次 retry 都会重新选择渠道。 +每次 attempt 独立解析 Flow Pool。 +每次 attempt 独立 acquire 和 release。 +Flow control rejection 默认不应继续 retry，除非后续明确设计。 +``` + +billing 语义： + +```text +排队前只做只读 precheck。 +排队中不扣费。 +acquire 成功后才 reserve/preconsume。 +如果等待后扣费失败，要 release slot，并返回清晰错误。 +``` + +schedule 语义： + +```text +always: enabled 即生效。 +datetime_range: start <= now < end。 +weekly: 按配置时区判断本地时间窗口。 +weekly 跨午夜: 开始日晚上到次日早上。 +窗口外: 绕过该 Flow Pool。 +``` + +## 9. 推荐生产配置 + +60 并发上游资源池建议初始配置： + +```text +max_inflight=60 +max_queue_size=240 +max_queue_per_user=2 或 3 +queue_timeout_ms=120000 +queue_policy=fifo +on_limit=queue +backend=redis +redis_failure_policy=fail_closed +``` + +策略建议： + +- 严格保护上游容量：`fail_closed`。 +- 开发环境或低风险场景：可用 `local_memory`。 +- 谨慎使用 `fail_open`，因为 Redis 故障时可能导致上游被打爆。 + +## 10. 待确认问题 + +后续需要产品或技术负责人确认： + +1. `fail_open` 是否允许在生产 UI 中直接选择，还是作为高级选项并加警告。 +2. 重复绑定是否必须增加 DB 约束，还是通过 service + health check 管控。 +3. Flow Pool event 保留多久：7 天、30 天，还是可配置。 +4. 每周计划是否只支持一个窗口，还是支持多个窗口。 +5. `channel + upstream_model` 绑定是在 Redis 稳定前做，还是稳定后做。 +6. 实时状态后续是否改成 SSE/WebSocket。 + +## 11. 接手检查清单 + +下一位同事开始前建议先做： + +```text +git status --short --branch +git log -1 --oneline +``` + +阅读： + +- `docs/channel-flow-control-queue-design-v3.md` +- `docs/channel-flow-next-phase-plan.md` + +跑基线测试： + +```text +GOCACHE=/private/tmp/new-api-go-build go test ./model ./service ./pkg/channel_flow_metrics -count=1 +cd web/default && bun run build +cd web/default && bun run i18n:sync +``` + +第一个实际任务建议从这里开始： + +```text +审计 controller/relay.go 和所有 streaming relay helper。 +证明 FlowGuard 是否持有到真实 stream 结束。 +如果不能证明，就补 lifecycle wrapper 或测试。 +``` diff --git a/main.go b/main.go index 3361b8ce933..85049f5d408 100644 --- a/main.go +++ b/main.go @@ -19,6 +19,7 @@ import ( "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/oauth" + channelflowmetrics "github.com/QuantumNous/new-api/pkg/channel_flow_metrics" perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics" "github.com/QuantumNous/new-api/relay" "github.com/QuantumNous/new-api/router" @@ -308,6 +309,8 @@ func InitResources() error { } perfmetrics.Init() + channelflowmetrics.Init() + service.StartChannelFlowStatusSampler() // 启动系统监控 common.StartSystemMonitor() diff --git a/model/channel_flow.go b/model/channel_flow.go new file mode 100644 index 00000000000..22ee0b79ec7 --- /dev/null +++ b/model/channel_flow.go @@ -0,0 +1,547 @@ +package model + +import ( + "fmt" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +const ( + ChannelFlowBackendMemory = "memory" + ChannelFlowBackendRedis = "redis" + + ChannelFlowQueuePolicyFIFO = "fifo" + + ChannelFlowOnLimitQueue = "queue" + ChannelFlowOnLimitReject = "reject" + ChannelFlowOnLimitFallback = "fallback" + + ChannelFlowScheduleAlways = "always" + ChannelFlowScheduleDateTimeRange = "datetime_range" + ChannelFlowScheduleWeekly = "weekly" + + ChannelFlowRedisFailureFailOpen = "fail_open" + ChannelFlowRedisFailureFailClosed = "fail_closed" + ChannelFlowRedisFailureLocalMemory = "local_memory" + + ChannelFlowMatchModeChannel = "channel" + ChannelFlowMatchModeChannelModel = "channel_model" + + ChannelFlowEventQueued = "queued" + ChannelFlowEventAcquired = "acquired" + ChannelFlowEventSucceeded = "succeeded" + ChannelFlowEventFailed = "failed" + ChannelFlowEventReleased = "released" + ChannelFlowEventRejected = "rejected" + ChannelFlowEventTimeout = "timeout" + ChannelFlowEventCancelled = "cancelled" + ChannelFlowEventBillingFailed = "billing_failed" + ChannelFlowEventLeaseRenewFailed = "lease_renew_failed" + ChannelFlowEventLeaseExpired = "lease_expired" + ChannelFlowEventStatusSample = "status_sample" +) + +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"` + Backend string `json:"backend" gorm:"type:varchar(32);default:'memory'"` + MaxInflight int `json:"max_inflight" gorm:"default:0"` + MaxInflightPerUser int `json:"max_inflight_per_user" 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"` + ScheduleMode string `json:"schedule_mode" gorm:"type:varchar(32);default:'always'"` + ScheduleTimezone string `json:"schedule_timezone" gorm:"type:varchar(64);default:'Asia/Shanghai'"` + EffectiveStartTime int64 `json:"effective_start_time" gorm:"bigint;default:0"` + EffectiveEndTime int64 `json:"effective_end_time" gorm:"bigint;default:0"` + ScheduleWindows string `json:"schedule_windows" gorm:"type:text"` + ConfigVersion int64 `json:"config_version" gorm:"bigint;default:1"` + CreatedTime int64 `json:"created_time" gorm:"bigint"` + UpdatedTime int64 `json:"updated_time" gorm:"bigint"` +} + +type ChannelFlowScheduleWindow struct { + Weekdays []int `json:"weekdays"` + StartMinute int `json:"start_minute"` + EndMinute int `json:"end_minute"` +} + +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"` + 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;uniqueIndex:idx_channel_flow_metric_bucket,priority:1;index"` + PoolKey string `json:"pool_key" gorm:"type:varchar(64);uniqueIndex:idx_channel_flow_metric_bucket,priority:2;index"` + ChannelId int `json:"channel_id" gorm:"uniqueIndex:idx_channel_flow_metric_bucket,priority:3;index"` + Model string `json:"model" gorm:"type:varchar(191);uniqueIndex:idx_channel_flow_metric_bucket,priority:4;index"` + SampleCount int64 `json:"-" gorm:"bigint;default:0"` + RunningSum int64 `json:"-" gorm:"bigint;default:0"` + RunningAvg float64 `json:"running_avg"` + RunningMax int `json:"running_max"` + QueuedSum int64 `json:"-" gorm:"bigint;default:0"` + QueuedAvg float64 `json:"queued_avg"` + QueuedMax int `json:"queued_max"` + AcquiredCount int `json:"acquired_count"` + QueuedCount int `json:"queued_count"` + SucceededCount int `json:"succeeded_count"` + FailedCount int `json:"failed_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"` + WaitMsSum int64 `json:"-" gorm:"bigint;default:0"` + WaitSampleCount int64 `json:"-" gorm:"bigint;default:0"` + WaitMsAvg int64 `json:"wait_ms_avg" gorm:"bigint"` + WaitMsMax int64 `json:"wait_ms_max" gorm:"bigint"` + ProcessMsSum int64 `json:"-" gorm:"bigint;default:0"` + ProcessSampleCount int64 `json:"-" gorm:"bigint;default:0"` + 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) + p.ScheduleMode = strings.TrimSpace(p.ScheduleMode) + p.ScheduleTimezone = strings.TrimSpace(p.ScheduleTimezone) + p.ScheduleWindows = strings.TrimSpace(p.ScheduleWindows) + 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.ScheduleMode == "" { + p.ScheduleMode = ChannelFlowScheduleAlways + } + if p.ScheduleTimezone == "" { + p.ScheduleTimezone = "Asia/Shanghai" + } + 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.MaxInflightPerUser < 0 || p.MaxQueueSize < 0 || p.MaxQueuePerUser < 0 { + return fmt.Errorf("flow pool limits cannot be negative") + } + if p.MaxInflightPerUser > 0 && p.MaxInflight > 0 && p.MaxInflightPerUser > p.MaxInflight { + return fmt.Errorf("max_inflight_per_user cannot exceed max_inflight") + } + 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) + } + switch p.ScheduleMode { + case ChannelFlowScheduleAlways: + case ChannelFlowScheduleDateTimeRange: + if p.EffectiveStartTime <= 0 || p.EffectiveEndTime <= 0 { + return fmt.Errorf("effective_start_time and effective_end_time are required for datetime_range schedule") + } + if p.EffectiveEndTime <= p.EffectiveStartTime { + return fmt.Errorf("effective_end_time must be after effective_start_time") + } + case ChannelFlowScheduleWeekly: + if _, err := p.ScheduleLocation(); err != nil { + return err + } + windows, err := p.ParseScheduleWindows() + if err != nil { + return err + } + if len(windows) == 0 { + return fmt.Errorf("schedule_windows is required for weekly schedule") + } + default: + return fmt.Errorf("invalid flow pool schedule_mode: %s", p.ScheduleMode) + } + return nil +} + +func (p *ChannelFlowPool) ScheduleLocation() (*time.Location, error) { + p.Normalize() + loc, err := time.LoadLocation(p.ScheduleTimezone) + if err != nil { + return nil, fmt.Errorf("invalid flow pool schedule_timezone: %s", p.ScheduleTimezone) + } + return loc, nil +} + +func (p *ChannelFlowPool) ParseScheduleWindows() ([]ChannelFlowScheduleWindow, error) { + p.Normalize() + if p.ScheduleWindows == "" { + return nil, nil + } + var windows []ChannelFlowScheduleWindow + if err := common.UnmarshalJsonStr(p.ScheduleWindows, &windows); err != nil { + return nil, fmt.Errorf("invalid flow pool schedule_windows: %w", err) + } + for _, window := range windows { + if len(window.Weekdays) == 0 { + return nil, fmt.Errorf("schedule window weekdays cannot be empty") + } + for _, weekday := range window.Weekdays { + if weekday < 0 || weekday > 6 { + return nil, fmt.Errorf("schedule window weekday must be between 0 and 6") + } + } + if window.StartMinute < 0 || window.StartMinute > 1439 { + return nil, fmt.Errorf("schedule window start_minute must be between 0 and 1439") + } + if window.EndMinute < 1 || window.EndMinute > 1440 { + return nil, fmt.Errorf("schedule window end_minute must be between 1 and 1440") + } + if window.StartMinute == window.EndMinute { + return nil, fmt.Errorf("schedule window start_minute and end_minute cannot be equal") + } + } + return windows, nil +} + +func (p *ChannelFlowPool) IsScheduleActiveAt(now time.Time) bool { + p.Normalize() + switch p.ScheduleMode { + case ChannelFlowScheduleAlways: + return true + case ChannelFlowScheduleDateTimeRange: + nowUnix := now.Unix() + return p.EffectiveStartTime <= nowUnix && nowUnix < p.EffectiveEndTime + case ChannelFlowScheduleWeekly: + loc, err := p.ScheduleLocation() + if err != nil { + return false + } + windows, err := p.ParseScheduleWindows() + if err != nil { + return false + } + localNow := now.In(loc) + for _, window := range windows { + if scheduleWindowContains(window, localNow) { + return true + } + } + } + return false +} + +func scheduleWindowContains(window ChannelFlowScheduleWindow, now time.Time) bool { + currentMinute := now.Hour()*60 + now.Minute() + currentWeekday := int(now.Weekday()) + if window.StartMinute < window.EndMinute { + return weekdayInSchedule(window.Weekdays, currentWeekday) && + currentMinute >= window.StartMinute && + currentMinute < window.EndMinute + } + previousWeekday := currentWeekday - 1 + if previousWeekday < 0 { + previousWeekday = 6 + } + return (weekdayInSchedule(window.Weekdays, currentWeekday) && currentMinute >= window.StartMinute) || + (weekdayInSchedule(window.Weekdays, previousWeekday) && currentMinute < window.EndMinute) +} + +func weekdayInSchedule(weekdays []int, target int) bool { + for _, weekday := range weekdays { + if weekday == target { + return true + } + } + return false +} + +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 ListEnabledChannelFlowPools() ([]*ChannelFlowPool, error) { + var pools []*ChannelFlowPool + err := DB.Where("enabled = ?", true).Order("id ASC").Find(&pools).Error + return pools, err +} + +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 +} + +func InsertChannelFlowEvent(event *ChannelFlowEvent) error { + if event == nil { + return nil + } + now := time.Now().Unix() + if event.CreatedTime == 0 { + event.CreatedTime = now + } + return DB.Create(event).Error +} + +func UpsertChannelFlowMetricMinute(delta *ChannelFlowMetricMinute) error { + if delta == nil || delta.PoolKey == "" || delta.BucketTs <= 0 { + return nil + } + now := time.Now().Unix() + if delta.CreatedTime == 0 { + delta.CreatedTime = now + } + delta.UpdatedTime = now + delta.recalculateAverages() + + table := "channel_flow_metric_minutes" + return DB.Clauses(clause.OnConflict{ + Columns: []clause.Column{ + {Name: "bucket_ts"}, + {Name: "pool_key"}, + {Name: "channel_id"}, + {Name: "model"}, + }, + DoUpdates: clause.Assignments(map[string]interface{}{ + "sample_count": gorm.Expr(table+".sample_count + ?", delta.SampleCount), + "running_sum": gorm.Expr(table+".running_sum + ?", delta.RunningSum), + "running_avg": gorm.Expr("CASE WHEN "+table+".sample_count + ? > 0 THEN ("+table+".running_sum + ?) * 1.0 / ("+table+".sample_count + ?) ELSE 0 END", delta.SampleCount, delta.RunningSum, delta.SampleCount), + "running_max": gorm.Expr("CASE WHEN "+table+".running_max > ? THEN "+table+".running_max ELSE ? END", delta.RunningMax, delta.RunningMax), + "queued_sum": gorm.Expr(table+".queued_sum + ?", delta.QueuedSum), + "queued_avg": gorm.Expr("CASE WHEN "+table+".sample_count + ? > 0 THEN ("+table+".queued_sum + ?) * 1.0 / ("+table+".sample_count + ?) ELSE 0 END", delta.SampleCount, delta.QueuedSum, delta.SampleCount), + "queued_max": gorm.Expr("CASE WHEN "+table+".queued_max > ? THEN "+table+".queued_max ELSE ? END", delta.QueuedMax, delta.QueuedMax), + "acquired_count": gorm.Expr(table+".acquired_count + ?", delta.AcquiredCount), + "queued_count": gorm.Expr(table+".queued_count + ?", delta.QueuedCount), + "succeeded_count": gorm.Expr(table+".succeeded_count + ?", delta.SucceededCount), + "failed_count": gorm.Expr(table+".failed_count + ?", delta.FailedCount), + "released_count": gorm.Expr(table+".released_count + ?", delta.ReleasedCount), + "rejected_count": gorm.Expr(table+".rejected_count + ?", delta.RejectedCount), + "timeout_count": gorm.Expr(table+".timeout_count + ?", delta.TimeoutCount), + "cancelled_count": gorm.Expr(table+".cancelled_count + ?", delta.CancelledCount), + "billing_failed_count": gorm.Expr(table+".billing_failed_count + ?", delta.BillingFailedCount), + "lease_renew_fail": gorm.Expr(table+".lease_renew_fail + ?", delta.LeaseRenewFail), + "lease_expired_count": gorm.Expr(table+".lease_expired_count + ?", delta.LeaseExpiredCount), + "wait_ms_sum": gorm.Expr(table+".wait_ms_sum + ?", delta.WaitMsSum), + "wait_sample_count": gorm.Expr(table+".wait_sample_count + ?", delta.WaitSampleCount), + "wait_ms_avg": gorm.Expr("CASE WHEN "+table+".wait_sample_count + ? > 0 THEN ("+table+".wait_ms_sum + ?) / ("+table+".wait_sample_count + ?) ELSE 0 END", delta.WaitSampleCount, delta.WaitMsSum, delta.WaitSampleCount), + "wait_ms_max": gorm.Expr("CASE WHEN "+table+".wait_ms_max > ? THEN "+table+".wait_ms_max ELSE ? END", delta.WaitMsMax, delta.WaitMsMax), + "process_ms_sum": gorm.Expr(table+".process_ms_sum + ?", delta.ProcessMsSum), + "process_sample_count": gorm.Expr(table+".process_sample_count + ?", delta.ProcessSampleCount), + "process_ms_avg": gorm.Expr("CASE WHEN "+table+".process_sample_count + ? > 0 THEN ("+table+".process_ms_sum + ?) / ("+table+".process_sample_count + ?) ELSE 0 END", delta.ProcessSampleCount, delta.ProcessMsSum, delta.ProcessSampleCount), + "process_ms_max": gorm.Expr("CASE WHEN "+table+".process_ms_max > ? THEN "+table+".process_ms_max ELSE ? END", delta.ProcessMsMax, delta.ProcessMsMax), + "updated_time": now, + }), + }).Create(delta).Error +} + +func GetChannelFlowMetricMinutes(poolKey string, startTs int64, endTs int64) ([]ChannelFlowMetricMinute, error) { + var metrics []ChannelFlowMetricMinute + err := DB.Model(&ChannelFlowMetricMinute{}). + Where("pool_key = ? AND bucket_ts >= ? AND bucket_ts <= ?", poolKey, startTs, endTs). + Order("bucket_ts ASC"). + Find(&metrics).Error + return metrics, err +} + +func DeleteChannelFlowMetricMinutesBefore(cutoffTs int64) error { + if cutoffTs <= 0 { + return nil + } + return DB.Where("bucket_ts < ?", cutoffTs).Delete(&ChannelFlowMetricMinute{}).Error +} + +func (m *ChannelFlowMetricMinute) recalculateAverages() { + if m == nil { + return + } + if m.SampleCount > 0 { + m.RunningAvg = float64(m.RunningSum) / float64(m.SampleCount) + m.QueuedAvg = float64(m.QueuedSum) / float64(m.SampleCount) + } + if m.WaitSampleCount > 0 { + m.WaitMsAvg = m.WaitMsSum / m.WaitSampleCount + } + if m.ProcessSampleCount > 0 { + m.ProcessMsAvg = m.ProcessMsSum / m.ProcessSampleCount + } +} diff --git a/model/channel_flow_schedule_test.go b/model/channel_flow_schedule_test.go new file mode 100644 index 00000000000..af62c84262d --- /dev/null +++ b/model/channel_flow_schedule_test.go @@ -0,0 +1,64 @@ +package model + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func testScheduledFlowPool() ChannelFlowPool { + return ChannelFlowPool{ + Name: "scheduled pool", + Enabled: true, + Backend: ChannelFlowBackendMemory, + MaxInflight: 1, + QueueTimeoutMs: 1000, + QueuePolicy: ChannelFlowQueuePolicyFIFO, + OnLimit: ChannelFlowOnLimitQueue, + ScheduleTimezone: "Asia/Shanghai", + } +} + +func TestChannelFlowPoolScheduleAlwaysActive(t *testing.T) { + pool := testScheduledFlowPool() + pool.ScheduleMode = ChannelFlowScheduleAlways + + require.True(t, pool.IsScheduleActiveAt(time.Date(2026, 6, 15, 12, 0, 0, 0, time.UTC)), + "always schedule should be active") +} + +func TestChannelFlowPoolScheduleDateTimeRange(t *testing.T) { + pool := testScheduledFlowPool() + pool.ScheduleMode = ChannelFlowScheduleDateTimeRange + pool.EffectiveStartTime = time.Date(2026, 6, 15, 10, 0, 0, 0, time.UTC).Unix() + pool.EffectiveEndTime = time.Date(2026, 6, 15, 11, 0, 0, 0, time.UTC).Unix() + + require.True(t, pool.IsScheduleActiveAt(time.Date(2026, 6, 15, 10, 30, 0, 0, time.UTC)), + "range schedule should be active inside the window") + require.False(t, pool.IsScheduleActiveAt(time.Date(2026, 6, 15, 11, 0, 0, 0, time.UTC)), + "range schedule should be inactive at the exclusive end") +} + +func TestChannelFlowPoolScheduleWeeklyCrossDay(t *testing.T) { + pool := testScheduledFlowPool() + pool.ScheduleMode = ChannelFlowScheduleWeekly + pool.ScheduleWindows = `[{"weekdays":[1],"start_minute":1320,"end_minute":120}]` + loc, err := time.LoadLocation("Asia/Shanghai") + require.NoError(t, err, "should load timezone") + + require.True(t, pool.IsScheduleActiveAt(time.Date(2026, 6, 15, 23, 0, 0, 0, loc)), + "weekly schedule should be active on the start day") + require.True(t, pool.IsScheduleActiveAt(time.Date(2026, 6, 16, 1, 30, 0, 0, loc)), + "weekly schedule should remain active after midnight") + require.False(t, pool.IsScheduleActiveAt(time.Date(2026, 6, 16, 3, 0, 0, 0, loc)), + "weekly schedule should be inactive after the cross-day end") +} + +func TestChannelFlowPoolScheduleWeeklyValidation(t *testing.T) { + pool := testScheduledFlowPool() + pool.ScheduleMode = ChannelFlowScheduleWeekly + pool.ScheduleWindows = `[{"weekdays":[7],"start_minute":60,"end_minute":120}]` + + require.Error(t, pool.Validate(), "invalid weekday should fail validation") +} diff --git a/model/main.go b/model/main.go index 6d900246287..87e6ef0b937 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/pkg/channel_flow_metrics/flush.go b/pkg/channel_flow_metrics/flush.go new file mode 100644 index 00000000000..e7c3d0671d7 --- /dev/null +++ b/pkg/channel_flow_metrics/flush.go @@ -0,0 +1,85 @@ +package channelflowmetrics + +import ( + "fmt" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" +) + +const flushInterval = time.Minute +const retentionDays = 8 + +func flushLoop() { + for { + time.Sleep(flushInterval) + flushCompletedBuckets() + cleanupExpiredMetrics() + } +} + +func flushCompletedBuckets() { + currentBucket := bucketStart(time.Now().Unix()) + hotBuckets.Range(func(key, value any) bool { + k := key.(bucketKey) + if k.bucketTs >= currentBucket { + return true + } + + bucket := value.(*atomicBucket) + drained := bucket.drain() + if !drained.hasData() { + hotBuckets.Delete(key) + return true + } + + if err := model.UpsertChannelFlowMetricMinute(metricFromCounters(k, drained)); err != nil { + bucket.addCounters(drained) + common.SysError(fmt.Sprintf("failed to flush channel flow metric pool=%s channel=%d model=%s bucket=%d: %s", k.poolKey, k.channelID, k.model, k.bucketTs, err.Error())) + return true + } + + hotBuckets.Delete(key) + return true + }) +} + +func metricFromCounters(k bucketKey, c counters) *model.ChannelFlowMetricMinute { + metric := &model.ChannelFlowMetricMinute{ + BucketTs: k.bucketTs, + PoolKey: k.poolKey, + ChannelId: k.channelID, + Model: k.model, + SampleCount: c.sampleCount, + RunningSum: c.runningSum, + RunningMax: safeInt(c.runningMax), + QueuedSum: c.queuedSum, + QueuedMax: safeInt(c.queuedMax), + AcquiredCount: safeInt(c.acquiredCount), + QueuedCount: safeInt(c.queuedCount), + SucceededCount: safeInt(c.succeededCount), + FailedCount: safeInt(c.failedCount), + ReleasedCount: safeInt(c.releasedCount), + RejectedCount: safeInt(c.rejectedCount), + TimeoutCount: safeInt(c.timeoutCount), + CancelledCount: safeInt(c.cancelledCount), + BillingFailedCount: safeInt(c.billingFailedCount), + LeaseRenewFail: safeInt(c.leaseRenewFail), + LeaseExpiredCount: safeInt(c.leaseExpiredCount), + WaitMsSum: c.waitMsSum, + WaitSampleCount: c.waitSampleCount, + WaitMsMax: c.waitMsMax, + ProcessMsSum: c.processMsSum, + ProcessSampleCount: c.processSampleCount, + ProcessMsMax: c.processMsMax, + } + return metric +} + +func cleanupExpiredMetrics() { + cutoff := time.Now().Add(-retentionDays * 24 * time.Hour).Unix() + if err := model.DeleteChannelFlowMetricMinutesBefore(cutoff); err != nil { + common.SysError("failed to cleanup expired channel flow metrics: " + err.Error()) + } +} diff --git a/pkg/channel_flow_metrics/metrics.go b/pkg/channel_flow_metrics/metrics.go new file mode 100644 index 00000000000..d5a266b5559 --- /dev/null +++ b/pkg/channel_flow_metrics/metrics.go @@ -0,0 +1,295 @@ +package channelflowmetrics + +import ( + "sync" + "time" + + "github.com/QuantumNous/new-api/model" +) + +const ( + defaultQueryHours = 6 + maxQueryHours = 24 * 7 + maxQueryMinutes = maxQueryHours * 60 + flowBucketSeconds = 60 +) + +var hotBuckets sync.Map + +func Init() { + go flushLoop() +} + +func Record(sample Sample) { + if sample.PoolKey == "" || sample.EventType == "" { + return + } + key := bucketKey{ + poolKey: sample.PoolKey, + channelID: sample.ChannelID, + model: sample.Model, + bucketTs: bucketStart(time.Now().Unix()), + } + actual, _ := hotBuckets.LoadOrStore(key, &atomicBucket{}) + actual.(*atomicBucket).add(sample) + recordRedis(key, sample) +} + +func Query(params QueryParams) (TrendResult, error) { + if params.Minutes <= 0 { + if params.Hours <= 0 { + params.Hours = defaultQueryHours + } + params.Minutes = params.Hours * 60 + } + if params.Minutes > maxQueryMinutes { + params.Minutes = maxQueryMinutes + } + endBucket := bucketStart(time.Now().Unix()) + bucketCount := params.Minutes * 60 / flowBucketSeconds + if bucketCount <= 0 { + bucketCount = 1 + } + startBucket := endBucket - int64(bucketCount-1)*flowBucketSeconds + + merged := map[int64]counters{} + rows, err := model.GetChannelFlowMetricMinutes(params.PoolKey, startBucket, endBucket) + if err != nil { + return TrendResult{}, err + } + for _, row := range rows { + mergeCounters(merged, row.BucketTs, metricToCounters(row)) + } + + redisActiveMerged := mergeRedisActiveBucket(merged, params.PoolKey, endBucket) + hotBuckets.Range(func(key, value any) bool { + k := key.(bucketKey) + if k.poolKey != params.PoolKey || k.bucketTs < startBucket || k.bucketTs > endBucket { + return true + } + if redisActiveMerged && k.bucketTs == endBucket { + return true + } + mergeCounters(merged, k.bucketTs, value.(*atomicBucket).snapshot()) + return true + }) + + return TrendResult{ + PoolKey: params.PoolKey, + Points: buildPoints(merged, startBucket, endBucket), + Totals: buildTotals(merged), + }, nil +} + +func bucketStart(ts int64) int64 { + return ts - (ts % flowBucketSeconds) +} + +func mergeCounters(merged map[int64]counters, bucketTs int64, value counters) { + if !value.hasData() { + return + } + current := merged[bucketTs] + current.sampleCount += value.sampleCount + current.runningSum += value.runningSum + if value.runningMax > current.runningMax { + current.runningMax = value.runningMax + } + current.queuedSum += value.queuedSum + if value.queuedMax > current.queuedMax { + current.queuedMax = value.queuedMax + } + current.acquiredCount += value.acquiredCount + current.queuedCount += value.queuedCount + current.succeededCount += value.succeededCount + current.failedCount += value.failedCount + current.releasedCount += value.releasedCount + current.rejectedCount += value.rejectedCount + current.timeoutCount += value.timeoutCount + current.cancelledCount += value.cancelledCount + current.billingFailedCount += value.billingFailedCount + current.leaseRenewFail += value.leaseRenewFail + current.leaseExpiredCount += value.leaseExpiredCount + current.waitMsSum += value.waitMsSum + current.waitSampleCount += value.waitSampleCount + if value.waitMsMax > current.waitMsMax { + current.waitMsMax = value.waitMsMax + } + current.processMsSum += value.processMsSum + current.processSampleCount += value.processSampleCount + if value.processMsMax > current.processMsMax { + current.processMsMax = value.processMsMax + } + merged[bucketTs] = current +} + +func buildPoints(merged map[int64]counters, startBucket int64, endBucket int64) []ChannelFlowTrendPoint { + if endBucket < startBucket { + return []ChannelFlowTrendPoint{} + } + points := make([]ChannelFlowTrendPoint, 0, int((endBucket-startBucket)/flowBucketSeconds)+1) + for ts := startBucket; ts <= endBucket; ts += flowBucketSeconds { + points = append(points, counterPoint(ts, merged[ts])) + } + return points +} + +func buildTotals(merged map[int64]counters) ChannelFlowTrendTotals { + total := counters{} + for _, value := range merged { + total = mergeTwoCounters(total, value) + } + return ChannelFlowTrendTotals{ + RequestCount: safeInt(total.requestCount()), + RunningAvg: safeInt(avg(total.runningSum, total.sampleCount)), + RunningMax: safeInt(total.runningMax), + QueuedAvg: safeInt(avg(total.queuedSum, total.sampleCount)), + QueuedMax: safeInt(total.queuedMax), + AcquiredCount: safeInt(total.acquiredCount), + QueuedCount: safeInt(total.queuedCount), + SucceededCount: safeInt(total.succeededCount), + FailedCount: safeInt(total.failedCount), + ReleasedCount: safeInt(total.releasedCount), + RejectedCount: safeInt(total.rejectedCount), + TimeoutCount: safeInt(total.timeoutCount), + CancelledCount: safeInt(total.cancelledCount), + BillingFailedCount: safeInt(total.billingFailedCount), + LeaseRenewFail: safeInt(total.leaseRenewFail), + LeaseExpiredCount: safeInt(total.leaseExpiredCount), + WaitMsAvg: avg(total.waitMsSum, total.waitSampleCount), + WaitMsMax: total.waitMsMax, + ProcessMsAvg: avg(total.processMsSum, total.processSampleCount), + ProcessMsMax: total.processMsMax, + } +} + +func mergeTwoCounters(current counters, value counters) counters { + current.sampleCount += value.sampleCount + current.runningSum += value.runningSum + if value.runningMax > current.runningMax { + current.runningMax = value.runningMax + } + current.queuedSum += value.queuedSum + if value.queuedMax > current.queuedMax { + current.queuedMax = value.queuedMax + } + current.acquiredCount += value.acquiredCount + current.queuedCount += value.queuedCount + current.succeededCount += value.succeededCount + current.failedCount += value.failedCount + current.releasedCount += value.releasedCount + current.rejectedCount += value.rejectedCount + current.timeoutCount += value.timeoutCount + current.cancelledCount += value.cancelledCount + current.billingFailedCount += value.billingFailedCount + current.leaseRenewFail += value.leaseRenewFail + current.leaseExpiredCount += value.leaseExpiredCount + current.waitMsSum += value.waitMsSum + current.waitSampleCount += value.waitSampleCount + if value.waitMsMax > current.waitMsMax { + current.waitMsMax = value.waitMsMax + } + current.processMsSum += value.processMsSum + current.processSampleCount += value.processSampleCount + if value.processMsMax > current.processMsMax { + current.processMsMax = value.processMsMax + } + return current +} + +func counterPoint(ts int64, value counters) ChannelFlowTrendPoint { + runningAvg := float64(0) + queuedAvg := float64(0) + if value.sampleCount > 0 { + runningAvg = float64(value.runningSum) / float64(value.sampleCount) + queuedAvg = float64(value.queuedSum) / float64(value.sampleCount) + } + return ChannelFlowTrendPoint{ + BucketTs: ts, + At: time.Unix(ts, 0).Format("15:04"), + Running: runningAvg, + RunningAvg: runningAvg, + RunningMax: safeInt(value.runningMax), + Queued: queuedAvg, + QueuedAvg: queuedAvg, + QueuedMax: safeInt(value.queuedMax), + RequestCount: safeInt(value.requestCount()), + AcquiredCount: safeInt(value.acquiredCount), + QueuedCount: safeInt(value.queuedCount), + SucceededCount: safeInt(value.succeededCount), + FailedCount: safeInt(value.failedCount), + ReleasedCount: safeInt(value.releasedCount), + RejectedCount: safeInt(value.rejectedCount), + TimeoutCount: safeInt(value.timeoutCount), + CancelledCount: safeInt(value.cancelledCount), + BillingFailedCount: safeInt(value.billingFailedCount), + LeaseRenewFail: safeInt(value.leaseRenewFail), + LeaseExpiredCount: safeInt(value.leaseExpiredCount), + WaitMsAvg: avg(value.waitMsSum, value.waitSampleCount), + WaitMsMax: value.waitMsMax, + ProcessMsAvg: avg(value.processMsSum, value.processSampleCount), + ProcessMsMax: value.processMsMax, + } +} + +func metricToCounters(metric model.ChannelFlowMetricMinute) counters { + return counters{ + sampleCount: metric.SampleCount, + runningSum: metric.RunningSum, + runningMax: int64(metric.RunningMax), + queuedSum: metric.QueuedSum, + queuedMax: int64(metric.QueuedMax), + acquiredCount: int64(metric.AcquiredCount), + queuedCount: int64(metric.QueuedCount), + succeededCount: int64(metric.SucceededCount), + failedCount: int64(metric.FailedCount), + releasedCount: int64(metric.ReleasedCount), + rejectedCount: int64(metric.RejectedCount), + timeoutCount: int64(metric.TimeoutCount), + cancelledCount: int64(metric.CancelledCount), + billingFailedCount: int64(metric.BillingFailedCount), + leaseRenewFail: int64(metric.LeaseRenewFail), + leaseExpiredCount: int64(metric.LeaseExpiredCount), + waitMsSum: metric.WaitMsSum, + waitSampleCount: metric.WaitSampleCount, + waitMsMax: metric.WaitMsMax, + processMsSum: metric.ProcessMsSum, + processSampleCount: metric.ProcessSampleCount, + processMsMax: metric.ProcessMsMax, + } +} + +func avg(sum int64, count int64) int64 { + if count <= 0 { + return 0 + } + return sum / count +} + +func safeInt(value int64) int { + if value <= 0 { + return 0 + } + return int(value) +} + +func (c counters) hasData() bool { + return c.sampleCount > 0 || + c.acquiredCount > 0 || + c.queuedCount > 0 || + c.succeededCount > 0 || + c.failedCount > 0 || + c.releasedCount > 0 || + c.rejectedCount > 0 || + c.timeoutCount > 0 || + c.cancelledCount > 0 || + c.billingFailedCount > 0 || + c.leaseRenewFail > 0 || + c.leaseExpiredCount > 0 || + c.waitSampleCount > 0 || + c.processSampleCount > 0 +} + +func (c counters) requestCount() int64 { + return c.acquiredCount + c.rejectedCount + c.timeoutCount + c.billingFailedCount +} diff --git a/pkg/channel_flow_metrics/metrics_test.go b/pkg/channel_flow_metrics/metrics_test.go new file mode 100644 index 00000000000..32c5a47dad2 --- /dev/null +++ b/pkg/channel_flow_metrics/metrics_test.go @@ -0,0 +1,143 @@ +package channelflowmetrics + +import ( + "sync" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupMetricDB(t *testing.T) { + t.Helper() + oldDB := model.DB + oldRedisEnabled := common.RedisEnabled + oldRDB := common.RDB + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.ChannelFlowMetricMinute{})) + model.DB = db + common.RedisEnabled = false + common.RDB = nil + hotBuckets = sync.Map{} + t.Cleanup(func() { + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, sqlDB.Close()) + model.DB = oldDB + common.RedisEnabled = oldRedisEnabled + common.RDB = oldRDB + hotBuckets = sync.Map{} + }) +} + +func TestQueryFillsEmptyMinuteBuckets(t *testing.T) { + setupMetricDB(t) + + result, err := Query(QueryParams{PoolKey: "flow_pool_empty_test", Hours: 1}) + require.NoError(t, err) + require.Len(t, result.Points, 60) + require.Equal(t, "flow_pool_empty_test", result.PoolKey) + for _, point := range result.Points { + require.Zero(t, point.RunningMax) + require.Zero(t, point.QueuedMax) + require.Zero(t, point.AcquiredCount) + } + require.Zero(t, result.Totals.AcquiredCount) + require.Zero(t, result.Totals.RejectedCount) +} + +func TestFlushCompletedBucketsUpsertsAndDeletesHotBucket(t *testing.T) { + setupMetricDB(t) + + bucketTs := bucketStart(time.Now().Add(-2 * time.Minute).Unix()) + key := bucketKey{ + poolKey: "flow_pool_metric_test", + channelID: 11, + model: "gpt-test", + bucketTs: bucketTs, + } + bucket := &atomicBucket{} + bucket.add(Sample{EventType: model.ChannelFlowEventAcquired, Running: 1, Queued: 2, WaitMs: 15}) + bucket.add(Sample{EventType: model.ChannelFlowEventSucceeded, Running: -1, Queued: -1}) + bucket.add(Sample{EventType: model.ChannelFlowEventReleased, Running: -1, Queued: -1, ProcessMs: 100}) + hotBuckets.Store(key, bucket) + + flushCompletedBuckets() + + _, ok := hotBuckets.Load(key) + require.False(t, ok, "flushed completed bucket should be removed from hot memory") + + rows, err := model.GetChannelFlowMetricMinutes(key.poolKey, bucketTs, bucketTs) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, int64(1), rows[0].SampleCount) + require.Equal(t, int64(1), rows[0].RunningSum) + require.Equal(t, 1.0, rows[0].RunningAvg) + require.Equal(t, int64(2), rows[0].QueuedSum) + require.Equal(t, 2.0, rows[0].QueuedAvg) + require.Equal(t, 1, rows[0].AcquiredCount) + require.Equal(t, 1, rows[0].SucceededCount) + require.Equal(t, 1, rows[0].ReleasedCount) + require.Equal(t, int64(15), rows[0].WaitMsAvg) + require.Equal(t, int64(100), rows[0].ProcessMsAvg) + + nextBucket := &atomicBucket{} + nextBucket.add(Sample{EventType: model.ChannelFlowEventAcquired, Running: 3, Queued: 0, WaitMs: 45}) + nextBucket.add(Sample{EventType: model.ChannelFlowEventFailed, Running: -1, Queued: -1}) + hotBuckets.Store(key, nextBucket) + + flushCompletedBuckets() + + rows, err = model.GetChannelFlowMetricMinutes(key.poolKey, bucketTs, bucketTs) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, int64(2), rows[0].SampleCount) + require.Equal(t, int64(4), rows[0].RunningSum) + require.Equal(t, 2.0, rows[0].RunningAvg) + require.Equal(t, 3, rows[0].RunningMax) + require.Equal(t, int64(2), rows[0].QueuedSum) + require.Equal(t, 1.0, rows[0].QueuedAvg) + require.Equal(t, 2, rows[0].AcquiredCount) + require.Equal(t, 1, rows[0].SucceededCount) + require.Equal(t, 1, rows[0].FailedCount) + require.Equal(t, int64(30), rows[0].WaitMsAvg) + require.Equal(t, int64(45), rows[0].WaitMsMax) +} + +func TestCleanupExpiredMetricsKeepsRetentionWindow(t *testing.T) { + setupMetricDB(t) + + poolKey := "flow_pool_cleanup_test" + oldTs := bucketStart(time.Now().Add(-9 * 24 * time.Hour).Unix()) + recentTs := bucketStart(time.Now().Add(-time.Hour).Unix()) + require.NoError(t, model.UpsertChannelFlowMetricMinute(&model.ChannelFlowMetricMinute{ + BucketTs: oldTs, + PoolKey: poolKey, + ChannelId: 1, + Model: "gpt-old", + SampleCount: 1, + RunningSum: 1, + AcquiredCount: 1, + })) + require.NoError(t, model.UpsertChannelFlowMetricMinute(&model.ChannelFlowMetricMinute{ + BucketTs: recentTs, + PoolKey: poolKey, + ChannelId: 1, + Model: "gpt-recent", + SampleCount: 1, + RunningSum: 1, + AcquiredCount: 1, + })) + + cleanupExpiredMetrics() + + rows, err := model.GetChannelFlowMetricMinutes(poolKey, oldTs, recentTs) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, recentTs, rows[0].BucketTs) +} diff --git a/pkg/channel_flow_metrics/redis.go b/pkg/channel_flow_metrics/redis.go new file mode 100644 index 00000000000..b1d868cc0d7 --- /dev/null +++ b/pkg/channel_flow_metrics/redis.go @@ -0,0 +1,188 @@ +package channelflowmetrics + +import ( + "context" + "encoding/base64" + "fmt" + "strconv" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/go-redis/redis/v8" +) + +const redisActiveBucketTTL = 2 * time.Hour + +const redisActiveBucketScript = ` +local metric_key = KEYS[1] +local index_key = KEYS[2] +local ttl = tonumber(ARGV[1]) +local member = ARGV[2] +redis.call("SADD", index_key, member) +redis.call("EXPIRE", index_key, ttl) + +local idx = 3 +local inc_count = tonumber(ARGV[idx]) +idx = idx + 1 +for i = 1, inc_count do + local field = ARGV[idx] + local value = tonumber(ARGV[idx + 1]) or 0 + if value ~= 0 then + redis.call("HINCRBY", metric_key, field, value) + end + idx = idx + 2 +end + +local max_count = tonumber(ARGV[idx]) +idx = idx + 1 +for i = 1, max_count do + local field = ARGV[idx] + local value = tonumber(ARGV[idx + 1]) or 0 + if value > 0 then + local current = tonumber(redis.call("HGET", metric_key, field) or "0") + if value > current then + redis.call("HSET", metric_key, field, value) + end + end + idx = idx + 2 +end + +redis.call("EXPIRE", metric_key, ttl) +return 1 +` + +var redisActiveBucketLua = redis.NewScript(redisActiveBucketScript) + +func recordRedis(key bucketKey, sample Sample) { + if !common.RedisEnabled || common.RDB == nil { + return + } + c := countersFromSample(sample) + if !c.hasData() { + return + } + + metricKey := redisMetricKey(key) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = redisActiveBucketLua.Run( + ctx, + common.RDB, + []string{metricKey, redisIndexKey(key.poolKey, key.bucketTs)}, + redisRecordArgs(metricKey, c)..., + ).Err() +} + +func mergeRedisActiveBucket(merged map[int64]counters, poolKey string, bucketTs int64) bool { + if !common.RedisEnabled || common.RDB == nil || poolKey == "" || bucketTs <= 0 { + return false + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + members, err := common.RDB.SMembers(ctx, redisIndexKey(poolKey, bucketTs)).Result() + if err != nil || len(members) == 0 { + return false + } + pipe := common.RDB.Pipeline() + cmds := make([]*redis.StringStringMapCmd, 0, len(members)) + for _, member := range members { + cmds = append(cmds, pipe.HGetAll(ctx, member)) + } + if _, err = pipe.Exec(ctx); err != nil && err != redis.Nil { + return false + } + + mergedAny := false + for _, cmd := range cmds { + value := redisCounters(cmd.Val()) + if !value.hasData() { + continue + } + mergeCounters(merged, bucketTs, value) + mergedAny = true + } + return mergedAny +} + +func redisRecordArgs(metricKey string, c counters) []interface{} { + increments := []interface{}{ + "sample_count", c.sampleCount, + "running_sum", c.runningSum, + "queued_sum", c.queuedSum, + "acquired_count", c.acquiredCount, + "queued_count", c.queuedCount, + "succeeded_count", c.succeededCount, + "failed_count", c.failedCount, + "released_count", c.releasedCount, + "rejected_count", c.rejectedCount, + "timeout_count", c.timeoutCount, + "cancelled_count", c.cancelledCount, + "billing_failed_count", c.billingFailedCount, + "lease_renew_fail", c.leaseRenewFail, + "lease_expired_count", c.leaseExpiredCount, + "wait_ms_sum", c.waitMsSum, + "wait_sample_count", c.waitSampleCount, + "process_ms_sum", c.processMsSum, + "process_sample_count", c.processSampleCount, + } + maxes := []interface{}{ + "running_max", c.runningMax, + "queued_max", c.queuedMax, + "wait_ms_max", c.waitMsMax, + "process_ms_max", c.processMsMax, + } + args := []interface{}{ + int(redisActiveBucketTTL.Seconds()), + metricKey, + len(increments) / 2, + } + args = append(args, increments...) + args = append(args, len(maxes)/2) + args = append(args, maxes...) + return args +} + +func redisCounters(values map[string]string) counters { + return counters{ + sampleCount: parseRedisInt(values["sample_count"]), + runningSum: parseRedisInt(values["running_sum"]), + runningMax: parseRedisInt(values["running_max"]), + queuedSum: parseRedisInt(values["queued_sum"]), + queuedMax: parseRedisInt(values["queued_max"]), + acquiredCount: parseRedisInt(values["acquired_count"]), + queuedCount: parseRedisInt(values["queued_count"]), + succeededCount: parseRedisInt(values["succeeded_count"]), + failedCount: parseRedisInt(values["failed_count"]), + releasedCount: parseRedisInt(values["released_count"]), + rejectedCount: parseRedisInt(values["rejected_count"]), + timeoutCount: parseRedisInt(values["timeout_count"]), + cancelledCount: parseRedisInt(values["cancelled_count"]), + billingFailedCount: parseRedisInt(values["billing_failed_count"]), + leaseRenewFail: parseRedisInt(values["lease_renew_fail"]), + leaseExpiredCount: parseRedisInt(values["lease_expired_count"]), + waitMsSum: parseRedisInt(values["wait_ms_sum"]), + waitSampleCount: parseRedisInt(values["wait_sample_count"]), + waitMsMax: parseRedisInt(values["wait_ms_max"]), + processMsSum: parseRedisInt(values["process_ms_sum"]), + processSampleCount: parseRedisInt(values["process_sample_count"]), + processMsMax: parseRedisInt(values["process_ms_max"]), + } +} + +func parseRedisInt(value string) int64 { + if value == "" { + return 0 + } + parsed, _ := strconv.ParseInt(value, 10, 64) + return parsed +} + +func redisIndexKey(poolKey string, bucketTs int64) string { + return fmt.Sprintf("channel_flow:metrics:%s:%d:index", poolKey, bucketTs) +} + +func redisMetricKey(key bucketKey) string { + modelKey := base64.RawURLEncoding.EncodeToString([]byte(key.model)) + return fmt.Sprintf("channel_flow:metrics:%s:%d:%s:%d", key.poolKey, key.channelID, modelKey, key.bucketTs) +} diff --git a/pkg/channel_flow_metrics/types.go b/pkg/channel_flow_metrics/types.go new file mode 100644 index 00000000000..cc811da5a2f --- /dev/null +++ b/pkg/channel_flow_metrics/types.go @@ -0,0 +1,274 @@ +package channelflowmetrics + +import "sync/atomic" + +type Sample struct { + PoolKey string + ChannelID int + Model string + EventType string + Running int + Queued int + WaitMs int64 + ProcessMs int64 +} + +type QueryParams struct { + PoolKey string + Hours int + Minutes int +} + +type TrendResult struct { + PoolKey string `json:"pool_key"` + Points []ChannelFlowTrendPoint `json:"points"` + Totals ChannelFlowTrendTotals `json:"totals"` +} + +type ChannelFlowTrendPoint struct { + BucketTs int64 `json:"bucket_ts"` + At string `json:"at"` + Running float64 `json:"running"` + RunningAvg float64 `json:"running_avg"` + RunningMax int `json:"running_max"` + Queued float64 `json:"queued"` + QueuedAvg float64 `json:"queued_avg"` + QueuedMax int `json:"queued_max"` + RequestCount int `json:"request_count"` + AcquiredCount int `json:"acquired_count"` + QueuedCount int `json:"queued_count"` + SucceededCount int `json:"succeeded_count"` + FailedCount int `json:"failed_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"` + WaitMsMax int64 `json:"wait_ms_max"` + ProcessMsAvg int64 `json:"process_ms_avg"` + ProcessMsMax int64 `json:"process_ms_max"` +} + +type ChannelFlowTrendTotals struct { + RequestCount int `json:"request_count"` + RunningAvg int `json:"running_avg"` + RunningMax int `json:"running_max"` + QueuedAvg int `json:"queued_avg"` + QueuedMax int `json:"queued_max"` + AcquiredCount int `json:"acquired_count"` + QueuedCount int `json:"queued_count"` + SucceededCount int `json:"succeeded_count"` + FailedCount int `json:"failed_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"` + WaitMsMax int64 `json:"wait_ms_max"` + ProcessMsAvg int64 `json:"process_ms_avg"` + ProcessMsMax int64 `json:"process_ms_max"` +} + +type bucketKey struct { + poolKey string + channelID int + model string + bucketTs int64 +} + +type counters struct { + sampleCount int64 + runningSum int64 + runningMax int64 + queuedSum int64 + queuedMax int64 + acquiredCount int64 + queuedCount int64 + succeededCount int64 + failedCount int64 + releasedCount int64 + rejectedCount int64 + timeoutCount int64 + cancelledCount int64 + billingFailedCount int64 + leaseRenewFail int64 + leaseExpiredCount int64 + waitMsSum int64 + waitSampleCount int64 + waitMsMax int64 + processMsSum int64 + processSampleCount int64 + processMsMax int64 +} + +type atomicBucket struct { + sampleCount atomic.Int64 + runningSum atomic.Int64 + runningMax atomic.Int64 + queuedSum atomic.Int64 + queuedMax atomic.Int64 + acquiredCount atomic.Int64 + queuedCount atomic.Int64 + succeededCount atomic.Int64 + failedCount atomic.Int64 + releasedCount atomic.Int64 + rejectedCount atomic.Int64 + timeoutCount atomic.Int64 + cancelledCount atomic.Int64 + billingFailedCount atomic.Int64 + leaseRenewFail atomic.Int64 + leaseExpiredCount atomic.Int64 + waitMsSum atomic.Int64 + waitSampleCount atomic.Int64 + waitMsMax atomic.Int64 + processMsSum atomic.Int64 + processSampleCount atomic.Int64 + processMsMax atomic.Int64 +} + +func (b *atomicBucket) add(sample Sample) { + b.addCounters(countersFromSample(sample)) +} + +func countersFromSample(sample Sample) counters { + c := counters{} + if sample.Running >= 0 && sample.Queued >= 0 { + c.sampleCount = 1 + c.runningSum = int64(sample.Running) + c.runningMax = int64(sample.Running) + c.queuedSum = int64(sample.Queued) + c.queuedMax = int64(sample.Queued) + } + switch sample.EventType { + case "acquired": + c.acquiredCount = 1 + case "queued": + c.queuedCount = 1 + case "succeeded": + c.succeededCount = 1 + case "failed": + c.failedCount = 1 + case "released": + c.releasedCount = 1 + case "rejected": + c.rejectedCount = 1 + case "timeout": + c.timeoutCount = 1 + case "cancelled": + c.cancelledCount = 1 + case "billing_failed": + c.billingFailedCount = 1 + case "lease_renew_failed": + c.leaseRenewFail = 1 + case "lease_expired": + c.leaseExpiredCount = 1 + } + if sample.WaitMs > 0 { + c.waitMsSum = sample.WaitMs + c.waitSampleCount = 1 + c.waitMsMax = sample.WaitMs + } + if sample.ProcessMs > 0 { + c.processMsSum = sample.ProcessMs + c.processSampleCount = 1 + c.processMsMax = sample.ProcessMs + } + return c +} + +func (b *atomicBucket) snapshot() counters { + return counters{ + sampleCount: b.sampleCount.Load(), + runningSum: b.runningSum.Load(), + runningMax: b.runningMax.Load(), + queuedSum: b.queuedSum.Load(), + queuedMax: b.queuedMax.Load(), + acquiredCount: b.acquiredCount.Load(), + queuedCount: b.queuedCount.Load(), + succeededCount: b.succeededCount.Load(), + failedCount: b.failedCount.Load(), + releasedCount: b.releasedCount.Load(), + rejectedCount: b.rejectedCount.Load(), + timeoutCount: b.timeoutCount.Load(), + cancelledCount: b.cancelledCount.Load(), + billingFailedCount: b.billingFailedCount.Load(), + leaseRenewFail: b.leaseRenewFail.Load(), + leaseExpiredCount: b.leaseExpiredCount.Load(), + waitMsSum: b.waitMsSum.Load(), + waitSampleCount: b.waitSampleCount.Load(), + waitMsMax: b.waitMsMax.Load(), + processMsSum: b.processMsSum.Load(), + processSampleCount: b.processSampleCount.Load(), + processMsMax: b.processMsMax.Load(), + } +} + +func (b *atomicBucket) drain() counters { + return counters{ + sampleCount: b.sampleCount.Swap(0), + runningSum: b.runningSum.Swap(0), + runningMax: b.runningMax.Swap(0), + queuedSum: b.queuedSum.Swap(0), + queuedMax: b.queuedMax.Swap(0), + acquiredCount: b.acquiredCount.Swap(0), + queuedCount: b.queuedCount.Swap(0), + succeededCount: b.succeededCount.Swap(0), + failedCount: b.failedCount.Swap(0), + releasedCount: b.releasedCount.Swap(0), + rejectedCount: b.rejectedCount.Swap(0), + timeoutCount: b.timeoutCount.Swap(0), + cancelledCount: b.cancelledCount.Swap(0), + billingFailedCount: b.billingFailedCount.Swap(0), + leaseRenewFail: b.leaseRenewFail.Swap(0), + leaseExpiredCount: b.leaseExpiredCount.Swap(0), + waitMsSum: b.waitMsSum.Swap(0), + waitSampleCount: b.waitSampleCount.Swap(0), + waitMsMax: b.waitMsMax.Swap(0), + processMsSum: b.processMsSum.Swap(0), + processSampleCount: b.processSampleCount.Swap(0), + processMsMax: b.processMsMax.Swap(0), + } +} + +func (b *atomicBucket) addCounters(c counters) { + b.sampleCount.Add(c.sampleCount) + b.runningSum.Add(c.runningSum) + updateAtomicMax(&b.runningMax, c.runningMax) + b.queuedSum.Add(c.queuedSum) + updateAtomicMax(&b.queuedMax, c.queuedMax) + b.acquiredCount.Add(c.acquiredCount) + b.queuedCount.Add(c.queuedCount) + b.succeededCount.Add(c.succeededCount) + b.failedCount.Add(c.failedCount) + b.releasedCount.Add(c.releasedCount) + b.rejectedCount.Add(c.rejectedCount) + b.timeoutCount.Add(c.timeoutCount) + b.cancelledCount.Add(c.cancelledCount) + b.billingFailedCount.Add(c.billingFailedCount) + b.leaseRenewFail.Add(c.leaseRenewFail) + b.leaseExpiredCount.Add(c.leaseExpiredCount) + b.waitMsSum.Add(c.waitMsSum) + b.waitSampleCount.Add(c.waitSampleCount) + updateAtomicMax(&b.waitMsMax, c.waitMsMax) + b.processMsSum.Add(c.processMsSum) + b.processSampleCount.Add(c.processSampleCount) + updateAtomicMax(&b.processMsMax, c.processMsMax) +} + +func updateAtomicMax(target *atomic.Int64, value int64) { + for { + current := target.Load() + if value <= current { + return + } + if target.CompareAndSwap(current, value) { + return + } + } +} diff --git a/router/api-router.go b/router/api-router.go index baf7cda2015..cc24abb5f41 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,20 @@ 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/trend", controller.GetChannelFlowPoolTrend) + 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 81daeed82c2..1168a26a78b 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 00000000000..bfca675fb76 --- /dev/null +++ b/service/channel_flow.go @@ -0,0 +1,1033 @@ +package service + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "sync" + "sync/atomic" + "time" + "unicode/utf8" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + channelflowmetrics "github.com/QuantumNous/new-api/pkg/channel_flow_metrics" + 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" + FlowDecisionRejectClientCancelled = "client_cancelled" + FlowDecisionRejectContextExceeded = "context_exceeded" + FlowDecisionRejectPerUserQueueFull = "per_user_queue_full" + FlowDecisionRejectPerUserInflightFull = "per_user_inflight_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"` + ScheduleActive bool `json:"schedule_active"` + 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"` + // WatchAttempts and TxConflicts are backend-global counters on the shared + // redisFlowBackend instance, not per-pool counters. + WatchAttempts int64 `json:"watch_attempts"` + TxConflicts int64 `json:"tx_conflicts"` +} + +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 || !pool.IsScheduleActiveAt(time.Now()) { + return binding, pool, false, nil + } + return binding, pool, true, nil +} + +func buildChannelFlowAcquireRequest(requestID string, pool model.ChannelFlowPool, channelID int, info *relaycommon.RelayInfo, now time.Time) AcquireRequest { + if now.IsZero() { + now = time.Now() + } + if requestID == "" { + requestID = common.GetUUID() + } + upstreamModel := "" + userID := 0 + tokenID := 0 + contextTokens := 0 + if info != nil { + upstreamModel = info.OriginModelName + if info.ChannelMeta != nil && info.UpstreamModelName != "" { + upstreamModel = info.UpstreamModelName + } + userID = info.UserId + tokenID = info.TokenId + contextTokens = info.GetEstimatePromptTokens() + } + return AcquireRequest{ + RequestID: requestID, + Pool: pool, + ChannelID: channelID, + UpstreamModel: upstreamModel, + UserID: userID, + TokenID: tokenID, + ContextTokens: contextTokens, + ContextChars: estimateChannelFlowContextChars(pool, info), + CreatedAtMs: now.UnixMilli(), + QueueTimeoutMs: pool.QueueTimeoutMs, + } +} + +func estimateChannelFlowContextChars(pool model.ChannelFlowPool, info *relaycommon.RelayInfo) int { + if pool.MaxContextChars <= 0 || info == nil || info.Request == nil { + return 0 + } + meta := info.Request.GetTokenCountMeta() + if meta == nil || meta.CombineText == "" { + return 0 + } + return utf8.RuneCountInString(meta.CombineText) +} + +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 + } + req := buildChannelFlowAcquireRequest(c.GetString(common.RequestIdKey), *pool, channelID, info, time.Now()) + guard, decision, acquireErr := GetChannelFlowController().Acquire(c.Request.Context(), req) + if acquireErr != nil { + if shouldPassThroughChannelFlowFallback(req.Pool, decision, acquireErr) { + return nil, nil, nil + } + if passThrough, fallbackPool, apiErr := handleRedisFlowAcquireError(c.Request.Context(), *pool, decision, acquireErr); apiErr != nil || passThrough { + if apiErr != nil { + recordChannelFlowMetric(req, channelFlowEventTypeFromDecision(decision), decision, true, decisionWaitMs(decision), 0) + } + return nil, decision, apiErr + } else if fallbackPool != nil { + req.Pool = *fallbackPool + guard, decision, acquireErr = GetChannelFlowController().Acquire(c.Request.Context(), req) + if acquireErr == nil { + bindChannelFlowGuardCallbacks(guard, req, decision) + return guard, decision, nil + } + } + recordChannelFlowMetric(req, channelFlowEventTypeFromDecision(decision), decision, true, decisionWaitMs(decision), 0) + return nil, decision, flowDecisionToAPIError(decision, acquireErr) + } + bindChannelFlowGuardCallbacks(guard, req, decision) + return guard, decision, nil +} + +func bindChannelFlowGuardCallbacks(guard FlowGuard, req AcquireRequest, decision *AcquireDecision) { + if guard == nil || decision == nil { + return + } + if decision.Queued { + recordChannelFlowMetric(req, model.ChannelFlowEventQueued, decision, false, 0, 0) + } + recordChannelFlowMetric(req, model.ChannelFlowEventAcquired, decision, true, decision.WaitedMs, 0) + + acquiredAt := time.Now() + stopRenew := startChannelFlowLeaseRenewer(guard, req) + guard.BindRelease(func() { + stopRenew() + recordChannelFlowMetric(req, model.ChannelFlowEventReleased, nil, false, 0, time.Since(acquiredAt).Milliseconds()) + }) +} + +func RecordChannelFlowOutcome(guard FlowGuard, channelID int, info *relaycommon.RelayInfo, success bool) { + if guard == nil || info == nil || guard.PoolKey() == "" { + return + } + upstreamModel := info.OriginModelName + if info.ChannelMeta != nil && info.UpstreamModelName != "" { + upstreamModel = info.UpstreamModelName + } + eventType := model.ChannelFlowEventFailed + if success { + eventType = model.ChannelFlowEventSucceeded + } + channelflowmetrics.Record(channelflowmetrics.Sample{ + PoolKey: guard.PoolKey(), + ChannelID: channelID, + Model: upstreamModel, + EventType: eventType, + Running: -1, + Queued: -1, + }) +} + +func startChannelFlowLeaseRenewer(guard FlowGuard, req AcquireRequest) func() { + if guard == nil || req.Pool.Backend != model.ChannelFlowBackendRedis { + return func() {} + } + interval := channelFlowRenewInterval(req.Pool) + ctx, cancel := context.WithCancel(context.Background()) + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if guard.IsReleased() { + return + } + if err := guard.RenewLease(context.Background()); err != nil { + recordChannelFlowMetric(req, model.ChannelFlowEventLeaseRenewFailed, nil, false, 0, 0) + } + } + } + }() + return cancel +} + +func channelFlowRenewInterval(pool model.ChannelFlowPool) time.Duration { + pool.Normalize() + interval := time.Duration(pool.RenewIntervalMs) * time.Millisecond + lease := time.Duration(pool.LeaseMs) * time.Millisecond + if lease > 0 && interval >= lease { + interval = lease / 2 + } + if interval < time.Second { + interval = time.Second + } + return interval +} + +func channelFlowEventTypeFromDecision(decision *AcquireDecision) string { + if decision == nil { + return model.ChannelFlowEventRejected + } + if decision.RejectCode == FlowDecisionRejectClientCancelled { + return model.ChannelFlowEventCancelled + } + if decision.RejectCode == FlowDecisionRejectQueueTimeout { + return model.ChannelFlowEventTimeout + } + return model.ChannelFlowEventRejected +} + +func decisionWaitMs(decision *AcquireDecision) int64 { + if decision == nil { + return 0 + } + return decision.WaitedMs +} + +func recordChannelFlowMetric(req AcquireRequest, eventType string, decision *AcquireDecision, includeCapacity bool, waitMs int64, processMs int64) { + if eventType == "" || req.Pool.PoolKey == "" { + return + } + running := -1 + queued := -1 + if includeCapacity && decision != nil { + running = decision.RunningNow + queued = decision.QueuedNow + } + channelflowmetrics.Record(channelflowmetrics.Sample{ + PoolKey: req.Pool.PoolKey, + ChannelID: req.ChannelID, + Model: req.UpstreamModel, + EventType: eventType, + Running: running, + Queued: queued, + WaitMs: waitMs, + ProcessMs: processMs, + }) +} + +func GetChannelFlowPoolStatus(ctx context.Context, pool model.ChannelFlowPool) (PoolStatus, error) { + if passThrough, fallbackPool, _ := resolveRedisFlowUnavailable(ctx, &pool); passThrough { + return withPoolStatusMetadata(degradedRedisFlowStatus(pool), pool), nil + } else if fallbackPool != nil { + status, err := GetChannelFlowController().Status(ctx, *fallbackPool) + return withPoolStatusMetadata(status, pool), err + } + status, err := GetChannelFlowController().Status(ctx, pool) + if err == nil { + return withPoolStatusMetadata(status, pool), nil + } + if errors.Is(err, ErrRedisFlowBackendUnavailable) { + switch pool.RedisFailurePolicy { + case model.ChannelFlowRedisFailureLocalMemory: + status, err := GetChannelFlowController().Status(ctx, localMemoryFallbackFlowPool(pool)) + return withPoolStatusMetadata(status, pool), err + default: + return withPoolStatusMetadata(degradedRedisFlowStatus(pool), pool), nil + } + } + return withPoolStatusMetadata(status, pool), err +} + +func withPoolStatusMetadata(status PoolStatus, pool model.ChannelFlowPool) PoolStatus { + pool.Normalize() + status.Name = pool.Name + status.ConfigVersion = pool.ConfigVersion + status.ScheduleActive = pool.Enabled && pool.IsScheduleActiveAt(time.Now()) + return status +} + +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", + ScheduleActive: pool.Enabled && pool.IsScheduleActiveAt(time.Now()), + 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 shouldPassThroughChannelFlowFallback(pool model.ChannelFlowPool, decision *AcquireDecision, err error) bool { + if err == nil || decision == nil || pool.OnLimit != model.ChannelFlowOnLimitFallback { + return false + } + switch decision.RejectCode { + case FlowDecisionRejectQueueFull, + FlowDecisionRejectPerUserQueueFull, + FlowDecisionRejectPerUserInflightFull: + return true + default: + return false + } +} + +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 FlowDecisionRejectClientCancelled: + errorCode = types.ErrorCodeChannelFlowClientCancelled + statusCode = http.StatusRequestTimeout + case FlowDecisionRejectContextExceeded: + errorCode = types.ErrorCodeChannelFlowContextExceeded + statusCode = http.StatusBadRequest + case FlowDecisionRejectPerUserQueueFull: + errorCode = types.ErrorCodeChannelFlowPerUserQueueFull + case FlowDecisionRejectPerUserInflightFull: + errorCode = types.ErrorCodeChannelFlowPerUserInflightFull + 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 req.Pool.MaxContextChars > 0 && req.ContextChars > req.Pool.MaxContextChars { + decision.RejectCode = FlowDecisionRejectContextExceeded + slot.mu.Unlock() + return nil, decision, fmt.Errorf("request context chars %d exceeds flow pool max_context_chars %d", req.ContextChars, req.Pool.MaxContextChars) + } + userInflightFull := req.Pool.MaxInflightPerUser > 0 && req.UserID > 0 && + slot.userRunningLocked(req.UserID) >= req.Pool.MaxInflightPerUser + if slot.hasCapacityLocked() && queued == 0 && !userInflightFull { + 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 { + if userInflightFull { + decision.RejectCode = FlowDecisionRejectPerUserInflightFull + } else { + 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 = FlowDecisionRejectClientCancelled + 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), + ScheduleActive: pool.Enabled && pool.IsScheduleActiveAt(time.Now()), + 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) userRunningLocked(userID int) int { + if userID <= 0 || s.config.MaxInflightPerUser <= 0 { + return 0 + } + count := 0 + for _, req := range s.queue { + if req.userID == userID && req.state == memoryFlowStateRunning && !req.cancelled { + count++ + } + } + return count +} + +func (s *memoryFlowSlot) dispatchLocked(now time.Time) { + for s.hasCapacityLocked() { + dispatched := false + for _, req := range s.queue { + if req.state != memoryFlowStateWaiting || req.cancelled { + continue + } + if s.config.MaxInflightPerUser > 0 && req.userID > 0 && + s.userRunningLocked(req.userID) >= s.config.MaxInflightPerUser { + 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, wasRunning bool) { + s.mu.Lock() + defer s.mu.Unlock() + now := time.Now() + for _, req := range s.queue { + if req.id == requestID && !req.cancelled { + switch req.state { + case memoryFlowStateWaiting: + req.cancelled = true + req.state = memoryFlowStateReleased + waitedMs = now.Sub(req.enqueuedAt).Milliseconds() + case memoryFlowStateRunning: + req.state = memoryFlowStateReleased + req.cancelled = true + wasRunning = true + } + break + } + } + s.compactIfNeededLocked() + s.dispatchLocked(now) + runningNow, queuedNow, _ = s.statsLocked(now) + return waitedMs, runningNow, queuedNow, wasRunning +} + +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 00000000000..b5e6274f24d --- /dev/null +++ b/service/channel_flow_redis.go @@ -0,0 +1,910 @@ +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 +) + +// redisFlowBackend implements Redis-backed channel flow state. +// +// WATCH/MULTI contention counters are backend-global: they accumulate across +// all pools served by this singleton backend. Pool-level counters would add +// work to every WATCH path; callers needing pool-level visibility should sample +// PoolStatus over time and compute deltas per pool. +type redisFlowBackend struct { + pollMin time.Duration + pollMax time.Duration + watchAttempts atomic.Int64 + txConflicts atomic.Int64 +} + +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 + userID int + 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) + } + if req.Pool.MaxContextChars > 0 && req.ContextChars > req.Pool.MaxContextChars { + decision.RejectCode = FlowDecisionRejectContextExceeded + return nil, decision, fmt.Errorf("request context chars %d exceeds flow pool max_context_chars %d", req.ContextChars, req.Pool.MaxContextChars) + } + + 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 = redisAcquireContextRejectCode(ctx, acquireCtx) + 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, redisAcquireContextError(decision.RejectCode, err) + } + + _ = b.cleanupExpired(acquireCtx, rdb, keys, req.Pool) + 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, + userID: req.UserID, + }, 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, pool) + + 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), + ScheduleActive: pool.Enabled && pool.IsScheduleActiveAt(time.Now()), + Running: int(running), + MaxInflight: pool.MaxInflight, + Queued: int(queued), + MaxQueueSize: pool.MaxQueueSize, + OldestWaitMs: oldestWaitMs, + ConfigVersion: pool.ConfigVersion, + WatchAttempts: b.watchAttempts.Load(), + TxConflicts: b.txConflicts.Load(), + }, 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)) + if req.Pool.MaxInflightPerUser > 0 { + watchKeys = append(watchKeys, keys.userRunning(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 { + userInflightFull := false + if req.Pool.MaxInflightPerUser > 0 && req.UserID > 0 { + userRunning, err := tx.ZCard(ctx, keys.userRunning(req.UserID)).Result() + if err != nil { + return err + } + userInflightFull = userRunning >= int64(req.Pool.MaxInflightPerUser) + } + if redisFlowHasCapacity(running, req.Pool.MaxInflight) && waiting == 0 && !userInflightFull { + dispatchedAt := time.Now() + expiresAtMs := dispatchedAt.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, + }) + if req.Pool.MaxInflightPerUser > 0 && req.UserID > 0 { + pipe.ZAdd(ctx, keys.userRunning(req.UserID), &redis.Z{ + Score: float64(expiresAtMs), + Member: req.RequestID, + }) + } + b.writeRequestMeta(ctx, pipe, keys, req, "running", 0, dispatchedAt.UnixMilli(), 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 + if userInflightFull { + attempt.decision.rejectCode = FlowDecisionRejectPerUserInflightFull + } else { + 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, 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 !redisFlowHasCapacity(running, req.Pool.MaxInflight) { + return nil + } + eligible, err := b.isEligibleWaitingRequest(ctx, tx, keys, req) + if err != nil { + return err + } + if !eligible { + return nil + } + if req.Pool.MaxInflightPerUser > 0 && req.UserID > 0 { + userRunning, err := tx.ZCard(ctx, keys.userRunning(req.UserID)).Result() + if err != nil { + return err + } + if userRunning >= int64(req.Pool.MaxInflightPerUser) { + return nil + } + } + dispatchedAt := time.Now() + expiresAtMs := dispatchedAt.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, + }) + if req.Pool.MaxInflightPerUser > 0 && req.UserID > 0 { + pipe.ZAdd(ctx, keys.userRunning(req.UserID), &redis.Z{ + Score: float64(expiresAtMs), + Member: req.RequestID, + }) + } + b.writeRequestMeta(ctx, pipe, keys, req, "running", 0, dispatchedAt.UnixMilli(), 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...) + b.watchAttempts.Add(1) + if errors.Is(err, redis.TxFailedErr) { + b.txConflicts.Add(1) + } + return attempt, err +} + +func (b *redisFlowBackend) isEligibleWaitingRequest(ctx context.Context, tx *redis.Tx, keys redisFlowKeys, req AcquireRequest) (bool, error) { + for start := int64(0); ; start += redisFlowCleanupBatch { + waiting, err := tx.ZRange(ctx, keys.Waiting, start, start+redisFlowCleanupBatch-1).Result() + if err != nil { + return false, err + } + for _, requestID := range waiting { + userID, err := b.requestIntFromTx(ctx, tx, keys, requestID, "user_id") + if err != nil { + return false, err + } + if req.Pool.MaxInflightPerUser > 0 && userID > 0 { + userRunning, err := tx.ZCard(ctx, keys.userRunning(userID)).Result() + if err != nil { + return false, err + } + if userRunning >= int64(req.Pool.MaxInflightPerUser) { + continue + } + } + return requestID == req.RequestID, nil + } + if len(waiting) < redisFlowCleanupBatch { + break + } + } + return false, nil +} + +func (b *redisFlowBackend) requestIntFromTx(ctx context.Context, tx *redis.Tx, keys redisFlowKeys, requestID string, field string) (int, error) { + value, err := tx.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.Atoi(value) + if err != nil { + return 0, nil + } + return parsed, nil +} + +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, userID int) error { + rdb, err := b.client() + if err != nil { + return err + } + keys := redisKeysForPool(pool) + pipe := rdb.TxPipeline() + pipe.ZRem(ctx, keys.Running, requestID) + if pool.MaxInflightPerUser > 0 && userID > 0 { + pipe.ZRem(ctx, keys.userRunning(userID), 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, userID int) 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, + }) + if pool.MaxInflightPerUser > 0 && userID > 0 { + pipe.ZAdd(ctx, keys.userRunning(userID), &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, pool model.ChannelFlowPool) error { + nowMs := time.Now().UnixMilli() + for { + expiredRunning, err := b.expiredRunningRequestIDs(ctx, rdb, keys, pool, nowMs) + if err != nil { + return err + } + if len(expiredRunning) == 0 { + break + } + if err := b.removeRunningRequests(ctx, rdb, keys, pool, expiredRunning); err != nil { + return err + } + } + for { + dirtyWaiting, err := b.dirtyWaitingRequestIDs(ctx, rdb, keys) + if err != nil { + return err + } + if len(dirtyWaiting) == 0 { + break + } + if err := b.removeWaitingRequests(ctx, rdb, keys, dirtyWaiting); err != nil { + return err + } + } + for { + 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 + } + if err := b.removeWaitingRequests(ctx, rdb, keys, expired); err != nil { + return err + } + } +} + +func (b *redisFlowBackend) expiredRunningRequestIDs(ctx context.Context, rdb *redis.Client, keys redisFlowKeys, pool model.ChannelFlowPool, nowMs int64) ([]string, error) { + seen := make(map[string]struct{}, redisFlowCleanupBatch) + expired, err := rdb.ZRangeByScore(ctx, keys.Running, &redis.ZRangeBy{ + Min: "-inf", + Max: strconv.FormatInt(nowMs, 10), + Offset: 0, + Count: redisFlowCleanupBatch, + }).Result() + if err != nil { + return nil, redisFlowUnavailable(err) + } + result := make([]string, 0, redisFlowCleanupBatch) + for _, requestID := range expired { + if _, ok := seen[requestID]; ok { + continue + } + seen[requestID] = struct{}{} + result = append(result, requestID) + } + if pool.MaxProcessingMs <= 0 || len(result) >= redisFlowCleanupBatch { + return result, nil + } + running, err := rdb.ZRange(ctx, keys.Running, 0, -1).Result() + if err != nil { + return nil, redisFlowUnavailable(err) + } + maxProcessingMs := pool.MaxProcessingMs + for _, requestID := range running { + if len(result) >= redisFlowCleanupBatch { + break + } + if _, ok := seen[requestID]; ok { + continue + } + dispatchedAtMs, err := b.requestInt64(ctx, rdb, keys, requestID, "dispatched_at_ms") + if err != nil { + return nil, err + } + if dispatchedAtMs > 0 && nowMs-dispatchedAtMs > maxProcessingMs { + seen[requestID] = struct{}{} + result = append(result, requestID) + } + } + return result, nil +} + +func (b *redisFlowBackend) dirtyWaitingRequestIDs(ctx context.Context, rdb *redis.Client, keys redisFlowKeys) ([]string, error) { + waiting, err := rdb.ZRange(ctx, keys.Waiting, 0, redisFlowCleanupBatch-1).Result() + if err != nil { + return nil, redisFlowUnavailable(err) + } + dirty := make([]string, 0, len(waiting)) + for _, requestID := range waiting { + exists, err := rdb.Exists(ctx, keys.request(requestID)).Result() + if err != nil { + return nil, redisFlowUnavailable(err) + } + if exists == 0 { + dirty = append(dirty, requestID) + continue + } + if _, err := rdb.ZScore(ctx, keys.Deadline, requestID).Result(); errors.Is(err, redis.Nil) { + dirty = append(dirty, requestID) + } else if err != nil { + return nil, redisFlowUnavailable(err) + } + } + return dirty, nil +} + +func (b *redisFlowBackend) removeRunningRequests(ctx context.Context, rdb *redis.Client, keys redisFlowKeys, pool model.ChannelFlowPool, requestIDs []string) error { + pipe := rdb.TxPipeline() + for _, requestID := range requestIDs { + pipe.ZRem(ctx, keys.Running, requestID) + userID, _ := b.requestInt(ctx, rdb, keys, requestID, "user_id") + if pool.MaxInflightPerUser > 0 && userID > 0 { + pipe.ZRem(ctx, keys.userRunning(userID), requestID) + } + pipe.Del(ctx, keys.request(requestID)) + } + _, err := pipe.Exec(ctx) + return redisFlowUnavailable(err) +} + +func (b *redisFlowBackend) removeWaitingRequests(ctx context.Context, rdb *redis.Client, keys redisFlowKeys, requestIDs []string) error { + pipe := rdb.TxPipeline() + for _, requestID := range requestIDs { + 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, dispatchedAtMs 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 dispatchedAtMs > 0 { + data["dispatched_at_ms"] = strconv.FormatInt(dispatchedAtMs, 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, g.userID) +} + +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, g.userID) +} + +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 (k redisFlowKeys) userRunning(userID int) string { + return fmt.Sprintf("%s:user:%d:running", 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 redisAcquireContextRejectCode(parent context.Context, acquireCtx context.Context) string { + if parent != nil && parent.Err() != nil { + return FlowDecisionRejectClientCancelled + } + if acquireCtx != nil && errors.Is(acquireCtx.Err(), context.Canceled) { + return FlowDecisionRejectClientCancelled + } + return FlowDecisionRejectQueueTimeout +} + +func redisAcquireContextError(rejectCode string, err error) error { + if rejectCode == FlowDecisionRejectClientCancelled { + return err + } + return fmt.Errorf("channel flow queue timeout") +} + +func redisRejectError(code string) error { + switch code { + case FlowDecisionRejectClientCancelled: + return context.Canceled + case FlowDecisionRejectPerUserInflightFull: + return fmt.Errorf("channel flow per-user inflight limit reached") + 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_status_sampler.go b/service/channel_flow_status_sampler.go new file mode 100644 index 00000000000..4e11bebfc2c --- /dev/null +++ b/service/channel_flow_status_sampler.go @@ -0,0 +1,78 @@ +package service + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + channelflowmetrics "github.com/QuantumNous/new-api/pkg/channel_flow_metrics" + + "github.com/bytedance/gopkg/util/gopool" +) + +const ( + channelFlowStatusSampleInterval = 30 * time.Second + channelFlowStatusSampleTimeout = 5 * time.Second +) + +var ( + channelFlowStatusSampleOnce sync.Once + channelFlowStatusSampleRunning atomic.Bool +) + +func StartChannelFlowStatusSampler() { + channelFlowStatusSampleOnce.Do(func() { + if !common.IsMasterNode { + return + } + gopool.Go(func() { + logger.LogInfo(context.Background(), fmt.Sprintf("channel flow status sampler started: tick=%s", channelFlowStatusSampleInterval)) + ticker := time.NewTicker(channelFlowStatusSampleInterval) + defer ticker.Stop() + + runChannelFlowStatusSampleOnce() + for range ticker.C { + runChannelFlowStatusSampleOnce() + } + }) + }) +} + +func runChannelFlowStatusSampleOnce() { + if !channelFlowStatusSampleRunning.CompareAndSwap(false, true) { + return + } + defer channelFlowStatusSampleRunning.Store(false) + + pools, err := model.ListEnabledChannelFlowPools() + if err != nil { + logger.LogWarn(context.Background(), fmt.Sprintf("channel flow status sampler: query pools failed: %v", err)) + return + } + for _, pool := range pools { + if pool == nil || pool.PoolKey == "" { + continue + } + ctx, cancel := context.WithTimeout(context.Background(), channelFlowStatusSampleTimeout) + status, err := GetChannelFlowPoolStatus(ctx, *pool) + cancel() + if err != nil { + logger.LogWarn(context.Background(), fmt.Sprintf("channel flow status sampler: pool=%s status failed: %v", pool.PoolKey, err)) + continue + } + if status.Running <= 0 && status.Queued <= 0 { + continue + } + channelflowmetrics.Record(channelflowmetrics.Sample{ + PoolKey: pool.PoolKey, + EventType: model.ChannelFlowEventStatusSample, + Running: status.Running, + Queued: status.Queued, + }) + } +} diff --git a/service/channel_flow_test.go b/service/channel_flow_test.go new file mode 100644 index 00000000000..e5ceee423ec --- /dev/null +++ b/service/channel_flow_test.go @@ -0,0 +1,1414 @@ +package service + +import ( + "context" + "fmt" + "os" + "sync" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/go-redis/redis/v8" + "github.com/stretchr/testify/require" +) + +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 TestBuildChannelFlowAcquireRequestIncludesRelayContextChars(t *testing.T) { + pool := testFlowPool() + info := &relaycommon.RelayInfo{ + UserId: 7, + TokenId: 11, + OriginModelName: "gpt-test", + Request: &dto.GeneralOpenAIRequest{ + Messages: []dto.Message{{Role: "user", Content: "hello"}}, + }, + } + info.SetEstimatePromptTokens(42) + pool.MaxContextChars = 100 + + req := buildChannelFlowAcquireRequest("req-context-chars", pool, 99, info, time.UnixMilli(1234)) + + require.Equal(t, "req-context-chars", req.RequestID) + require.Equal(t, 99, req.ChannelID) + require.Equal(t, 42, req.ContextTokens) + require.Equal(t, len([]rune("user\nhello")), req.ContextChars) + require.Equal(t, int64(1234), req.CreatedAtMs) + require.Equal(t, pool.QueueTimeoutMs, req.QueueTimeoutMs) +} + +func TestChannelFlowFallbackOnlyPassesCapacityRejections(t *testing.T) { + pool := testFlowPool() + pool.OnLimit = model.ChannelFlowOnLimitFallback + + require.True(t, shouldPassThroughChannelFlowFallback(pool, &AcquireDecision{RejectCode: FlowDecisionRejectQueueFull}, fmt.Errorf("queue full"))) + require.True(t, shouldPassThroughChannelFlowFallback(pool, &AcquireDecision{RejectCode: FlowDecisionRejectPerUserQueueFull}, fmt.Errorf("per-user queue full"))) + require.True(t, shouldPassThroughChannelFlowFallback(pool, &AcquireDecision{RejectCode: FlowDecisionRejectPerUserInflightFull}, fmt.Errorf("per-user inflight full"))) + require.False(t, shouldPassThroughChannelFlowFallback(pool, &AcquireDecision{RejectCode: FlowDecisionRejectContextExceeded}, fmt.Errorf("context exceeded"))) + require.False(t, shouldPassThroughChannelFlowFallback(pool, &AcquireDecision{RejectCode: FlowDecisionRejectBackendDisabled}, fmt.Errorf("backend disabled"))) + require.False(t, shouldPassThroughChannelFlowFallback(pool, nil, fmt.Errorf("unknown acquire failure"))) + + pool.OnLimit = model.ChannelFlowOnLimitReject + require.False(t, shouldPassThroughChannelFlowFallback(pool, &AcquireDecision{RejectCode: FlowDecisionRejectQueueFull}, fmt.Errorf("queue full"))) +} + +func TestMemoryFlowBackendReleaseDispatchesWaitingRequest(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err, "first acquire failed") + require.NotNil(t, guard1, "first acquire should be admitted immediately") + + 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: + require.NoError(t, err, "waiting acquire failed") + 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, + }) + require.NoError(t, err, "first acquire failed") + 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, + }) + require.Error(t, err, "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, + }) + require.NoError(t, err, "first acquire failed") + 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, + }) + require.Error(t, err, "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, + }) + require.NoError(t, err, "first acquire failed") + 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, + }) + require.Error(t, err, "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, + }) + require.NoError(t, err, "fallback acquire failed") + defer guard.Release(context.Background()) + + status, err := GetChannelFlowPoolStatus(context.Background(), pool) + require.NoError(t, err, "status failed") + 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, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err, "first redis acquire failed") + require.NotNil(t, guard1, "first redis acquire should be admitted") + + 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: + require.NoError(t, err, "waiting redis acquire failed") + 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, + }) + require.NoError(t, err, "first redis acquire failed") + 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, + }) + require.Error(t, err, "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: + require.NoError(t, 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, + }) + require.NoError(t, err, "first redis acquire failed") + 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, + }) + require.Error(t, err, "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() + if os.Getenv("CHANNEL_FLOW_REDIS_TEST") != "1" { + t.Skip("CHANNEL_FLOW_REDIS_TEST=1 is not set") + } + redisURL := os.Getenv("REDIS_CONN_STRING") + if redisURL == "" { + t.Skip("REDIS_CONN_STRING is not set") + } + opt, err := redis.ParseURL(redisURL) + require.NoError(t, err, "parse redis url") + 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() + require.NoError(t, err, "scan redis flow keys") + 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) +} + +func assertFlowStatus(t *testing.T, backend FlowBackend, pool model.ChannelFlowPool, wantRunning, wantQueued int) { + t.Helper() + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err) + require.Equal(t, wantRunning, status.Running) + require.Equal(t, wantQueued, status.Queued) +} + +// ── Lifecycle Consistency Tests ───────────────────────────────────────── +// +// These tests verify Phase 1/P0 lifecycle guarantees: +// - Guard.Release() is idempotent (safe to call multiple times) +// - Client abort during wait properly cleans up and releases capacity +// - max_inflight_per_user limits are enforced (Memory backend) + +func TestMemoryFlowGuardReleaseIdempotent(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + + guard, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "idempotent-req", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err, "acquire failed") + if !decision.Admitted { + t.Fatalf("should be admitted immediately") + } + + // First Release must succeed and free capacity + if err := guard.Release(context.Background()); err != nil { + t.Fatalf("first release failed: %v", err) + } + + // Second release must be a no-op (not panic, not error) + if err := guard.Release(context.Background()); err != nil { + t.Fatalf("second release should be no-op: %v", err) + } + + // Third release via BindRelease callback — also no-op + if err := guard.Release(context.Background()); err != nil { + t.Fatalf("third release should be no-op: %v", err) + } + + // Capacity must be restored after first release + guard2, decision2, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "idempotent-req-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err, "acquire after idempotent releases failed") + if !decision2.Admitted { + t.Fatalf("capacity should be available after release") + } + guard2.Release(context.Background()) +} + +func TestMemoryFlowBackendClientAbortReleasesCapacity(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + + // Fill inflight to capacity + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "abort-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err, "first acquire failed") + defer guard1.Release(context.Background()) + + // Create cancellable context to simulate client abort + abortCtx, abortCancel := context.WithCancel(context.Background()) + + resultCh := make(chan *AcquireDecision, 1) + go func() { + _, decision, _ := backend.Acquire(abortCtx, AcquireRequest{ + RequestID: "abort-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: 5000, // long timeout so abort is the trigger + }) + resultCh <- decision + }() + + time.Sleep(50 * time.Millisecond) + + // Verify queued + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err, "status failed") + if status.Queued != 1 { + t.Fatalf("expected 1 queued, got %d", status.Queued) + } + + // Simulate client abort + abortCancel() + + select { + case decision := <-resultCh: + require.NotNil(t, decision, "acquire should return decision on client abort") + require.Equal(t, FlowDecisionRejectClientCancelled, decision.RejectCode) + case <-time.After(time.Second): + t.Fatal("acquire did not return after client abort") + } + + // After abort, queued count should be 0 + status, err = backend.Status(context.Background(), pool) + require.NoError(t, err, "status after abort failed") + if status.Queued != 0 { + t.Fatalf("expected 0 queued after abort, got %d", status.Queued) + } + if status.Running != 1 { + t.Fatalf("running count should remain 1, got %d", status.Running) + } +} + +func TestMemoryFlowBackendQueueTimeoutRejectCode(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + pool.QueueTimeoutMs = 30 + + guard, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "timeout-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + defer guard.Release(context.Background()) + + _, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "timeout-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.Error(t, err) + require.NotNil(t, decision) + require.Equal(t, FlowDecisionRejectQueueTimeout, decision.RejectCode) + + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err) + require.Equal(t, 1, status.Running) + require.Equal(t, 0, status.Queued) +} + +func TestMemoryFlowBackendCleanupExpiredRunning(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + pool.MaxInflight = 2 + pool.MaxProcessingMs = 30 + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "cleanup-expired-running-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + guard2, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "cleanup-expired-running-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + + time.Sleep(60 * time.Millisecond) + + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err) + require.Equal(t, 0, status.Running) + require.Equal(t, 0, status.Queued) + require.NoError(t, guard1.Release(context.Background())) + require.NoError(t, guard2.Release(context.Background())) +} + +func TestMemoryFlowBackendDispatchAfterCleanup(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + pool.MaxInflight = 1 + pool.MaxQueueSize = 4 + pool.MaxProcessingMs = 40 + + type guardResult struct { + guard FlowGuard + decision *AcquireDecision + err error + } + resultCh := make(chan guardResult, 2) + + guard1, decision1, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "dispatch-cleanup-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: 10000, + }) + require.NoError(t, err) + require.True(t, decision1.Admitted) + assertFlowStatus(t, backend, pool, 1, 0) + + go func() { + guard, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "dispatch-cleanup-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: 10000, + }) + resultCh <- guardResult{guard: guard, decision: decision, err: err} + }() + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 1 + }) + + time.Sleep(70 * time.Millisecond) + assertFlowStatus(t, backend, pool, 0, 1) + + ctx3, cancel3 := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel3() + go func() { + guard, decision, err := backend.Acquire(ctx3, AcquireRequest{ + RequestID: "dispatch-cleanup-3", + Pool: pool, + UserID: 3, + QueueTimeoutMs: 10000, + }) + resultCh <- guardResult{guard: guard, decision: decision, err: err} + }() + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 1 + }) + + var guard2 FlowGuard + select { + case result := <-resultCh: + require.NoError(t, result.err) + require.NotNil(t, result.guard) + require.NotNil(t, result.decision) + require.True(t, result.decision.Admitted) + require.True(t, result.decision.Queued) + guard2 = result.guard + case <-time.After(2 * time.Second): + t.Fatal("queued request was not dispatched after cleanup freed capacity") + } + + require.NoError(t, guard2.Release(context.Background())) + assertFlowStatus(t, backend, pool, 1, 0) + + select { + case result := <-resultCh: + require.NoError(t, result.err) + require.NotNil(t, result.guard) + require.NotNil(t, result.decision) + require.True(t, result.decision.Admitted) + require.NoError(t, result.guard.Release(context.Background())) + case <-time.After(2 * time.Second): + t.Fatal("third request was not dispatched after releasing promoted guard") + } + + assertFlowStatus(t, backend, pool, 0, 0) + require.NoError(t, guard1.Release(context.Background())) +} + +func TestMemoryFlowBackendMaxInflightPerUser(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + pool.MaxInflight = 5 + pool.MaxInflightPerUser = 2 + pool.MaxQueueSize = 2 + + // User 1: 2 requests should be admitted (hits max_inflight_per_user) + guard1a, d1a, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "user1-req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil || !d1a.Admitted { + t.Fatalf("user1 req1 should be admitted: %v, decision=%+v", err, d1a) + } + defer guard1a.Release(context.Background()) + + guard1b, d1b, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "user1-req-2", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil || !d1b.Admitted { + t.Fatalf("user1 req2 should be admitted: %v, decision=%+v", err, d1b) + } + defer guard1b.Release(context.Background()) + + user1Third := make(chan error, 1) + user1ThirdCtx, cancelUser1Third := context.WithCancel(context.Background()) + defer cancelUser1Third() + go func() { + guard, decision, err := backend.Acquire(user1ThirdCtx, AcquireRequest{ + RequestID: "user1-req-3", + Pool: pool, + UserID: 1, + QueueTimeoutMs: 5000, + }) + if err == nil { + if guard == nil || decision == nil || !decision.Admitted || !decision.Queued { + user1Third <- fmt.Errorf("expected queued admission after release, decision=%+v guard=%v", decision, guard) + return + } + _ = guard.Release(context.Background()) + } + user1Third <- err + }() + + time.Sleep(50 * time.Millisecond) + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err, "status failed") + if status.Queued != 1 { + t.Fatalf("expected user1 third request to queue at per-user inflight limit, got queued=%d", status.Queued) + } + + // User 2: should still be admitted (different user, pool has capacity) + guard2, d2, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "user2-req-1", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil || !d2.Admitted { + t.Fatalf("user2 req1 should be admitted: %v, decision=%+v", err, d2) + } + defer guard2.Release(context.Background()) + + if err := guard1a.Release(context.Background()); err != nil { + t.Fatalf("release user1 req1 failed: %v", err) + } + select { + case err := <-user1Third: + require.NoError(t, err, "user1 queued request should be admitted after release") + case <-time.After(time.Second): + t.Fatal("user1 queued request was not admitted after release") + } +} + +func TestMemoryFlowBackendDispatchRespectsMaxInflightPerUser(t *testing.T) { + backend := NewMemoryFlowBackend() + pool := testFlowPool() + pool.MaxInflight = 2 + pool.MaxInflightPerUser = 1 + pool.MaxQueueSize = 5 + + // User 1: fill inflight slot + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "dispatch-user1-req-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: 5000, + }) + require.NoError(t, err, "user1 req1 acquire failed") + + // User 2: fill another inflight slot + guard2, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "dispatch-user2-req-1", + Pool: pool, + UserID: 2, + QueueTimeoutMs: 5000, + }) + require.NoError(t, err, "user2 req1 acquire failed") + + // User 1: queued (2nd request, user1 already has 1 running) + ch1 := make(chan error, 1) + go func() { + _, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "dispatch-user1-req-2", + Pool: pool, + UserID: 1, + QueueTimeoutMs: 5000, + }) + ch1 <- err + }() + time.Sleep(50 * time.Millisecond) + + // User 3: queued (user3 has 0 running, should be dispatchable) + ch3 := make(chan error, 1) + go func() { + g3, d3, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "dispatch-user3-req-1", + Pool: pool, + UserID: 3, + QueueTimeoutMs: 5000, + }) + if err == nil && d3.Admitted { + g3.Release(context.Background()) + } + ch3 <- err + }() + time.Sleep(50 * time.Millisecond) + + // Release user2's slot — user3 should be dispatched (not user1, since user1 already at max_inflight_per_user) + if err := guard2.Release(context.Background()); err != nil { + t.Fatalf("guard2 release failed: %v", err) + } + + select { + case err := <-ch3: + require.NoError(t, err, "user3 should be admitted after release") + case <-time.After(time.Second): + t.Fatal("user3 was not dispatched after release — dispatch may be blocked by user1's per-user limit") + } + + // Cleanup + guard1.Release(context.Background()) + + // user1's queued request should timeout or complete + select { + case <-ch1: + case <-time.After(2 * time.Second): + } +} + +// ── Redis Lifecycle Tests ─────────────────────────────────────────────── +// These are only run when REDIS_CONN_STRING is set. + +func TestRedisFlowGuardReleaseIdempotent(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + + guard, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-idempotent-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err, "redis acquire failed") + if !decision.Admitted { + t.Fatalf("should be admitted immediately") + } + + // First Release + if err := guard.Release(context.Background()); err != nil { + t.Fatalf("first release failed: %v", err) + } + + // Second Release must be no-op + if err := guard.Release(context.Background()); err != nil { + t.Fatalf("second release should be no-op: %v", err) + } + + // Third Release also no-op + if err := guard.Release(context.Background()); err != nil { + t.Fatalf("third release should be no-op: %v", err) + } + + // Capacity restored + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 0 + }) +} + +func TestRedisFlowBackendMaxInflightPerUser(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxInflight = 3 + pool.MaxInflightPerUser = 2 + pool.MaxQueueSize = 0 + + // User 1: 2 requests admitted + guard1a, d1a, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-max-inflight-user1-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil || !d1a.Admitted { + t.Fatalf("user1 req1 should be admitted: %v, decision=%+v", err, d1a) + } + defer guard1a.Release(context.Background()) + + guard1b, d1b, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-max-inflight-user1-2", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil || !d1b.Admitted { + t.Fatalf("user1 req2 should be admitted: %v, decision=%+v", err, d1b) + } + defer guard1b.Release(context.Background()) + + // User 2: 1 request admitted (different user) + guard2, d2, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-max-inflight-user2-1", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err != nil || !d2.Admitted { + t.Fatalf("user2 req1 should be admitted: %v, decision=%+v", err, d2) + } + defer guard2.Release(context.Background()) + + // User 1: 3rd request cannot acquire (per-user inflight limit already hit) + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 3 + }) + + // User1's 3rd request can queue, but after release it must not be promoted + // because user1 already has 2 running + waitCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + queuedCh := make(chan error, 1) + go func() { + _, _, err := backend.Acquire(waitCtx, AcquireRequest{ + RequestID: "redis-max-inflight-user1-3", + Pool: pool, + UserID: 1, + QueueTimeoutMs: 10000, + }) + queuedCh <- err + }() + + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Queued >= 1 + }) + + // Also enqueue User3 which should be promoted when a slot opens + waitCtx2, cancel2 := context.WithCancel(context.Background()) + defer cancel2() + user3Ch := make(chan error, 1) + go func() { + g, d, err := backend.Acquire(waitCtx2, AcquireRequest{ + RequestID: "redis-max-inflight-user3-1", + Pool: pool, + UserID: 3, + QueueTimeoutMs: 10000, + }) + if err == nil && d.Admitted { + g.Release(context.Background()) + } + user3Ch <- err + }() + + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Queued >= 2 + }) + + // Release user2's slot — user3 should get it (user1 already at per-user limit) + guard2.Release(context.Background()) + + select { + case err := <-user3Ch: + require.NoError(t, err, "user3 should be admitted after release") + case <-time.After(3 * time.Second): + t.Fatal("user3 was not promoted — dispatch may be blocked by per-user inflight limit") + } + + cancel2() + cancel() + select { + case <-queuedCh: + case <-time.After(2 * time.Second): + } +} + +func TestRedisFlowBackendClientAbortRejectCode(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + + guard, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-abort-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + defer guard.Release(context.Background()) + + abortCtx, abortCancel := context.WithCancel(context.Background()) + resultCh := make(chan *AcquireDecision, 1) + go func() { + _, decision, _ := backend.Acquire(abortCtx, AcquireRequest{ + RequestID: "redis-abort-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: 5000, + }) + resultCh <- decision + }() + + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 1 + }) + + abortCancel() + select { + case decision := <-resultCh: + require.NotNil(t, decision) + require.Equal(t, FlowDecisionRejectClientCancelled, decision.RejectCode) + case <-time.After(2 * time.Second): + t.Fatal("redis acquire did not return after client abort") + } + + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 0 + }) +} + +func TestRedisFlowBackendMaxProcessingCleanupIgnoresRenewedLease(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxProcessingMs = 80 + pool.LeaseMs = 1000 + + guard, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-max-processing-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + + time.Sleep(40 * time.Millisecond) + require.NoError(t, guard.RenewLease(context.Background())) + time.Sleep(70 * time.Millisecond) + + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err) + require.Equal(t, 0, status.Running, "max_processing_ms should release running request even when lease was renewed") +} + +func TestRedisFlowBackendCleanupDrainsExpiredRunningBatch(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxInflight = redisFlowCleanupBatch + 5 + pool.LeaseMs = 20 + + for i := 0; i < redisFlowCleanupBatch+5; i++ { + guard, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: fmt.Sprintf("redis-expired-running-%d", i), + Pool: pool, + UserID: i + 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + require.NotNil(t, guard) + } + + time.Sleep(60 * time.Millisecond) + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err) + require.Equal(t, 0, status.Running, "cleanup should drain more than one expired running batch") +} + +func TestRedisFlowBackendStatusReportsWatchContention(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxInflight = 1 + pool.MaxQueueSize = 100 + pool.QueueTimeoutMs = 15000 + pool.LeaseMs = 30000 + + const workers = 30 + resultCh := make(chan error, workers) + start := make(chan struct{}) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + i := i + wg.Add(1) + go func() { + defer wg.Done() + <-start + guard, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: fmt.Sprintf("redis-contention-%d", i), + Pool: pool, + UserID: i + 1, + QueueTimeoutMs: 15000, + }) + if guard != nil { + time.Sleep(5 * time.Millisecond) + _ = guard.Release(context.Background()) + } + resultCh <- err + }() + } + close(start) + wg.Wait() + for i := 0; i < workers; i++ { + select { + case err := <-resultCh: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("contention acquire did not finish") + } + } + + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err) + require.Greater(t, status.WatchAttempts, int64(0)) + require.GreaterOrEqual(t, status.TxConflicts, int64(0)) + if status.TxConflicts > 0 { + t.Logf("WATCH/MULTI contention confirmed: WatchAttempts=%d TxConflicts=%d conflict_rate=%.2f%%", + status.WatchAttempts, + status.TxConflicts, + float64(status.TxConflicts)/float64(status.WatchAttempts)*100) + } else { + t.Logf("No WATCH/MULTI conflicts observed in this run (WatchAttempts=%d). "+ + "This can happen when local Redis completes WATCH/EXEC faster than competing goroutines overlap; "+ + "acceptable follow-up observations are high-concurrency spike runs, multi-instance E2E, or "+ + "production PoolStatus deltas for TxConflicts.", status.WatchAttempts) + } +} + +func TestRedisFlowBackendDirtyHeadCleanup(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxQueueSize = 5 + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-dirty-head-running", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + + keys := redisKeysForPool(pool) + require.NoError(t, common.RDB.ZAdd(context.Background(), keys.Waiting, &redis.Z{ + Score: 1, + Member: "redis-dirty-head-stale", + }).Err()) + require.NoError(t, common.RDB.Set(context.Background(), keys.Seq, 1, 0).Err()) + + resultCh := make(chan error, 1) + go func() { + guard2, decision2, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-dirty-head-valid", + 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("valid request was not admitted after dirty head cleanup: 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 + }) + require.NoError(t, guard1.Release(context.Background())) + + select { + case err := <-resultCh: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("valid request was blocked behind stale waiting head") + } +} + +func TestRedisFlowBackendLeaseRenewal(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.LeaseMs = 80 + + guard, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-lease-renewal", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + time.Sleep(50 * time.Millisecond) + require.NoError(t, guard.RenewLease(context.Background())) + time.Sleep(50 * time.Millisecond) + + status, err := backend.Status(context.Background(), pool) + require.NoError(t, err) + require.Equal(t, 1, status.Running, "renewed lease should keep request running") + + time.Sleep(60 * time.Millisecond) + status, err = backend.Status(context.Background(), pool) + require.NoError(t, err) + require.Equal(t, 0, status.Running, "request should expire after renewed lease elapses") +} + +func TestRedisFlowBackendFIFOOrdering(t *testing.T) { + backend, pool, cleanup := newRedisFlowBackendForTest(t) + defer cleanup() + pool.MaxQueueSize = 5 + + guard1, _, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-fifo-1", + Pool: pool, + UserID: 1, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + require.NoError(t, err) + + admittedCh := make(chan string, 2) + go func() { + guard, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-fifo-2", + Pool: pool, + UserID: 2, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err == nil && guard != nil && decision != nil && decision.Admitted { + admittedCh <- "redis-fifo-2" + _ = guard.Release(context.Background()) + return + } + admittedCh <- "error-2" + }() + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 1 + }) + go func() { + guard, decision, err := backend.Acquire(context.Background(), AcquireRequest{ + RequestID: "redis-fifo-3", + Pool: pool, + UserID: 3, + QueueTimeoutMs: pool.QueueTimeoutMs, + }) + if err == nil && guard != nil && decision != nil && decision.Admitted { + admittedCh <- "redis-fifo-3" + _ = guard.Release(context.Background()) + return + } + admittedCh <- "error-3" + }() + eventuallyFlowStatus(t, backend, pool, func(status PoolStatus) bool { + return status.Running == 1 && status.Queued == 2 + }) + require.NoError(t, guard1.Release(context.Background())) + + select { + case requestID := <-admittedCh: + require.Equal(t, "redis-fifo-2", requestID) + case <-time.After(2 * time.Second): + t.Fatal("first queued request was not admitted") + } +} + +func TestRedisFlowOutagePolicyFailClosed(t *testing.T) { + oldRedisEnabled := common.RedisEnabled + oldRDB := common.RDB + common.RedisEnabled = false + common.RDB = nil + defer func() { + common.RedisEnabled = oldRedisEnabled + common.RDB = oldRDB + }() + + pool := testFlowPool() + pool.Backend = model.ChannelFlowBackendRedis + pool.RedisFailurePolicy = model.ChannelFlowRedisFailureFailClosed + + passThrough, fallbackPool, apiErr := resolveRedisFlowUnavailable(context.Background(), &pool) + require.False(t, passThrough) + require.Nil(t, fallbackPool) + require.NotNil(t, apiErr) +} + +func TestRedisFlowOutagePolicyFailOpen(t *testing.T) { + oldRedisEnabled := common.RedisEnabled + oldRDB := common.RDB + common.RedisEnabled = false + common.RDB = nil + defer func() { + common.RedisEnabled = oldRedisEnabled + common.RDB = oldRDB + }() + + pool := testFlowPool() + pool.Backend = model.ChannelFlowBackendRedis + pool.RedisFailurePolicy = model.ChannelFlowRedisFailureFailOpen + + passThrough, fallbackPool, apiErr := resolveRedisFlowUnavailable(context.Background(), &pool) + require.True(t, passThrough) + require.Nil(t, fallbackPool) + require.Nil(t, apiErr) +} + +func TestRedisFlowOutagePolicyLocalMemory(t *testing.T) { + oldRedisEnabled := common.RedisEnabled + oldRDB := common.RDB + common.RedisEnabled = false + common.RDB = nil + defer func() { + common.RedisEnabled = oldRedisEnabled + common.RDB = oldRDB + }() + + pool := testFlowPool() + pool.Backend = model.ChannelFlowBackendRedis + pool.RedisFailurePolicy = model.ChannelFlowRedisFailureLocalMemory + + passThrough, fallbackPool, apiErr := resolveRedisFlowUnavailable(context.Background(), &pool) + require.False(t, passThrough) + require.NotNil(t, fallbackPool) + require.Nil(t, apiErr) + require.Equal(t, model.ChannelFlowBackendMemory, fallbackPool.Backend) +} diff --git a/tools/channel-flow-spike/README.md b/tools/channel-flow-spike/README.md new file mode 100644 index 00000000000..336c9070a19 --- /dev/null +++ b/tools/channel-flow-spike/README.md @@ -0,0 +1,31 @@ +# 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 00000000000..bf0a0e84169 --- /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 +} diff --git a/types/error.go b/types/error.go index 9717401ae7b..a7246df8c1f 100644 --- a/types/error.go +++ b/types/error.go @@ -85,6 +85,18 @@ 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" + ErrorCodeChannelFlowClientCancelled ErrorCode = "channel_flow_client_cancelled" + 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" + ErrorCodeChannelFlowPerUserInflightFull ErrorCode = "channel_flow_per_user_inflight_full" ) type NewAPIError struct { diff --git a/web/default/scripts/sync-i18n.mjs b/web/default/scripts/sync-i18n.mjs index c962f7d220b..b5210ab8b53 100644 --- a/web/default/scripts/sync-i18n.mjs +++ b/web/default/scripts/sync-i18n.mjs @@ -313,8 +313,6 @@ async function main() { } main().catch((err) => { - console.error(err) process.exitCode = 1 }) - 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 00000000000..572ffc34976 --- /dev/null +++ b/web/default/src/features/channel-flow/api.ts @@ -0,0 +1,130 @@ +/* +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, + ChannelFlowTrend, + 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 getChannelFlowPoolTrend( + poolId: number, + minutes = 60 +): Promise> { + const res = await api.get(`/api/channel_flow/pools/${poolId}/trend`, { + params: { minutes }, + 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 00000000000..6c2eea96a61 --- /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 = 1000 + +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 availableChannels = useMemo(() => { + const rawChannels = channelsQuery.data?.data?.items ?? [] + const boundChannelIds = new Set( + props.bindings + .filter((binding) => binding.enabled) + .map((binding) => binding.channel_id) + ) + return rawChannels.filter((channel) => !boundChannelIds.has(channel.id)) + }, [channelsQuery.data, 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 00000000000..0da5f8930ef --- /dev/null +++ b/web/default/src/features/channel-flow/components/pool-bindings-panel.tsx @@ -0,0 +1,146 @@ +/* +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 { useMemo } from 'react' +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() + const { pool, bindings, loading, deletingBindingId, onAddBinding, onDeleteBinding } = props + + const columns = useMemo( + () => [ + { + id: 'channel', + header: t('Channel'), + cell: (binding: ChannelFlowPoolBinding) => ( +
+
+ #{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: ChannelFlowPoolBinding) => ( + + {binding.match_mode === 'channel_model' + ? t('Channel and model') + : t('Channel')} + + ), + }, + { + id: 'enabled', + header: t('Status'), + cell: (binding: ChannelFlowPoolBinding) => ( + + {binding.enabled ? t('Enabled') : t('Disabled')} + + ), + }, + { + id: 'actions', + header: '', + className: 'w-16 text-right', + cellClassName: 'text-right', + cell: (binding: ChannelFlowPoolBinding) => ( + + ), + }, + ], + [t, deletingBindingId, onDeleteBinding] + ) + const emptyContent = useMemo( + () => ( + + {pool + ? t('No channels bound to this Flow Pool') + : t('Select a Flow Pool to view bindings')} + + ), + [pool, t] + ) + + return ( +
+
+
+

{t('Channel bindings')}

+

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

+
+ +
+ + {loading ? ( +
+ {Array.from({ length: 3 }).map((_, index) => ( + + ))} +
+ ) : ( + binding.id} + columns={columns} + emptyContent={emptyContent} + /> + )} +
+ ) +} 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 00000000000..3c5d0797789 --- /dev/null +++ b/web/default/src/features/channel-flow/components/pool-form-sheet.tsx @@ -0,0 +1,673 @@ +/* +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 { Loader2 } from 'lucide-react' +import { useTranslation } from 'react-i18next' +import { formatTimestampForInput, parseTimestampFromInput } from '@/lib/format' +import { Button } from '@/components/ui/button' +import { Checkbox } from '@/components/ui/checkbox' +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 +} + +// eslint-disable-next-line @typescript-eslint/no-unused-vars +const numberFields = [ + 'max_inflight', + 'max_inflight_per_user', + '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 scheduleMode = form.watch('schedule_mode') + 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') }, + ] + const scheduleModeOptions: SelectOption< + ChannelFlowPoolFormValues['schedule_mode'] + >[] = [ + { value: 'always', label: t('Always active') }, + { value: 'datetime_range', label: t('Date range') }, + { value: 'weekly', label: t('Weekly schedule') }, + ] + const weekdayOptions = [ + { value: 0, label: t('Sun') }, + { value: 1, label: t('Mon') }, + { value: 2, label: t('Tue') }, + { value: 3, label: t('Wed') }, + { value: 4, label: t('Thu') }, + { value: 5, label: t('Fri') }, + { value: 6, label: t('Sat') }, + ] + + 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')} + +