Skip to content

feat: worker-side per-tenant rate limiting - #551

Merged
beinan merged 13 commits into
milvus-io:mainfrom
somfornot:feat/tenant-id
Aug 19, 2026
Merged

feat: worker-side per-tenant rate limiting#551
beinan merged 13 commits into
milvus-io:mainfrom
somfornot:feat/tenant-id

Conversation

@somfornot

Copy link
Copy Markdown
Collaborator

Summary

Worker-side per-tenant rate limiting for the SDK-direct data plane: when a
client declares a tenant, the worker admits or throttles the read against a
static per-tenant policy. Supersedes #550 (opened from a different fork; same
feature, now with the requested read_throughput / burst_ratio config shape).

What it does

  • Tenant identity (talon-core): TenantId::Named (SDK-declared) /
    Unattributed (undeclared, e.g. FUSE). The direct path is unauthenticated, so
    a declared tenant is trusted for fairness, not as a security boundary.
  • Wire (talon-transport): new MsgType::GetRangeTenant carrying a
    TenantScopedRange; the reply is an ordinary GetRange frame. Fail-closed
    on older workers
    (a distinct opcode, like GetCachedRange), and conformance
    vectors are unchanged so the other-language clients are unaffected.
  • Policy ([rate_limits] table): enabled, a default MetricLimits, and
    per-tenant overrides. Each metric is RateLimit { rate, burst_ratio } where
    the burst is rate * burst_ratio. Metrics: read_iops (req/s) and
    read_throughput (bytes/s served to the client). origin_read_bytes is a
    defined metric for a later stage.
  • Enforcement (talon-worker): a lock-free TenantRateLimiter (one GCRA
    cell per (tenant, metric), single-atomic TAT, Arc-shared across io_uring
    rings). Admission before serve in both the io_uring and Tokio handlers;
    charges read_iops + read_throughput. A throttle returns
    DataErrorCode::RateLimited.
  • Client (talon-cache-client): WorkerClient::with_tenant sends
    GetRangeTenant; RateLimited maps to a non-origin-fallback
    CacheReadError (a throttle can't be evaded via origin bypass).
  • Gateway: RateLimited → S3 503 SlowDown / Azure 503 ServerBusy.
  • Telemetry: bounded talon_worker_rate_limited_total{metric}; no raw
    tenant labels (per-tenant drill-down is a management-API concern).

Config example

[rate_limits]
enabled = true
[rate_limits.default]
read_iops = { rate = 5000, burst_ratio = 2.0 }
read_throughput = { rate = 1073741824 }        # 1 GiB/s; burst_ratio defaults to 1.0
[rate_limits.overrides.acme]
read_iops = { rate = 20000, burst_ratio = 2.0 }

Testing

cargo test green: talon-core (71) and talon-worker (253). Unit coverage:
GCRA burst/refill/burst-ratio, limiter admission (per-tenant, overrides,
throughput), [rate_limits] TOML parse, and WorkerClient tenant-scoped vs
plain-GetRange send.

Deliberate follow-ups (noted in code)

  • origin_read_bytes enforcement (needs the tenant threaded into the
    cache-miss path).
  • Per-SDK tenant knobs (BlockReader / talon-client CLI / C / Python / Java).
  • A socket-level worker end-to-end throttle test.
  • Later stages: coordinator policy distribution → consistent-hash owner/peer
    layer → Tier-2 monitoring (docs/explanation/tenant-traffic-observability.md).

🤖 Generated with Claude Code

@beinan

beinan commented Aug 18, 2026

Copy link
Copy Markdown
Collaborator

Great feature! Thank you for the contribution! the cache industry and oss community were waiting for this for years!

could you fix the ci? thanks a lot!

@beinan beinan left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The overall structure is coherent, but the inline findings below leave rate limiting bypassable or capable of permanently rejecting valid reads. Please address these before merge. The dependency-advisories check also currently fails on RUSTSEC-2026-0258 (h2 0.4.15; patched in 0.4.16).

Comment thread crates/talon-core/src/rate_limit.rs Outdated
}
// Total TAT advance for this cost, saturating rather than wrapping for
// pathological (huge cost, tiny rate) combinations.
let increment = (cost as u128 * self.interval_nanos as u128).min(u64::MAX as u128) as u64;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A charge larger than the burst capacity can never be admitted. For example, at 10 MiB/s with burst_ratio = 1.0, a 16 MiB read always has a 1.6s increment against a 1s tolerance. Retrying after the returned 0.6s does not help because base advances to the new now, so every retry returns the same throttle forever. Please either charge/police throughput in bounded transfer chunks (as the proposal describes) or explicitly handle costs larger than the bucket capacity instead of returning a retry time that can never succeed. Add a regression test where cost > rate * burst_ratio.

Comment thread crates/talon-core/src/rate_limit.rs Outdated
pub fn new(limit: RateLimit) -> Self {
// Nanoseconds of credit per unit. Clamped to at least 1ns so a single
// unit always advances the TAT and an admitted cost is never a no-op.
let interval_nanos = NANOS_PER_SEC

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This representation caps every configured rate above 1,000,000,000 units/s at an effective 1,000,000,000 units/s because the per-unit interval is clamped to 1ns. That already makes the PR's 1 GiB/s example about 6.9% too restrictive, and larger byte rates diverge further. Please retain fractional precision (for example with fixed-point TAT units or by computing the cost-scaled increment before division), and cover a rate above 1e9 in tests.

Comment thread crates/talon-worker/src/rate_limit.rs Outdated
// Slow path (first time this tenant is seen): insert, then charge. The
// brief write lock is paid once per tenant, not per request.
let mut map = self.cells.write().unwrap();
map.entry(tenant.clone())

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This map grows permanently for every tenant name observed on the unauthenticated data plane. Neither handler calls TenantId::validate(), and there is no cardinality bound or expiry, so a client can continuously send unique (up to control-frame-sized) tenant names and grow worker memory without limit; even tenants whose resolved MetricLimits are empty get an entry. Please validate at decode/admission and introduce bounded lifecycle/cardinality semantics for tenant cells, with a test that exercises the bound.

/// (`Unattributed`) sends an ordinary `GetRange`. Send a named tenant only
/// to workers known to understand the type — an older worker fails the
/// request closed rather than misreading it.
pub fn with_tenant(mut self, tenant: TenantId) -> Self {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

with_tenant does not actually attribute every fetch: fetch_cached_range() still emits an unscoped GetCachedRange, and both worker cached-range handlers bypass TenantRateLimiter. A tenant can therefore avoid both read_iops and read_throughput whenever the requested version is resident. Please add a tenant-scoped cache-only operation (or fail closed for this combination) and enforce the same admission in both worker handlers.

somfornot and others added 12 commits August 19, 2026 12:22
Introduce TenantId, the authenticated resource owner a request is
attributed to for per-tenant rate limiting. An authenticated tenant is a
gateway provider account (TenantId::Principal); traffic that bypasses the
gateway (direct FUSE / native clients) is the reserved
TenantId::Unattributed tenant, so it is still counted and bounded rather
than escaping per-tenant policy.

This is the shared vocabulary the rate-limit hot path, control-plane
policy, and bounded telemetry build on; wiring follows in later PRs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add the tier-agnostic rate-limit building blocks: RateMetric (read_iops,
client_egress_bytes, origin_read_bytes), RateLimit (rate, burst), and a
lock-free single-atomic GCRA cell (Gcra) that admits or throttles one
(tenant, metric) stream without a lock. Time is caller-supplied, so the
cell is deterministic and unit-tested.

These are shared by whichever tier enforces the limit (the gateway for
authenticated tenants, the worker for direct clients).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The tenant reaches the worker either declared by the SDK on the direct
data plane or resolved from the gateway principal, so name the variant
TenantId::Named rather than Principal, and rename the accessor to name().
Unattributed is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a distinct MsgType::GetRangeTenant whose payload is a bincode
TenantScopedRange (tenant + RangeRequest); the reply is an ordinary
GetRange frame. A worker that predates per-tenant QoS rejects the unknown
type (fail-closed) instead of misreading the tenant prefix, mirroring how
GetCachedRange was introduced. Conformance vectors are unchanged, so other
language clients are unaffected.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add RateMetric (read_iops, read_throughput, origin_read_bytes), RateLimit
{rate, burst_ratio} where burst = rate * burst_ratio, and RateLimitPolicy /
MetricLimits (a default plus per-tenant overrides) deserialized from a
worker [rate_limits] table. read_throughput is the client read/egress
byte rate (clearer than "client_egress_bytes"); origin_read_bytes is
defined for a later stage.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a lock-free TenantRateLimiter (one GCRA cell per (tenant, metric),
Arc-shared across io_uring rings) and admit before serve in both the
io_uring and Tokio data-plane handlers. A throttle returns a typed
DataErrorCode::RateLimited; the client maps it to a non-origin-fallback
CacheReadError (so a throttle cannot be evaded via origin bypass) and the
gateway maps it to S3 503 SlowDown / Azure 503 ServerBusy. Bounded metric
talon_worker_rate_limited_total{metric}; no raw tenant labels.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
WorkerClient::with_tenant binds a client to a tenant; fetch_range then
sends a GetRangeTenant frame, so the worker attributes and rate-limits the
read. Without a tenant it sends an ordinary GetRange (backward compatible).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Charge a warm tenant's GCRA cells in place under the map read lock instead
of cloning an Arc per request, and skip the clock read when the limiter is
disabled. A divan benchmark (benches/rate_limit_benches.rs) drove this:
admit's warm full-path median drops ~361ns -> ~90ns and the disabled path
~48ns -> ~6.5ns. Also add a deterministic capability test proving the
limiter caps a 10x-overload offered rate to the burst plus one second of
the sustained rate (~2000 admitted of 10000 offered).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two defects in the GCRA cell, both on the enforced read_throughput path
(cost = bytes), found in review of milvus-io#551:

1. A single charge whose cost exceeds the burst capacity
   (cost > rate * burst_ratio) could never be admitted: the increment
   always exceeded the tolerance, so every retry re-threw the same
   throttle. Worse, the returned retry_after was non-zero but never
   sufficient — a client that slept and retried it looped forever. A read
   larger than a tenant's configured throughput burst (e.g. a 16 MiB read
   under a 10 MiB/s, burst 1.0 limit) was thus permanently rejected. Fix:
   the admission ceiling for a charge is now max(tolerance, its own
   increment), so an oversized request is admitted from an idle cell and
   then paces the tenant for its full cost, and retry_after is truthful.

2. `interval = 1e9 / rate` used integer division, flooring to 0 (clamped
   to 1ns) for any rate above 1e9 units/s and silently capping the
   effective rate at 1e9 units/s. Byte-rate limits above ~0.93 GiB/s were
   therefore under-enforced (1 GiB/s ran ~6.9% slow; 10 GiB/s ran ~11x
   slow). Fix: compute the increment as cost * 1e9 / rate with the
   multiply first (u128), storing the rate rather than a pre-divided
   per-unit interval, so high rates keep full precision.

Adds regression tests for cost > burst and for a rate above 1e9.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The declared tenant on the direct data plane is unauthenticated and
attacker-controlled, but the per-tenant cell map grew one permanent entry
per distinct tenant name ever seen, with no validation and no ceiling — a
memory-exhaustion vector (raised in review of milvus-io#551). Even tenants whose
resolved limits were empty got an entry.

- Both range handlers now call `TenantId::validate()` after decode and
  reject an empty or over-long name with InvalidRequest, before the name
  is used as a cell key or telemetry label.
- The limiter skips allocating a cell entirely when a tenant's resolved
  limits are empty (an empty default therefore never grows the map).
- Distinct tenant cells are capped at MAX_TENANT_CELLS; tenants first seen
  beyond the cap share one overflow cell governed by the default limit, so
  memory stays bounded while the excess is still paced. Configured
  overrides are pre-created and exempt from the cap so they are always
  honored regardless of arrival order.

Adds tests for the empty-limit skip, the cardinality cap + overflow, and
override exemption.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The cache-only read path bypassed per-tenant rate limiting entirely: both
worker handlers served GetCachedRange without calling the limiter, and
`WorkerClient::with_tenant` did not attribute cache-only fetches. A tenant
could therefore escape both read_iops and read_throughput whenever the
requested version was resident (raised in review of milvus-io#551).

Close the gap end to end:
- New wire op MsgType::GetCachedRangeTenant (= 9) carrying a
  TenantScopedCachedRange, mirroring GetRangeTenant: reply is an ordinary
  GetRange frame and an older worker rejects the distinct type fail-closed
  rather than misreading the tenant prefix. Tightly payload-capped.
- Both worker handlers now decode the declared tenant on the cached path,
  validate it, and charge the limiter before serving — a plain
  GetCachedRange carries no tenant and is metered as Unattributed under the
  default, so the opcode itself is no longer an escape hatch.
- WorkerClient::fetch_cached_range sends GetCachedRangeTenant when the
  client is bound to a named tenant, else a plain GetCachedRange.

Adds transport round-trip + fail-closed tests and a client test that a
tenant-bound client sends the scoped cached frame.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add two capability checks to the worker rate-limit suite that confirm
each tenant is throttled to its own limit when several tenants issue
requests at once, not just in isolation:

- caps_each_tenant_independently_under_concurrent_overload: three tenants
  with different read_iops limits (default 1000, vip 5000, tiny 100) all
  offer overload at every simulated instant; each admits ~its own rate,
  proving a greedy tenant cannot borrow another's budget.
- isolates_tenant_rates_under_real_thread_contention: one OS thread per
  tenant hammers the shared limiter for a fixed window, exercising the
  lock-free cells under real contention; each is capped near its own rate
  and the 8000/s tier admits far more than the 2000/s tier.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Re-wrap the transport re-export and worker-client import lists and the
multi-tenant test overrides to satisfy rustfmt; no behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@beinan
beinan merged commit 43add72 into milvus-io:main Aug 19, 2026
58 of 60 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants