test: count deallocations in hot path guards

This commit is contained in:
floor-licker
2026-06-22 19:42:14 -04:00
parent c821262af5
commit 543329f4d6
4 changed files with 93 additions and 57 deletions
+68 -50
View File
@@ -11,28 +11,29 @@ use polyfill_rs::{
use rust_decimal::Decimal;
thread_local! {
static ALLOCATIONS: Cell<usize> = const { Cell::new(0) };
static HEAP_OPERATIONS: Cell<usize> = const { Cell::new(0) };
}
struct CountingAllocator;
unsafe impl GlobalAlloc for CountingAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
ALLOCATIONS.with(|count| count.set(count.get() + 1));
HEAP_OPERATIONS.with(|count| count.set(count.get() + 1));
System.alloc(layout)
}
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
ALLOCATIONS.with(|count| count.set(count.get() + 1));
HEAP_OPERATIONS.with(|count| count.set(count.get() + 1));
System.alloc_zeroed(layout)
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
ALLOCATIONS.with(|count| count.set(count.get() + 1));
HEAP_OPERATIONS.with(|count| count.set(count.get() + 1));
System.realloc(ptr, layout, new_size)
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
HEAP_OPERATIONS.with(|count| count.set(count.get() + 1));
System.dealloc(ptr, layout)
}
}
@@ -40,32 +41,45 @@ unsafe impl GlobalAlloc for CountingAllocator {
#[global_allocator]
static GLOBAL: CountingAllocator = CountingAllocator;
fn allocation_count() -> usize {
ALLOCATIONS.with(|count| count.get())
fn heap_operation_count() -> usize {
HEAP_OPERATIONS.with(|count| count.get())
}
struct NoAllocGuard {
struct NoHeapTrafficGuard {
before: usize,
}
impl NoAllocGuard {
impl NoHeapTrafficGuard {
fn new() -> Self {
Self {
before: allocation_count(),
before: heap_operation_count(),
}
}
fn assert_no_allocations(self) {
let after = allocation_count();
fn assert_no_heap_traffic(self) {
let after = heap_operation_count();
assert_eq!(
after,
self.before,
"expected no heap allocations, but saw {} allocation(s)",
"expected no heap traffic, but saw {} allocator operation(s)",
after - self.before
);
}
}
#[test]
fn allocator_counter_tracks_deallocations() {
let vec = Vec::<u8>::with_capacity(1024);
let before = heap_operation_count();
drop(vec);
let after = heap_operation_count();
assert!(
after > before,
"expected dropping an allocated Vec to count as heap traffic"
);
}
fn token_id_hash(token_id: &str) -> u64 {
let mut hasher = DefaultHasher::new();
token_id.hash(&mut hasher);
@@ -101,15 +115,15 @@ fn no_alloc_mid_and_spread_fast() {
book.apply_delta_fast(mk_delta(token_hash, Side::SELL, 7600, 1_000_000, 2))
.unwrap();
// Warm up TLS access before measuring (defensive).
let _ = allocation_count();
// Warm up allocator-counter TLS access before measuring (defensive).
let _ = heap_operation_count();
let guard = NoAllocGuard::new();
let guard = NoHeapTrafficGuard::new();
assert!(book.best_bid_fast().is_some());
assert!(book.best_ask_fast().is_some());
assert!(book.spread_fast().is_some());
assert!(book.mid_price_fast().is_some());
guard.assert_no_allocations();
guard.assert_no_heap_traffic();
}
#[test]
@@ -134,9 +148,9 @@ fn no_alloc_book_analysis_fast_paths() {
let expected_buy_liquidity = Decimal::from_str("200.0").unwrap();
let expected_sell_liquidity = Decimal::from_str("150.0").unwrap();
let _ = allocation_count();
let _ = heap_operation_count();
let guard = NoAllocGuard::new();
let guard = NoHeapTrafficGuard::new();
let impact = book
.calculate_market_impact(Side::BUY, impact_size)
.unwrap();
@@ -150,7 +164,7 @@ fn no_alloc_book_analysis_fast_paths() {
expected_sell_liquidity
);
assert!(book.is_valid());
guard.assert_no_allocations();
guard.assert_no_heap_traffic();
}
#[test]
@@ -163,14 +177,14 @@ fn no_alloc_apply_delta_fast_existing_level_update() {
book.apply_delta_fast(mk_delta(token_hash, Side::BUY, 7500, 1_000_000, 1))
.unwrap();
// Warm up TLS access before measuring (defensive).
let _ = allocation_count();
// Warm up allocator-counter TLS access before measuring (defensive).
let _ = heap_operation_count();
let guard = NoAllocGuard::new();
// Updating an existing level should not require heap allocation.
let guard = NoHeapTrafficGuard::new();
// Updating an existing level should not touch the heap allocator.
book.apply_delta_fast(mk_delta(token_hash, Side::BUY, 7500, 2_000_000, 2))
.unwrap();
guard.assert_no_allocations();
guard.assert_no_heap_traffic();
}
#[test]
@@ -200,12 +214,12 @@ fn no_alloc_apply_book_update_existing_levels() {
hash: None,
};
// Warm up TLS access before measuring (defensive).
let _ = allocation_count();
// Warm up allocator-counter TLS access before measuring (defensive).
let _ = heap_operation_count();
let guard = NoAllocGuard::new();
let guard = NoHeapTrafficGuard::new();
book.apply_book_update(&update).unwrap();
guard.assert_no_allocations();
guard.assert_no_heap_traffic();
}
#[test]
@@ -214,7 +228,7 @@ fn no_alloc_book_manager_apply_book_update_existing_levels() {
let manager = OrderBookManager::new(100);
manager.get_or_create_book(asset_id).unwrap();
// Warm up the internal book with initial levels (allocations allowed).
// Warm up the internal book with initial levels (allocator traffic allowed).
manager
.apply_delta(polyfill_rs::types::OrderDelta {
token_id: asset_id.to_string(),
@@ -251,12 +265,12 @@ fn no_alloc_book_manager_apply_book_update_existing_levels() {
hash: None,
};
// Warm up TLS access before measuring (defensive).
let _ = allocation_count();
// Warm up allocator-counter TLS access before measuring (defensive).
let _ = heap_operation_count();
let guard = NoAllocGuard::new();
let guard = NoHeapTrafficGuard::new();
manager.apply_book_update(&update).unwrap();
guard.assert_no_allocations();
guard.assert_no_heap_traffic();
}
#[test]
@@ -265,7 +279,7 @@ fn no_alloc_ws_book_update_processor_apply_existing_levels() {
let manager = OrderBookManager::new(100);
manager.get_or_create_book(asset_id).unwrap();
// Warm up the internal book with initial levels (allocations allowed).
// Warm up the internal book with initial levels (allocator traffic allowed).
manager
.apply_delta(polyfill_rs::types::OrderDelta {
token_id: asset_id.to_string(),
@@ -303,23 +317,23 @@ fn no_alloc_ws_book_update_processor_apply_existing_levels() {
)
.into_bytes();
// Warm up TLS access before measuring (defensive).
let _ = allocation_count();
// Warm up allocator-counter TLS access before measuring (defensive).
let _ = heap_operation_count();
let guard = NoAllocGuard::new();
let guard = NoHeapTrafficGuard::new();
processor
.process_bytes(msg.as_mut_slice(), &manager)
.unwrap();
guard.assert_no_allocations();
guard.assert_no_heap_traffic();
}
#[test]
fn no_alloc_websocket_book_applier_apply_text_message_existing_levels() {
fn no_alloc_websocket_book_applier_apply_bytes_message_existing_levels() {
let asset_id = "test_asset_id";
let manager = OrderBookManager::new(100);
manager.get_or_create_book(asset_id).unwrap();
// Warm up the internal book with initial levels (allocations allowed).
// Warm up the internal book with initial levels (allocator traffic allowed).
manager
.apply_delta(polyfill_rs::types::OrderDelta {
token_id: asset_id.to_string(),
@@ -346,19 +360,23 @@ fn no_alloc_websocket_book_applier_apply_text_message_existing_levels() {
let mut applier = stream.into_book_applier(&manager, processor);
// Warm up simd-json buffers/tape outside the guarded section.
let warmup_msg = format!(
let mut warmup_msg = format!(
"{{\"event_type\":\"book\",\"asset_id\":\"{asset_id}\",\"market\":\"0xabc\",\"timestamp\":10,\"bids\":[{{\"price\":\"0.75\",\"size\":\"200.0\"}}],\"asks\":[{{\"price\":\"0.76\",\"size\":\"50.0\"}}]}}"
);
applier.apply_text_message(warmup_msg).unwrap();
)
.into_bytes();
applier
.apply_bytes_message(warmup_msg.as_mut_slice())
.unwrap();
let msg = format!(
let mut msg = format!(
"{{\"event_type\":\"book\",\"asset_id\":\"{asset_id}\",\"market\":\"0xabc\",\"timestamp\":11,\"bids\":[{{\"price\":\"0.75\",\"size\":\"150.0\"}}],\"asks\":[{{\"price\":\"0.76\",\"size\":\"75.0\"}}]}}"
);
)
.into_bytes();
// Warm up TLS access before measuring (defensive).
let _ = allocation_count();
// Warm up allocator-counter TLS access before measuring (defensive).
let _ = heap_operation_count();
let guard = NoAllocGuard::new();
applier.apply_text_message(msg).unwrap();
guard.assert_no_allocations();
let guard = NoHeapTrafficGuard::new();
applier.apply_bytes_message(msg.as_mut_slice()).unwrap();
guard.assert_no_heap_traffic();
}