use std::{
num::NonZeroU32,
sync::atomic::{AtomicU64, Ordering},
time::Duration,
};
use dashmap::DashMap;
#[cfg(test)]
use nautilus_network::ratelimiter::clock::FakeRelativeClock;
use nautilus_network::ratelimiter::clock::{Clock, MonotonicClock, Reference};
use ustr::Ustr;
pub const DERIVE_MATCHING_RATE_KEY: &str = "derive:matching";
pub const DERIVE_NON_MATCHING_RATE_KEY: &str = "derive:non-matching";
pub const DERIVE_CANCEL_ALL_RATE_KEY: &str = "derive:cancel-all";
pub const DERIVE_CANCEL_BY_LABEL_RATE_KEY: &str = "derive:cancel-by-label";
const DERIVE_PER_INSTRUMENT_RATE_KEY_PREFIX: &str = "derive:matching:instrument:";
pub const DERIVE_DEFAULT_MATCHING_TPS: u32 = 1;
pub const DERIVE_DEFAULT_PER_INSTRUMENT_MATCHING_TPS: u32 = 1;
pub const DERIVE_NON_MATCHING_TPS: u32 = 10;
pub const DERIVE_WEBSOCKET_NON_MATCHING_TPS: u32 = 5;
pub const DERIVE_CANCEL_ALL_TPS: u32 = 1;
pub const DERIVE_CANCEL_BY_LABEL_TPS: u32 = 10;
pub const DERIVE_RATE_WINDOW_SECS: u64 = 5;
pub const DERIVE_RATE_BURST_MULTIPLIER: u32 = 5;
const RATE_WINDOW_NANOS: u64 = DERIVE_RATE_WINDOW_SECS * 1_000_000_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RateClass {
NonMatching,
Matching,
CancelAll,
CancelByLabel,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RateBucket<'a> {
NonMatching,
Matching,
PerInstrument(&'a Ustr),
CancelAll,
CancelByLabel,
}
#[must_use]
pub(crate) fn rate_class_for_method(method: &str) -> RateClass {
match method.trim_start_matches('/') {
"private/order"
| "private/trigger_order"
| "private/replace"
| "private/cancel"
| "private/cancel_trigger_order" => RateClass::Matching,
"private/cancel_all" => RateClass::CancelAll,
"private/cancel_by_label" => RateClass::CancelByLabel,
_ => RateClass::NonMatching,
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct FixedWindowLimits {
pub(crate) non_matching: NonZeroU32,
pub(crate) matching: NonZeroU32,
pub(crate) per_instrument_matching: NonZeroU32,
pub(crate) cancel_all: NonZeroU32,
pub(crate) cancel_by_label: NonZeroU32,
}
impl FixedWindowLimits {
#[must_use]
pub(crate) fn rest(
matching_tps: Option<u32>,
per_instrument_matching_tps: Option<u32>,
) -> Self {
Self {
non_matching: window_limit(DERIVE_NON_MATCHING_TPS),
matching: window_limit(resolve_tps(matching_tps, DERIVE_DEFAULT_MATCHING_TPS)),
per_instrument_matching: window_limit(resolve_tps(
per_instrument_matching_tps,
DERIVE_DEFAULT_PER_INSTRUMENT_MATCHING_TPS,
)),
cancel_all: window_limit(DERIVE_CANCEL_ALL_TPS),
cancel_by_label: window_limit(DERIVE_CANCEL_BY_LABEL_TPS),
}
}
#[must_use]
pub(crate) fn websocket(
matching_tps: Option<u32>,
per_instrument_matching_tps: Option<u32>,
) -> Self {
Self {
non_matching: window_limit(DERIVE_WEBSOCKET_NON_MATCHING_TPS),
..Self::rest(matching_tps, per_instrument_matching_tps)
}
}
#[must_use]
pub(crate) fn limit_for(&self, bucket: RateBucket<'_>) -> NonZeroU32 {
match bucket {
RateBucket::NonMatching => self.non_matching,
RateBucket::Matching => self.matching,
RateBucket::PerInstrument(_) => self.per_instrument_matching,
RateBucket::CancelAll => self.cancel_all,
RateBucket::CancelByLabel => self.cancel_by_label,
}
}
}
pub(crate) struct FixedWindowLimiter<C: Clock> {
limits: FixedWindowLimits,
cells: DashMap<Ustr, AtomicU64>,
clock: C,
start: C::Instant,
}
impl<C: Clock> FixedWindowLimiter<C> {
pub(crate) fn new(limits: FixedWindowLimits, clock: C) -> Self {
let start = clock.now();
Self {
limits,
cells: DashMap::new(),
clock,
start,
}
}
#[cfg(test)]
pub(crate) fn check_bucket(&self, bucket: RateBucket<'_>) -> Result<(), Duration> {
loop {
let elapsed = self.elapsed_nanos();
let window = window_index(elapsed);
let limit = self.limits.limit_for(bucket).get();
let key = bucket_key(bucket);
let cell = self.cells.entry(key).or_default();
match consume_cell_fixed_window(cell.value(), limit, window) {
CellOutcome::Consumed => return Ok(()),
CellOutcome::Exhausted => {
let window_end_nanos = (u64::from(window) + 1) * RATE_WINDOW_NANOS;
return Err(Duration::from_nanos(
window_end_nanos.saturating_sub(elapsed),
));
}
CellOutcome::Advanced => {}
}
}
}
pub(crate) async fn await_buckets_ready(&self, buckets: &[RateBucket<'_>]) -> u32 {
loop {
let elapsed = self.elapsed_nanos();
let window = window_index(elapsed);
let mut acquired: Vec<Ustr> = Vec::with_capacity(buckets.len());
let mut denial = None;
for bucket in buckets {
let limit = self.limits.limit_for(*bucket).get();
let key = bucket_key(*bucket);
let cell = self.cells.entry(key).or_default();
match consume_cell_fixed_window(cell.value(), limit, window) {
CellOutcome::Consumed => acquired.push(key),
CellOutcome::Exhausted => {
let window_end_nanos = (u64::from(window) + 1) * RATE_WINDOW_NANOS;
denial = Some(Duration::from_nanos(
window_end_nanos.saturating_sub(elapsed),
));
break;
}
CellOutcome::Advanced => break,
}
}
match denial {
None if acquired.len() == buckets.len() => return window,
None => {
self.rollback_window(acquired, window);
}
Some(wait) => {
self.rollback_window(acquired, window);
self.clock.sleep(wait).await;
}
}
}
}
pub(crate) async fn await_class_ready(
&self,
class: RateClass,
instrument_name: Option<&Ustr>,
) -> u32 {
match class {
RateClass::Matching if instrument_name.is_some() => {
let instrument = instrument_name.expect("checked above");
self.await_buckets_ready(&[
RateBucket::Matching,
RateBucket::PerInstrument(instrument),
])
.await
}
RateClass::Matching => self.await_buckets_ready(&[RateBucket::Matching]).await,
RateClass::NonMatching => self.await_buckets_ready(&[RateBucket::NonMatching]).await,
RateClass::CancelAll => self.await_buckets_ready(&[RateBucket::CancelAll]).await,
RateClass::CancelByLabel => {
self.await_buckets_ready(&[RateBucket::CancelByLabel]).await
}
}
}
fn rollback_window(&self, keys: Vec<Ustr>, window: u32) {
for key in keys {
if let Some(cell) = self.cells.get(&key) {
rollback_cell_fixed_window(cell.value(), window);
}
}
}
pub(crate) async fn ensure_window_current(
&self,
class: RateClass,
instrument_name: Option<&Ustr>,
reserved_window: u32,
) {
if window_index(self.elapsed_nanos()) != reserved_window {
self.await_class_ready(class, instrument_name).await;
}
}
fn elapsed_nanos(&self) -> u64 {
self.clock.now().duration_since(self.start).as_u64()
}
}
#[cfg(test)]
impl FixedWindowLimiter<FakeRelativeClock> {
pub(crate) fn advance_clock(&self, by: Duration) {
self.clock.advance(by);
}
}
impl<C: Clock> std::fmt::Debug for FixedWindowLimiter<C> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct(stringify!(FixedWindowLimiter)).finish()
}
}
pub(crate) type DeriveRateLimiter = FixedWindowLimiter<MonotonicClock>;
fn bucket_key(bucket: RateBucket<'_>) -> Ustr {
match bucket {
RateBucket::NonMatching => Ustr::from(DERIVE_NON_MATCHING_RATE_KEY),
RateBucket::Matching => Ustr::from(DERIVE_MATCHING_RATE_KEY),
RateBucket::PerInstrument(instrument_name) => Ustr::from(
format!(
"{DERIVE_PER_INSTRUMENT_RATE_KEY_PREFIX}{}",
instrument_name.as_str(),
)
.as_str(),
),
RateBucket::CancelAll => Ustr::from(DERIVE_CANCEL_ALL_RATE_KEY),
RateBucket::CancelByLabel => Ustr::from(DERIVE_CANCEL_BY_LABEL_RATE_KEY),
}
}
fn resolve_tps(configured: Option<u32>, default_tps: u32) -> u32 {
configured.filter(|&v| v > 0).unwrap_or(default_tps)
}
fn window_limit(tps: u32) -> NonZeroU32 {
NonZeroU32::new(tps.saturating_mul(DERIVE_RATE_BURST_MULTIPLIER))
.expect("window limit must be non-zero")
}
fn window_index(elapsed_nanos: u64) -> u32 {
u32::try_from(elapsed_nanos / RATE_WINDOW_NANOS).expect("window index fits u32")
}
fn pack(window: u32, consumed: u32) -> u64 {
(u64::from(window) << 32) | u64::from(consumed)
}
fn unpack(packed: u64) -> (u32, u32) {
(
u32::try_from(packed >> 32).expect("window index fits u32"),
packed as u32,
)
}
enum CellOutcome {
Consumed,
Exhausted,
Advanced,
}
fn consume_cell_fixed_window(cell: &AtomicU64, limit: u32, window: u32) -> CellOutcome {
let mut prev = cell.load(Ordering::Acquire);
loop {
let (prev_window, prev_consumed) = unpack(prev);
if prev_window > window {
return CellOutcome::Advanced;
}
let next = if prev_window < window {
pack(window, 1)
} else if prev_consumed < limit {
pack(window, prev_consumed + 1)
} else {
return CellOutcome::Exhausted;
};
match cell.compare_exchange_weak(prev, next, Ordering::Release, Ordering::Relaxed) {
Ok(_) => return CellOutcome::Consumed,
Err(contended) => prev = contended,
}
}
}
fn rollback_cell_fixed_window(cell: &AtomicU64, window: u32) {
let mut prev = cell.load(Ordering::Acquire);
loop {
let (prev_window, prev_consumed) = unpack(prev);
if prev_window != window || prev_consumed == 0 {
return;
}
let next = pack(prev_window, prev_consumed - 1);
match cell.compare_exchange_weak(prev, next, Ordering::Release, Ordering::Relaxed) {
Ok(_) => return,
Err(contended) => prev = contended,
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn instrument(name: &str) -> Ustr {
Ustr::from(name)
}
fn trader_limits() -> FixedWindowLimits {
FixedWindowLimits::websocket(None, None)
}
fn limiter() -> FixedWindowLimiter<FakeRelativeClock> {
FixedWindowLimiter::new(trader_limits(), FakeRelativeClock::default())
}
#[rstest]
fn test_rest_limits_match_documented_trader_contract() {
let limits = FixedWindowLimits::rest(None, None);
assert_eq!(limits.non_matching.get(), 50); assert_eq!(limits.matching.get(), 5); assert_eq!(limits.per_instrument_matching.get(), 5);
assert_eq!(limits.cancel_all.get(), 5); assert_eq!(limits.cancel_by_label.get(), 50); }
#[rstest]
fn test_websocket_limits_match_documented_trader_contract() {
let limits = FixedWindowLimits::websocket(None, None);
assert_eq!(limits.non_matching.get(), 25); assert_eq!(limits.matching.get(), 5);
assert_eq!(limits.per_instrument_matching.get(), 5);
}
#[rstest]
fn test_matching_overrides_do_not_leak_into_per_instrument_allowance() {
let limits = FixedWindowLimits::websocket(Some(500), None);
assert_eq!(limits.matching.get(), 2_500);
assert_eq!(limits.per_instrument_matching.get(), 5);
let limits = FixedWindowLimits::websocket(None, Some(10));
assert_eq!(limits.matching.get(), 5);
assert_eq!(limits.per_instrument_matching.get(), 50);
}
#[rstest]
fn test_matching_overrides_treat_zero_as_unset() {
let limits = FixedWindowLimits::websocket(Some(0), Some(0));
assert_eq!(limits.matching.get(), 5);
assert_eq!(limits.per_instrument_matching.get(), 5);
}
#[rstest]
#[case("private/order", RateClass::Matching)]
#[case("/private/order", RateClass::Matching)]
#[case("private/trigger_order", RateClass::Matching)]
#[case("private/replace", RateClass::Matching)]
#[case("private/cancel", RateClass::Matching)]
#[case("private/cancel_trigger_order", RateClass::Matching)]
#[case("private/cancel_all", RateClass::CancelAll)]
#[case("private/cancel_by_label", RateClass::CancelByLabel)]
#[case("private/get_subaccount", RateClass::NonMatching)]
#[case("private/get_open_orders", RateClass::NonMatching)]
#[case("public/get_instruments", RateClass::NonMatching)]
#[case("public/login", RateClass::NonMatching)]
#[case("subscribe", RateClass::NonMatching)]
fn test_rate_class_for_method(#[case] method: &str, #[case] expected: RateClass) {
assert_eq!(rate_class_for_method(method), expected);
}
#[rstest]
fn test_full_matching_burst_denies_sixth_request_until_window_reset() {
let limiter = limiter();
for _ in 0..5 {
assert!(
limiter.check_bucket(RateBucket::Matching).is_ok(),
"Trader matching burst is five requests",
);
}
assert!(
limiter.check_bucket(RateBucket::Matching).is_err(),
"sixth matching request must wait for the window reset",
);
}
#[rstest]
fn test_allowance_refills_discretely_at_window_boundary() {
let limiter = limiter();
for _ in 0..5 {
limiter.check_bucket(RateBucket::Matching).expect("burst");
}
limiter.advance_clock(Duration::from_millis(4_999));
assert!(
limiter.check_bucket(RateBucket::Matching).is_err(),
"window has not rolled: nothing refills before the boundary",
);
limiter.advance_clock(Duration::from_millis(1));
for sequence in 0..5 {
assert!(
limiter.check_bucket(RateBucket::Matching).is_ok(),
"full allowance must refill at the boundary, request {sequence}",
);
}
assert!(
limiter.check_bucket(RateBucket::Matching).is_err(),
"only one window's allowance refills",
);
}
#[rstest]
fn test_window_reset_does_not_refill_one_token_at_a_time() {
let limiter = limiter();
for _ in 0..5 {
limiter.check_bucket(RateBucket::Matching).expect("burst");
}
for _ in 0..4 {
limiter.advance_clock(Duration::from_secs(1));
assert!(
limiter.check_bucket(RateBucket::Matching).is_err(),
"sustained-rate refill must not apply inside a window",
);
}
limiter.advance_clock(Duration::from_secs(1));
assert!(
limiter.check_bucket(RateBucket::Matching).is_ok(),
"full refill lands exactly at the five-second boundary",
);
}
#[rstest]
#[tokio::test]
async fn test_await_buckets_ready_waits_for_window_reset_and_consumes() {
let limiter = limiter();
for _ in 0..5 {
limiter.check_bucket(RateBucket::Matching).expect("burst");
}
limiter.await_buckets_ready(&[RateBucket::Matching]).await;
for _ in 0..4 {
limiter
.check_bucket(RateBucket::Matching)
.expect("fresh window minus the awaited cell");
}
assert!(
limiter.check_bucket(RateBucket::Matching).is_err(),
"await must consume from the fresh window",
);
}
#[rstest]
fn test_per_instrument_buckets_are_independent() {
let limiter = FixedWindowLimiter::new(
FixedWindowLimits::websocket(Some(10), None),
FakeRelativeClock::default(),
);
for _ in 0..5 {
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("ETH-PERP")))
.expect("ETH-PERP burst");
}
assert!(
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("ETH-PERP")))
.is_err(),
"ETH-PERP allowance is exhausted",
);
assert!(
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("BTC-PERP")))
.is_ok(),
"BTC-PERP has an independent allowance",
);
assert!(
limiter.check_bucket(RateBucket::Matching).is_ok(),
"account-wide matching still has headroom (10 TPS)",
);
}
#[rstest]
#[tokio::test]
async fn test_global_matching_bucket_enforced_alongside_per_instrument() {
let clock = FakeRelativeClock::default();
let limiter =
FixedWindowLimiter::new(FixedWindowLimits::websocket(None, Some(10)), clock.clone());
for _ in 0..5 {
limiter
.await_class_ready(RateClass::Matching, Some(&instrument("ETH-PERP")))
.await;
}
limiter
.await_class_ready(RateClass::Matching, Some(&instrument("BTC-PERP")))
.await;
assert_eq!(
clock.now().as_u64(),
RATE_WINDOW_NANOS,
"BTC-PERP write must wait for the global window reset",
);
}
#[rstest]
#[tokio::test]
async fn test_matching_write_consumes_global_and_per_instrument_buckets() {
let limiter = FixedWindowLimiter::new(
FixedWindowLimits::websocket(Some(2), None),
FakeRelativeClock::default(),
);
for _ in 0..5 {
limiter
.await_class_ready(RateClass::Matching, Some(&instrument("ETH-PERP")))
.await;
}
assert!(
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("ETH-PERP")))
.is_err(),
"each write consumes the instrument bucket",
);
assert!(
limiter.check_bucket(RateBucket::Matching).is_ok(),
"five of the global override's ten window cells remain",
);
assert!(
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("BTC-PERP")))
.is_ok(),
"other instruments are unaffected",
);
}
#[rstest]
#[tokio::test]
async fn test_multi_bucket_wait_consumes_both_buckets_from_one_window() {
let clock = FakeRelativeClock::default();
let limiter =
FixedWindowLimiter::new(FixedWindowLimits::websocket(Some(10), None), clock.clone());
for _ in 0..5 {
limiter
.await_class_ready(RateClass::Matching, Some(&instrument("ETH-PERP")))
.await;
}
limiter
.await_class_ready(RateClass::Matching, Some(&instrument("ETH-PERP")))
.await;
assert_eq!(
clock.now().as_u64(),
RATE_WINDOW_NANOS,
"the denied write must wait for the window boundary",
);
let mut remaining = 0;
while limiter.check_bucket(RateBucket::Matching).is_ok() {
remaining += 1;
}
assert_eq!(remaining, 49, "global cell must come from window 1");
for _ in 0..4 {
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("ETH-PERP")))
.expect("window 1 holds one consumed cell of five");
}
assert!(
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("ETH-PERP")))
.is_err(),
"the awaited write consumed the fifth ETH-PERP cell of window 1",
);
}
#[rstest]
#[tokio::test]
async fn test_ensure_window_current_reacquires_only_after_rollover() {
let clock = FakeRelativeClock::default();
let limiter =
FixedWindowLimiter::new(FixedWindowLimits::websocket(None, None), clock.clone());
let reserved_window = limiter
.await_class_ready(RateClass::Matching, Some(&instrument("ETH-PERP")))
.await;
assert_eq!(reserved_window, 0);
limiter
.ensure_window_current(
RateClass::Matching,
Some(&instrument("ETH-PERP")),
reserved_window,
)
.await;
assert!(
limiter.check_bucket(RateBucket::Matching).is_ok(),
"same-window refresh consumes nothing",
);
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("ETH-PERP")))
.expect("same-window refresh consumes nothing");
clock.advance(Duration::from_secs(5));
limiter
.ensure_window_current(
RateClass::Matching,
Some(&instrument("ETH-PERP")),
reserved_window,
)
.await;
for _ in 0..4 {
limiter
.check_bucket(RateBucket::Matching)
.expect("window 1 global has 4 cells left of 5");
}
assert!(
limiter.check_bucket(RateBucket::Matching).is_err(),
"rolled-window refresh consumed a window-1 global cell",
);
for _ in 0..4 {
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("ETH-PERP")))
.expect("window 1 instrument has 4 cells left of 5");
}
assert!(
limiter
.check_bucket(RateBucket::PerInstrument(&instrument("ETH-PERP")))
.is_err(),
"rolled-window refresh consumed a window-1 instrument cell",
);
}
#[rstest]
fn test_custom_cancel_all_quota_is_one_tps_burst() {
let limiter = limiter();
for _ in 0..5 {
limiter.check_bucket(RateBucket::CancelAll).expect("burst");
}
assert!(
limiter.check_bucket(RateBucket::CancelAll).is_err(),
"custom cancel_all allowance is 5 per window",
);
limiter.advance_clock(Duration::from_secs(5));
assert!(limiter.check_bucket(RateBucket::CancelAll).is_ok());
}
#[rstest]
fn test_custom_unscoped_cancel_by_label_quota_is_ten_tps_burst() {
let limiter = limiter();
for _ in 0..50 {
limiter
.check_bucket(RateBucket::CancelByLabel)
.expect("burst");
}
assert!(
limiter.check_bucket(RateBucket::CancelByLabel).is_err(),
"unscoped cancel_by_label allowance is 50 per window",
);
limiter.advance_clock(Duration::from_secs(5));
assert!(limiter.check_bucket(RateBucket::CancelByLabel).is_ok());
}
#[rstest]
fn test_rest_non_matching_quota_is_fifty_per_window() {
let limiter = FixedWindowLimiter::new(
FixedWindowLimits::rest(None, None),
FakeRelativeClock::default(),
);
for _ in 0..50 {
limiter
.check_bucket(RateBucket::NonMatching)
.expect("burst");
}
assert!(
limiter.check_bucket(RateBucket::NonMatching).is_err(),
"REST non-matching allowance is 50 per window",
);
}
#[rstest]
fn test_websocket_non_matching_quota_is_twenty_five_per_window() {
let limiter = limiter();
for _ in 0..25 {
limiter
.check_bucket(RateBucket::NonMatching)
.expect("burst");
}
assert!(
limiter.check_bucket(RateBucket::NonMatching).is_err(),
"WebSocket non-matching allowance is 25 per window",
);
}
#[rstest]
fn test_window_limit_and_index_arithmetic() {
assert_eq!(window_limit(1).get(), 5);
assert_eq!(window_index(0), 0);
assert_eq!(window_index(RATE_WINDOW_NANOS - 1), 0);
assert_eq!(window_index(RATE_WINDOW_NANOS), 1);
assert_eq!(unpack(pack(7, 3)), (7, 3));
assert_eq!(unpack(0), (0, 0));
}
}