diff --git a/charts/kthena/charts/networking/README.md b/charts/kthena/charts/networking/README.md
index c7a7a0328..30aa3cd63 100644
--- a/charts/kthena/charts/networking/README.md
+++ b/charts/kthena/charts/networking/README.md
@@ -52,18 +52,19 @@ kthenaRouter:
#### Configuration Parameters
-| Parameter | Type | Default | Description |
-| ------------------------------------------ | ------- | ---------------- | ------------------------------------------------------------------ |
-| `kthenaRouter.fairness.enabled` | boolean | `false` | Enable user-fairness scheduling (mutually exclusive with boost) |
-| `kthenaRouter.fairness.windowSize` | string | `"1h"` | Fairness: sliding window duration (1m-1h) |
-| `kthenaRouter.fairness.inputTokenWeight` | float | `1.0` | Fairness: weight for input tokens (≥0) |
-| `kthenaRouter.fairness.outputTokenWeight` | float | `2.0` | Fairness: weight for output tokens (≥0) |
-| `kthenaRouter.fairness.maxConcurrent` | int | `0` | Fairness: global inflight limit (`0` falls back to QPS mode) |
-| `kthenaRouter.sessionBoost.enabled` | boolean | `false` | Enable session-boost scheduling (mutually exclusive with fairness) |
-| `kthenaRouter.sessionBoost.header` | string | `"X-Session-ID"` | HTTP header used to identify conversation sessions |
-| `kthenaRouter.sessionBoost.maxSessions` | int | `4096` | Max recently-completed sessions kept warm (LRU-evicted) |
-| `kthenaRouter.sessionBoost.inflightPerPod` | int | `16` | Inflight requests per backend pod; total = perPod x pod count |
-| `kthenaRouter.sessionBoost.gracePeriod` | string | `"0s"` | Wait time for a same-session follow-up (disabled by default) |
+| Parameter | Type | Default | Description |
+| ------------------------------------------ | ------- | ---------------- | ----------------------------------------------------------------------------- |
+| `kthenaRouter.fairness.enabled` | boolean | `false` | Enable user-fairness scheduling (mutually exclusive with boost) |
+| `kthenaRouter.fairness.windowSize` | string | `"1h"` | Fairness: sliding window duration (1m-1h) |
+| `kthenaRouter.fairness.inputTokenWeight` | float | `1.0` | Fairness: weight for input tokens (≥0) |
+| `kthenaRouter.fairness.outputTokenWeight` | float | `2.0` | Fairness: weight for output tokens (≥0) |
+| `kthenaRouter.fairness.maxConcurrent` | int | `0` | Fairness: global inflight limit (`0` falls back to QPS mode) |
+| `kthenaRouter.sessionBoost.enabled` | boolean | `false` | Enable session-boost scheduling (mutually exclusive with fairness) |
+| `kthenaRouter.sessionBoost.header` | string | `"X-Session-ID"` | HTTP header used to identify conversation sessions |
+| `kthenaRouter.sessionBoost.maxSessions` | int | `4096` | Max recently-completed sessions kept warm (LRU-evicted) |
+| `kthenaRouter.sessionBoost.inflightPerPod` | int | `16` | Inflight requests per backend pod; total = perPod x pod count |
+| `kthenaRouter.sessionBoost.gracePeriod` | string | `"0s"` | Wait time for a same-session follow-up (disabled by default) |
+| `kthenaRouter.sessionBoost.timeout` | string | `"30s"` | Max queue wait before 504; set a non-positive duration (e.g. `0s`) to disable |
#### Session Boost Configuration
diff --git a/charts/kthena/charts/networking/templates/kthena-router/component/deployment.yaml b/charts/kthena/charts/networking/templates/kthena-router/component/deployment.yaml
index 58f345282..e4087762f 100644
--- a/charts/kthena/charts/networking/templates/kthena-router/component/deployment.yaml
+++ b/charts/kthena/charts/networking/templates/kthena-router/component/deployment.yaml
@@ -106,6 +106,10 @@ spec:
{{- end }}
- name: SESSION_BOOST_GRACE_PERIOD
value: {{ .gracePeriod | quote }}
+ {{- if .timeout }}
+ - name: SESSION_BOOST_TIMEOUT
+ value: {{ .timeout | quote }}
+ {{- end }}
{{- end }}
{{- end }}
# Access log configuration
diff --git a/charts/kthena/charts/networking/values.yaml b/charts/kthena/charts/networking/values.yaml
index 54b158891..28ac1edcd 100644
--- a/charts/kthena/charts/networking/values.yaml
+++ b/charts/kthena/charts/networking/values.yaml
@@ -70,6 +70,11 @@ kthenaRouter:
# follow-up to arrive. Disabled by default (0s); enable only if you understand
# the latency trade-off for non-boosted requests.
gracePeriod: "0s"
+ # timeout is the maximum time a request may wait in the session-boost queue
+ # before it is rejected with HTTP 504. Defaults to 30s; set it to a non-positive
+ # duration (e.g. "0s") to disable the timeout, in which case a request is bounded
+ # only by client disconnect.
+ timeout: "30s"
# accessLog configuration for request logging
accessLog:
# enabled controls whether access logging is active
diff --git a/charts/kthena/values.yaml b/charts/kthena/values.yaml
index 8943e3737..8ba840225 100644
--- a/charts/kthena/values.yaml
+++ b/charts/kthena/values.yaml
@@ -100,6 +100,10 @@ networking:
# -- Wait time after a request completes for a same-session follow-up.
# Disabled by default (`0s`).
gracePeriod: "0s"
+ # -- Maximum time a request may wait in the session-boost queue before it is
+ # rejected with HTTP 504. Defaults to `30s`; set a non-positive duration
+ # (e.g. `0s`) to disable it (bounded only by client disconnect).
+ timeout: "30s"
gatewayAPI:
# -- Enable Gateway API related features.
enabled: false
diff --git a/docs/kthena/docs/reference/helm-chart-values.md b/docs/kthena/docs/reference/helm-chart-values.md
index 3c9c86a7b..abc3935f6 100644
--- a/docs/kthena/docs/reference/helm-chart-values.md
+++ b/docs/kthena/docs/reference/helm-chart-values.md
@@ -37,6 +37,7 @@ A Helm chart for deploying Kthena
| networking.kthenaRouter.sessionBoost.header | string | `"X-Session-ID"` | HTTP header used to identify conversation sessions. |
| networking.kthenaRouter.sessionBoost.inflightPerPod | int | `16` | Maximum inflight requests admitted per backend pod.
Total inflight limit is this value times the number of backend pods.
`0` or unset uses the router default (16). |
| networking.kthenaRouter.sessionBoost.maxSessions | int | `4096` | Maximum number of recently-completed sessions kept warm for boosting.
Bounds an LRU cache; the least-recently-used session is evicted automatically.
Size it by the number of concurrent conversations to keep boosted. |
+| networking.kthenaRouter.sessionBoost.timeout | string | `"30s"` | Maximum time a request may wait in the session-boost queue before it is
rejected with HTTP 504. Defaults to `30s`; set a non-positive duration
(e.g. `0s`) to disable it (bounded only by client disconnect). |
| networking.kthenaRouter.terminationGracePeriodSeconds | int | `330` | The router will drain all in-flight requests before forcefully closing connections. |
| networking.kthenaRouter.tls.dnsName | string | `"your-domain.com"` | DNS name to use for the certificate. |
| networking.kthenaRouter.tls.enabled | bool | `false` | Enable TLS for Kthena Router server. |
diff --git a/docs/kthena/docs/user-guide/fairness-scheduling.md b/docs/kthena/docs/user-guide/fairness-scheduling.md
index 73ce21f76..1cad45885 100644
--- a/docs/kthena/docs/user-guide/fairness-scheduling.md
+++ b/docs/kthena/docs/user-guide/fairness-scheduling.md
@@ -121,13 +121,13 @@ env:
### Core Settings
-| Environment Variable | Purpose | Default | Notes |
-| ------------------------------ | ----------------------------------------------------- | -------------------- | --------------------------------------------------------------------- |
-| `ENABLE_FAIRNESS_SCHEDULING` | Enables fairness scheduling in the router | `false` | Global feature switch. Mutually exclusive with `ENABLE_SESSION_BOOST` |
-| `FAIRNESS_WINDOW_SIZE` | Sliding window used to track recent usage | runtime default `5m` | The Helm chart default sets this to `1h` when fairness is enabled |
-| `FAIRNESS_INPUT_TOKEN_WEIGHT` | Weight applied to input tokens when recording usage | `1.0` | Used by the token tracker |
-| `FAIRNESS_OUTPUT_TOKEN_WEIGHT` | Weight applied to output tokens when recording usage | `2.0` | Used by the token tracker |
-| `FAIRNESS_QUEUE_TIMEOUT` | Maximum time a request may wait in the fairness queue | `60s` | Waiting longer returns a timeout to the client |
+| Environment Variable | Purpose | Default | Notes |
+| ------------------------------ | ----------------------------------------------------- | -------------------- | ---------------------------------------------------------------------------------------------------------------- |
+| `ENABLE_FAIRNESS_SCHEDULING` | Enables fairness scheduling in the router | `false` | Global feature switch. Mutually exclusive with `ENABLE_SESSION_BOOST` |
+| `FAIRNESS_WINDOW_SIZE` | Sliding window used to track recent usage | runtime default `5m` | The Helm chart default sets this to `1h` when fairness is enabled |
+| `FAIRNESS_INPUT_TOKEN_WEIGHT` | Weight applied to input tokens when recording usage | `1.0` | Used by the token tracker |
+| `FAIRNESS_OUTPUT_TOKEN_WEIGHT` | Weight applied to output tokens when recording usage | `2.0` | Used by the token tracker |
+| `FAIRNESS_QUEUE_TIMEOUT` | Maximum time a request may wait in the fairness queue | `60s` | Waiting longer returns a timeout (504) to the client. Applies only to the user-fairness queue, not session boost |
### Queue Policy Settings
diff --git a/docs/kthena/docs/user-guide/session-boost.md b/docs/kthena/docs/user-guide/session-boost.md
index 60e717069..d12b2dc44 100644
--- a/docs/kthena/docs/user-guide/session-boost.md
+++ b/docs/kthena/docs/user-guide/session-boost.md
@@ -30,7 +30,9 @@ flowchart TD
E --> F
F -- Yes --> G[Dispatch to backend]
F -- No --> H[Wait / backpressure]
- H --> F
+ H --> K{Waited longer than
SESSION_BOOST_TIMEOUT?}
+ K -- Yes --> R[Reject with HTTP 504]
+ K -- No --> F
G --> I[Request completes]
I --> J[Mark session completed
promote in LRU cache]
J --> M[Dequeue next request]
@@ -73,6 +75,7 @@ networking:
header: "X-Session-ID"
maxSessions: 4096 # LRU cache of recently-completed sessions kept warm
inflightPerPod: 16 # total inflight = inflightPerPod x backend pod count
+ timeout: "30s" # reject requests that wait longer than this with HTTP 504 (set "0s" to disable)
```
Apply with Helm:
@@ -98,16 +101,19 @@ env:
value: "4096"
- name: SESSION_BOOST_INFLIGHT_PER_POD
value: "16" # total inflight = perPod x backend pod count
+- name: SESSION_BOOST_TIMEOUT
+ value: "30s" # reject requests that wait longer than this with HTTP 504 (set "0s" to disable)
```
## Configuration Reference
-| Environment Variable | Purpose | Default | Notes |
-| -------------------------------- | -------------------------------------------------------------------------------- | -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| `ENABLE_SESSION_BOOST` | Enable the session-boost scheduling strategy | `false` | Mutually exclusive with `ENABLE_FAIRNESS_SCHEDULING` |
-| `SESSION_BOOST_HEADER` | HTTP header used to identify conversation sessions | `X-Session-ID` | Must match what your clients send |
-| `SESSION_BOOST_MAX_SESSIONS` | Max number of recently-completed sessions kept "warm" for boosting (LRU-bounded) | `4096` | When the cache is full, the least-recently-used session is evicted automatically. Size it by the number of concurrent conversations you want to keep boosted—no time-based tuning required |
-| `SESSION_BOOST_INFLIGHT_PER_POD` | Inflight requests admitted per backend pod | `16` | The total inflight limit is this value multiplied by the number of backend pods. Size it from the per-pod concurrency (e.g., vLLM's `--max-num-seqs`) |
+| Environment Variable | Purpose | Default | Notes |
+| -------------------------------- | -------------------------------------------------------------------------------- | -------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `ENABLE_SESSION_BOOST` | Enable the session-boost scheduling strategy | `false` | Mutually exclusive with `ENABLE_FAIRNESS_SCHEDULING` |
+| `SESSION_BOOST_HEADER` | HTTP header used to identify conversation sessions | `X-Session-ID` | Must match what your clients send |
+| `SESSION_BOOST_MAX_SESSIONS` | Max number of recently-completed sessions kept "warm" for boosting (LRU-bounded) | `4096` | When the cache is full, the least-recently-used session is evicted automatically. Size it by the number of concurrent conversations you want to keep boosted—no time-based tuning required |
+| `SESSION_BOOST_INFLIGHT_PER_POD` | Inflight requests admitted per backend pod | `16` | The total inflight limit is this value multiplied by the number of backend pods. Size it from the per-pod concurrency (e.g., vLLM's `--max-num-seqs`) |
+| `SESSION_BOOST_TIMEOUT` | Maximum time a request may wait in the queue before it is rejected with 504 | `30s` | Enabled by default. A request queued longer than this is rejected with `504 Gateway Timeout`. Set a non-positive duration (e.g. `0s`) to disable it. This is the only server-side queue-wait bound in session-boost mode; `FAIRNESS_QUEUE_TIMEOUT` does not apply |
> `SESSION_BOOST_GRACE_PERIOD` is intentionally omitted from the table above. It is an advanced, scenario-specific knob that is disabled by default; see [Advanced: Grace Period](#advanced-grace-period-use-with-caution).
@@ -174,6 +180,19 @@ The queue uses two-level admission control to avoid flooding backends:
When a request completes, the queue immediately attempts to dequeue the next request (release-driven dequeue) rather than waiting for the next metrics refresh.
+### Queue Wait Timeout (504 Rejection)
+
+Under heavy load a request may sit in the queue for a long time before backend capacity frees up. Session boost does **not** apply the fairness queue timeout (`FAIRNESS_QUEUE_TIMEOUT` only governs the user-fairness queue). Instead, session boost bounds the wait with its own **queue wait timeout** so latency-sensitive front ends **fail fast** and can retry or shed load rather than waiting indefinitely.
+
+The wait timeout rejects over-queued requests with `504 Gateway Timeout`:
+
+- It is controlled by `SESSION_BOOST_TIMEOUT` and **enabled by default** at `30s`.
+- Set a non-positive duration (e.g. `0s`) to disable it; a session-boost request is then bounded only by client disconnect.
+
+If a request waits in the queue longer than `SESSION_BOOST_TIMEOUT`, the router removes it from the queue and responds with `504 Gateway Timeout`. The 504 signals that the request *timed out* waiting for backend capacity (the queue is saturated); the client can retry later or shed load.
+
+> `SESSION_BOOST_TIMEOUT` is the only server-side queue-wait bound in session-boost mode. `FAIRNESS_QUEUE_TIMEOUT` has no effect here—it applies exclusively to the user-fairness strategy. When `SESSION_BOOST_TIMEOUT` is disabled (a non-positive value), a session-boost request is bounded only by client disconnect.
+
## Session Boost vs User Fairness
Session boost and user fairness are two mutually exclusive scheduling strategies for the per-model request queue; they configure that queue for different goals:
@@ -249,6 +268,14 @@ Session boost only controls queue ordering. If the boosted request is routed to
Each tracked session consumes minimal memory (just a session ID). The tracker is a bounded LRU cache, so total memory is capped at `SESSION_BOOST_MAX_SESSIONS` entries regardless of how much traffic flows through the router. If memory is a concern, reduce `SESSION_BOOST_MAX_SESSIONS`.
+### Clients receive HTTP 504 responses
+
+When `SESSION_BOOST_TIMEOUT` is set (enabled by default at `30s`), the router returns `504 Gateway Timeout` for requests that wait in the queue longer than it. This is expected load-shedding under sustained overload—it means the backends cannot keep up with demand. If you see more 504s than desired:
+
+1. Confirm the backends are healthy and scaled appropriately; add capacity if they are saturated.
+2. Increase `SESSION_BOOST_TIMEOUT` to allow requests to wait longer before rejection.
+3. Disable the timeout by setting `SESSION_BOOST_TIMEOUT` to a non-positive value (e.g. `0s`) so requests are no longer rejected on a wait timeout (they then wait until backend capacity frees up or the client disconnects).
+
## Advanced: Grace Period (Use With Caution)
:::warning
diff --git a/docs/proposal/session-boost-strategy.md b/docs/proposal/session-boost-strategy.md
index 159b83c16..05b9b12eb 100644
--- a/docs/proposal/session-boost-strategy.md
+++ b/docs/proposal/session-boost-strategy.md
@@ -39,6 +39,7 @@ Enabling session boost therefore reconfigures the same per-model priority queue
3. **KV cache optimization**: Prioritize follow-up requests from recently completed sessions to maximize warm cache hits.
4. **Grace period (advanced, off by default)**: An optional, tricky tuning knob that, after a request completes, briefly holds the dequeue slot for a potential follow-up from the same session before dispatching unrelated requests. It is **disabled by default** and should only be enabled by operators who fully understand that it deliberately delays unrelated requests in exchange for a higher same-session prefix-cache hit rate.
5. **Backpressure-aware**: Respect backend pod capacity to avoid flooding, using two-level admission control (inflight limit + backend metrics).
+6. **Fail-fast queue wait timeout (optional, off by default)**: Optionally reject requests that wait in the queue longer than a configurable threshold with HTTP 504, so latency-sensitive clients shed load or retry instead of waiting indefinitely for backend capacity. `FAIRNESS_QUEUE_TIMEOUT` governs only the user-fairness queue and does not apply in session-boost mode.
#### Non-Goals
@@ -218,13 +219,14 @@ The net effect is a strict precedence: **grace timing → inflight gate → back
#### Configuration
-| Environment Variable | Default | Description |
-| -------------------------------- | -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `ENABLE_SESSION_BOOST` | `false` | Enable the session-boost scheduling strategy (mutually exclusive with `ENABLE_FAIRNESS_SCHEDULING`) |
-| `SESSION_BOOST_HEADER` | `X-Session-ID` | HTTP header used to identify conversation sessions |
-| `SESSION_BOOST_MAX_SESSIONS` | `4096` | Maximum number of recently-completed sessions kept warm for boosting. Bounds an LRU cache; the least-recently-used session is evicted when exceeded. Sized by session count, not time |
-| `SESSION_BOOST_GRACE_PERIOD` | `0` | Wait time after release for same-session follow-up. Disabled by default; enable only when you understand the latency trade-off |
-| `SESSION_BOOST_INFLIGHT_PER_POD` | `16` | Inflight requests admitted per backend pod; total inflight = perPod x backend pod count. Size it from the estimated per-pod concurrency (e.g. vLLM's --max-num-seqs) |
+| Environment Variable | Default | Description |
+| -------------------------------- | -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| `ENABLE_SESSION_BOOST` | `false` | Enable the session-boost scheduling strategy (mutually exclusive with `ENABLE_FAIRNESS_SCHEDULING`) |
+| `SESSION_BOOST_HEADER` | `X-Session-ID` | HTTP header used to identify conversation sessions |
+| `SESSION_BOOST_MAX_SESSIONS` | `4096` | Maximum number of recently-completed sessions kept warm for boosting. Bounds an LRU cache; the least-recently-used session is evicted when exceeded. Sized by session count, not time |
+| `SESSION_BOOST_GRACE_PERIOD` | `0` | Wait time after release for same-session follow-up. Disabled by default; enable only when you understand the latency trade-off |
+| `SESSION_BOOST_INFLIGHT_PER_POD` | `16` | Inflight requests admitted per backend pod; total inflight = perPod x backend pod count. Size it from the estimated per-pod concurrency (e.g. vLLM's --max-num-seqs) |
+| `SESSION_BOOST_TIMEOUT` | `30s` | Maximum time a request may wait in the queue before it is rejected with HTTP 504. Enabled by default; set a non-positive duration (e.g. `0s`) to disable it. It is the only server-side queue-wait bound in session-boost mode (`FAIRNESS_QUEUE_TIMEOUT` does not apply) |
### Design Details
@@ -297,6 +299,23 @@ The queue uses two-level admission control:
When a request completes (Release), the queue immediately attempts to dequeue the next request (release-driven dequeue), ensuring minimal latency between sequential requests. The loop is fully event-driven — there is no independent polling timer. In single-router operation every moment a backend frees capacity coincides with one of our own requests completing (a release), so release and arrival events alone cover every dequeue opportunity; the capacity check simply reads the pod metrics already scraped by the store (`METRICS_SCRAPE_INTERVAL`).
+#### Queue Wait Timeout (504 Rejection)
+
+Under sustained overload the two-level admission control legitimately holds requests in the queue until backend capacity frees up. `FAIRNESS_QUEUE_TIMEOUT` governs **only** the user-fairness queue and does not apply in session-boost mode. Instead, session boost bounds the wait with its own timeout so latency-sensitive front ends are told *quickly* that the system is saturated and can retry elsewhere or shed load rather than waiting open-endedly (until the client disconnects, which returns `503`).
+
+The queue wait timeout provides this fail-fast behavior:
+
+- It is controlled by `SESSION_BOOST_TIMEOUT` and enabled by default at `30s`. Setting it to a non-positive duration (e.g. `0s`) disables the timeout, leaving a session-boost request bounded only by client disconnect.
+- When a request has been waiting in the queue longer than `SESSION_BOOST_TIMEOUT`, the router stops waiting, removes the request from the queue, and responds with `504 Gateway Timeout`.
+
+It is implemented in the router's request-handling path rather than in the queue's dequeue loop, mirroring how the fairness queue's own `FAIRNESS_QUEUE_TIMEOUT` (504) is handled. Both queue-wait timeouts therefore return the same `504` status; they differ only in what arms them (`FAIRNESS_QUEUE_TIMEOUT` vs `SESSION_BOOST_TIMEOUT`). In session-boost mode the request context carries the `SESSION_BOOST_TIMEOUT` deadline (via `context.WithTimeout`) instead of the fairness deadline, so `FAIRNESS_QUEUE_TIMEOUT` has no effect. The handler waits on the admission (`NotifyChan`) and cancellation/timeout (`reqCtx.Done()`) signals:
+
+- **Admitted first**: the request proceeds to the backend as usual (no rejection).
+- **Timeout fires first**: the request context's deadline is exceeded, which makes the queue drop the request from the heap via its existing cancellation check (`isCancelled` / `drainCancelledLocked`), reusing the same cleanup path as client disconnects; the handler abandons the request, releases any permit that may have been granted concurrently, and returns `504 Gateway Timeout`.
+- **Client disconnect fires first**: the existing behavior is unchanged (`503`).
+
+Because it reuses the queue's cancellation cleanup, the wait timeout adds no new state to the queue and no extra polling. It is orthogonal to the inflight and backend capacity gates: those gates decide *whether* a request may run, while the wait timeout only bounds *how long* a request is willing to wait before giving up. `SESSION_BOOST_TIMEOUT` is the only server-side queue-wait bound in session-boost mode; the feature only applies in session-boost mode.
+
### Multi-Turn Conversation Advantages
#### 1. Prefix Cache Hit Rate Improvement
diff --git a/pkg/kthena-router/datastore/fairness_queue.go b/pkg/kthena-router/datastore/fairness_queue.go
index 02c3f0c78..099887fc2 100644
--- a/pkg/kthena-router/datastore/fairness_queue.go
+++ b/pkg/kthena-router/datastore/fairness_queue.go
@@ -152,6 +152,55 @@ type Request struct {
NotifyChan chan struct{}
CancelCh <-chan struct{} // Request-scoped cancellation signal
Release func() // Set by the queue when a permit is acquired
+
+ // admitMu serializes admission (by the dequeue loop) against abandonment (by
+ // the waiting caller on timeout/cancel), closing the race where the loop has
+ // popped the request and passed its cancellation check but has not yet
+ // installed Release / closed NotifyChan. Guards admitted and abandoned.
+ admitMu sync.Mutex
+ admitted bool // admission committed: Release is set and about to be signalled
+ abandoned bool // caller gave up before admission; the loop must not admit
+}
+
+// commitAdmission runs fn under the request lock, but only if the caller has not
+// already abandoned the request. fn performs the admission side effects that must
+// be fully visible before the request is considered admitted: acquiring the
+// inflight permit, installing Release, and incrementing the inflight metric.
+// It returns true if admission was committed. When it returns false the request
+// was abandoned first, so the dequeue loop must not mark it inflight or signal it.
+//
+// Because fn (including installing Release) completes before admitted is set, any
+// caller that observes admitted via abandon() is guaranteed to see a non-nil
+// Release, so the permit can always be returned.
+func (r *Request) commitAdmission(fn func()) bool {
+ r.admitMu.Lock()
+ defer r.admitMu.Unlock()
+ if r.abandoned {
+ return false
+ }
+ fn()
+ r.admitted = true
+ return true
+}
+
+// Abandon marks the request as given up by the waiting caller (queue timeout,
+// wait-reject, or client cancellation). If the request had already been admitted,
+// the caller owned the inflight permit, so Abandon releases it to avoid leaking
+// capacity. Otherwise it marks the request abandoned so admission is guaranteed
+// not to proceed (commitAdmission observes abandoned and skips) and no permit can
+// leak.
+func (r *Request) Abandon() {
+ r.admitMu.Lock()
+ if !r.admitted {
+ r.abandoned = true
+ r.admitMu.Unlock()
+ return
+ }
+ r.admitMu.Unlock()
+ // Admission raced in first, so we own the inflight permit; release it here.
+ // Release is guaranteed non-nil once admitted is set (installed by fn before
+ // commitAdmission sets admitted).
+ r.Release()
}
// RequestPriorityQueue implements the heap.Interface
@@ -485,21 +534,27 @@ func (pq *RequestPriorityQueue) runSemaphoreMode(ctx context.Context) {
// Permit acquired
}
- releaseOnce := sync.Once{}
- trackedInflight := false
- req.Release = func() {
- releaseOnce.Do(func() {
- <-pq.sem
- if trackedInflight && pq.metrics != nil {
- pq.metrics.DecFairnessQueueInflight(req.ModelName)
- }
- })
- }
-
- if pq.metrics != nil {
- pq.metrics.IncFairnessQueueInflight(req.ModelName)
+ // Commit admission atomically with respect to a concurrent Abandon(): install
+ // Release and increment the inflight metric under the request lock so the
+ // waiting caller either observes the admission (and releases the permit on
+ // timeout) or blocks admission entirely. Increment the metric inside the
+ // committed section so a racing Release always has a matching increment.
+ admitted := req.commitAdmission(func() {
+ releaseOnce := sync.Once{}
+ req.Release = func() {
+ releaseOnce.Do(func() {
+ <-pq.sem
+ pq.metricDecInflight(req.ModelName)
+ })
+ }
+ pq.metricIncInflight(req.ModelName)
+ })
+ if !admitted {
+ // Caller abandoned before admission; return the just-acquired permit and
+ // drop the request without signalling it.
+ <-pq.sem
+ continue
}
- trackedInflight = true
close(req.NotifyChan)
}
}
diff --git a/pkg/kthena-router/datastore/session_boost_queue.go b/pkg/kthena-router/datastore/session_boost_queue.go
index e6aa13bc6..95ee4d5dd 100644
--- a/pkg/kthena-router/datastore/session_boost_queue.go
+++ b/pkg/kthena-router/datastore/session_boost_queue.go
@@ -124,27 +124,40 @@ func (pq *RequestPriorityQueue) GetInflightCount() int64 {
}
// admitSessionBoost marks a request as inflight, installs its release callback and
-// unblocks the waiting caller by closing its NotifyChan.
-func (pq *RequestPriorityQueue) admitSessionBoost(req *Request) {
- pq.inflightCount.Add(1)
- // req.Release returns the inflight permit this admission consumed. The request
- // handler invokes it (via defer) once the backend response is fully proxied or
- // the request fails/times out. It decrements the inflight count, signals
- // releaseCh so the dequeue loop can immediately admit the next waiting request,
- // and updates the inflight metric. sync.Once makes it idempotent so capacity is
- // never released twice on overlapping exit paths.
- releaseOnce := sync.Once{}
- req.Release = func() {
- releaseOnce.Do(func() {
- pq.inflightCount.Add(-1)
- select {
- case pq.releaseCh <- struct{}{}:
- default:
- }
- pq.metricDecInflight(req.ModelName)
- })
+// unblocks the waiting caller by closing its NotifyChan. It returns false without
+// admitting when the caller has already abandoned the request (timed out or
+// cancelled): in that case the request has already left the queue and admitting it
+// would leak an inflight permit that no one will release.
+func (pq *RequestPriorityQueue) admitSessionBoost(req *Request) bool {
+ // Commit the inflight accounting and Release installation atomically with
+ // respect to a concurrent Abandon(). If the caller abandoned first, skip
+ // admission entirely so the inflight permit is never consumed.
+ admitted := req.commitAdmission(func() {
+ pq.inflightCount.Add(1)
+ // req.Release returns the inflight permit this admission consumed. The request
+ // handler invokes it (via defer) once the backend response is fully proxied or
+ // the request fails/times out. It decrements the inflight count, signals
+ // releaseCh so the dequeue loop can immediately admit the next waiting request,
+ // and updates the inflight metric. sync.Once makes it idempotent so capacity is
+ // never released twice on overlapping exit paths.
+ releaseOnce := sync.Once{}
+ req.Release = func() {
+ releaseOnce.Do(func() {
+ pq.inflightCount.Add(-1)
+ select {
+ case pq.releaseCh <- struct{}{}:
+ default:
+ }
+ pq.metricDecInflight(req.ModelName)
+ })
+ }
+ pq.metricIncInflight(req.ModelName)
+ })
+ if !admitted {
+ klog.V(4).Infof("[SessionBoost] admission skipped, request abandoned before admission: user=%s model=%s",
+ req.UserID, req.ModelName)
+ return false
}
- pq.metricIncInflight(req.ModelName)
// Closing NotifyChan is the admission signal: the caller blocked in Enqueue is
// waiting on this channel and proceeds to the backend only once it is closed.
// We notify here because admission (a free inflight slot plus backend capacity)
@@ -153,6 +166,7 @@ func (pq *RequestPriorityQueue) admitSessionBoost(req *Request) {
if req.NotifyChan != nil {
close(req.NotifyChan)
}
+ return true
}
// runSessionBoostMode is the session-boost dequeue loop. It dequeues requests only
@@ -347,7 +361,12 @@ func (pq *RequestPriorityQueue) tryBackpressureDequeue(ctx context.Context) {
return
}
- pq.admitSessionBoost(req)
+ if !pq.admitSessionBoost(req) {
+ // The request was abandoned (timed out / cancelled) between the
+ // cancellation check and admission; it consumed no inflight slot, so
+ // continue admitting the next queued request.
+ continue
+ }
klog.V(4).Infof("[SessionBoost] backpressure dequeue: user=%s model=%s sessionBoost=%v inflight=%d/%d",
req.UserID, req.ModelName, req.SessionBoost, pq.inflightCount.Load(), maxInflight)
diff --git a/pkg/kthena-router/datastore/session_boost_queue_test.go b/pkg/kthena-router/datastore/session_boost_queue_test.go
index 36faacb3c..30859b3e2 100644
--- a/pkg/kthena-router/datastore/session_boost_queue_test.go
+++ b/pkg/kthena-router/datastore/session_boost_queue_test.go
@@ -18,6 +18,7 @@ package datastore
import (
"context"
+ "sync"
"testing"
"time"
)
@@ -556,3 +557,113 @@ func TestSessionBoostQueue_Len(t *testing.T) {
t.Errorf("Expected len=5, got %d", q.Len())
}
}
+
+// TestSessionBoost_AbandonBeforeAdmissionBlocksAdmission verifies that a request
+// abandoned by the caller before the dequeue loop admits it consumes no inflight
+// slot: admitSessionBoost must observe the abandonment and skip.
+func TestSessionBoost_AbandonBeforeAdmissionBlocksAdmission(t *testing.T) {
+ cfg := sessionBoostConfig()
+ q := newSessionBoostQueue(cfg, nil)
+ defer q.Close()
+
+ req := &Request{
+ UserID: "user-A",
+ ModelName: "model-1",
+ NotifyChan: make(chan struct{}),
+ }
+
+ // Caller gives up first.
+ req.Abandon()
+
+ // The dequeue loop then tries to admit; it must skip.
+ if q.admitSessionBoost(req) {
+ t.Fatal("admitSessionBoost should skip an already-abandoned request")
+ }
+ if got := q.GetInflightCount(); got != 0 {
+ t.Fatalf("inflight count should stay 0 after skipped admission, got %d", got)
+ }
+ select {
+ case <-req.NotifyChan:
+ t.Fatal("NotifyChan must not be closed for a skipped admission")
+ default:
+ }
+}
+
+// TestSessionBoost_AdmissionBeforeAbandonReleases verifies that when admission wins
+// the race, the caller's Abandon() releases the permit it owns, returning the
+// inflight count to zero.
+func TestSessionBoost_AdmissionBeforeAbandonReleases(t *testing.T) {
+ cfg := sessionBoostConfig()
+ q := newSessionBoostQueue(cfg, nil)
+ defer q.Close()
+
+ req := &Request{
+ UserID: "user-A",
+ ModelName: "model-1",
+ NotifyChan: make(chan struct{}),
+ }
+
+ if !q.admitSessionBoost(req) {
+ t.Fatal("admitSessionBoost should admit a live request")
+ }
+ select {
+ case <-req.NotifyChan:
+ default:
+ t.Fatal("NotifyChan should be closed after admission")
+ }
+ if got := q.GetInflightCount(); got != 1 {
+ t.Fatalf("inflight count should be 1 after admission, got %d", got)
+ }
+
+ // Caller times out after admission raced in: Abandon() must release the permit
+ // it owns.
+ req.Abandon()
+ if got := q.GetInflightCount(); got != 0 {
+ t.Fatalf("inflight count should return to 0 after Abandon, got %d", got)
+ }
+}
+
+// TestSessionBoost_ConcurrentAdmitAbandonNoLeak stresses the admission/abandonment
+// race directly: for each request one goroutine admits while another abandons. No
+// matter which wins, the inflight count must return to zero, proving the permit is
+// never leaked and never double-released. Run with -race to catch data races.
+func TestSessionBoost_ConcurrentAdmitAbandonNoLeak(t *testing.T) {
+ cfg := sessionBoostConfig()
+ q := newSessionBoostQueue(cfg, nil)
+ defer q.Close()
+
+ const iterations = 2000
+ for i := 0; i < iterations; i++ {
+ req := &Request{
+ UserID: "user-A",
+ ModelName: "model-1",
+ NotifyChan: make(chan struct{}),
+ }
+
+ var wg sync.WaitGroup
+ wg.Add(2)
+ // Dequeue-loop side: attempt admission.
+ go func() {
+ defer wg.Done()
+ q.admitSessionBoost(req)
+ }()
+ // Handler side: time out and abandon; Abandon releases if admission raced in.
+ // This mirrors the router's queue-wait timeout handling.
+ go func() {
+ defer wg.Done()
+ req.Abandon()
+ }()
+ wg.Wait()
+
+ // After both settle, drain any releaseCh signal the release may have posted so
+ // it does not affect the next iteration's assertions.
+ select {
+ case <-q.releaseCh:
+ default:
+ }
+
+ if got := q.GetInflightCount(); got != 0 {
+ t.Fatalf("iteration %d: inflight leaked, count=%d", i, got)
+ }
+ }
+}
diff --git a/pkg/kthena-router/router/router.go b/pkg/kthena-router/router/router.go
index 469fef76c..1cc668e0e 100644
--- a/pkg/kthena-router/router/router.go
+++ b/pkg/kthena-router/router/router.go
@@ -92,6 +92,12 @@ type Router struct {
queueTimeout time.Duration
tokenWeight float64 // Weight for token-based priority in the fairness strategy (default 1.0)
requestNumWeight float64 // Weight for request-count-based priority in the fairness strategy (default 0.0)
+
+ // Session-boost queue-wait timeout. A request that waits in the session-boost
+ // queue longer than sessionBoostTimeout is rejected with HTTP 504 instead of
+ // waiting indefinitely for backend capacity. It defaults to 30s; a non-positive
+ // value disables the timeout (the request is bounded only by client disconnect).
+ sessionBoostTimeout time.Duration
}
// ActiveRequestCount returns the number of requests currently being handled by the router.
@@ -182,6 +188,8 @@ func NewRouter(store datastore.Store, routerConfigPath string) *Router {
queueTimeout: parseQueueTimeout(),
tokenWeight: parseEnvFloat("FAIRNESS_PRIORITY_TOKEN_WEIGHT", 1.0),
requestNumWeight: parseEnvFloat("FAIRNESS_PRIORITY_REQUEST_NUM_WEIGHT", 0.0),
+
+ sessionBoostTimeout: parseSessionBoostTimeout(),
}
}
@@ -197,6 +205,29 @@ func parseQueueTimeout() time.Duration {
return defaultQueueTimeout
}
+// defaultSessionBoostTimeout is the session-boost queue-wait timeout applied
+// when SESSION_BOOST_TIMEOUT is not set. It is enabled by default so that a
+// session-boost request does not wait indefinitely for backend capacity.
+const defaultSessionBoostTimeout = 30 * time.Second
+
+// parseSessionBoostTimeout reads the session-boost queue-wait timeout from the
+// SESSION_BOOST_TIMEOUT environment variable. A request that waits in the
+// session-boost queue longer than the timeout is rejected with HTTP 504. The
+// timeout defaults to defaultSessionBoostTimeout (30s) when the variable is
+// unset. Setting it to a non-positive duration (e.g. "0s") disables the timeout,
+// in which case a session-boost request is bounded only by client disconnect. An
+// invalid value falls back to the default.
+func parseSessionBoostTimeout() time.Duration {
+ if s, ok := os.LookupEnv("SESSION_BOOST_TIMEOUT"); ok {
+ if d, err := time.ParseDuration(s); err == nil {
+ // A non-positive duration explicitly disables the timeout.
+ return d
+ }
+ klog.Warningf("Invalid SESSION_BOOST_TIMEOUT %q, using default %v", s, defaultSessionBoostTimeout)
+ }
+ return defaultSessionBoostTimeout
+}
+
func parseEnvFloat(key string, fallback float64) float64 {
if s, ok := os.LookupEnv(key); ok {
if v, err := strconv.ParseFloat(s, 64); err == nil && !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 {
@@ -1077,11 +1108,35 @@ func (r *Router) handleFairnessScheduling(c *gin.Context, modelRequest ModelRequ
klog.Warningf("user ID not found in request %s", requestID)
}
- klog.V(4).Infof("[FairnessScheduling] incoming request: reqID=%s user=%s model=%s",
- requestID, userId, modelName)
+ // logPrefix reflects the active scheduling strategy so log lines emitted from
+ // this shared handler are attributed to the right queue (session boost vs
+ // user fairness).
+ logPrefix := "[FairnessScheduling]"
+ if EnableSessionBoost {
+ logPrefix = "[SessionBoost]"
+ }
- // Create request-scoped context that unifies client disconnect and server timeout
- reqCtx, cancel := context.WithTimeout(c.Request.Context(), r.queueTimeout)
+ klog.V(4).Infof("%s incoming request: reqID=%s user=%s model=%s",
+ logPrefix, requestID, userId, modelName)
+
+ // Create the request-scoped context that also drives the queue's cancellation
+ // cleanup (CancelCh). The queue-wait deadline differs by strategy:
+ // - Fairness mode: bounded by FAIRNESS_QUEUE_TIMEOUT; exceeding it returns 504.
+ // - Session-boost mode: FAIRNESS_QUEUE_TIMEOUT does NOT apply. SESSION_BOOST_TIMEOUT
+ // (default 30s) bounds the wait and exceeding it returns 504; setting it to a
+ // non-positive value disables the timeout, leaving the request bounded only by
+ // client disconnect.
+ var reqCtx context.Context
+ var cancel context.CancelFunc
+ if EnableSessionBoost {
+ if r.sessionBoostTimeout > 0 {
+ reqCtx, cancel = context.WithTimeout(c.Request.Context(), r.sessionBoostTimeout)
+ } else {
+ reqCtx, cancel = context.WithCancel(c.Request.Context())
+ }
+ } else {
+ reqCtx, cancel = context.WithTimeout(c.Request.Context(), r.queueTimeout)
+ }
defer cancel()
var pri float64
@@ -1107,8 +1162,8 @@ func (r *Router) handleFairnessScheduling(c *gin.Context, modelRequest ModelRequ
}
if err := r.store.Enqueue(queueReq); err != nil {
- klog.Errorf("[FairnessScheduling] failed to enqueue: reqID=%s sessionID=%s user=%s model=%s err=%v",
- requestID, sessionID, userId, modelName, err)
+ klog.Errorf("%s failed to enqueue: reqID=%s sessionID=%s user=%s model=%s err=%v",
+ logPrefix, requestID, sessionID, userId, modelName, err)
c.AbortWithStatusJSON(http.StatusInternalServerError, fmt.Sprintf("failed to enqueue request: %v", err))
return fmt.Errorf("failed to enqueue request: %v", err)
}
@@ -1118,8 +1173,8 @@ func (r *Router) handleFairnessScheduling(c *gin.Context, modelRequest ModelRequ
if queueReq.Release != nil {
defer queueReq.Release()
}
- klog.V(4).Infof("[FairnessScheduling] request dequeued: reqID=%s user=%s model=%s sessionBoost=%v waitTime=%v",
- requestID, userId, modelName, queueReq.SessionBoost, time.Since(queueReq.RequestTime))
+ klog.V(4).Infof("%s request dequeued: reqID=%s user=%s model=%s sessionBoost=%v waitTime=%v",
+ logPrefix, requestID, userId, modelName, queueReq.SessionBoost, time.Since(queueReq.RequestTime))
lbErr := r.doLoadbalance(c, modelRequest)
// After a successful proxy, mark the session request as completed so follow-up
@@ -1130,18 +1185,29 @@ func (r *Router) handleFairnessScheduling(c *gin.Context, modelRequest ModelRequ
}
return nil
case <-reqCtx.Done():
- if queueReq.Release != nil {
- queueReq.Release()
- }
+ // Abandon() atomically coordinates with the dequeue loop: if admission raced
+ // in first it releases the inflight permit we own; otherwise it marks the
+ // request abandoned so the loop skips admission and no permit can leak.
+ queueReq.Abandon()
if errors.Is(reqCtx.Err(), context.DeadlineExceeded) {
- klog.Errorf("[FairnessScheduling] request timed out in queue: reqID=%s sessionID=%s user=%s model=%s timeout=%v",
- requestID, sessionID, userId, modelName, r.queueTimeout)
- c.AbortWithStatusJSON(http.StatusGatewayTimeout, "Request processing timed out")
- return fmt.Errorf("request processing timed out in fairness queue")
+ // Exceeded the queue-wait timeout. In session-boost mode this is expected
+ // load-shedding when SESSION_BOOST_TIMEOUT is set, and under sustained
+ // overload it can fire for many requests, so log at a verbose level to avoid
+ // flooding the logs. In fairness mode the FAIRNESS_QUEUE_TIMEOUT is unexpected
+ // and logs at error level.
+ if EnableSessionBoost {
+ klog.V(2).Infof("%s request rejected after exceeding queue-wait timeout: reqID=%s sessionID=%s user=%s model=%s timeout=%v",
+ logPrefix, requestID, sessionID, userId, modelName, r.sessionBoostTimeout)
+ } else {
+ klog.Errorf("%s request timed out in queue: reqID=%s sessionID=%s user=%s model=%s timeout=%v",
+ logPrefix, requestID, sessionID, userId, modelName, r.queueTimeout)
+ }
+ c.AbortWithStatusJSON(http.StatusGatewayTimeout, "Request processing timed out in queue")
+ return fmt.Errorf("request processing timed out in queue")
}
- klog.V(4).Infof("[FairnessScheduling] request cancelled (client disconnected): reqID=%s sessionID=%s user=%s model=%s",
- requestID, sessionID, userId, modelName)
- c.AbortWithStatusJSON(http.StatusServiceUnavailable, "Client disconnected while waiting in fairness queue")
- return fmt.Errorf("client disconnected while waiting in fairness queue")
+ klog.V(4).Infof("%s request cancelled (client disconnected): reqID=%s sessionID=%s user=%s model=%s",
+ logPrefix, requestID, sessionID, userId, modelName)
+ c.AbortWithStatusJSON(http.StatusServiceUnavailable, "Client disconnected while waiting in queue")
+ return fmt.Errorf("client disconnected while waiting in queue")
}
}
diff --git a/pkg/kthena-router/router/router_test.go b/pkg/kthena-router/router/router_test.go
index 60a84cd57..d9fb50f81 100644
--- a/pkg/kthena-router/router/router_test.go
+++ b/pkg/kthena-router/router/router_test.go
@@ -1363,6 +1363,9 @@ func TestHandleFairnessScheduling(t *testing.T) {
wantErrMsg string
wantHTTPStatus int
wantBodyContains string
+ // Session-boost queue-wait timeout configuration for the test case.
+ enableSessionBoost bool
+ sessionBoostTimeout time.Duration
}{
{
name: "happy path with userId",
@@ -1410,6 +1413,17 @@ func TestHandleFairnessScheduling(t *testing.T) {
wantErrMsg: "failed to enqueue request",
wantHTTPStatus: http.StatusInternalServerError,
},
+ {
+ name: "session boost queue-wait timeout returns 504",
+ fairnessTimeout: 10 * time.Second,
+ setUserID: true,
+ storeWrapper: func(real datastore.Store) datastore.Store { return &blockingEnqueueStore{Store: real} },
+ enableSessionBoost: true,
+ sessionBoostTimeout: 50 * time.Millisecond,
+ wantErr: true,
+ wantErrMsg: "timed out",
+ wantHTTPStatus: http.StatusGatewayTimeout,
+ },
}
for _, tt := range tests {
@@ -1418,6 +1432,12 @@ func TestHandleFairnessScheduling(t *testing.T) {
defer backend.Close()
router.queueTimeout = tt.fairnessTimeout
+ router.sessionBoostTimeout = tt.sessionBoostTimeout
+ // Set the package-level flag explicitly for every case (and restore it)
+ // so subtests stay isolated regardless of execution order.
+ prevEnableSessionBoost := EnableSessionBoost
+ EnableSessionBoost = tt.enableSessionBoost
+ defer func() { EnableSessionBoost = prevEnableSessionBoost }()
if tt.storeWrapper != nil {
router.store = tt.storeWrapper(store)
}