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