Add pool topology, unified gate interface, labels, and tier-priority dispatch

[joeltg]

Motivation

We want to support tier-priority dispatch across a shared inference fleet: strict priority across tiers (interactive, async, batch), per-team reservation quotas that are floors not ceilings so idle capacity isn't stranded, and interactive overflow that fails fast under saturation rather than waiting through redelivery delays. This PR implements the framework primitives that make this possible, and the policy itself.

A fundamental issue is that the existing gate interface — a float "dispatch budget" consulted at pull time — works against this:

• When the backend saturates, the budget drops to zero, the streaming pull connection to Pub/Sub is torn down, and pre-fetched messages return to the broker.
• When capacity recovers and the gate restarts the pull, there's a 1–2 second connection-reestablishment gap during which workers are idle.
• Under sustained load this oscillates: a sawtooth where throughput collapses on every saturation cycle while waiting for the new pull connection. In our setup, we saw drops from ~4,900 tok/s down to ~400 tok/s at 200 concurrency.

Instead, we want to pull aggressively and always, and gate admission at the worker rather than at the pull loop. Messages wait in the in-process prefetch buffer at zero cost and dispatch the instant capacity opens, with no nack-and-redeliver round-trip.

To support this architecture, we implement the following scoped refactoring:

1. Explicit pool topology
2. Unified gate interface, with two gate evaluation sites (per-subscription and per-pool)
3. Gate-assigned labels for request classification
4. Tier-priority RequestMergePolicy and admission gates

---

Pool topology

A named Pool declares its gateway URL, worker count, and gate chain. Subscriptions reference a pool by ID. The merge policy now emits one buffered channel per pool; each pool gets its own dedicated worker pool so backpressure on one pool's downstream endpoint stalls only that pool's prefetch — not the entire pod's.

An example FlowConfig schema illustrating this (and the rest of the changes):

pools:
  - id: my-pool
    gateway_url: http://igw.svc:8000
    request_path: /v1/chat/completions
    model_name_override: llama-3.1-8b-instant
    workers: 32
    gates:
      - type: local-max-concurrency
        params: { max_concurrency: "64" }

subscriptions:
  - subscriber_id: inference-requests-teamA-async
    pool: my-pool
    labels: { team: teamA, tier: async }
    gates:
      - type: deadline-drop

Unified gate interface

The DispatchGate.Budget float interface is replaced by Gate.Apply, which returns a Verdict:

var Continue = Verdict{}
func Drop(r *ResultMessage) Verdict  // ack; publish r if non-nil (e.g. 429 fail-fast)
func Refuse() Verdict                // nack; transport redelivers

Gates run at two distinct sites with different blocking rules:

• Subscription gates run in the Pub/Sub receive callback, before the merge policy. They must not block (blocking stalls prefetch). Their primary job is classification: they read and mutate msg.Labels to stamp decisions like class=reserved or class=overflow that the merge policy and pool gates act on downstream. Gate-written labels are the key mechanism for runtime classification — the subscription config supplies static identity labels (team, tier, model); gates supply dynamic classification labels whose values depend on system state at dispatch time.

• Pool gates run in the per-pool worker pool, after the merge policy routes the message, before HTTP dispatch. They may block — because the message is already prefetched and in memory, parking the worker until capacity opens dispatches it with zero extra latency and no nack round-trip.

In both sites, multiple gates are provided in an array, are evaluated in order, and short-circuit on the first non-Continue result.

Labels

Subscriptions declare static labels (team, tier, model) in their config; the Flow stamps these onto every pulled message. Labels are entirely operator-controlled: transport attributes from producers round-trip via body.Metadata / result.Metadata instead, which is a separate namespace. Labels are a plain map[string]string with last-write-wins semantics: gates may overwrite any key, and each gate in the chain sees the mutations made by earlier ones. This is what makes runtime classification composable — a deadline gate can run first and drop the message; a reservation classifier runs next and stamps class; the merge policy buckets on that value.

Tier-priority RMP and admission gate

Two pieces that compose to implement the priority-tier dispatch:

• The tier-priority RMP buckets messages by (pool, tier, class) and dispatches in strict priority across tiers (e.g. interactive > async > batch) with round-robin across (team, model) source channels within each tier. Here tier is a static subscription label set by the operator; class is a gate-written label (reserved or overflow) stamped at subscription-gate time based on whether the request is within the team's reservation cap — this is the reason the gate-labels pipeline exists. The RMP stamps the engine-side priority header and is pure ordering: it never consults a gate.

• The tier-priority-admission pool gate reads a pool saturation signal and returns a three-way verdict based on labels: block if class=reserved (don't nack work that legitimately deserves a slot; park and dispatch the instant capacity opens); Drop with 429 if tier=interactive, class=overflow (interactive's SLA is sub-second; fail fast so the producer can fall back); Refuse for all other overflow (nacking async/batch overflow is correct — blocking it would tie up workers and head-of-line-block reserved messages behind it in the per-pool FIFO channel). The saturation signal is pluggable: a Prometheus-backed source for self-hosted IGW pools, a redis-counter source for external providers with hard concurrency limits.

                              ┌──────────┐
                              │ Producer │
                              └────┬─────┘
                                   │ publish
                                   ▼
                              ┌─────────┐
                              │ Pub/Sub │  request topic
                              │  topic  │
                              └────┬────┘
                                   │ pull
  ╔════════════════════════════════▼═════════════════════════════════╗
  ║                       async-processor pod                        ║
  ║                                                                  ║
  ║         ┌─────────────────────────────────────────┐              ║
  ║         │ Flow receive callback                   │              ║
  ║         │ ├ register tracker, spawn ack goroutine │              ║
  ║         │ ├ subscription gate chain (non-blocking)│              ║
  ║         │ │   • deadline-drop                     │              ║
  ║         │ │   • reservation-classifier            │              ║
  ║         │ └ → per-subscription channel            │              ║
  ║         └────────────────────┬────────────────────┘              ║
  ║                              ▼                                   ║
  ║         ┌─────────────────────────────────────────┐              ║
  ║         │ Merge policy (tier-priority RMP)        │              ║
  ║         │ ├ bucket by (pool, tier, class)         │              ║
  ║         │ ├ reserved before overflow (strict)     │              ║
  ║         │ ├ within each class: tier order (strict)│              ║
  ║         │ ├ round-robin across (team,model) at    │              ║
  ║         │ │   the leaves                          │              ║
  ║         │ └ → per-pool channel                    │              ║
  ║         └────────────────────┬────────────────────┘              ║
  ║                              ▼                                   ║
  ║         ┌─────────────────────────────────────────┐              ║
  ║         │ Per-pool worker (1 of K)                │              ║
  ║         │ ├ pool gate chain (may block)           │              ║
  ║         │ │   • tier-priority-admission           │              ║
  ║         │ │   • local-rate-limit (optional)       │              ║
  ║         │ └ HTTP dispatch                         │              ║
  ║         └────────────────────┬────────────────────┘              ║
  ║                              │                                   ║
  ╚══════════════════════════════│═══════════════════════════════════╝
                                 │ HTTP
                                 ▼
                            ┌─────────┐
                            │   IGW   │
                            └────┬────┘
                                 ▼
                            ┌─────────┐
                            │ Engine  │  (SGLang / vLLM / external)
                            └────┬────┘
                                 │ response
                                 ▼
                            ┌─────────┐
                            │ Pub/Sub │  result topic
                            │  topic  │
                            └─────────┘

  ack/nack flows back to the request topic from the ack goroutine
  spawned in the receive callback (above): every Verdict outcome routes
  through resultsChannel → ack goroutine → msg.Ack() / msg.Nack().

[github-actions]

Unsigned commits detected! Please sign your commits.

For instructions on how to set up GPG/SSH signing and verify your commits, please see GitHub Documentation.

[lioraron]

Review: 8 inline comments covering correctness, safety, and test coverage.

[lioraron]

Worker goroutines fire-and-forget: Workers are spawned with bare go asyncworker.Worker(...) and nothing waits for them to finish — the old WaitGroup was removed. After <-ctx.Done() returns (line 224), main() exits immediately, so in-flight requests are abandoned mid-dispatch. This can lose results or leave partial state in Redis.

Add a sync.WaitGroup (or similar) so the shutdown path waits for all workers to drain before exiting.

[lioraron]

Start is not idempotent despite the doc claiming it is: Each call launches a new goroutine with a new ticker. If called twice you get duplicate concurrent refresh loops hammering Prometheus.

Add a sync.Once guard (or an atomic flag) to match the documented contract.

[lioraron]

time.After timer leak: time.After(wait) inside this for loop allocates a timer per iteration that cannot be GC'd until it fires. Under sustained load with many goroutines blocked here, this leaks memory.

Replace with time.NewTimer + defer timer.Stop() (or reset the same timer across iterations).

[lioraron]

context.Context stored in struct: This is a well-documented Go anti-pattern (style guide). The context here drives background goroutines (PromQL refresh loops).

Consider passing it through a Start(ctx) method or a dedicated Run(ctx) lifecycle method instead of baking it into the struct.

[lioraron]

panic on duplicate registration: Combined with init()-based registration, a misconfigured import crashes the process at startup with no recovery path. Return an error instead — callers (e.g. NewMergePolicy) already have error paths.

[lioraron]

Set panics on nil receiver: Get and Has are explicitly nil-safe, but Set (and Merge) panic on nil. Since external gate authors call msg.Labels.Set(...), this is a footgun if Labels was never initialized.

Either auto-initialize in Set/Merge (return a new map), or add a prominent doc warning that callers must check for nil.

[lioraron]

Mutable package-level sentinel: var Continue = Verdict{} can be accidentally mutated by any gate (e.g. pipeline.Continue.Terminate = true), corrupting it globally for all callers.

Make it a function func Continue() Verdict { return Verdict{} }, or keep the var but make it unexported with a public accessor that returns a copy.

[lioraron]

No test coverage: This file has no _test.go counterpart. Token-bucket logic with timing (refill rates, burst capacity, edge cases around zero/negative inputs) is easy to get wrong and should have unit tests before merge.

[lioraron]

Thanks @joeltg for this extensive work!

This PR covers a lot of ground — gate interface migration, labels, pool topology, merge policy registry, flow control gates, and tier-priority dispatch. Each of these is a meaningful feature with its own design surface, and bundling them makes review harder and riskier (e.g. the 8 inline comments I left are spread across unrelated subsystems).

Would it be possible to split this into smaller, independently mergeable PRs? A reasonable split might be:

1. Gate interface migration — Verdict-based Gate replacing DispatchGate.Budget(), ApplyChain, AlwaysContinue
2. Labels + merge policy registry — Labels type, RegisterMergePolicy, MergePolicyFactory
3. Pool topology + per-pool workers — Pool type, per-pool worker spawning, MergeRequestChannels
4. Flow control gates + tier-priority dispatch — LocalRateLimitGate, PromQLLoadSource, tier-priority admission, reservation classifier

I'd also suggest opening a corresponding issue for each PR that includes at least:

• Problem description: what's missing or broken today, and why it matters
• High-level design: the approach, key types/interfaces, and how it interacts with existing code

This makes each change independently reviewable, gives reviewers the context to evaluate design decisions, and creates a traceable record of why each piece exists.

[evacchi]

I would go for an alternative design here, instead of the double bool; the two-bool encoding allows a fourth state (Terminate=false, Redeliver=true) and forces every consumer to reason about their combination; I'd advise going with something like:

type VerdictKind int

const (
    VerdictContinue VerdictKind = iota
    VerdictDrop
    VerdictRefuse
)

type Verdict struct {
   Kind   VerdictKind
   Result *api.ResultMessage 
}

[github-actions]

This PR is marked as stale after 21d of inactivity. After an additional 14d of inactivity (7d to become rotten, then 7d more), it will be closed. To prevent this PR from being closed, add a comment or remove the lifecycle/stale label.