polyfill_rs/
ws_hot_path.rs

1//! Zero-allocation-ish WebSocket hot-path processing.
2//!
3//! This module is focused on the "decode + apply" path for WS `book` events:
4//! after warmup, existing-level happy-path processing should not perform heap
5//! allocation, reallocation, or deallocation.
6//!
7//! Important: using the current tokio-tungstenite transport, the *network layer*
8//! may still allocate when producing `Message::Text(String)`. This module aims to
9//! make the *processing* layer allocation-free so we can enforce it with tests.
10
11use crate::book::{OrderBookManager, ParsedBookLevel};
12use crate::errors::{PolyfillError, Result};
13use crate::types::{Price, Qty, Side, MAX_PRICE_TICKS, MAX_QTY, MIN_PRICE_TICKS, SCALE_FACTOR};
14use simd_json::prelude::*;
15
16/// Summary of what happened while processing a WS payload.
17#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
18pub struct WsBookApplyStats {
19    pub book_messages: usize,
20    pub book_levels_applied: usize,
21}
22
23/// In-place WS `book` message processor built on `simd-json`'s tape API.
24///
25/// This avoids building a DOM (which allocates for arrays/objects) by decoding into a
26/// reusable tape, then traversing it to extract the fields needed for order book updates.
27pub struct WsBookUpdateProcessor {
28    buffers: simd_json::Buffers,
29    tape: Option<simd_json::Tape<'static>>,
30    parsed_levels: Vec<ParsedBookLevel>,
31}
32
33impl WsBookUpdateProcessor {
34    /// Create a new processor.
35    ///
36    /// `input_len_hint` should be set to the typical WS message size to reduce warmup reallocs.
37    pub fn new(input_len_hint: usize) -> Self {
38        Self {
39            buffers: simd_json::Buffers::new(input_len_hint),
40            // Store an empty tape with a `'static` lifetime so we can reuse its allocation.
41            tape: Some(simd_json::Tape::null().reset()),
42            parsed_levels: Vec::with_capacity((input_len_hint / 32).max(8)),
43        }
44    }
45
46    /// Process a WS payload in-place (bytes will be mutated by the JSON parser).
47    pub fn process_bytes(
48        &mut self,
49        bytes: &mut [u8],
50        books: &OrderBookManager,
51    ) -> Result<WsBookApplyStats> {
52        let mut tape = self
53            .tape
54            .take()
55            .expect("WsBookUpdateProcessor tape must be present")
56            .reset();
57
58        let result = match simd_json::fill_tape(bytes, &mut self.buffers, &mut tape) {
59            Ok(()) => {
60                let root = tape.as_value();
61                process_root_value(root, books, &mut self.parsed_levels)
62            },
63            Err(e) => Err(PolyfillError::parse(
64                "Failed to parse WebSocket JSON",
65                Some(Box::new(e)),
66            )),
67        };
68
69        // Reset the tape to detach lifetimes and keep capacity for reuse.
70        self.tape = Some(tape.reset());
71        result
72    }
73
74    /// Convenience: process an owned text message without allocating an additional buffer.
75    pub fn process_text(
76        &mut self,
77        text: String,
78        books: &OrderBookManager,
79    ) -> Result<WsBookApplyStats> {
80        let mut bytes = text.into_bytes();
81        self.process_bytes(bytes.as_mut_slice(), books)
82    }
83}
84
85fn process_root_value<'tape, 'input>(
86    value: simd_json::tape::Value<'tape, 'input>,
87    books: &OrderBookManager,
88    parsed_levels: &mut Vec<ParsedBookLevel>,
89) -> Result<WsBookApplyStats> {
90    if let Some(obj) = value.as_object() {
91        return process_stream_object(obj, books, parsed_levels);
92    }
93
94    let Some(arr) = value.as_array() else {
95        return Ok(WsBookApplyStats::default());
96    };
97
98    let mut total = WsBookApplyStats::default();
99    for elem in arr.iter() {
100        let Some(obj) = elem.as_object() else {
101            continue;
102        };
103        let stats = process_stream_object(obj, books, parsed_levels)?;
104        total.book_messages += stats.book_messages;
105        total.book_levels_applied += stats.book_levels_applied;
106    }
107
108    Ok(total)
109}
110
111fn process_stream_object<'tape, 'input>(
112    obj: simd_json::tape::Object<'tape, 'input>,
113    books: &OrderBookManager,
114    parsed_levels: &mut Vec<ParsedBookLevel>,
115) -> Result<WsBookApplyStats> {
116    let Some(event_type) = obj.get("event_type").and_then(|v| v.into_string()) else {
117        return Ok(WsBookApplyStats::default());
118    };
119
120    if event_type != "book" {
121        return Ok(WsBookApplyStats::default());
122    }
123
124    let asset_id = obj
125        .get("asset_id")
126        .and_then(|v| v.into_string())
127        .ok_or_else(|| PolyfillError::parse("Missing asset_id", None))?;
128
129    let timestamp_value = obj
130        .get("timestamp")
131        .ok_or_else(|| PolyfillError::parse("Missing timestamp", None))?;
132    let timestamp = parse_u64(timestamp_value)
133        .ok_or_else(|| PolyfillError::parse("Invalid timestamp", None))?;
134    let hash = obj.get("hash").and_then(|v| v.into_string());
135
136    let bids = obj
137        .get("bids")
138        .ok_or_else(|| PolyfillError::parse("Missing bids", None))?
139        .as_array()
140        .ok_or_else(|| PolyfillError::parse("Invalid bids", None))?;
141    let asks = obj
142        .get("asks")
143        .ok_or_else(|| PolyfillError::parse("Missing asks", None))?
144        .as_array()
145        .ok_or_else(|| PolyfillError::parse("Invalid asks", None))?;
146
147    parsed_levels.clear();
148
149    let should_parse = books.with_book_mut(asset_id, |book| {
150        book.should_apply_ws_book_update(asset_id, timestamp, hash)
151    })?;
152    if !should_parse {
153        return Ok(WsBookApplyStats {
154            book_messages: 1,
155            book_levels_applied: 0,
156        });
157    }
158
159    let parsed_count = match (|| {
160        collect_levels(Side::BUY, bids, parsed_levels)?;
161        collect_levels(Side::SELL, asks, parsed_levels)?;
162        Ok(parsed_levels.len())
163    })() {
164        Ok(parsed_count) => parsed_count,
165        Err(error) => {
166            parsed_levels.clear();
167            return Err(error);
168        },
169    };
170
171    let result = books.with_book_mut(asset_id, |book| {
172        if book.apply_ws_book_snapshot_fast(asset_id, timestamp, hash, parsed_levels)? {
173            Ok(parsed_count)
174        } else {
175            Ok(0)
176        }
177    });
178    parsed_levels.clear();
179    let levels_applied = result?;
180
181    Ok(WsBookApplyStats {
182        book_messages: 1,
183        book_levels_applied: levels_applied,
184    })
185}
186
187fn parse_u64<'tape, 'input>(value: simd_json::tape::Value<'tape, 'input>) -> Option<u64> {
188    value
189        .as_u64()
190        .or_else(|| value.into_string().and_then(|s| s.parse::<u64>().ok()))
191}
192
193fn collect_levels<'tape, 'input>(
194    side: Side,
195    levels: simd_json::tape::Array<'tape, 'input>,
196    parsed_levels: &mut Vec<ParsedBookLevel>,
197) -> Result<usize> {
198    let mut applied = 0usize;
199    for level in levels.iter() {
200        let Some(obj) = level.as_object() else {
201            continue;
202        };
203
204        let price_str = obj
205            .get("price")
206            .and_then(|v| v.into_string())
207            .ok_or_else(|| PolyfillError::parse("Missing price", None))?;
208        let size_str = obj
209            .get("size")
210            .and_then(|v| v.into_string())
211            .ok_or_else(|| PolyfillError::parse("Missing size", None))?;
212
213        let price_ticks = parse_price_ticks_4dp(price_str)?;
214        let size_units = parse_qty_scaled_4dp(size_str)?;
215
216        parsed_levels.push(ParsedBookLevel {
217            side,
218            price_ticks,
219            size_units,
220        });
221        applied += 1;
222    }
223
224    Ok(applied)
225}
226
227#[inline]
228fn parse_price_ticks_4dp(value: &str) -> Result<Price> {
229    let scaled = parse_scaled_4_u64(value)?;
230    if scaled < MIN_PRICE_TICKS as u64 {
231        return Err(PolyfillError::validation("Invalid price"));
232    }
233    if scaled > MAX_PRICE_TICKS as u64 {
234        return Err(PolyfillError::validation("Invalid price"));
235    }
236
237    Ok(scaled as Price)
238}
239
240#[inline]
241fn parse_qty_scaled_4dp(value: &str) -> Result<Qty> {
242    let scaled = parse_scaled_4_u64(value)?;
243    if scaled > MAX_QTY as u64 {
244        return Err(PolyfillError::validation("Invalid size"));
245    }
246
247    Ok(scaled as Qty)
248}
249
250#[inline]
251fn parse_scaled_4_u64(value: &str) -> Result<u64> {
252    if value.is_empty() {
253        return Err(PolyfillError::parse("invalid decimal", None));
254    }
255
256    let mut whole = 0u64;
257    let mut frac = 0u64;
258    let mut frac_digits = 0u8;
259    let mut seen_dot = false;
260    let mut seen_digit = false;
261
262    for &byte in value.as_bytes() {
263        match byte {
264            b'0'..=b'9' => {
265                seen_digit = true;
266                let digit = (byte - b'0') as u64;
267                if seen_dot {
268                    if frac_digits >= 4 {
269                        if digit != 0 {
270                            return Err(PolyfillError::parse("too many decimal places", None));
271                        }
272                    } else {
273                        frac = frac
274                            .checked_mul(10)
275                            .and_then(|x| x.checked_add(digit))
276                            .ok_or_else(|| PolyfillError::parse("scaled value overflow", None))?;
277                        frac_digits += 1;
278                    }
279                } else {
280                    whole = whole
281                        .checked_mul(10)
282                        .and_then(|x| x.checked_add(digit))
283                        .ok_or_else(|| PolyfillError::parse("scaled value overflow", None))?;
284                }
285            },
286            b'.' if !seen_dot => {
287                seen_dot = true;
288            },
289            _ => return Err(PolyfillError::parse("invalid decimal", None)),
290        }
291    }
292
293    if !seen_digit {
294        return Err(PolyfillError::parse("invalid decimal", None));
295    }
296
297    while frac_digits < 4 {
298        frac *= 10;
299        frac_digits += 1;
300    }
301
302    whole
303        .checked_mul(SCALE_FACTOR as u64)
304        .and_then(|x| x.checked_add(frac))
305        .ok_or_else(|| PolyfillError::parse("scaled value overflow", None))
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311    use crate::types::{BookUpdate, OrderSummary};
312    use rust_decimal_macros::dec;
313
314    #[test]
315    fn fixed_point_parser_matches_expected_price_ticks() {
316        assert_eq!(parse_price_ticks_4dp("0.6543").unwrap(), 6543);
317        assert_eq!(parse_price_ticks_4dp("1.0000").unwrap(), 10_000);
318        assert_eq!(parse_price_ticks_4dp("1.000000").unwrap(), 10_000);
319        assert!(parse_price_ticks_4dp("0.00005").is_err());
320        assert!(parse_price_ticks_4dp("0").is_err());
321        assert!(parse_price_ticks_4dp("-0.1").is_err());
322    }
323
324    #[test]
325    fn fixed_point_parser_matches_expected_qty_units() {
326        assert_eq!(parse_qty_scaled_4dp("100.0").unwrap(), 1_000_000);
327        assert_eq!(parse_qty_scaled_4dp("0.0000").unwrap(), 0);
328        assert_eq!(parse_qty_scaled_4dp("1.234500").unwrap(), 12_345);
329        assert!(parse_qty_scaled_4dp("-50.5").is_err());
330        assert!(parse_qty_scaled_4dp("0.00004").is_err());
331        assert!(parse_qty_scaled_4dp("0.00005").is_err());
332    }
333
334    #[test]
335    fn processor_recovers_after_parse_and_validation_errors() {
336        let books = OrderBookManager::new(10);
337        books.get_or_create_book("test_asset_id").unwrap();
338        let mut processor = WsBookUpdateProcessor::new(1024);
339
340        let mut malformed_json = br#"{"event_type":"#.to_vec();
341        assert!(processor
342            .process_bytes(malformed_json.as_mut_slice(), &books)
343            .is_err());
344
345        let mut invalid_level = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1000,"bids":[{"price":"0.75001","size":"1.0000"}],"asks":[]}"#.to_vec();
346        assert!(processor
347            .process_bytes(invalid_level.as_mut_slice(), &books)
348            .is_err());
349
350        let mut valid_update = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1001,"bids":[{"price":"0.7500","size":"1.0000"}],"asks":[]}"#.to_vec();
351        let stats = processor
352            .process_bytes(valid_update.as_mut_slice(), &books)
353            .unwrap();
354
355        assert_eq!(stats.book_messages, 1);
356        assert_eq!(stats.book_levels_applied, 1);
357    }
358
359    #[test]
360    fn processor_error_keeps_existing_snapshot() {
361        let books = OrderBookManager::new(10);
362        books.get_or_create_book("test_asset_id").unwrap();
363        books
364            .apply_book_update(&BookUpdate {
365                asset_id: "test_asset_id".to_string(),
366                market: "0xabc".to_string(),
367                timestamp: 1000,
368                bids: vec![OrderSummary {
369                    price: dec!(0.50),
370                    size: dec!(10),
371                }],
372                asks: vec![OrderSummary {
373                    price: dec!(0.60),
374                    size: dec!(20),
375                }],
376                hash: None,
377            })
378            .unwrap();
379
380        let mut processor = WsBookUpdateProcessor::new(1024);
381        let mut invalid_snapshot = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1001,"bids":[{"price":"0.5100","size":"11.0000"},{"price":"0.51001","size":"12.0000"}],"asks":[{"price":"0.6100","size":"21.0000"}]}"#.to_vec();
382        assert!(processor
383            .process_bytes(invalid_snapshot.as_mut_slice(), &books)
384            .is_err());
385
386        let snapshot = books.get_book("test_asset_id").unwrap();
387        assert_eq!(snapshot.sequence, 0);
388        assert_eq!(snapshot.last_delta_sequence, 0);
389        assert_eq!(snapshot.last_snapshot_timestamp_ms, 1000);
390        assert_eq!(snapshot.timestamp.timestamp_millis(), 1000);
391        assert_eq!(snapshot.bids.len(), 1);
392        assert_eq!(snapshot.asks.len(), 1);
393        assert_eq!(snapshot.bids[0].price, dec!(0.50));
394        assert_eq!(snapshot.bids[0].size, dec!(10));
395        assert_eq!(snapshot.asks[0].price, dec!(0.60));
396        assert_eq!(snapshot.asks[0].size, dec!(20));
397    }
398
399    #[test]
400    fn processor_missing_side_keeps_existing_snapshot() {
401        let books = OrderBookManager::new(10);
402        books.get_or_create_book("test_asset_id").unwrap();
403        books
404            .apply_book_update(&BookUpdate {
405                asset_id: "test_asset_id".to_string(),
406                market: "0xabc".to_string(),
407                timestamp: 1000,
408                bids: vec![OrderSummary {
409                    price: dec!(0.50),
410                    size: dec!(10),
411                }],
412                asks: vec![OrderSummary {
413                    price: dec!(0.60),
414                    size: dec!(20),
415                }],
416                hash: None,
417            })
418            .unwrap();
419
420        let mut processor = WsBookUpdateProcessor::new(1024);
421        let mut missing_asks = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1001,"bids":[{"price":"0.5100","size":"11.0000"}]}"#.to_vec();
422        assert!(processor
423            .process_bytes(missing_asks.as_mut_slice(), &books)
424            .is_err());
425
426        let snapshot = books.get_book("test_asset_id").unwrap();
427        assert_eq!(snapshot.sequence, 0);
428        assert_eq!(snapshot.last_delta_sequence, 0);
429        assert_eq!(snapshot.last_snapshot_timestamp_ms, 1000);
430        assert_eq!(snapshot.timestamp.timestamp_millis(), 1000);
431        assert_eq!(snapshot.bids.len(), 1);
432        assert_eq!(snapshot.asks.len(), 1);
433        assert_eq!(snapshot.bids[0].price, dec!(0.50));
434        assert_eq!(snapshot.asks[0].price, dec!(0.60));
435    }
436
437    #[test]
438    fn processor_allows_same_timestamp_with_different_hash() {
439        let books = OrderBookManager::new(10);
440        books.get_or_create_book("test_asset_id").unwrap();
441        let mut processor = WsBookUpdateProcessor::new(1024);
442
443        let mut first = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1000,"hash":"hash_a","bids":[{"price":"0.5000","size":"10.0000"}],"asks":[{"price":"0.6000","size":"20.0000"}]}"#.to_vec();
444        let first_stats = processor
445            .process_bytes(first.as_mut_slice(), &books)
446            .unwrap();
447        assert_eq!(first_stats.book_levels_applied, 2);
448
449        let mut second = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1000,"hash":"hash_b","bids":[{"price":"0.5100","size":"11.0000"}],"asks":[{"price":"0.6100","size":"21.0000"}]}"#.to_vec();
450        let second_stats = processor
451            .process_bytes(second.as_mut_slice(), &books)
452            .unwrap();
453        assert_eq!(second_stats.book_levels_applied, 2);
454
455        let snapshot = books.get_book("test_asset_id").unwrap();
456        assert_eq!(snapshot.sequence, 0);
457        assert_eq!(snapshot.last_delta_sequence, 0);
458        assert_eq!(snapshot.last_snapshot_timestamp_ms, 1000);
459        assert_eq!(snapshot.timestamp.timestamp_millis(), 1000);
460        assert_eq!(snapshot.bids[0].price, dec!(0.51));
461        assert_eq!(snapshot.bids[0].size, dec!(11));
462        assert_eq!(snapshot.asks[0].price, dec!(0.61));
463        assert_eq!(snapshot.asks[0].size, dec!(21));
464
465        let mut duplicate = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1000,"hash":"hash_b","bids":[{"price":"0.5200","size":"12.0000"}],"asks":[{"price":"0.6200","size":"22.0000"}]}"#.to_vec();
466        let duplicate_stats = processor
467            .process_bytes(duplicate.as_mut_slice(), &books)
468            .unwrap();
469        assert_eq!(duplicate_stats.book_levels_applied, 0);
470
471        let snapshot = books.get_book("test_asset_id").unwrap();
472        assert_eq!(snapshot.bids[0].price, dec!(0.51));
473        assert_eq!(snapshot.asks[0].price, dec!(0.61));
474    }
475}
polyfill_rs/ws_hot_path.rs

polyfill_rs/
ws_hot_path.rs
