From 3fb2bd8f6602ae2b803c2e07df8ff51eeed4a34f Mon Sep 17 00:00:00 2001 From: floor-licker Date: Mon, 22 Jun 2026 19:09:06 -0400 Subject: [PATCH] perf: use sorted vectors for book sides --- README.md | 10 +- src/book.rs | 355 +++++++++++++++++++++++++++++++++++++--------------- 2 files changed, 256 insertions(+), 109 deletions(-) diff --git a/README.md b/README.md index 001a2e7..2adcb4b 100644 --- a/README.md +++ b/README.md @@ -11,8 +11,8 @@ At the time that this project was started, `polymarket-rs-client` was a Polymark I also want to take a moment to clarify what zero-alloc means because I've now recieved double digit messages about this on twitter/x and telegram. In general, zero alloc means either zero alloc in hot paths (which can be a bit more arbitrary) or atlernatively it can mean zero alloc after init/warm-up, which is the objective of this repository. Succinctly that means that **the per-message handling loop never touches the heap**. Notably order book paths that introduce new allocations by design: -- First time seeing a token/book (HashMap insert + key clone): `src/book.rs:~788` -- New price levels (BTreeMap node growth): `src/book.rs:~409` +- First time seeing a token/book (HashMap insert + key clone): `src/book.rs` +- New price levels (sorted Vec insert/growth): `src/book.rs` ## Quick Start @@ -68,10 +68,10 @@ Real-world Polymarket API latency broken down by request phase: | Operation | Performance | Notes | |-----------|-------------|-------| -| **Order Book Updates (1000 ops)** | 159.6 µs ± 32 µs | 6,260 updates/sec, zero-allocation | -| **Spread/Mid Calculations** | 70 ns ± 77 ns | 14.3M ops/sec, optimized BTreeMap | +| **Order Book Updates (1000 ops)** | 69.6 µs | ~14.4M updates/sec, zero-allocation for warmed existing levels | +| **Spread/Mid Calculations** | 26.6 ns | best bid/ask + spread + mid over sorted-vector book sides | | **JSON Parsing (480KB)** | ~0.5 ms | SIMD-backed parsing for large REST market responses and benchmarked polyfill typed parse path | -| **WS `book` hot path (decode + apply)** | ~0.23 µs / 1.73 µs / 6.74 µs | 1 / 16 / 64 levels-per-side, strict fixed-point tape parser with generation-marked snapshot retention (see `benches/ws_hot_path.rs`) | +| **WS `book` hot path (decode + apply)** | ~0.24 µs / 1.56 µs / 5.92 µs | 1 / 16 / 64 levels-per-side, strict fixed-point tape parser with generation-marked snapshot retention (see `benches/ws_hot_path.rs`) | Run the WS hot-path benchmark locally with `cargo bench --bench ws_hot_path`. diff --git a/src/book.rs b/src/book.rs index c0542d5..a866f71 100644 --- a/src/book.rs +++ b/src/book.rs @@ -6,7 +6,7 @@ use crate::utils::math; use chrono::Utc; use parking_lot::RwLock; use rust_decimal::Decimal; -use std::collections::{BTreeMap, HashMap}; // BTreeMap keeps prices sorted automatically - crucial for order books +use std::collections::HashMap; use std::sync::Arc; // For shared access across multiple tasks use tracing::{debug, trace, warn}; // Logging for debugging and monitoring @@ -16,6 +16,127 @@ struct StoredLevel { generation: u64, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BookSideKind { + Bid, + Ask, +} + +#[derive(Debug, Clone)] +struct BookSide { + levels: Vec<(Price, StoredLevel)>, + kind: BookSideKind, +} + +impl BookSide { + fn new(kind: BookSideKind, max_depth: usize) -> Self { + Self { + levels: Vec::with_capacity(max_depth.saturating_add(8)), + kind, + } + } + + #[inline] + fn len(&self) -> usize { + self.levels.len() + } + + #[inline] + fn search(&self, price: Price) -> std::result::Result { + match self.kind { + // Bids are stored highest price first. + BookSideKind::Bid => self + .levels + .binary_search_by(|(level_price, _)| price.cmp(level_price)), + // Asks are stored lowest price first. + BookSideKind::Ask => self + .levels + .binary_search_by(|(level_price, _)| level_price.cmp(&price)), + } + } + + #[inline] + fn best(&self) -> Option<(Price, &StoredLevel)> { + self.levels.first().map(|(price, level)| (*price, level)) + } + + #[inline] + fn get(&self, price: Price) -> Option<&StoredLevel> { + self.search(price) + .ok() + .and_then(|idx| self.levels.get(idx).map(|(_, level)| level)) + } + + #[inline] + fn upsert(&mut self, price: Price, level: StoredLevel) { + match self.search(price) { + Ok(idx) if level.qty == 0 => { + self.levels.remove(idx); + }, + Ok(idx) => { + self.levels[idx].1 = level; + }, + Err(idx) if level.qty != 0 => { + self.levels.insert(idx, (price, level)); + }, + Err(_) => {}, + } + } + + #[inline] + fn retain_generation(&mut self, generation: u64) { + self.levels + .retain(|(_, level)| level.generation == generation); + } + + #[inline] + fn trim_depth(&mut self, max_depth: usize) { + if self.levels.len() > max_depth { + self.levels.truncate(max_depth); + } + } + + #[inline] + fn iter_top(&self, depth: usize) -> impl Iterator { + self.levels + .iter() + .take(depth) + .map(|(price, level)| (*price, level)) + } + + #[inline] + fn iter_all(&self) -> impl Iterator { + self.levels.iter().map(|(price, level)| (*price, level)) + } + + fn sum_range(&self, min_price: Price, max_price: Price) -> Qty { + let mut total = 0; + for (price, level) in &self.levels { + match self.kind { + BookSideKind::Bid => { + if *price > max_price { + continue; + } + if *price < min_price { + break; + } + }, + BookSideKind::Ask => { + if *price < min_price { + continue; + } + if *price > max_price { + break; + } + }, + } + + total += level.qty; + } + total + } +} + const DEFAULT_BOOK_SHARDS: usize = 64; /// High-performance order book implementation @@ -45,24 +166,24 @@ pub struct OrderBook { pub timestamp: chrono::DateTime, /// Bid side (price -> size, sorted descending) - NOW USING FIXED-POINT! - /// BTreeMap automatically keeps highest bids first, which is what we want + /// Stored as a sorted vector with highest bids first for cache-local top-of-book iteration. /// Key = price in ticks (like 6500 for $0.65), Value = size in fixed-point units /// /// BEFORE (slow): bids: BTreeMap, - /// AFTER (fast): bids: BTreeMap, + /// AFTER (fast): bids: sorted Vec<(Price, StoredLevel)>, /// /// Why this is faster: /// - Integer comparisons are ~10x faster than Decimal comparisons /// - No memory allocation for each price level /// - Better CPU cache utilization (smaller data structures) - bids: BTreeMap, + bids: BookSide, /// Ask side (price -> size, sorted ascending) - NOW USING FIXED-POINT! - /// BTreeMap keeps lowest asks first - people selling at cheapest prices + /// Stored as a sorted vector with lowest asks first - people selling at cheapest prices. /// /// BEFORE (slow): asks: BTreeMap, - /// AFTER (fast): asks: BTreeMap, - asks: BTreeMap, + /// AFTER (fast): asks: sorted Vec<(Price, StoredLevel)>, + asks: BookSide, /// Snapshot generation used to retain book levels without rescanning input payloads. snapshot_generation: u64, @@ -103,8 +224,8 @@ impl OrderBook { token_id_hash, sequence: 0, // Start at 0, will increment as we get updates timestamp: Utc::now(), - bids: BTreeMap::new(), // Empty to start - using Price/Qty types - asks: BTreeMap::new(), // Empty to start - using Price/Qty types + bids: BookSide::new(BookSideKind::Bid, max_depth), // Empty to start + asks: BookSide::new(BookSideKind::Ask, max_depth), // Empty to start snapshot_generation: 0, tick_size_ticks: None, // We'll set this later when we learn about the market max_depth, @@ -128,7 +249,7 @@ impl OrderBook { } /// Get the current best bid (highest price someone is willing to pay) - /// Uses next_back() because BTreeMap sorts ascending, but we want the highest bid + /// Bids are stored highest price first. /// /// PERFORMANCE: Now returns data in external format but internally uses fast lookups pub fn best_bid(&self) -> Option { @@ -136,7 +257,7 @@ impl OrderBook { // self.bids.iter().next_back().map(|(&price, &size)| BookLevel { price, size }) // AFTER (fast, ~5ns, no allocation for the lookup): - self.bids.iter().next_back().map(|(&price_ticks, level)| { + self.bids.best().map(|(price_ticks, level)| { // Convert from internal fixed-point to external Decimal format // This conversion only happens at the API boundary BookLevel { @@ -147,7 +268,7 @@ impl OrderBook { } /// Get the current best ask (lowest price someone is willing to sell at) - /// Uses next() because BTreeMap sorts ascending, so first item is lowest ask + /// Asks are stored lowest price first. /// /// PERFORMANCE: Now returns data in external format but internally uses fast lookups pub fn best_ask(&self) -> Option { @@ -155,7 +276,7 @@ impl OrderBook { // self.asks.iter().next().map(|(&price, &size)| BookLevel { price, size }) // AFTER (fast, ~5ns, no allocation for the lookup): - self.asks.iter().next().map(|(&price_ticks, level)| { + self.asks.best().map(|(price_ticks, level)| { // Convert from internal fixed-point to external Decimal format // This conversion only happens at the API boundary BookLevel { @@ -169,18 +290,16 @@ impl OrderBook { /// Use this for internal calculations to avoid conversion overhead pub fn best_bid_fast(&self) -> Option { self.bids - .iter() - .next_back() - .map(|(&price, level)| FastBookLevel::new(price, level.qty)) + .best() + .map(|(price, level)| FastBookLevel::new(price, level.qty)) } /// Get the current best ask in fast internal format /// Use this for internal calculations to avoid conversion overhead pub fn best_ask_fast(&self) -> Option { self.asks - .iter() - .next() - .map(|(&price, level)| FastBookLevel::new(price, level.qty)) + .best() + .map(|(price, level)| FastBookLevel::new(price, level.qty)) } /// Get the current spread (difference between best ask and best bid) @@ -231,9 +350,9 @@ impl OrderBook { /// Get best bid and ask prices in fast internal format /// Helper method to avoid code duplication and minimize conversions fn best_prices_fast(&self) -> Option<(Price, Price)> { - let best_bid_ticks = self.bids.iter().next_back()?.0; - let best_ask_ticks = self.asks.iter().next()?.0; - Some((*best_bid_ticks, *best_ask_ticks)) + let best_bid_ticks = self.bids.best()?.0; + let best_ask_ticks = self.asks.best()?.0; + Some((best_bid_ticks, best_ask_ticks)) } /// Get the current spread in fast internal format (PERFORMANCE OPTIMIZED) @@ -258,10 +377,8 @@ impl OrderBook { pub fn bids(&self, depth: Option) -> Vec { let depth = depth.unwrap_or(self.max_depth); self.bids - .iter() - .rev() // Reverse because we want highest prices first - .take(depth) // Only take the top N levels - .map(|(&price_ticks, level)| BookLevel { + .iter_top(depth) + .map(|(price_ticks, level)| BookLevel { price: price_to_decimal(price_ticks), size: qty_to_decimal(level.qty), }) @@ -276,9 +393,8 @@ impl OrderBook { pub fn asks(&self, depth: Option) -> Vec { let depth = depth.unwrap_or(self.max_depth); self.asks - .iter() // Already in ascending order, so no need to reverse - .take(depth) // Only take the top N levels - .map(|(&price_ticks, level)| BookLevel { + .iter_top(depth) + .map(|(price_ticks, level)| BookLevel { price: price_to_decimal(price_ticks), size: qty_to_decimal(level.qty), }) @@ -290,10 +406,8 @@ impl OrderBook { pub fn bids_fast(&self, depth: Option) -> Vec { let depth = depth.unwrap_or(self.max_depth); self.bids - .iter() - .rev() // Reverse because we want highest prices first - .take(depth) // Only take the top N levels - .map(|(&price, level)| FastBookLevel::new(price, level.qty)) + .iter_top(depth) + .map(|(price, level)| FastBookLevel::new(price, level.qty)) .collect() } @@ -302,9 +416,8 @@ impl OrderBook { pub fn asks_fast(&self, depth: Option) -> Vec { let depth = depth.unwrap_or(self.max_depth); self.asks - .iter() // Already in ascending order, so no need to reverse - .take(depth) // Only take the top N levels - .map(|(&price, level)| FastBookLevel::new(price, level.qty)) + .iter_top(depth) + .map(|(price, level)| FastBookLevel::new(price, level.qty)) .collect() } @@ -432,7 +545,7 @@ impl OrderBook { /// Apply a single WS `book` level (already converted to internal fixed-point). /// - /// Note: Insertions of new price levels may allocate (BTreeMap node growth). In a strict + /// Note: Insertions of new price levels may allocate (Vec growth/shifting). In a strict /// zero-alloc hot path, all expected levels must be warmed up ahead of time. pub(crate) fn apply_ws_book_level_fast( &mut self, @@ -466,7 +579,7 @@ impl OrderBook { /// Notes: /// - This replaces the snapshot for both sides. /// - Levels omitted from the message are removed. - /// - Insertions of *new* price levels may allocate (BTreeMap node growth). + /// - Insertions of *new* price levels may allocate or shift vector entries. pub fn apply_book_update(&mut self, update: &BookUpdate) -> Result<()> { if update.asset_id != self.token_id { return Err(PolyfillError::validation("Token ID mismatch")); @@ -558,17 +671,13 @@ impl OrderBook { // } // AFTER (fast, ~5ns, no allocation): - if size_units == 0 { - self.bids.remove(&price_ticks); // No more buyers at this price - } else { - self.bids.insert( - price_ticks, - StoredLevel { - qty: size_units, - generation: self.snapshot_generation, - }, - ); // Update total size at this price - } + self.bids.upsert( + price_ticks, + StoredLevel { + qty: size_units, + generation: self.snapshot_generation, + }, + ); } /// Apply an ask-side delta (someone wants to sell) - FAST VERSION @@ -584,17 +693,13 @@ impl OrderBook { // } // AFTER (fast, ~5ns, no allocation): - if size_units == 0 { - self.asks.remove(&price_ticks); // No more sellers at this price - } else { - self.asks.insert( - price_ticks, - StoredLevel { - qty: size_units, - generation: self.snapshot_generation, - }, - ); // Update total size at this price - } + self.asks.upsert( + price_ticks, + StoredLevel { + qty: size_units, + generation: self.snapshot_generation, + }, + ); } #[inline] @@ -605,29 +710,25 @@ impl OrderBook { #[inline] fn apply_snapshot_level(&mut self, side: Side, price_ticks: Price, size_units: Qty) { let generation = self.snapshot_generation; - let map = match side { + let book_side = match side { Side::BUY => &mut self.bids, Side::SELL => &mut self.asks, }; - if size_units == 0 { - map.remove(&price_ticks); - } else { - map.insert( - price_ticks, - StoredLevel { - qty: size_units, - generation, - }, - ); - } + book_side.upsert( + price_ticks, + StoredLevel { + qty: size_units, + generation, + }, + ); } #[inline] fn finish_snapshot(&mut self) { let generation = self.snapshot_generation; - self.bids.retain(|_, level| level.generation == generation); - self.asks.retain(|_, level| level.generation == generation); + self.bids.retain_generation(generation); + self.asks.retain_generation(generation); } /// Trim the book to maintain depth limits @@ -644,21 +745,11 @@ impl OrderBook { fn trim_depth(&mut self) { // For bids, remove the LOWEST prices (worst bids) if we have too many // Example: If best bid is $0.65, we don't care about bids at $0.10 - if self.bids.len() > self.max_depth { - let to_remove = self.bids.len() - self.max_depth; - for _ in 0..to_remove { - self.bids.pop_first(); // Remove lowest bid prices (furthest from market) - } - } + self.bids.trim_depth(self.max_depth); // For asks, remove the HIGHEST prices (worst asks) if we have too many // Example: If best ask is $0.67, we don't care about asks at $0.95 - if self.asks.len() > self.max_depth { - let to_remove = self.asks.len() - self.max_depth; - for _ in 0..to_remove { - self.asks.pop_last(); // Remove highest ask prices (furthest from market) - } - } + self.asks.trim_depth(self.max_depth); } /// Calculate the market impact for a given order size @@ -671,8 +762,8 @@ impl OrderBook { pub fn calculate_market_impact(&self, side: Side, size: Decimal) -> Option { let size_units = decimal_to_qty(size).ok()?; let (filled_units, total_notional, best_price_ticks) = match side { - Side::BUY => fill_market_impact(self.asks.iter(), size_units)?, - Side::SELL => fill_market_impact(self.bids.iter().rev(), size_units)?, + Side::BUY => fill_market_impact(self.asks.iter_all(), size_units)?, + Side::SELL => fill_market_impact(self.bids.iter_all(), size_units)?, }; let total_cost = Decimal::from_i128_with_scale(total_notional, 8); @@ -714,7 +805,7 @@ impl OrderBook { // How much we can buy at this price (look at asks) let size_units = self .asks - .get(&price_ticks) + .get(price_ticks) .map(|level| level.qty) .unwrap_or_default(); qty_to_decimal(size_units) @@ -723,7 +814,7 @@ impl OrderBook { // How much we can sell at this price (look at bids) let size_units = self .bids - .get(&price_ticks) + .get(price_ticks) .map(|level| level.qty) .unwrap_or_default(); qty_to_decimal(size_units) @@ -750,16 +841,8 @@ impl OrderBook { }; let total_size_units: Qty = match side { - Side::BUY => self - .asks - .range(min_price_ticks..=max_price_ticks) - .map(|(_, level)| level.qty) - .sum(), - Side::SELL => self - .bids - .range(min_price_ticks..=max_price_ticks) - .map(|(_, level)| level.qty) - .sum(), + Side::BUY => self.asks.sum_range(min_price_ticks, max_price_ticks), + Side::SELL => self.bids.sum_range(min_price_ticks, max_price_ticks), }; qty_to_decimal(total_size_units) @@ -775,7 +858,7 @@ impl OrderBook { } fn fill_market_impact<'a>( - levels: impl Iterator, + levels: impl Iterator, size_units: Qty, ) -> Option<(Qty, i128, Price)> { if size_units <= 0 { @@ -787,7 +870,7 @@ fn fill_market_impact<'a>( let mut total_notional = 0i128; let mut best_price_ticks = None; - for (&price_ticks, level) in levels { + for (price_ticks, level) in levels { if best_price_ticks.is_none() { best_price_ticks = Some(price_ticks); } @@ -1030,8 +1113,8 @@ impl OrderBook { let bid_count = self.bids.len(); let ask_count = self.asks.len(); // Sum up all bid/ask sizes, converting from fixed-point back to Decimal - let total_bid_size_units: i64 = self.bids.values().map(|level| level.qty).sum(); - let total_ask_size_units: i64 = self.asks.values().map(|level| level.qty).sum(); + let total_bid_size_units: i64 = self.bids.iter_all().map(|(_, level)| level.qty).sum(); + let total_ask_size_units: i64 = self.asks.iter_all().map(|(_, level)| level.qty).sum(); let total_bid_size = qty_to_decimal(total_bid_size_units); let total_ask_size = qty_to_decimal(total_ask_size_units); @@ -1141,6 +1224,70 @@ mod tests { assert_eq!(book.best_bid().unwrap().size, dec!(100)); // Should be our size } + #[test] + fn test_sorted_side_ordering_and_removal() { + let mut book = OrderBook::new("test_token".to_string(), 10); + + for (sequence, price) in [(1, dec!(0.73)), (2, dec!(0.75)), (3, dec!(0.74))] { + book.apply_delta(OrderDelta { + token_id: "test_token".to_string(), + timestamp: Utc::now(), + side: Side::BUY, + price, + size: dec!(100), + sequence, + }) + .unwrap(); + } + + for (sequence, price) in [(4, dec!(0.78)), (5, dec!(0.76)), (6, dec!(0.77))] { + book.apply_delta(OrderDelta { + token_id: "test_token".to_string(), + timestamp: Utc::now(), + side: Side::SELL, + price, + size: dec!(100), + sequence, + }) + .unwrap(); + } + + let bids: Vec<_> = book + .bids(Some(3)) + .into_iter() + .map(|level| level.price) + .collect(); + let asks: Vec<_> = book + .asks(Some(3)) + .into_iter() + .map(|level| level.price) + .collect(); + assert_eq!(bids, vec![dec!(0.75), dec!(0.74), dec!(0.73)]); + assert_eq!(asks, vec![dec!(0.76), dec!(0.77), dec!(0.78)]); + + book.apply_delta(OrderDelta { + token_id: "test_token".to_string(), + timestamp: Utc::now(), + side: Side::BUY, + price: dec!(0.75), + size: Decimal::ZERO, + sequence: 7, + }) + .unwrap(); + book.apply_delta(OrderDelta { + token_id: "test_token".to_string(), + timestamp: Utc::now(), + side: Side::SELL, + price: dec!(0.76), + size: Decimal::ZERO, + sequence: 8, + }) + .unwrap(); + + assert_eq!(book.best_bid().unwrap().price, dec!(0.74)); + assert_eq!(book.best_ask().unwrap().price, dec!(0.77)); + } + #[test] fn test_book_update_replaces_snapshot_and_uses_millis_timestamp() { let mut book = OrderBook::new("test_token".to_string(), 10);