1#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17pub enum TradingBucket {
18 Order,
20 Cancel,
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
30pub enum Tier {
31 #[default]
33 Standard,
34 Copper,
36 Bronze,
38 Silver,
40 Gold,
42 Platinum,
44 Diamond,
46 Elite,
48}
49
50impl Tier {
51 pub fn rate(self, bucket: TradingBucket) -> u32 {
53 let (order_rate, _, cancel_rate, _) = self.allowances();
54 match bucket {
55 TradingBucket::Order => order_rate,
56 TradingBucket::Cancel => cancel_rate,
57 }
58 }
59
60 pub fn burst(self, bucket: TradingBucket) -> u32 {
65 let (_, order_burst, _, cancel_burst) = self.allowances();
66 match bucket {
67 TradingBucket::Order => order_burst,
68 TradingBucket::Cancel => cancel_burst,
69 }
70 }
71
72 fn allowances(self) -> (u32, u32, u32, u32) {
79 match self {
80 Tier::Standard => (40, 60, 80, 120),
81 Tier::Copper => (60, 90, 120, 180),
82 Tier::Bronze => (80, 120, 160, 240),
83 Tier::Silver => (200, 300, 400, 600),
84 Tier::Gold => (400, 600, 800, 1_200),
85 Tier::Platinum => (450, 675, 900, 1_350),
86 Tier::Diamond => (525, 787, 1_050, 1_575),
87 Tier::Elite => (600, 900, 1_200, 1_800),
88 }
89 }
90
91 pub fn from_header(value: &str) -> Option<Self> {
96 match value.trim().to_ascii_lowercase().as_str() {
97 "standard" => Some(Tier::Standard),
98 "copper" => Some(Tier::Copper),
99 "bronze" => Some(Tier::Bronze),
100 "silver" => Some(Tier::Silver),
101 "gold" => Some(Tier::Gold),
102 "platinum" => Some(Tier::Platinum),
103 "diamond" => Some(Tier::Diamond),
104 "elite" => Some(Tier::Elite),
105 _ => None,
106 }
107 }
108}
109
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
112pub enum TradingRequest {
113 PostOrder,
115 PostOrders {
117 count: u32,
119 },
120 CancelOrder,
122 CancelOrders {
124 count: u32,
126 },
127 CancelAll,
131 CancelMarketOrders,
135}
136
137impl TradingRequest {
138 pub fn bucket(self) -> TradingBucket {
140 match self {
141 TradingRequest::PostOrder | TradingRequest::PostOrders { .. } => TradingBucket::Order,
142 TradingRequest::CancelOrder
143 | TradingRequest::CancelOrders { .. }
144 | TradingRequest::CancelAll
145 | TradingRequest::CancelMarketOrders => TradingBucket::Cancel,
146 }
147 }
148
149 pub fn cost(self) -> u32 {
158 match self {
159 TradingRequest::PostOrder | TradingRequest::CancelOrder => 1,
160 TradingRequest::PostOrders { count } | TradingRequest::CancelOrders { count } => count,
161 TradingRequest::CancelAll | TradingRequest::CancelMarketOrders => 1,
163 }
164 }
165
166 pub fn cost_is_exact(self) -> bool {
171 !matches!(
172 self,
173 TradingRequest::CancelAll | TradingRequest::CancelMarketOrders
174 )
175 }
176}
177
178#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
188pub struct RateLimitStatus {
189 pub remaining: Option<u32>,
191 pub reset: Option<u64>,
193 pub tier: Option<Tier>,
195 pub warning: bool,
197}
198
199impl RateLimitStatus {
200 pub fn from_headers(headers: &reqwest::header::HeaderMap) -> Self {
206 let get = |name: &str| headers.get(name).and_then(|v| v.to_str().ok());
207
208 Self {
209 remaining: get("poly-ratelimit-remaining").and_then(|v| v.trim().parse().ok()),
210 reset: get("poly-ratelimit-reset").and_then(|v| v.trim().parse().ok()),
211 tier: get("poly-ratelimit-tier").and_then(Tier::from_header),
212 warning: get("poly-ratelimit-warning")
213 .is_some_and(|v| v.trim().eq_ignore_ascii_case("true")),
214 }
215 }
216
217 pub fn is_empty(&self) -> bool {
219 self.remaining.is_none() && self.reset.is_none() && self.tier.is_none() && !self.warning
220 }
221}
222
223#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
230#[error(
231 "batch costs {cost} tokens but the {bucket:?} bucket at tier {tier:?} holds at most \
232 {capacity}; this can never succeed — split it into batches of {capacity} or fewer"
233)]
234pub struct BurstCapacityExceeded {
235 pub cost: u32,
237 pub capacity: u32,
239 pub tier: Tier,
241 pub bucket: TradingBucket,
243}
244
245type DirectLimiter = governor::RateLimiter<
246 governor::state::NotKeyed,
247 governor::state::InMemoryState,
248 governor::clock::DefaultClock,
249>;
250
251struct Buckets {
252 tier: Tier,
253 order: std::sync::Arc<DirectLimiter>,
254 cancel: std::sync::Arc<DirectLimiter>,
255}
256
257impl Buckets {
258 fn for_tier(tier: Tier) -> Self {
259 let build = |bucket: TradingBucket| {
260 let rate = tier.rate(bucket).max(1);
261 let burst = tier.burst(bucket).max(1);
262 let quota = governor::Quota::with_period(std::time::Duration::from_secs(1) / rate)
263 .expect("per-token interval is non-zero")
264 .allow_burst(std::num::NonZeroU32::new(burst).expect("burst is non-zero"));
265 std::sync::Arc::new(DirectLimiter::direct(quota))
266 };
267 Self {
268 tier,
269 order: build(TradingBucket::Order),
270 cancel: build(TradingBucket::Cancel),
271 }
272 }
273}
274
275#[derive(Clone)]
289pub struct SignerLimiter {
290 inner: std::sync::Arc<SignerLimiterInner>,
291}
292
293struct SignerLimiterInner {
294 buckets: std::sync::RwLock<Buckets>,
295 status: std::sync::RwLock<RateLimitStatus>,
296}
297
298impl std::fmt::Debug for SignerLimiter {
299 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
300 f.debug_struct("SignerLimiter")
301 .field("tier", &self.tier())
302 .finish()
303 }
304}
305
306impl Default for SignerLimiter {
307 fn default() -> Self {
308 Self::new()
309 }
310}
311
312impl SignerLimiter {
313 pub fn new() -> Self {
315 Self {
316 inner: std::sync::Arc::new(SignerLimiterInner {
317 buckets: std::sync::RwLock::new(Buckets::for_tier(Tier::default())),
318 status: std::sync::RwLock::new(RateLimitStatus::default()),
319 }),
320 }
321 }
322
323 pub fn tier(&self) -> Tier {
325 self.inner
326 .buckets
327 .read()
328 .expect("lock is never poisoned")
329 .tier
330 }
331
332 pub fn last_status(&self) -> RateLimitStatus {
334 *self.inner.status.read().expect("lock is never poisoned")
335 }
336
337 pub fn observe(&self, headers: &reqwest::header::HeaderMap) {
342 let status = RateLimitStatus::from_headers(headers);
343 if status.is_empty() {
344 return;
345 }
346 *self.inner.status.write().expect("lock is never poisoned") = status;
347
348 if let Some(tier) = status.tier {
351 let mut buckets = self.inner.buckets.write().expect("lock is never poisoned");
352 if buckets.tier != tier {
353 tracing::debug!("adopting rate limit tier {tier:?} (was {:?})", buckets.tier);
354 *buckets = Buckets::for_tier(tier);
355 }
356 }
357 }
358
359 pub async fn acquire(&self, request: TradingRequest) -> Result<(), BurstCapacityExceeded> {
366 let bucket = request.bucket();
367
368 let (tier, limiter) = {
371 let buckets = self.inner.buckets.read().expect("lock is never poisoned");
372 let limiter = match bucket {
373 TradingBucket::Order => buckets.order.clone(),
374 TradingBucket::Cancel => buckets.cancel.clone(),
375 };
376 (buckets.tier, limiter)
377 };
378
379 let cost = request.cost().max(1);
380 let n = std::num::NonZeroU32::new(cost).expect("cost floor is 1");
381
382 limiter
386 .until_n_ready(n)
387 .await
388 .map_err(|_| BurstCapacityExceeded {
389 cost,
390 capacity: tier.burst(bucket),
391 tier,
392 bucket,
393 })
394 }
395}
396
397#[cfg(test)]
398mod limiter_tests {
399 use super::*;
400 use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
401 use std::time::Duration;
402
403 fn tier_header(tier: &str) -> HeaderMap {
404 let mut map = HeaderMap::new();
405 map.insert(
406 HeaderName::from_static("poly-ratelimit-tier"),
407 HeaderValue::from_str(tier).unwrap(),
408 );
409 map
410 }
411
412 #[test]
413 fn starts_at_the_tightest_tier() {
414 assert_eq!(SignerLimiter::new().tier(), Tier::Standard);
415 }
416
417 #[test]
418 fn observing_a_tier_header_adopts_it() {
419 let limiter = SignerLimiter::new();
420 limiter.observe(&tier_header("gold"));
421 assert_eq!(limiter.tier(), Tier::Gold);
422 assert_eq!(limiter.last_status().tier, Some(Tier::Gold));
423 }
424
425 #[test]
426 fn an_unrecognised_tier_leaves_the_current_one_in_force() {
427 let limiter = SignerLimiter::new();
430 limiter.observe(&tier_header("silver"));
431 limiter.observe(&tier_header("titanium"));
432 assert_eq!(limiter.tier(), Tier::Silver);
433 }
434
435 #[test]
436 fn a_response_without_the_headers_does_not_clear_telemetry() {
437 let limiter = SignerLimiter::new();
438 limiter.observe(&tier_header("silver"));
439 limiter.observe(&HeaderMap::new());
440 assert_eq!(limiter.tier(), Tier::Silver);
441 assert_eq!(limiter.last_status().tier, Some(Tier::Silver));
442 }
443
444 #[tokio::test]
445 async fn an_over_capacity_batch_is_rejected_immediately_not_queued() {
446 let limiter = SignerLimiter::new();
449 let request = TradingRequest::CancelOrders { count: 2_000 };
450
451 let result = tokio::time::timeout(Duration::from_millis(100), limiter.acquire(request))
452 .await
453 .expect("must not hang waiting for capacity that can never exist");
454
455 let err = result.expect_err("2,000 tokens exceeds Standard's 120 cancel burst");
456 assert_eq!(err.cost, 2_000);
457 assert_eq!(err.capacity, 120);
458 assert_eq!(err.bucket, TradingBucket::Cancel);
459 }
460
461 #[tokio::test]
462 async fn a_batch_within_capacity_is_admitted() {
463 let limiter = SignerLimiter::new();
464 limiter
465 .acquire(TradingRequest::CancelOrders { count: 100 })
466 .await
467 .expect("100 fits Standard's 120 cancel burst");
468 }
469
470 #[tokio::test]
471 async fn adopting_a_higher_tier_admits_a_batch_that_was_impossible() {
472 let limiter = SignerLimiter::new();
476 let batch = TradingRequest::CancelOrders { count: 500 };
477 assert!(limiter.acquire(batch).await.is_err());
478
479 limiter.observe(&tier_header("gold"));
480 limiter
481 .acquire(batch)
482 .await
483 .expect("500 fits Gold's 1,200 cancel burst");
484 }
485
486 #[tokio::test]
487 async fn the_order_and_cancel_buckets_are_independent() {
488 let limiter = SignerLimiter::new();
492 limiter
493 .acquire(TradingRequest::PostOrders { count: 60 })
494 .await
495 .expect("60 fills Standard's order burst exactly");
496
497 let start = std::time::Instant::now();
498 limiter
499 .acquire(TradingRequest::CancelOrder)
500 .await
501 .expect("cancel bucket is untouched");
502 assert!(
503 start.elapsed() < Duration::from_millis(25),
504 "cancelling was throttled by order placement"
505 );
506 }
507
508 #[tokio::test]
509 async fn batch_cost_is_charged_in_full_not_as_one_request() {
510 let limiter = SignerLimiter::new();
514 limiter
515 .acquire(TradingRequest::PostOrders { count: 60 })
516 .await
517 .unwrap();
518
519 let start = std::time::Instant::now();
520 limiter.acquire(TradingRequest::PostOrder).await.unwrap();
521 assert!(
522 start.elapsed() >= Duration::from_millis(10),
523 "a 60-order batch was charged as a single token"
524 );
525 }
526}
527
528#[cfg(test)]
529mod status_tests {
530 use super::*;
531 use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
532
533 fn headers(pairs: &[(&str, &str)]) -> HeaderMap {
534 let mut map = HeaderMap::new();
535 for (k, v) in pairs {
536 map.insert(
537 HeaderName::from_bytes(k.as_bytes()).unwrap(),
538 HeaderValue::from_str(v).unwrap(),
539 );
540 }
541 map
542 }
543
544 #[test]
545 fn reads_the_full_header_family() {
546 let status = RateLimitStatus::from_headers(&headers(&[
547 ("poly-ratelimit-remaining", "57"),
548 ("poly-ratelimit-reset", "1767225660"),
549 ("poly-ratelimit-tier", "silver"),
550 ("poly-ratelimit-warning", "true"),
551 ]));
552
553 assert_eq!(status.remaining, Some(57));
554 assert_eq!(status.reset, Some(1_767_225_660));
555 assert_eq!(status.tier, Some(Tier::Silver));
556 assert!(status.warning);
557 assert!(!status.is_empty());
558 }
559
560 #[test]
561 fn header_names_are_matched_case_insensitively() {
562 let status = RateLimitStatus::from_headers(&headers(&[
565 ("Poly-RateLimit-Tier", "GOLD"),
566 ("POLY-RATELIMIT-REMAINING", "3"),
567 ]));
568 assert_eq!(status.tier, Some(Tier::Gold));
569 assert_eq!(status.remaining, Some(3));
570 }
571
572 #[test]
573 fn absent_headers_yield_an_empty_status_not_an_error() {
574 let status = RateLimitStatus::from_headers(&HeaderMap::new());
575 assert!(status.is_empty());
576 assert_eq!(status, RateLimitStatus::default());
577 }
578
579 #[test]
580 fn malformed_values_are_dropped_rather_than_guessed() {
581 let status = RateLimitStatus::from_headers(&headers(&[
585 ("poly-ratelimit-remaining", "not-a-number"),
586 ("poly-ratelimit-reset", ""),
587 ("poly-ratelimit-tier", "titanium"),
588 ("poly-ratelimit-warning", "false"),
589 ]));
590
591 assert_eq!(status.remaining, None);
592 assert_eq!(status.reset, None);
593 assert_eq!(status.tier, None);
594 assert!(!status.warning);
595 }
596}
597
598#[cfg(test)]
599mod tests {
600 use super::*;
601
602 fn published() -> Vec<(Tier, u32, u32, u32, u32)> {
606 vec![
607 (Tier::Standard, 40, 60, 80, 120),
608 (Tier::Copper, 60, 90, 120, 180),
609 (Tier::Bronze, 80, 120, 160, 240),
610 (Tier::Silver, 200, 300, 400, 600),
611 (Tier::Gold, 400, 600, 800, 1_200),
612 (Tier::Platinum, 450, 675, 900, 1_350),
613 (Tier::Diamond, 525, 787, 1_050, 1_575),
614 (Tier::Elite, 600, 900, 1_200, 1_800),
615 ]
616 }
617
618 #[test]
619 fn every_tier_matches_the_published_table() {
620 for (tier, o_rate, o_burst, c_rate, c_burst) in published() {
621 assert_eq!(
622 tier.rate(TradingBucket::Order),
623 o_rate,
624 "{tier:?} order rate"
625 );
626 assert_eq!(
627 tier.burst(TradingBucket::Order),
628 o_burst,
629 "{tier:?} order burst"
630 );
631 assert_eq!(
632 tier.rate(TradingBucket::Cancel),
633 c_rate,
634 "{tier:?} cancel rate"
635 );
636 assert_eq!(
637 tier.burst(TradingBucket::Cancel),
638 c_burst,
639 "{tier:?} cancel burst"
640 );
641 }
642 }
643
644 #[test]
645 fn the_default_tier_is_the_tightest_one() {
646 let default = Tier::default();
648 for (tier, ..) in published() {
649 assert!(
650 default.rate(TradingBucket::Order) <= tier.rate(TradingBucket::Order),
651 "default tier {default:?} is looser than {tier:?}"
652 );
653 assert!(
654 default.burst(TradingBucket::Cancel) <= tier.burst(TradingBucket::Cancel),
655 "default tier {default:?} bursts higher than {tier:?}"
656 );
657 }
658 }
659
660 #[test]
661 fn tier_headers_parse_case_insensitively() {
662 assert_eq!(Tier::from_header("standard"), Some(Tier::Standard));
663 assert_eq!(Tier::from_header("Silver"), Some(Tier::Silver));
664 assert_eq!(Tier::from_header("ELITE"), Some(Tier::Elite));
665 }
666
667 #[test]
668 fn an_unknown_tier_header_is_not_guessed() {
669 assert_eq!(Tier::from_header("titanium"), None);
672 assert_eq!(Tier::from_header(""), None);
673 }
674
675 #[test]
676 fn batch_costs_scale_with_the_payload() {
677 assert_eq!(TradingRequest::PostOrder.cost(), 1);
678 assert_eq!(TradingRequest::PostOrders { count: 40 }.cost(), 40);
679 assert_eq!(TradingRequest::CancelOrder.cost(), 1);
680 assert_eq!(TradingRequest::CancelOrders { count: 250 }.cost(), 250);
681 }
682
683 #[test]
684 fn requests_draw_from_the_right_bucket() {
685 assert_eq!(TradingRequest::PostOrder.bucket(), TradingBucket::Order);
686 assert_eq!(
687 TradingRequest::PostOrders { count: 2 }.bucket(),
688 TradingBucket::Order
689 );
690 assert_eq!(TradingRequest::CancelOrder.bucket(), TradingBucket::Cancel);
691 assert_eq!(
692 TradingRequest::CancelOrders { count: 2 }.bucket(),
693 TradingBucket::Cancel
694 );
695 assert_eq!(TradingRequest::CancelAll.bucket(), TradingBucket::Cancel);
696 assert_eq!(
697 TradingRequest::CancelMarketOrders.bucket(),
698 TradingBucket::Cancel
699 );
700 }
701
702 #[test]
703 fn cancel_all_reports_a_floor_cost_and_says_so() {
704 assert_eq!(TradingRequest::CancelAll.cost(), 1);
708 assert!(!TradingRequest::CancelAll.cost_is_exact());
709 assert_eq!(TradingRequest::CancelMarketOrders.cost(), 1);
710 assert!(!TradingRequest::CancelMarketOrders.cost_is_exact());
711
712 for exact in [
713 TradingRequest::PostOrder,
714 TradingRequest::PostOrders { count: 3 },
715 TradingRequest::CancelOrder,
716 TradingRequest::CancelOrders { count: 3 },
717 ] {
718 assert!(exact.cost_is_exact(), "{exact:?} cost is computable");
719 }
720 }
721
722 #[test]
723 fn a_batch_larger_than_elite_burst_is_impossible_on_every_tier() {
724 let batch = TradingRequest::CancelOrders { count: 2_000 };
728 for (tier, ..) in published() {
729 assert!(
730 batch.cost() > tier.burst(batch.bucket()),
731 "{tier:?} could absorb a 2,000-ID batch — check the published table"
732 );
733 }
734 }
735}