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 = 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
69const MIN_RATE: f64 = 1.0;
71
72const FAST_RAMP_FACTOR: f64 = 2.0;
75const PROBE_FRACTION: f64 = 0.1;
77const PROBE_MIN: f64 = 5.0;
78const BETA: f64 = 0.7;
80const BETA_NO_SIGNAL: f64 = 0.5;
82
83const DEFAULT_RATE_HORIZON_MS: u64 = 200;
86const DEFAULT_RATE_MAX_WAIT_MS: u64 = 10_000;
96const TICK_INTERVAL: Duration = Duration::from_secs(1);
98const REFRACTORY: Duration = Duration::from_secs(1);
100const PROBE_QUIET_PERIOD: Duration = Duration::from_secs(2);
102
103const SATURATION_THRESHOLD: f64 = 0.8;
106const SUCCESS_SMOOTH: f64 = 0.8;
108const IDLE_DECAY: f64 = 0.9;
110
111pub(crate) const HTTP_RATE_LIMIT_WAKE_TICK: Duration = Duration::from_millis(1);
114
115pub type RateCell = Arc<AtomicU64>;
118
119#[derive(Debug, Clone, Copy, PartialEq, Eq)]
122#[allow(unused)]
123pub(crate) enum InitPolicy {
124 SetToInit,
126 SetToFloor,
128 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#[derive(Debug, Clone)]
144pub(crate) struct DirectionalRateLimitConfig {
145 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#[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#[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#[derive(Debug, Clone)]
281pub(crate) struct RateState {
282 pub rate_bits: RateCell,
283 pub max_bits: RateCell, }
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 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#[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 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 pub(crate) fn apply_init_policy(&self) {
366 self.state.apply_init_policy(&self.config);
367 }
368
369 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 #[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 Search,
393 Recover { until: Instant, anchor: f64 },
395 Track { anchor: f64 },
397}
398
399#[derive(Debug)]
403struct AimdState {
404 regime: Regime,
406
407 shared: RateState,
410 success_rate: Option<f64>,
412 window_start: Option<Instant>,
414 last_cut_time: Option<Instant>,
416
417 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 fn update_success_rate(&mut self, successes: u64, elapsed_s: f64) {
447 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#[derive(Debug)]
466pub(crate) struct PacerSignal {
467 window_end_ns: AtomicU64,
469 last_cut_ns: AtomicU64,
471 resp_succeeded: AtomicU64,
473 resp_throttled: AtomicU64,
474}
475
476#[derive(Debug)]
477pub(crate) struct AdaptiveRateController {
478 epoch: Instant,
479 pacer: Arc<Pacer>,
482 signal: PacerSignal,
484 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 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 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 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 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 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 state.regime = Regime::Recover {
597 until: now.checked_add(PROBE_QUIET_PERIOD).unwrap(),
598 anchor,
599 };
600
601 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 if self.now_ns() >= self.signal.window_end_ns.load(Relaxed) {
632 self.settle_lazy_tick();
633 }
634 }
635
636 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 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 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 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 state.update_success_rate(win_succeeded, elapsed_s);
673
674 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 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 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 "decay (idle)"
707 } else {
708 match state.regime {
709 Regime::Search => {
710 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 "hold (recover)"
722 },
723 Regime::Track { anchor } => {
724 if rate < anchor && util_bound {
726 state.set_rate((rate * FAST_RAMP_FACTOR).min(anchor).min(ceiling_rate));
727 "increase (reclaim)"
728 } else
729 if rate < anchor {
731 "hold (app-limited)"
732 } else
733 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 if queued {
741 "hold (draining)"
742 } else {
743 "hold (app-limited)"
745 }
746 },
747 }
748 }
749 };
750
751 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 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#[derive(Debug)]
782pub(crate) struct RateController {
783 pub read: AdaptiveRateController,
785 pub write: AdaptiveRateController,
787}
788
789impl RateController {
790 pub(crate) fn new(rate_limiter: &RateLimiter) -> Arc<Self> {
791 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#[derive(Debug, Default)]
839struct WaiterQueue {
840 parked: Mutex<VecDeque<oneshot::Sender<()>>>,
841 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 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; }
865 }
867 false
868 }
869
870 fn depth(&self) -> u64 {
871 self.depth.load(Relaxed)
872 }
873}
874
875#[derive(Debug)]
877pub struct Pacer {
878 bucket: Arc<TokenBucket>,
879 queue: Arc<WaiterQueue>,
880 max_wait: Duration,
881 admitted: AtomicU64,
883 denied: AtomicU64,
885}
886
887impl Pacer {
888 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 pub async fn admit(&self) -> Result<(), HttpError> {
904 if self.queue.depth() == 0 && self.bucket.try_acquire().is_ok() {
906 self.admitted.fetch_add(1, Relaxed);
907 return Ok(());
908 }
909
910 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, PacerBusy {
918 est_wait_ms: est_wait.as_millis() as u64,
919 },
920 ));
921 }
922
923 let rx = self.queue.park();
925
926 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 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 break;
963 }
964 }
965 }
966 });
967 }
968}
969
970#[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 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 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}