1use std::sync::{Arc, Mutex, MutexGuard};
11use std::time::{Duration, SystemTime, UNIX_EPOCH};
12
13use tokio::time::Instant;
14
15use crate::usdm::types::DepthLimit;
16
17pub const PUBLISHED_WEIGHT_PER_MINUTE: u32 = 2400;
19
20pub const PUBLISHED_FUNDING_PER_FIVE_MINUTES: u32 = 500;
23
24pub const DEFAULT_BAN: Duration = Duration::from_secs(120);
27
28pub const MAX_COOLDOWN: Duration = Duration::from_secs(3 * 24 * 60 * 60);
31
32const RESERVED_FRACTION: u32 = 10;
36
37const MINUTE_MS: u64 = 60_000;
38const FUNDING_PERIOD: Duration = Duration::from_secs(300);
39
40fn after_reserve(published: u32) -> u32 {
41 published - published.div_ceil(RESERVED_FRACTION)
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum Cost {
47 Weight(u32),
49 Funding,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55#[non_exhaustive]
56pub enum Route {
57 Ping,
59 Time,
61 ExchangeInfo,
63 FundingInfo,
65 Ticker24h {
67 all: bool,
69 },
70 PremiumIndex {
72 all: bool,
74 },
75 Klines {
77 limit: Option<u32>,
79 },
80 FundingRate,
82 OpenInterest,
84 AggTrades,
86 Depth {
88 limit: Option<DepthLimit>,
90 },
91}
92
93impl Route {
94 pub fn path(self) -> &'static str {
96 match self {
97 Self::Ping => "/fapi/v1/ping",
98 Self::Time => "/fapi/v1/time",
99 Self::ExchangeInfo => "/fapi/v1/exchangeInfo",
100 Self::FundingInfo => "/fapi/v1/fundingInfo",
101 Self::Ticker24h { .. } => "/fapi/v1/ticker/24hr",
102 Self::PremiumIndex { .. } => "/fapi/v1/premiumIndex",
103 Self::Klines { .. } => "/fapi/v1/klines",
104 Self::FundingRate => "/fapi/v1/fundingRate",
105 Self::OpenInterest => "/fapi/v1/openInterest",
106 Self::AggTrades => "/fapi/v1/aggTrades",
107 Self::Depth { .. } => "/fapi/v1/depth",
108 }
109 }
110
111 pub fn cost(self) -> Cost {
119 match self {
120 Self::Ping | Self::Time | Self::ExchangeInfo | Self::OpenInterest => Cost::Weight(1),
121 Self::FundingInfo | Self::FundingRate => Cost::Funding,
122 Self::Ticker24h { all } => Cost::Weight(if all { 40 } else { 1 }),
123 Self::PremiumIndex { all } => Cost::Weight(if all { 10 } else { 1 }),
124 Self::Klines { limit: None } => Cost::Weight(5),
125 Self::Klines { limit: Some(limit) } => Cost::Weight(match limit {
126 0..=100 => 1,
127 101..=500 => 2,
128 501..=1000 => 5,
129 _ => 10,
130 }),
131 Self::AggTrades => Cost::Weight(20),
132 Self::Depth { limit: None } => Cost::Weight(1),
133 Self::Depth { limit: Some(limit) } => Cost::Weight(match limit {
134 DepthLimit::Five | DepthLimit::Ten | DepthLimit::Twenty | DepthLimit::Fifty => 2,
135 DepthLimit::Hundred => 5,
136 DepthLimit::FiveHundred => 10,
137 DepthLimit::Thousand => 20,
138 }),
139 }
140 }
141}
142
143#[derive(Debug, Clone)]
172pub struct WeightBudget {
173 inner: Arc<Inner>,
174}
175
176#[derive(Debug)]
177struct Inner {
178 per_minute: u32,
179 funding_interval: Duration,
180 clock: Clock,
181 state: Mutex<State>,
182}
183
184#[derive(Debug)]
185struct State {
186 minute: u64,
188 used: u32,
189 cooldown_until: Option<Instant>,
190 next_funding: Option<Instant>,
191}
192
193#[derive(Debug, Clone, Copy, PartialEq, Eq)]
197pub(crate) struct Charge {
198 minute: Option<u64>,
199}
200
201#[derive(Debug, Clone)]
203enum Clock {
204 System,
206 #[cfg(test)]
208 Tokio { start: Instant, start_ms: u64 },
209 #[cfg(test)]
211 Manual(Arc<std::sync::atomic::AtomicU64>),
212}
213
214impl Clock {
215 fn now_ms(&self) -> u64 {
216 match self {
217 Self::System => SystemTime::now()
218 .duration_since(UNIX_EPOCH)
219 .map_or(0, |since| since.as_millis() as u64),
220 #[cfg(test)]
221 Self::Tokio { start, start_ms } => *start_ms + start.elapsed().as_millis() as u64,
222 #[cfg(test)]
223 Self::Manual(ms) => ms.load(std::sync::atomic::Ordering::SeqCst),
224 }
225 }
226}
227
228impl Default for WeightBudget {
229 fn default() -> Self {
230 Self::new()
231 }
232}
233
234impl WeightBudget {
235 pub fn new() -> Self {
237 Self::with_clock(
238 PUBLISHED_WEIGHT_PER_MINUTE,
239 PUBLISHED_FUNDING_PER_FIVE_MINUTES,
240 Clock::System,
241 )
242 }
243
244 fn with_clock(weight_per_minute: u32, funding_per_five_minutes: u32, clock: Clock) -> Self {
245 let funding_slots = after_reserve(funding_per_five_minutes).max(2) - 1;
248 Self {
249 inner: Arc::new(Inner {
250 per_minute: after_reserve(weight_per_minute),
251 funding_interval: FUNDING_PERIOD / funding_slots,
252 clock,
253 state: Mutex::new(State {
254 minute: 0,
255 used: 0,
256 cooldown_until: None,
257 next_funding: None,
258 }),
259 }),
260 }
261 }
262
263 #[cfg(test)]
265 pub(crate) fn is_shared_with(&self, other: &Self) -> bool {
266 Arc::ptr_eq(&self.inner, &other.inner)
267 }
268
269 pub fn per_minute(&self) -> u32 {
271 self.inner.per_minute
272 }
273
274 pub fn used(&self) -> u32 {
276 let minute = self.inner.clock.now_ms() / MINUTE_MS;
277 let state = self.lock();
278 if minute > state.minute {
279 0
280 } else {
281 state.used
282 }
283 }
284
285 pub(crate) async fn acquire(&self, cost: Cost) -> Charge {
287 match cost {
288 Cost::Funding => loop {
289 self.await_cooldown().await;
293 let slot = self.reserve_funding_slot();
294 tokio::time::sleep_until(slot).await;
295 if !self.in_cooldown() {
296 return Charge { minute: None };
297 }
298 },
301 Cost::Weight(weight) => loop {
302 self.await_cooldown().await;
303 match self.try_charge(weight) {
304 Ok(charge) => return charge,
305 Err(until_next_minute) => tokio::time::sleep(until_next_minute).await,
306 }
307 },
308 }
309 }
310
311 fn try_charge(&self, weight: u32) -> Result<Charge, Duration> {
316 let now_ms = self.inner.clock.now_ms();
317 let mut state = self.lock();
318 let minute = now_ms / MINUTE_MS;
322 if minute > state.minute {
323 state.minute = minute;
324 state.used = 0;
325 }
326 if state.used == 0 || state.used.saturating_add(weight) <= self.inner.per_minute {
327 state.used = state.used.saturating_add(weight);
328 Ok(Charge {
329 minute: Some(state.minute),
330 })
331 } else {
332 Err(Duration::from_millis(MINUTE_MS - now_ms % MINUTE_MS))
333 }
334 }
335
336 pub(crate) fn record_used(&self, charge: Charge, used: u32) {
338 let Some(minute) = charge.minute else { return };
339 let mut state = self.lock();
340 if state.minute == minute && used > state.used {
341 state.used = used;
342 }
343 }
344
345 pub(crate) fn begin_cooldown(&self, delay: Duration) {
349 let until = Instant::now() + delay.min(MAX_COOLDOWN);
350 let mut state = self.lock();
351 if state.cooldown_until.is_none_or(|current| until > current) {
352 state.cooldown_until = Some(until);
353 }
354 }
355
356 pub(crate) fn hold_until_next_minute(&self) -> Duration {
360 let now_ms = self.inner.clock.now_ms();
361 let hold = Duration::from_millis(MINUTE_MS - now_ms % MINUTE_MS);
362 self.begin_cooldown(hold);
363 hold
364 }
365
366 async fn await_cooldown(&self) {
367 loop {
368 let until = self.lock().cooldown_until;
370 match until {
371 Some(until) if until > Instant::now() => tokio::time::sleep_until(until).await,
374 _ => return,
375 }
376 }
377 }
378
379 fn in_cooldown(&self) -> bool {
380 self.lock()
381 .cooldown_until
382 .is_some_and(|until| until > Instant::now())
383 }
384
385 fn reserve_funding_slot(&self) -> Instant {
386 let now = Instant::now();
387 let mut state = self.lock();
388 let slot = state.next_funding.map_or(now, |next| next.max(now));
389 state.next_funding = Some(slot + self.inner.funding_interval);
390 slot
391 }
392
393 fn lock(&self) -> MutexGuard<'_, State> {
396 self.inner
397 .state
398 .lock()
399 .unwrap_or_else(|poisoned| poisoned.into_inner())
400 }
401}
402
403#[cfg(test)]
404mod tests {
405 use super::*;
406 use std::sync::atomic::{AtomicU64, Ordering};
407
408 const ON_A_MINUTE: u64 = 1_791_362_100_000;
410
411 fn budget(weight_per_minute: u32, start_ms: u64) -> WeightBudget {
412 WeightBudget::with_clock(
413 weight_per_minute,
414 PUBLISHED_FUNDING_PER_FIVE_MINUTES,
415 Clock::Tokio {
416 start: Instant::now(),
417 start_ms,
418 },
419 )
420 }
421
422 fn manual_budget(start_ms: u64) -> (WeightBudget, Arc<AtomicU64>) {
423 let clock = Arc::new(AtomicU64::new(start_ms));
424 let budget = WeightBudget::with_clock(
425 PUBLISHED_WEIGHT_PER_MINUTE,
426 PUBLISHED_FUNDING_PER_FIVE_MINUTES,
427 Clock::Manual(Arc::clone(&clock)),
428 );
429 (budget, clock)
430 }
431
432 #[test]
433 fn documented_weights() {
434 use DepthLimit::*;
435 let w = Cost::Weight;
436 let table = [
437 (Route::Ping, w(1)),
438 (Route::Time, w(1)),
439 (Route::ExchangeInfo, w(1)),
440 (Route::FundingInfo, Cost::Funding),
441 (Route::Ticker24h { all: false }, w(1)),
442 (Route::Ticker24h { all: true }, w(40)),
443 (Route::PremiumIndex { all: false }, w(1)),
444 (Route::PremiumIndex { all: true }, w(10)),
445 (Route::Klines { limit: None }, w(5)),
446 (Route::Klines { limit: Some(1) }, w(1)),
447 (Route::Klines { limit: Some(100) }, w(1)),
448 (Route::Klines { limit: Some(101) }, w(2)),
449 (Route::Klines { limit: Some(500) }, w(2)),
450 (Route::Klines { limit: Some(501) }, w(5)),
451 (Route::Klines { limit: Some(1000) }, w(5)),
452 (Route::Klines { limit: Some(1001) }, w(10)),
453 (Route::Klines { limit: Some(1500) }, w(10)),
454 (Route::FundingRate, Cost::Funding),
455 (Route::OpenInterest, w(1)),
456 (Route::AggTrades, w(20)),
457 (Route::Depth { limit: None }, w(1)),
458 (Route::Depth { limit: Some(Five) }, w(2)),
459 (Route::Depth { limit: Some(Ten) }, w(2)),
460 (
461 Route::Depth {
462 limit: Some(Twenty),
463 },
464 w(2),
465 ),
466 (Route::Depth { limit: Some(Fifty) }, w(2)),
467 (
468 Route::Depth {
469 limit: Some(Hundred),
470 },
471 w(5),
472 ),
473 (
474 Route::Depth {
475 limit: Some(FiveHundred),
476 },
477 w(10),
478 ),
479 (
480 Route::Depth {
481 limit: Some(Thousand),
482 },
483 w(20),
484 ),
485 ];
486 for (route, cost) in table {
487 assert_eq!(route.cost(), cost, "{route:?}");
488 }
489 }
490
491 #[test]
492 fn the_budget_aims_a_tenth_below_the_published_limit() {
493 assert_eq!(WeightBudget::new().per_minute(), 2160);
494 }
495
496 #[tokio::test(start_paused = true)]
497 async fn a_charge_past_the_budget_waits_for_the_next_minute() {
498 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE + 15_000);
499 for _ in 0..54 {
500 budget.acquire(Cost::Weight(40)).await;
501 }
502 assert_eq!(budget.used(), 2160);
503
504 let start = Instant::now();
505 budget.acquire(Cost::Weight(1)).await;
506 assert_eq!(
507 start.elapsed(),
508 Duration::from_secs(45),
509 "held to the boundary"
510 );
511 assert_eq!(budget.used(), 1, "the new minute starts from zero");
512 }
513
514 #[tokio::test(start_paused = true)]
515 async fn a_header_raises_the_count_and_never_lowers_it() {
516 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE);
517 let charge = budget.acquire(Cost::Weight(1)).await;
518 budget.record_used(charge, 2000);
519 assert_eq!(budget.used(), 2000, "another process on this IP spent 1999");
520 budget.record_used(charge, 3);
521 assert_eq!(budget.used(), 2000);
522
523 let start = Instant::now();
524 budget.acquire(Cost::Weight(200)).await;
525 assert_eq!(
526 start.elapsed(),
527 Duration::from_secs(60),
528 "2000 + 200 passes 2160"
529 );
530 }
531
532 #[tokio::test(start_paused = true)]
533 async fn a_header_is_applied_only_to_the_minute_its_request_was_charged_in() {
534 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE + 59_900);
535 let in_flight = budget.acquire(Cost::Weight(5)).await;
536 tokio::time::advance(Duration::from_millis(200)).await;
537 budget.acquire(Cost::Weight(5)).await;
538 budget.record_used(in_flight, 2150);
541 assert_eq!(budget.used(), 5);
542 }
543
544 #[tokio::test(start_paused = true)]
545 async fn a_request_heavier_than_the_budget_is_sent_alone() {
546 let budget = budget(20, ON_A_MINUTE);
547 let start = Instant::now();
548 budget.acquire(Cost::Weight(40)).await;
549 assert_eq!(start.elapsed(), Duration::ZERO);
550 budget.acquire(Cost::Weight(1)).await;
551 assert_eq!(start.elapsed(), Duration::from_secs(60));
552 }
553
554 #[tokio::test(start_paused = true)]
555 async fn a_cooldown_holds_both_limits_and_is_only_extended() {
556 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE);
557 budget.begin_cooldown(Duration::from_secs(10));
558 budget.begin_cooldown(Duration::from_secs(2));
559
560 let start = Instant::now();
561 budget.acquire(Cost::Weight(1)).await;
562 assert_eq!(start.elapsed(), Duration::from_secs(10));
563
564 budget.begin_cooldown(Duration::from_secs(3));
565 let start = Instant::now();
566 budget.acquire(Cost::Funding).await;
567 assert_eq!(start.elapsed(), Duration::from_secs(3));
568 }
569
570 #[tokio::test(start_paused = true)]
571 async fn a_cooldown_is_clamped_to_the_longest_documented_ban() {
572 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE);
573 budget.begin_cooldown(Duration::MAX);
574 let start = Instant::now();
575 budget.acquire(Cost::Weight(1)).await;
576 assert_eq!(start.elapsed(), MAX_COOLDOWN);
577 }
578
579 #[tokio::test(start_paused = true)]
580 async fn a_hold_lasts_until_the_next_minute_for_both_limits() {
581 for cost in [Cost::Weight(1), Cost::Funding] {
582 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE + 15_000);
583 assert_eq!(budget.hold_until_next_minute(), Duration::from_secs(45));
584 let start = Instant::now();
585 budget.acquire(cost).await;
586 assert_eq!(start.elapsed(), Duration::from_secs(45), "{cost:?}");
587 }
588 }
589
590 #[tokio::test(start_paused = true)]
591 async fn the_funding_bucket_admits_450_per_five_minutes() {
592 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE);
593 let start = Instant::now();
594 let mut admitted = 0;
595 while start.elapsed() <= FUNDING_PERIOD {
596 budget.acquire(Cost::Funding).await;
597 if start.elapsed() <= FUNDING_PERIOD {
598 admitted += 1;
599 }
600 }
601 assert_eq!(admitted, 450);
602 }
603
604 #[tokio::test(start_paused = true)]
605 async fn the_funding_bucket_never_bursts() {
606 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE);
607 let start = Instant::now();
608 budget.acquire(Cost::Funding).await;
609 budget.acquire(Cost::Funding).await;
610 assert!(
611 start.elapsed() >= Duration::from_millis(668),
612 "{:?}",
613 start.elapsed()
614 );
615 }
616
617 #[tokio::test(start_paused = true)]
618 async fn a_cooldown_extended_mid_wait_is_waited_out_in_full() {
619 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE);
620 budget.begin_cooldown(Duration::from_secs(2));
621 let extender = budget.clone();
622 tokio::spawn(async move {
623 tokio::time::sleep(Duration::from_secs(1)).await;
624 extender.begin_cooldown(Duration::from_secs(10));
625 });
626 let start = Instant::now();
627 budget.acquire(Cost::Weight(1)).await;
628 assert_eq!(start.elapsed(), Duration::from_secs(11));
629 }
630
631 #[tokio::test(start_paused = true)]
632 async fn funding_requests_parked_by_a_cooldown_are_paced_when_it_ends() {
633 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE);
634 budget.begin_cooldown(Duration::from_secs(60));
635 let start = Instant::now();
636 let tasks: Vec<_> = (0..10)
637 .map(|_| {
638 let budget = budget.clone();
639 tokio::spawn(async move {
640 budget.acquire(Cost::Funding).await;
641 start.elapsed()
642 })
643 })
644 .collect();
645 let mut released = Vec::new();
646 for task in tasks {
647 released.push(task.await.unwrap());
648 }
649 released.sort();
650 assert_eq!(released[0], Duration::from_secs(60));
651 for pair in released.windows(2) {
652 assert!(
653 pair[1] - pair[0] >= Duration::from_millis(668),
654 "released together: {released:?}"
655 );
656 }
657 }
658
659 #[tokio::test(start_paused = true)]
660 async fn the_counted_minute_never_moves_backwards() {
661 let (budget, clock) = manual_budget(ON_A_MINUTE + 1_000);
662 budget.acquire(Cost::Weight(2000)).await;
663 clock.store(ON_A_MINUTE - 1_000, Ordering::SeqCst);
667 assert_eq!(budget.used(), 2000);
668 let second =
669 tokio::time::timeout(Duration::from_secs(5), budget.acquire(Cost::Weight(2000))).await;
670 assert!(second.is_err(), "charged into a minute already counted");
671 let stale = tokio::time::timeout(Duration::from_secs(5), budget.acquire(Cost::Weight(1)))
674 .await
675 .expect("2001 fits under 2160");
676 budget.record_used(stale, 2100);
677 assert_eq!(budget.used(), 2100);
678 }
679
680 #[tokio::test(start_paused = true)]
681 async fn a_cooldown_begun_during_the_slot_wait_holds_the_funding_request() {
682 let budget = budget(PUBLISHED_WEIGHT_PER_MINUTE, ON_A_MINUTE);
684 budget.acquire(Cost::Funding).await; let start = Instant::now();
686 let waiter = {
687 let budget = budget.clone();
688 tokio::spawn(async move {
689 budget.acquire(Cost::Funding).await;
690 start.elapsed()
691 })
692 };
693 tokio::time::sleep(Duration::from_millis(100)).await;
694 budget.begin_cooldown(Duration::from_secs(10));
695 assert_eq!(waiter.await.unwrap(), Duration::from_millis(10_100));
696 }
697}