feat: worker-side per-tenant rate limiting - #551
Conversation
|
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
left a comment
There was a problem hiding this comment.
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).
| } | ||
| // 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; |
There was a problem hiding this comment.
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.
| 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 |
There was a problem hiding this comment.
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.
| // 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()) |
There was a problem hiding this comment.
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 { |
There was a problem hiding this comment.
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.
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>
be38cfe to
bc1bb35
Compare
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>
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_ratioconfig shape).What it does
talon-core):TenantId::Named(SDK-declared) /Unattributed(undeclared, e.g. FUSE). The direct path is unauthenticated, soa declared tenant is trusted for fairness, not as a security boundary.
talon-transport): newMsgType::GetRangeTenantcarrying aTenantScopedRange; the reply is an ordinaryGetRangeframe. Fail-closedon older workers (a distinct opcode, like
GetCachedRange), and conformancevectors are unchanged so the other-language clients are unaffected.
[rate_limits]table):enabled, adefaultMetricLimits, andper-tenant
overrides. Each metric isRateLimit { rate, burst_ratio }wherethe burst is
rate * burst_ratio. Metrics:read_iops(req/s) andread_throughput(bytes/s served to the client).origin_read_bytesis adefined metric for a later stage.
talon-worker): a lock-freeTenantRateLimiter(one GCRAcell per
(tenant, metric), single-atomic TAT,Arc-shared across io_uringrings). Admission before serve in both the io_uring and Tokio handlers;
charges
read_iops+read_throughput. A throttle returnsDataErrorCode::RateLimited.talon-cache-client):WorkerClient::with_tenantsendsGetRangeTenant;RateLimitedmaps to a non-origin-fallbackCacheReadError(a throttle can't be evaded via origin bypass).RateLimited→ S3503 SlowDown/ Azure503 ServerBusy.talon_worker_rate_limited_total{metric}; no rawtenant labels (per-tenant drill-down is a management-API concern).
Config example
Testing
cargo testgreen:talon-core(71) andtalon-worker(253). Unit coverage:GCRA burst/refill/burst-ratio, limiter admission (per-tenant, overrides,
throughput),
[rate_limits]TOML parse, andWorkerClienttenant-scoped vsplain-
GetRangesend.Deliberate follow-ups (noted in code)
origin_read_bytesenforcement (needs the tenant threaded into thecache-miss path).
talon-clientCLI / C / Python / Java).layer → Tier-2 monitoring (
docs/explanation/tenant-traffic-observability.md).🤖 Generated with Claude Code