1use 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
51static CONTROLLER_ID: AtomicU32 = AtomicU32::new(0);
53
54static LOG_HTTP_RATE_LIMIT: LazyLock<bool> =
56 LazyLock::new(|| std::env::var("POLARS_LOG_HTTP_RATE_LIMIT").is_ok());
57
58const 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
65const MIN_RATE: f64 = 1.0;
67
68const FAST_RAMP_FACTOR: f64 = 2.0;
71const PROBE_FRACTION: f64 = 0.1;
73const PROBE_MIN: f64 = 5.0;
74const BETA: f64 = 0.7;
76const BETA_NO_SIGNAL: f64 = 0.5;
78
79const DEFAULT_RATE_HORIZON_MS: u64 = 200;
82const DEFAULT_RATE_MAX_WAIT_MS: u64 = 10_000;
92const TICK_INTERVAL: Duration = Duration::from_secs(1);
94const REFRACTORY: Duration = Duration::from_secs(1);
96const PROBE_QUIET_PERIOD: Duration = Duration::from_secs(2);
98
99const SATURATION_THRESHOLD: f64 = 0.8;
102const SUCCESS_SMOOTH: f64 = 0.8;
104const IDLE_DECAY: f64 = 0.9;
106
107pub(crate) const HTTP_RATE_LIMIT_WAKE_TICK: Duration = Duration::from_millis(1);
110
111pub type RateCell = Arc<AtomicU64>;
114
115#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118#[allow(unused)]
119pub(crate) enum InitPolicy {
120 SetToInit,
122 SetToFloor,
124 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#[derive(Debug, Clone)]
140pub(crate) struct DirectionalRateLimitConfig {
141 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#[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#[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#[derive(Debug, Clone)]
277pub(crate) struct RateState {
278 pub rate_bits: RateCell,
279 pub max_bits: RateCell, }
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 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#[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 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 pub(crate) fn apply_init_policy(&self) {
362 self.state.apply_init_policy(&self.config);
363 }
364
365 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 #[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 Search,
389 Recover { until: Instant, anchor: f64 },
391 Track { anchor: f64 },
393}
394
395#[derive(Debug)]
399struct AimdState {
400 regime: Regime,
402
403 shared: RateState,
406 success_rate: Option<f64>,
408 window_start: Option<Instant>,
410 last_cut_time: Option<Instant>,
412
413 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 fn update_success_rate(&mut self, successes: u64, elapsed_s: f64) {
443 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#[derive(Debug)]
462pub(crate) struct PacerSignal {
463 window_end_ns: AtomicU64,
465 last_cut_ns: AtomicU64,
467 resp_succeeded: AtomicU64,
469 resp_throttled: AtomicU64,
470}
471
472#[derive(Debug)]
473pub(crate) struct AdaptiveRateController {
474 epoch: Instant,
475 pacer: Arc<Pacer>,
478 signal: PacerSignal,
480 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 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 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 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 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 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 state.regime = Regime::Recover {
593 until: now.checked_add(PROBE_QUIET_PERIOD).unwrap(),
594 anchor,
595 };
596
597 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 if self.now_ns() >= self.signal.window_end_ns.load(Relaxed) {
628 self.settle_lazy_tick();
629 }
630 }
631
632 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 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 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 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 state.update_success_rate(win_succeeded, elapsed_s);
669
670 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 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 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 "decay (idle)"
703 } else {
704 match state.regime {
705 Regime::Search => {
706 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 "hold (recover)"
718 },
719 Regime::Track { anchor } => {
720 if rate < anchor && util_bound {
722 state.set_rate((rate * FAST_RAMP_FACTOR).min(anchor).min(ceiling_rate));
723 "increase (reclaim)"
724 } else
725 if rate < anchor {
727 "hold (app-limited)"
728 } else
729 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 if queued {
737 "hold (draining)"
738 } else {
739 "hold (app-limited)"
741 }
742 },
743 }
744 }
745 };
746
747 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 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#[derive(Debug)]
778pub(crate) struct RateController {
779 pub read: AdaptiveRateController,
781 pub write: AdaptiveRateController,
783}
784
785impl RateController {
786 pub(crate) fn new(rate_limiter: &RateLimiter) -> Arc<Self> {
787 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#[derive(Debug, Default)]
835struct WaiterQueue {
836 parked: Mutex<VecDeque<oneshot::Sender<()>>>,
837 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 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; }
861 }
863 false
864 }
865
866 fn depth(&self) -> u64 {
867 self.depth.load(Relaxed)
868 }
869}
870
871#[derive(Debug)]
873pub struct Pacer {
874 bucket: Arc<TokenBucket>,
875 queue: Arc<WaiterQueue>,
876 max_wait: Duration,
877 admitted: AtomicU64,
879 denied: AtomicU64,
881}
882
883impl Pacer {
884 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 pub async fn admit(&self) -> Result<(), HttpError> {
900 if self.queue.depth() == 0 && self.bucket.try_acquire().is_ok() {
902 self.admitted.fetch_add(1, Relaxed);
903 return Ok(());
904 }
905
906 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, PacerBusy {
914 est_wait_ms: est_wait.as_millis() as u64,
915 },
916 ));
917 }
918
919 let rx = self.queue.park();
921
922 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 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 break;
959 }
960 }
961 }
962 });
963 }
964}
965
966#[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 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 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}