Skip to main content

polars_io/cloud/
http_rate_limit.rs

1// Additive Increase / Multiplicative Decrease (AIMD) adaptive rate-limiter at
2// the HttpService with JIT-pricing.
3//
4// Components:
5// - Rate-limiter: responsible for config-based rate-limiting, including adapting
6//   to the observed success rate.
7// - Pacer (lock-free hot path): responsible for pacing the requests on `admit()`.
8//   The pacer will deny requests that have an estimated wait that is too far out.
9//   Holds an AtomicU64 f64 bits representation of the learned rate.
10// - PacerSignal (lock-free warm path): collects metrics from the pacer and holds
11//   time window boundaries to protect the cold path Mutex.
12// - AimdState (cold path): actuator, adapts the rate based on the observed
13//   HTTP success or failure rate. Authoritative for the atomic (learned) rate,
14//   shared real-time lock-free with the pacer and inflight concurrency controller.
15//
16// Integration points:
17// - The inflight concurrency controller adjusts its concurrency based on the
18//   learned rate. The `RateCell` contains the shared information.
19// - The object_store is responsible for retries, acting as a shock absorber
20//   for when the pipeline pushes above the rate-limit.
21// - The `InitPolicy` facilitates persistent learning when an object_store
22//   is rebuilt on error.
23//
24// TODO:
25// - Rate-limit is one per object_store per verb. AWS S3 scaled out its
26//   rate-limit per-prefix when hot. This can be used to refine the rate-limiter.
27// - The bridge between rate-limiter and inflight concurrency controller allows
28//   for a more reliable RTT_min signal than the clamped one currently in use,
29//   since it sits below the retry layer.
30
31use std::collections::VecDeque;
32use std::fmt;
33use std::str::FromStr;
34use std::sync::atomic::Ordering::Relaxed;
35use std::sync::atomic::{AtomicU32, AtomicU64};
36use std::sync::{Arc, LazyLock, Mutex};
37use std::time::{Duration, Instant};
38
39use async_trait::async_trait;
40use chrono::Utc;
41use object_store::ClientOptions;
42use object_store::client::{
43    HttpClient, HttpConnector, HttpError, HttpErrorKind, HttpRequest, HttpResponse, HttpService,
44};
45use polars_core::runtime::ASYNC;
46use tokio::sync::oneshot;
47
48use super::token_bucket::TokenBucket;
49use crate::cloud::{CloudDirectionalRateLimitConfig, CloudRateLimitConfig};
50
51// For tracing purposes.
52static CONTROLLER_ID: AtomicU32 = AtomicU32::new(0);
53
54// Verbose.
55static LOG_HTTP_RATE_LIMIT: LazyLock<bool> =
56    LazyLock::new(|| std::env::var("POLARS_LOG_HTTP_RATE_LIMIT").is_ok());
57
58// Request/s rate init and boundaries.
59const DEFAULT_READ_INIT_RATE: f64 = 2500.0;
60const DEFAULT_WRITE_INIT_RATE: f64 = 1000.0;
61const DEFAULT_READ_FLOOR_RATE: f64 = 50.0;
62const DEFAULT_WRITE_FLOOR_RATE: f64 = 10.0;
63const DEFAULT_CEILING_RATE: f64 = 50_000.0;
64
65// Growth out of a cut is multiplicative, so a zero rate cannot recover.
66const MIN_RATE: f64 = 1.0;
67
68// Increase / decrease parameters.
69// Cold-start multiplicative increase, per tick.
70const FAST_RAMP_FACTOR: f64 = 2.0;
71// Additive step increase (relative to max), per tick.
72const PROBE_FRACTION: f64 = 0.1;
73const PROBE_MIN: f64 = 5.0;
74// Multiplicative decrease (on cut), with signal
75const BETA: f64 = 0.7;
76// Multiplicative decrease (on cut), no signal available (e.g. cold start)
77const BETA_NO_SIGNAL: f64 = 0.5;
78
79// Timing parameters.
80// Target queue depth as communicated to the concurrency controller.
81const DEFAULT_RATE_HORIZON_MS: u64 = 200;
82// Maximum queue depth, or admission denial ceiling: fail-fast rather than park
83// when the estimated wait exceeds this.
84// Note. Under nominal conditions the concurrency-driven population cap
85// sets depth below this value. Kicks in as buffer on rate collapse.
86// TODO: Dynamically size based on CloudRetryConfig and Init/Floor values.
87// Indicative sizing: to avoid object store erroring out, the rate-limiter must have
88// sufficient capacity to handle a drop-off from Init to Floor. This comes down
89// (MAX_WAIT * FLOOR) >= (INIT * HORIZON). In addition, the combined system
90// must be able to absorb a retry-storm, the size of which is unknown.
91const DEFAULT_RATE_MAX_WAIT_MS: u64 = 10_000;
92// Settle frequency, where AIMD updates its state based on the observed traffic.
93const TICK_INTERVAL: Duration = Duration::from_secs(1);
94// Period during which additional cuts will be suppressed.
95const REFRACTORY: Duration = Duration::from_secs(1);
96// Period after last cut where growth beyond last_max is deferred.
97const PROBE_QUIET_PERIOD: Duration = Duration::from_secs(2);
98
99// Load and freshness parameters.
100// Utilization bound for growth.
101const SATURATION_THRESHOLD: f64 = 0.8;
102// EWMA weight of fresh goodput sample.
103const SUCCESS_SMOOTH: f64 = 0.8;
104// Decay multiplier towards init on idle, per tick.
105const IDLE_DECAY: f64 = 0.9;
106
107// Wake-latency tick. Bounds how long a parked waiter oversleeps past its
108// arithmetic due-time; does NOT affect pacing accuracy.
109pub(crate) const HTTP_RATE_LIMIT_WAKE_TICK: Duration = Duration::from_millis(1);
110
111// Used for sharing values between concurrency controller and rate-limiter.
112// Captures the f64 value in bits.
113pub type RateCell = Arc<AtomicU64>;
114
115/// What happens to the learned rate state when a store inits. Important
116/// when an object store rebuilds on error.
117#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118#[allow(unused)]
119pub(crate) enum InitPolicy {
120    /// Reset to init values: a rebuild is a deliberate clean slate.
121    SetToInit,
122    /// Reset to floor values: a rebuild falls back to a conservative start.
123    SetToFloor,
124    /// Leave the rate cells untouched: inherit the learned rate, if any.
125    Inherit,
126}
127
128impl InitPolicy {
129    fn target_rate(self, config: &DirectionalRateLimitConfig) -> Option<f64> {
130        match self {
131            Self::SetToInit => Some(config.init_rate),
132            Self::SetToFloor => Some(config.floor_rate),
133            Self::Inherit => None,
134        }
135    }
136}
137
138/// Rate-limit parameters.
139#[derive(Debug, Clone)]
140pub(crate) struct DirectionalRateLimitConfig {
141    // HTTP requests per second (rps).
142    pub(crate) init_rate: f64,
143    pub(crate) floor_rate: f64,
144    pub(crate) ceiling_rate: f64,
145    pub(crate) horizon: Duration,
146    pub(crate) max_wait: Duration,
147    pub(crate) init_policy: InitPolicy,
148}
149
150/// Rate-limit parameters.
151#[derive(Debug, Clone)]
152pub(crate) struct RateLimitConfig {
153    pub(crate) read: DirectionalRateLimitConfig,
154    pub(crate) write: DirectionalRateLimitConfig,
155}
156
157impl From<CloudRateLimitConfig> for RateLimitConfig {
158    fn from(value: CloudRateLimitConfig) -> Self {
159        fn to_rate_limit_config(
160            config: &CloudDirectionalRateLimitConfig,
161            defaults: &DirectionalRateLimitConfig,
162        ) -> DirectionalRateLimitConfig {
163            let ceiling_rate = config
164                .ceiling_rate
165                .map_or(defaults.ceiling_rate, |r| r as f64)
166                .max(MIN_RATE);
167
168            let floor_rate = config
169                .floor_rate
170                .map_or(defaults.floor_rate, |r| r as f64)
171                .min(ceiling_rate)
172                .max(MIN_RATE);
173
174            let init_rate = config
175                .init_rate
176                .map_or(defaults.init_rate, |r| r as f64)
177                .clamp(floor_rate, ceiling_rate);
178
179            DirectionalRateLimitConfig {
180                init_rate,
181                floor_rate,
182                ceiling_rate,
183                horizon: defaults.horizon,
184                max_wait: defaults.max_wait,
185                init_policy: defaults.init_policy,
186            }
187        }
188
189        let read_config = to_rate_limit_config(&value.read, &READ_DEFAULTS);
190        let write_config = to_rate_limit_config(&value.write, &WRITE_DEFAULTS);
191
192        return RateLimitConfig {
193            read: read_config,
194            write: write_config,
195        };
196
197        static READ_DEFAULTS: LazyLock<DirectionalRateLimitConfig> = LazyLock::new(|| {
198            directional_defaults(
199                DEFAULT_READ_INIT_RATE,
200                DEFAULT_READ_FLOOR_RATE,
201                "POLARS_CLOUD_READ_INIT_RATE",
202                "POLARS_CLOUD_READ_FLOOR_RATE",
203                "POLARS_CLOUD_READ_CEILING_RATE",
204            )
205        });
206
207        static WRITE_DEFAULTS: LazyLock<DirectionalRateLimitConfig> = LazyLock::new(|| {
208            directional_defaults(
209                DEFAULT_WRITE_INIT_RATE,
210                DEFAULT_WRITE_FLOOR_RATE,
211                "POLARS_CLOUD_WRITE_INIT_RATE",
212                "POLARS_CLOUD_WRITE_FLOOR_RATE",
213                "POLARS_CLOUD_WRITE_CEILING_RATE",
214            )
215        });
216
217        fn directional_defaults(
218            default_init_rate: f64,
219            default_floor_rate: f64,
220            init_rate_env: &'static str,
221            floor_rate_env: &'static str,
222            ceiling_rate_env: &'static str,
223        ) -> DirectionalRateLimitConfig {
224            DirectionalRateLimitConfig {
225                init_rate: parse_env_var(default_init_rate, init_rate_env),
226                floor_rate: parse_env_var(default_floor_rate, floor_rate_env),
227                ceiling_rate: parse_env_var(DEFAULT_CEILING_RATE, ceiling_rate_env),
228                horizon: Duration::from_millis(parse_env_var(
229                    DEFAULT_RATE_HORIZON_MS,
230                    "POLARS_CLOUD_RATE_HORIZON_MS",
231                )),
232                max_wait: Duration::from_millis(parse_env_var(
233                    DEFAULT_RATE_MAX_WAIT_MS,
234                    "POLARS_CLOUD_RATE_MAX_WAIT_MS",
235                )),
236                init_policy: InitPolicy::Inherit,
237            }
238        }
239
240        fn parse_env_var<T: FromStr>(default: T, name: &'static str) -> T {
241            std::env::var(name).map_or(default, |x| {
242                x.parse::<T>()
243                    .ok()
244                    .unwrap_or_else(|| panic!("invalid value for {name}: {x}"))
245            })
246        }
247    }
248}
249
250/// Read-only view of the pacing budget including the learned rate-limit signal
251/// for internal consumers (e.g., ConcurrencyController).
252/// The cells contains learned rates (f64 bits in an AtomicU64), and are updated
253/// by the AIMD loop only.
254/// Valid for the process lifetime: object store rebuilds can rotate behind the
255/// cells rather than replacing them.
256#[derive(Debug, Clone)]
257pub struct PacingBudget {
258    rate_bits: RateCell,
259    horizon: Duration,
260}
261
262impl PacingBudget {
263    pub fn rate(&self) -> f64 {
264        f64::from_bits(self.rate_bits.load(Relaxed))
265    }
266
267    pub fn horizon(&self) -> Duration {
268        self.horizon
269    }
270
271    pub fn request_budget(&self, bdp: f64) -> f64 {
272        bdp.min(self.rate() * self.horizon().as_secs_f64())
273    }
274}
275// Persist rate-limit state when rebuilding an object_store with InitPolicy::Inherit.
276#[derive(Debug, Clone)]
277pub(crate) struct RateState {
278    pub rate_bits: RateCell,
279    pub max_bits: RateCell, // NaN represents None
280}
281
282impl RateState {
283    pub fn new(init_rate: f64) -> Self {
284        Self {
285            rate_bits: Arc::new(AtomicU64::new(init_rate.to_bits())),
286            max_bits: Arc::new(AtomicU64::new(f64::NAN.to_bits())),
287        }
288    }
289
290    #[inline]
291    pub fn get_rate(&self) -> f64 {
292        f64::from_bits(self.rate_bits.load(Relaxed))
293    }
294
295    #[inline]
296    pub fn set_rate(&self, rate: f64) {
297        self.rate_bits.store(rate.to_bits(), Relaxed);
298    }
299
300    #[inline]
301    pub fn get_last_max(&self) -> Option<f64> {
302        // NaN means no value.
303        let value = f64::from_bits(self.max_bits.load(Relaxed));
304        (!value.is_nan()).then_some(value)
305    }
306
307    #[inline]
308    pub fn set_last_max(&self, max: Option<f64>) {
309        let bits = max.unwrap_or(f64::NAN).to_bits();
310        self.max_bits.store(bits, Relaxed);
311    }
312
313    pub fn reset(&self, init_rate: f64) {
314        self.set_rate(init_rate);
315        self.set_last_max(None);
316    }
317}
318
319/// Builder-owned. State that may survive object store rebuilds.
320#[derive(Debug)]
321struct RateLimitState {
322    read_state: RateState,
323    write_state: RateState,
324}
325
326impl RateLimitState {
327    pub(crate) fn new(config: &RateLimitConfig) -> Self {
328        Self {
329            read_state: RateState::new(config.read.init_rate),
330            write_state: RateState::new(config.write.init_rate),
331        }
332    }
333
334    /// Initialize the cell values from prior state on object store rebuild.
335    pub(crate) fn apply_init_policy(&self, config: &RateLimitConfig) {
336        fn apply_one(state: &RateState, config: &DirectionalRateLimitConfig) {
337            if let Some(rate) = config.init_policy.target_rate(config) {
338                state.reset(rate);
339            }
340        }
341
342        apply_one(&self.read_state, &config.read);
343        apply_one(&self.write_state, &config.write);
344    }
345}
346
347#[derive(Debug)]
348pub(crate) struct RateLimiter {
349    pub(crate) config: RateLimitConfig,
350    state: RateLimitState,
351}
352
353impl RateLimiter {
354    pub(crate) fn new(config: RateLimitConfig) -> Self {
355        let state = RateLimitState::new(&config);
356        Self { config, state }
357    }
358
359    /// Initialize the cell values. Called at every store (re)build, before
360    /// constructing the new PacedHttpConnector.
361    pub(crate) fn apply_init_policy(&self) {
362        self.state.apply_init_policy(&self.config);
363    }
364
365    // Read-only view of pacing budget for 'read' based on learned rate.
366    // Targeted at upstream consumers.
367    pub(crate) fn read_budget(&self) -> PacingBudget {
368        PacingBudget {
369            rate_bits: Arc::clone(&self.state.read_state.rate_bits),
370            horizon: self.config.read.horizon,
371        }
372    }
373
374    // Read-only view of pacing budget for 'write' based on learned rate.
375    // Targeted at upstream consumers.
376    #[allow(unused)]
377    pub(crate) fn write_budget(&self) -> PacingBudget {
378        PacingBudget {
379            rate_bits: Arc::clone(&self.state.write_state.rate_bits),
380            horizon: self.config.write.horizon,
381        }
382    }
383}
384
385#[derive(Debug, Copy, Clone)]
386enum Regime {
387    // No knowledge about ceiling or cut. Fast-ramp.
388    Search,
389    // Recover from a congestion event.
390    Recover { until: Instant, anchor: f64 },
391    // Track towards a learned ceiling and probe above.
392    Track { anchor: f64 },
393}
394
395// Adaptive Increase Multiplicative Decrease (AIMD) state.
396// Note: the initial fast_ramp is multiplicative, not additive.
397// Cold path: state gets updated every tick interval.
398#[derive(Debug)]
399struct AimdState {
400    // How to drive the rate based on what has been observed.
401    regime: Regime,
402
403    // Source of truth for the learned request rates in requests per second (rps).
404    // Represented as f64 in bits and shared as Atomic cells.
405    shared: RateState,
406    // Exponentially weighted moving average (EWMA) of the success rate.
407    success_rate: Option<f64>,
408    // Start of window, where window is the interval between 2 ticks.
409    window_start: Option<Instant>,
410    // Last cut time.
411    last_cut_time: Option<Instant>,
412
413    // Pacer statistics at last tick.
414    prev_admitted: u64,
415    prev_denied: u64,
416    prev_resp_succeeded: u64,
417    prev_resp_throttled: u64,
418}
419
420impl AimdState {
421    #[inline]
422    pub fn rate(&self) -> f64 {
423        self.shared.get_rate()
424    }
425
426    #[inline]
427    pub fn set_rate(&mut self, rate: f64) {
428        self.shared.set_rate(rate);
429    }
430
431    #[inline]
432    pub fn last_max(&self) -> Option<f64> {
433        self.shared.get_last_max()
434    }
435
436    #[inline]
437    pub fn set_last_max(&mut self, max: Option<f64>) {
438        self.shared.set_last_max(max);
439    }
440
441    /// The success_rate is the observed 'goodput' signal.
442    fn update_success_rate(&mut self, successes: u64, elapsed_s: f64) {
443        // Too short to be statistically meaningful (also guards div-by-zero).
444        if elapsed_s < 0.5 * TICK_INTERVAL.as_secs_f64() {
445            return;
446        }
447
448        if successes == 0 {
449            return;
450        }
451
452        let success_rate = successes as f64 / elapsed_s;
453        self.success_rate = Some(match self.success_rate {
454            None => success_rate,
455            Some(rate) => SUCCESS_SMOOTH * success_rate + (1.0 - SUCCESS_SMOOTH) * rate,
456        });
457    }
458}
459
460// Lock-free metrics and time window filter to guard the Mutex, warm path.
461#[derive(Debug)]
462pub(crate) struct PacerSignal {
463    // Fast-path view of window end.
464    window_end_ns: AtomicU64,
465    // Advisory refractory pre-check (0 = never).
466    last_cut_ns: AtomicU64,
467    // Cumulative HTTP response counters, since epoch.
468    resp_succeeded: AtomicU64,
469    resp_throttled: AtomicU64,
470}
471
472#[derive(Debug)]
473pub(crate) struct AdaptiveRateController {
474    epoch: Instant,
475    // Hot path, lock-free atomics.
476    //  TBD - do we need an Arc?
477    pacer: Arc<Pacer>,
478    // Warm signal path, lock-free atomics.
479    signal: PacerSignal,
480    // Cold path, change state and rate.
481    state: Mutex<AimdState>,
482    label: &'static str,
483    id: u32,
484    config: DirectionalRateLimitConfig,
485}
486
487impl AdaptiveRateController {
488    fn new(
489        label: &'static str,
490        id: u32,
491        shared: RateState,
492        config: DirectionalRateLimitConfig,
493    ) -> Self {
494        let epoch = Instant::now();
495
496        let regime = match shared.get_last_max() {
497            Some(anchor) => Regime::Track { anchor },
498            None => Regime::Search,
499        };
500
501        let token_bucket = Arc::new(TokenBucket::new(shared.rate_bits.clone()));
502        let pacer = Pacer::start(token_bucket, config.max_wait);
503
504        let signal = PacerSignal {
505            window_end_ns: AtomicU64::new(u64::MAX),
506            last_cut_ns: AtomicU64::new(0),
507            resp_succeeded: AtomicU64::new(0),
508            resp_throttled: AtomicU64::new(0),
509        };
510
511        let state = Mutex::new(AimdState {
512            regime,
513            shared,
514            success_rate: None,
515            window_start: None,
516            last_cut_time: None,
517            prev_admitted: 0,
518            prev_denied: 0,
519            prev_resp_succeeded: 0,
520            prev_resp_throttled: 0,
521        });
522
523        Self {
524            epoch,
525            pacer,
526            signal,
527            state,
528            label,
529            id,
530            config,
531        }
532    }
533
534    #[inline]
535    fn now_ns(&self) -> u64 {
536        self.epoch.elapsed().as_nanos() as u64
537    }
538
539    #[inline]
540    fn mark_first_traffic(&self) {
541        // The window origin is FIRST TRAFFIC, not construction: the pipeline
542        // takes an unknown time to spin up, and a window that spans the idle
543        // prologue measures emptiness (4 successes / 533ms at cold start).
544        if self.signal.window_end_ns.load(Relaxed) == u64::MAX {
545            let now_ns = self.now_ns();
546            self.signal
547                .window_end_ns
548                .store(now_ns + TICK_INTERVAL.as_nanos() as u64, Relaxed);
549            self.state.lock().unwrap().window_start = Some(Instant::now());
550        }
551    }
552
553    fn on_congestion(&self) {
554        self.signal.resp_throttled.fetch_add(1, Relaxed);
555
556        self.mark_first_traffic();
557        self.maybe_settle();
558
559        // Lock-free fast-path (1): advisory refractory pre-check, avoid Mutex storm.
560        let now_ns = self.now_ns();
561        let last = self.signal.last_cut_ns.load(Relaxed);
562        if last != 0 && now_ns.saturating_sub(last) < REFRACTORY.as_nanos() as u64 {
563            return;
564        }
565
566        // Locking fast-path (2): check the authoritative `last_cut` in AimdState.
567        let now = Instant::now();
568        let mut state = self.state.lock().unwrap();
569        if state
570            .last_cut_time
571            .is_some_and(|t| now.duration_since(t) < REFRACTORY)
572        {
573            return;
574        }
575
576        // Calculate anchor rate, which is our (conservative) estimate for the unknown rate-limit enforced
577        // by the back-end.
578        let (anchor, beta) = match state.success_rate {
579            None => (state.rate(), BETA_NO_SIGNAL),
580            Some(success_rate) => (
581                state.rate().min(success_rate.max(self.config.floor_rate)),
582                BETA,
583            ),
584        };
585
586        // Activate new rate and update state.
587        state.set_rate((anchor * beta).max(self.config.floor_rate));
588        state.last_cut_time = Some(now);
589        self.signal.last_cut_ns.store(now_ns, Relaxed);
590
591        // Unconditionally move into Recover on every cut.
592        state.regime = Regime::Recover {
593            until: now.checked_add(PROBE_QUIET_PERIOD).unwrap(),
594            anchor,
595        };
596
597        // Log.
598        if *LOG_HTTP_RATE_LIMIT {
599            eprintln!(
600                "[http rate_limit #{}_{} {}] ..cut (anchored): rate: {:.1}, success_rate: {:.1}, last_max: {:.1}",
601                self.id,
602                self.label,
603                Utc::now(),
604                state.rate(),
605                state.success_rate.unwrap_or_default(),
606                state.last_max().unwrap_or_default(),
607            );
608        }
609    }
610
611    fn on_success(&self) {
612        self.signal.resp_succeeded.fetch_add(1, Relaxed);
613
614        self.mark_first_traffic();
615        self.maybe_settle();
616    }
617
618    fn on_other(&self) {
619        self.mark_first_traffic();
620        self.maybe_settle();
621    }
622
623    #[inline]
624    fn maybe_settle(&self) {
625        // Lock-free fast path when settlement is not due.
626
627        if self.now_ns() >= self.signal.window_end_ns.load(Relaxed) {
628            self.settle_lazy_tick();
629        }
630    }
631
632    // Settle the rate-limiter after every tick interval, on the first success response.
633    fn settle_lazy_tick(&self) {
634        let mut state = self.state.lock().unwrap();
635        let now = Instant::now();
636        let now_ns = self.now_ns();
637
638        // Another success event may have settled while we waited on the lock.
639        let elapsed = now.duration_since(state.window_start.unwrap());
640        let ticks = (elapsed.as_secs_f64() / TICK_INTERVAL.as_secs_f64()).floor();
641        if ticks < 1.0 {
642            return;
643        }
644
645        // Update AIMD stats.
646        let admitted = self.pacer.admitted();
647        let win_admitted = admitted.saturating_sub(state.prev_admitted);
648        state.prev_admitted = admitted;
649
650        let denied = self.pacer.denied.load(Relaxed);
651        let win_denied = denied.saturating_sub(state.prev_denied);
652        state.prev_denied = denied;
653
654        let succeeded = self.signal.resp_succeeded.load(Relaxed);
655        let win_succeeded = succeeded.saturating_sub(state.prev_resp_succeeded);
656        state.prev_resp_succeeded = succeeded;
657
658        let throttled = self.signal.resp_throttled.load(Relaxed);
659        let win_throttled = throttled.saturating_sub(state.prev_resp_throttled);
660        state.prev_resp_throttled = throttled;
661
662        // Growth conditions.
663        let elapsed_s = elapsed.as_secs_f64();
664        let queued = self.pacer.queue_depth() > 0;
665        let util_bound = win_admitted as f64 >= SATURATION_THRESHOLD * state.rate() * elapsed_s;
666
667        // Update signal, if any.
668        state.update_success_rate(win_succeeded, elapsed_s);
669
670        // Goodput is a always a lower bound on capacity, so an observation may always
671        // increase. It may only decrease when the system was pressure-tested (i.e., demand saturated).
672        if let Some(observed) = state.success_rate {
673            let pressure_tested =
674                win_admitted as f64 / elapsed_s >= state.last_max().unwrap_or(0.0);
675            let new_max = match state.last_max() {
676                Some(prev) if !pressure_tested => prev.max(observed),
677                _ => observed,
678            };
679            state.set_last_max(Some(new_max.max(self.config.floor_rate)));
680        }
681
682        // Evaluate and apply rate and regime changes.
683        let verdict = {
684            let rate = state.rate();
685            let init_rate = self.config.init_rate;
686            let floor_rate = self.config.floor_rate;
687            let ceiling_rate = self.config.ceiling_rate;
688
689            if let Regime::Recover { until, anchor } = state.regime
690                && now > until
691            {
692                state.regime = Regime::Track { anchor }
693            };
694
695            if win_admitted + win_denied == 0 {
696                // No demand - decay toward seed only if we are currently above it.
697                if rate > init_rate {
698                    let rate = init_rate + (state.rate() - init_rate) * IDLE_DECAY.powf(ticks);
699                    state.set_rate(rate.clamp(floor_rate, ceiling_rate));
700                }
701                // Note: we leave last_max as-is which is not always right.
702                "decay (idle)"
703            } else {
704                match state.regime {
705                    Regime::Search => {
706                        // Cold start -> aggressive exponential growth, if earned
707                        // Note. One may expect `.powf(ticks)`; this is deliberately omitted.
708                        if util_bound && win_throttled == 0 {
709                            state.set_rate((rate * FAST_RAMP_FACTOR).min(ceiling_rate));
710                            "increase (fast_ramp)"
711                        } else {
712                            "hold (app-limited)"
713                        }
714                    },
715                    Regime::Recover { .. } => {
716                        // No-op, rate changes were handled in `on_congestion`
717                        "hold (recover)"
718                    },
719                    Regime::Track { anchor } => {
720                        // Below anchor + backlogged/saturated -> fast reclaim up to our known-good ceiling
721                        if rate < anchor && util_bound {
722                            state.set_rate((rate * FAST_RAMP_FACTOR).min(anchor).min(ceiling_rate));
723                            "increase (reclaim)"
724                        } else
725                        // Below anchor but not hitting capacity bounds -> wait for more traffic
726                        if rate < anchor {
727                            "hold (app-limited)"
728                        } else
729                        // At or above anchor + saturated + past quiet period -> cautious additive probing
730                        if util_bound {
731                            let delta = (PROBE_FRACTION * anchor).max(PROBE_MIN);
732                            state.set_rate((rate + delta).min(ceiling_rate));
733                            "increase (probe)"
734                        } else
735                        // At/above anchor and backlog is draining out.
736                        if queued {
737                            "hold (draining)"
738                        } else {
739                            // Fallback for unutilized capacity above the anchor.
740                            "hold (app-limited)"
741                        }
742                    },
743                }
744            }
745        };
746
747        // Update state.
748        state.window_start = Some(now);
749        self.signal
750            .window_end_ns
751            .store(now_ns + TICK_INTERVAL.as_nanos() as u64, Relaxed);
752
753        // Logging.
754        if *LOG_HTTP_RATE_LIMIT {
755            eprintln!(
756                "[http rate_limit #{}_{} {}] {}: rate: {:.1}, success_rate_ewma: {:.1}, last_max: {:.1}, \
757                    elapsed: {:.1}s, win_admit: {}, win_deny: {}, win_success: {}, win_throttle: {},  q_depth: {}",
758                self.id,
759                self.label,
760                Utc::now(),
761                verdict,
762                state.rate(),
763                state.success_rate.unwrap_or_default(),
764                state.last_max().unwrap_or_default(),
765                elapsed_s,
766                win_admitted,
767                win_denied,
768                win_succeeded,
769                win_throttled,
770                self.pacer.queue_depth(),
771            );
772        }
773    }
774}
775
776/// Grouped rate-limit controller as used by a logical object_store instance.
777#[derive(Debug)]
778pub(crate) struct RateController {
779    // GET, HEAD
780    pub read: AdaptiveRateController,
781    // PUT, POST, DELETE
782    pub write: AdaptiveRateController,
783}
784
785impl RateController {
786    pub(crate) fn new(rate_limiter: &RateLimiter) -> Arc<Self> {
787        // NOTE: A rebuilt object_store can share one rate state.
788        // Acceptable transient situation.
789        let id = CONTROLLER_ID.fetch_add(1, Relaxed);
790        Arc::new(Self {
791            read: AdaptiveRateController::new(
792                "read",
793                id,
794                rate_limiter.state.read_state.clone(),
795                rate_limiter.config.read.clone(),
796            ),
797            write: AdaptiveRateController::new(
798                "write",
799                id,
800                rate_limiter.state.write_state.clone(),
801                rate_limiter.config.write.clone(),
802            ),
803        })
804    }
805
806    fn class(&self, req: &HttpRequest) -> &AdaptiveRateController {
807        match req.method().as_str() {
808            "GET" | "HEAD" => &self.read,
809            _ => &self.write,
810        }
811    }
812}
813
814#[derive(Debug)]
815struct PacerBusy {
816    est_wait_ms: u64,
817}
818impl fmt::Display for PacerBusy {
819    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
820        write!(
821            f,
822            "internal pacer denied: estimated wait {}ms exceeds bound",
823            self.est_wait_ms
824        )
825    }
826}
827impl std::error::Error for PacerBusy {}
828
829// PACER INCL WAITER_QUEUE
830
831// The pacer is responsible for admissin/denial into the `TokenBucket` annex
832// `WaiterQueue`. Once admitted, it paces the requests to the prescribed rate.
833
834#[derive(Debug, Default)]
835struct WaiterQueue {
836    parked: Mutex<VecDeque<oneshot::Sender<()>>>,
837    /// Lock-free mirror of parked.len(): read by the fast path (anti-barge),
838    /// the population coupling, and telemetry.
839    depth: AtomicU64,
840}
841
842impl WaiterQueue {
843    fn park(&self) -> oneshot::Receiver<()> {
844        let (tx, rx) = oneshot::channel();
845        let mut q = self.parked.lock().unwrap();
846        q.push_back(tx);
847        self.depth.store(q.len() as u64, Relaxed);
848        rx
849    }
850
851    /// Pop the head and deliver a grant. Returns false when the queue is
852    /// empty. A dead (cancelled) head consumes no grant: we skip it and keep
853    /// the token for the next live waiter.
854    fn grant_one(&self) -> bool {
855        let mut q = self.parked.lock().unwrap();
856        while let Some(tx) = q.pop_front() {
857            self.depth.store(q.len() as u64, Relaxed);
858            if tx.send(()).is_ok() {
859                return true; // grant delivered
860            }
861            // Cancelled waiter: token still in hand, try the next.
862        }
863        false
864    }
865
866    fn depth(&self) -> u64 {
867        self.depth.load(Relaxed)
868    }
869}
870
871/// Lock-free hot path enforcing the rate-limit by pacing requests through the token bucket.
872#[derive(Debug)]
873pub struct Pacer {
874    bucket: Arc<TokenBucket>,
875    queue: Arc<WaiterQueue>,
876    max_wait: Duration,
877    // Decided to admit.
878    admitted: AtomicU64,
879    // Decided not to park.
880    denied: AtomicU64,
881}
882
883impl Pacer {
884    /// Construct and spawn the wake tick. The tick is per-pacer.
885    pub fn start(bucket: Arc<TokenBucket>, max_wait: Duration) -> Arc<Self> {
886        let pacer = Arc::new(Self {
887            bucket,
888            queue: Arc::new(WaiterQueue::default()),
889            max_wait,
890            admitted: AtomicU64::new(0),
891            denied: AtomicU64::new(0),
892        });
893        Self::spawn_wake_tick(Arc::downgrade(&pacer));
894        pacer
895    }
896
897    /// Admit a request to the pacer, which may get queued internally. Returns a
898    /// PacerBusy Error when the estimated wait exceeds the wait bound.
899    pub async fn admit(&self) -> Result<(), HttpError> {
900        // Anti-barge plus fast path.
901        if self.queue.depth() == 0 && self.bucket.try_acquire().is_ok() {
902            self.admitted.fetch_add(1, Relaxed);
903            return Ok(());
904        }
905
906        // JIT rejection: price the line, before parking.
907        let est_wait =
908            Duration::from_secs_f64((self.queue.depth() + 1) as f64 / self.bucket.rate());
909        if est_wait > self.max_wait {
910            self.denied.fetch_add(1, Relaxed);
911            return Err(HttpError::new(
912                HttpErrorKind::Timeout, // retryable by object_store
913                PacerBusy {
914                    est_wait_ms: est_wait.as_millis() as u64,
915                },
916            ));
917        }
918
919        // Park. Wake when granted.
920        let rx = self.queue.park();
921
922        // Wait. On Err, allow through unpaced when WaiterQueue/Pacer is torn down.
923        let _ = rx.await;
924        self.admitted.fetch_add(1, Relaxed);
925        Ok(())
926    }
927
928    pub fn queue_depth(&self) -> u64 {
929        self.queue.depth()
930    }
931
932    pub fn admitted(&self) -> u64 {
933        self.admitted.load(Relaxed)
934    }
935
936    pub fn bucket(&self) -> &Arc<TokenBucket> {
937        &self.bucket
938    }
939
940    /// The single timer in the design. Each tick: while waiters exist AND the
941    /// bucket yields a token, hand grants to the FIFO head. Weak handle so a
942    /// dropped pacer (store rebuild) tears its tick down with it.
943    fn spawn_wake_tick(gate: std::sync::Weak<Self>) {
944        ASYNC.spawn(async move {
945            let mut tick = tokio::time::interval(HTTP_RATE_LIMIT_WAKE_TICK);
946            tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
947            loop {
948                tick.tick().await;
949                let Some(g) = gate.upgrade() else { return };
950
951                while g.queue.depth() > 0 {
952                    if g.bucket.try_acquire().is_err() {
953                        break;
954                    }
955                    if !g.queue.grant_one() {
956                        // Queue drained between the depth check and the grant
957                        // (or all-cancelled): one token over-taken.
958                        break;
959                    }
960                }
961            }
962        });
963    }
964}
965
966// OBJECT_STORE HTTP_SERVICE MIDDLEWARE
967
968#[derive(Debug)]
969pub(crate) struct PacedHttpConnector {
970    inner: Box<dyn HttpConnector>,
971    controller: Arc<RateController>,
972}
973
974impl PacedHttpConnector {
975    pub(crate) fn new(inner: Box<dyn HttpConnector>, rate_limiter: &RateLimiter) -> Self {
976        // The HTTP Connector may re-use its learned request rate.
977        rate_limiter.apply_init_policy();
978
979        Self {
980            inner,
981            controller: RateController::new(rate_limiter),
982        }
983    }
984}
985
986impl HttpConnector for PacedHttpConnector {
987    fn connect(&self, options: &ClientOptions) -> object_store::Result<HttpClient> {
988        let client = self.inner.connect(options)?;
989        Ok(HttpClient::new(PacedHttpService {
990            inner: client,
991            controller: Arc::clone(&self.controller),
992        }))
993    }
994}
995
996#[derive(Debug)]
997pub(crate) struct PacedHttpService {
998    inner: HttpClient,
999    controller: Arc<RateController>,
1000}
1001
1002#[async_trait]
1003impl HttpService for PacedHttpService {
1004    async fn call(&self, req: HttpRequest) -> Result<HttpResponse, HttpError> {
1005        let verb_pacer = self.controller.class(&req);
1006
1007        // Enforce pacing: on admission, wait; on denied, raise an error.
1008        // Requests get denied when the estimated wait is too high (aka 'shed').
1009        // The object_store is responsible for retry.
1010        verb_pacer.pacer.admit().await?;
1011
1012        let response = self.inner.execute(req).await;
1013        match &response {
1014            Ok(r) if r.status().as_u16() == 429 || r.status().as_u16() == 503 => {
1015                verb_pacer.on_congestion()
1016            },
1017            Ok(r) if r.status().is_success() => verb_pacer.on_success(),
1018            _ => verb_pacer.on_other(),
1019        }
1020
1021        response
1022    }
1023}