1use std::{
32 fmt::{Debug, Display},
33 path::PathBuf,
34 str::FromStr,
35 sync::Arc,
36 time::Duration,
37};
38
39use ahash::{AHashMap, HashSet, HashSetExt};
40use databento::{
41 dbn::{self, PitSymbolMap, Record, SymbolIndex},
42 live::Subscription,
43};
44use indexmap::IndexMap;
45use nautilus_core::{
46 AtomicMap, UnixNanos, consts::NAUTILUS_USER_AGENT, time::get_atomic_clock_realtime,
47};
48use nautilus_model::{
49 data::{Data, InstrumentStatus, OrderBookDelta, OrderBookDeltas},
50 enums::RecordFlag,
51 identifiers::{InstrumentId, Symbol, Venue},
52 instruments::{Instrument, InstrumentAny},
53 types::Currency,
54};
55use nautilus_network::backoff::ExponentialBackoff;
56use time::OffsetDateTime;
57
58use super::{
59 decode::{
60 decode_imbalance_msg, decode_statistics_msg, decode_status_msg, is_supported_stat_type,
61 },
62 types::{DatabentoImbalance, DatabentoStatistics, SubscriptionAckEvent},
63};
64use crate::{
65 common::{Credential, build_publisher_venue_map, load_publishers},
66 decode::{decode_instrument_def_msg, decode_record},
67 symbology::{check_consistent_symbology, infer_symbology_type},
68 types::PublisherId,
69};
70
71#[derive(Debug)]
72pub enum HandlerCommand {
73 Subscribe(Subscription),
74 SetPricePrecision(Symbol, u8),
75 Start,
76 Close,
77}
78
79#[derive(Debug)]
80pub enum DatabentoMessage {
81 Data(Data),
82 Instrument(Box<InstrumentAny>),
83 Status(InstrumentStatus),
84 Imbalance(DatabentoImbalance),
85 Statistics(DatabentoStatistics),
86 SubscriptionAck(SubscriptionAckEvent),
87 Error(anyhow::Error),
88 Close,
89}
90
91#[cfg_attr(
92 feature = "python",
93 pyo3::pyclass(module = "nautilus_trader.adapters.databento")
94)]
95#[cfg_attr(
96 feature = "python",
97 pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.adapters.databento")
98)]
99pub struct DatabentoLiveClient {
100 credential: Credential,
101 pub dataset: String,
102 is_running: bool,
103 is_closed: bool,
104 cmd_tx: tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
105 cmd_rx: Option<tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>>,
106 publisher_venue_map: IndexMap<PublisherId, Venue>,
107 symbol_venue_map: Arc<AtomicMap<Symbol, Venue>>,
108 use_exchange_as_venue: bool,
109 bars_timestamp_on_close: bool,
110 reconnect_timeout_mins: Option<u64>,
111}
112
113impl Debug for DatabentoLiveClient {
114 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115 f.debug_struct(stringify!(DatabentoLiveClient))
116 .field("credential", &self.credential)
117 .field("dataset", &self.dataset)
118 .field("is_running", &self.is_running)
119 .field("is_closed", &self.is_closed)
120 .finish()
121 }
122}
123
124impl DatabentoLiveClient {
125 pub fn new(
131 key: String,
132 dataset: String,
133 publishers_filepath: PathBuf,
134 use_exchange_as_venue: bool,
135 bars_timestamp_on_close: Option<bool>,
136 reconnect_timeout_mins: Option<i64>,
137 ) -> anyhow::Result<Self> {
138 let publishers = load_publishers(publishers_filepath)?;
139 let publisher_venue_map = build_publisher_venue_map(&publishers);
140
141 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
142
143 let reconnect_timeout_mins = reconnect_timeout_mins
144 .and_then(|mins| if mins >= 0 { Some(mins as u64) } else { None });
145
146 Ok(Self {
147 credential: Credential::new(key),
148 dataset,
149 cmd_tx,
150 cmd_rx: Some(cmd_rx),
151 is_running: false,
152 is_closed: false,
153 publisher_venue_map,
154 symbol_venue_map: Arc::new(AtomicMap::new()),
155 use_exchange_as_venue,
156 bars_timestamp_on_close: bars_timestamp_on_close.unwrap_or(true),
157 reconnect_timeout_mins,
158 })
159 }
160
161 #[must_use]
162 pub const fn is_running(&self) -> bool {
163 self.is_running
164 }
165
166 #[must_use]
167 pub const fn is_closed(&self) -> bool {
168 self.is_closed
169 }
170
171 #[expect(clippy::needless_pass_by_value)]
178 pub fn subscribe(
179 &mut self,
180 schema: String,
181 instrument_ids: Vec<InstrumentId>,
182 start: Option<u64>,
183 snapshot: Option<bool>,
184 price_precisions: Option<Vec<Option<u8>>>,
185 stype_in: Option<String>,
186 ) -> anyhow::Result<()> {
187 if let Some(precisions) = &price_precisions
188 && precisions.len() != instrument_ids.len()
189 {
190 anyhow::bail!(
191 "`price_precisions` length ({}) must match `instrument_ids` length ({})",
192 precisions.len(),
193 instrument_ids.len()
194 );
195 }
196
197 let symbols: Vec<String> = instrument_ids
198 .iter()
199 .map(|id| id.symbol.to_string())
200 .collect();
201 let first_symbol = symbols
202 .first()
203 .ok_or_else(|| anyhow::anyhow!("No symbols provided"))?;
204 let stype_in = match stype_in {
205 Some(stype_in) => dbn::SType::from_str(&stype_in)?,
206 None => infer_symbology_type(first_symbol),
207 };
208 let symbols: Vec<&str> = symbols.iter().map(String::as_str).collect();
209 check_consistent_symbology(symbols.as_slice())?;
210 let mut sub = Subscription::builder()
211 .symbols(symbols)
212 .schema(dbn::Schema::from_str(&schema)?)
213 .stype_in(stype_in)
214 .build();
215
216 if let Some(start) = start {
217 sub.start = Some(OffsetDateTime::from_unix_timestamp_nanos(i128::from(
218 start,
219 ))?);
220 }
221 sub.use_snapshot = snapshot.unwrap_or(false);
222
223 self.symbol_venue_map.rcu(|m| {
224 for id in &instrument_ids {
225 m.entry(id.symbol).or_insert(id.venue);
226 }
227 });
228
229 if let Some(precisions) = price_precisions {
230 for (instrument_id, precision) in instrument_ids.iter().zip(precisions) {
231 if let Some(precision) = precision {
232 self.send_command(HandlerCommand::SetPricePrecision(
233 instrument_id.symbol,
234 precision,
235 ))?;
236 }
237 }
238 }
239
240 self.send_command(HandlerCommand::Subscribe(sub))
241 }
242
243 pub fn start(
249 &mut self,
250 ) -> anyhow::Result<(
251 DatabentoFeedHandler,
252 tokio::sync::mpsc::UnboundedReceiver<DatabentoMessage>,
253 )> {
254 if self.is_closed {
255 anyhow::bail!("Client already closed");
256 }
257
258 if self.is_running {
259 anyhow::bail!("Client already running");
260 }
261
262 log::debug!("Starting client");
263
264 let (msg_tx, msg_rx) = tokio::sync::mpsc::unbounded_channel::<DatabentoMessage>();
265 let cmd_rx = self
266 .cmd_rx
267 .take()
268 .ok_or_else(|| anyhow::anyhow!("Command receiver already taken"))?;
269
270 let feed_handler = DatabentoFeedHandler::new(
271 self.credential.clone(),
272 self.dataset.clone(),
273 cmd_rx,
274 msg_tx,
275 self.publisher_venue_map.clone(),
276 self.symbol_venue_map.clone(),
277 self.use_exchange_as_venue,
278 self.bars_timestamp_on_close,
279 self.reconnect_timeout_mins,
280 );
281
282 self.send_command(HandlerCommand::Start)?;
283 self.is_running = true;
284
285 Ok((feed_handler, msg_rx))
286 }
287
288 pub fn close(&mut self) -> anyhow::Result<()> {
295 if !self.is_running {
296 anyhow::bail!("Client never started");
297 }
298
299 if self.is_closed {
300 anyhow::bail!("Client already closed");
301 }
302
303 log::debug!("Closing client");
304
305 if !self.cmd_tx.is_closed() {
306 self.send_command(HandlerCommand::Close)?;
307 }
308
309 self.is_running = false;
310 self.is_closed = true;
311
312 Ok(())
313 }
314
315 fn send_command(&self, cmd: HandlerCommand) -> anyhow::Result<()> {
316 self.cmd_tx.send(cmd).map_err(|e| {
317 anyhow::Error::new(CommandSendError {
318 details: e.to_string(),
319 })
320 })
321 }
322}
323
324#[cfg(any(test, feature = "python"))]
325pub(crate) fn is_command_send_error(error: &anyhow::Error) -> bool {
326 error.is::<CommandSendError>()
327}
328
329#[derive(Debug)]
330struct CommandSendError {
331 details: String,
332}
333
334impl Display for CommandSendError {
335 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
336 write!(
337 f,
338 "Failed to send command to Databento feed handler: {}",
339 self.details
340 )
341 }
342}
343
344impl std::error::Error for CommandSendError {}
345
346pub struct DatabentoFeedHandler {
358 credential: Credential,
359 dataset: String,
360 cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
361 msg_tx: tokio::sync::mpsc::UnboundedSender<DatabentoMessage>,
362 publisher_venue_map: IndexMap<PublisherId, Venue>,
363 symbol_venue_map: Arc<AtomicMap<Symbol, Venue>>,
364 replay: bool,
365 use_exchange_as_venue: bool,
366 bars_timestamp_on_close: bool,
367 reconnect_timeout_mins: Option<u64>,
368 backoff: ExponentialBackoff,
369 subscriptions: Vec<Subscription>,
370 buffered_commands: Vec<HandlerCommand>,
371 price_precision_overrides: AHashMap<Symbol, u8>,
372 gateway_addr: Option<String>,
373 success_threshold: Duration,
374}
375
376impl Debug for DatabentoFeedHandler {
377 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
378 f.debug_struct(stringify!(DatabentoFeedHandler))
379 .field("credential", &self.credential)
380 .field("dataset", &self.dataset)
381 .field("replay", &self.replay)
382 .field("reconnect_timeout_mins", &self.reconnect_timeout_mins)
383 .field("subscriptions", &self.subscriptions.len())
384 .finish()
385 }
386}
387
388impl DatabentoFeedHandler {
389 #[must_use]
395 #[expect(clippy::too_many_arguments)]
396 pub fn new(
397 credential: Credential,
398 dataset: String,
399 rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
400 tx: tokio::sync::mpsc::UnboundedSender<DatabentoMessage>,
401 publisher_venue_map: IndexMap<PublisherId, Venue>,
402 symbol_venue_map: Arc<AtomicMap<Symbol, Venue>>,
403 use_exchange_as_venue: bool,
404 bars_timestamp_on_close: bool,
405 reconnect_timeout_mins: Option<u64>,
406 ) -> Self {
407 let delay_max = if reconnect_timeout_mins.is_some() {
411 Duration::from_secs(60)
412 } else {
413 Duration::from_secs(600)
414 };
415
416 let backoff = ExponentialBackoff::new(Duration::from_secs(1), delay_max, 2.0, 1000, false)
417 .expect("hardcoded backoff parameters are valid");
418
419 Self {
420 credential,
421 dataset,
422 cmd_rx: rx,
423 msg_tx: tx,
424 publisher_venue_map,
425 symbol_venue_map,
426 replay: false,
427 use_exchange_as_venue,
428 bars_timestamp_on_close,
429 reconnect_timeout_mins,
430 backoff,
431 subscriptions: Vec::new(),
432 buffered_commands: Vec::new(),
433 price_precision_overrides: AHashMap::new(),
434 gateway_addr: None,
435 success_threshold: Duration::from_secs(60),
436 }
437 }
438
439 #[must_use]
441 pub fn with_gateway_addr(mut self, addr: String) -> Self {
442 self.gateway_addr = Some(addr);
443 self
444 }
445
446 #[must_use]
451 pub fn with_success_threshold(mut self, threshold: Duration) -> Self {
452 self.success_threshold = threshold;
453 self
454 }
455
456 pub async fn run(&mut self) -> anyhow::Result<()> {
469 log::debug!("Running feed handler");
470
471 let mut reconnect_start: Option<tokio::time::Instant> = None;
472 let mut attempt = 0;
473
474 loop {
475 attempt += 1;
476
477 match self.run_session(attempt).await {
478 Ok(ran_successfully) => {
479 if ran_successfully {
480 log::info!("Resetting reconnection cycle after successful session");
481 reconnect_start = None;
482 attempt = 0;
483 self.backoff.reset();
484 } else {
485 log::debug!("Session ended normally");
486 break Ok(());
487 }
488 }
489 Err(e) => {
490 let cycle_start = reconnect_start.get_or_insert_with(tokio::time::Instant::now);
491
492 if let Some(timeout_mins) = self.reconnect_timeout_mins {
493 let elapsed = cycle_start.elapsed();
494 let timeout = Duration::from_mins(timeout_mins);
495
496 if elapsed >= timeout {
497 log::error!("Giving up reconnection after {timeout_mins} minutes");
498 self.send_msg(DatabentoMessage::Error(anyhow::anyhow!(
499 "Reconnection timeout after {timeout_mins} minutes: {e}"
500 )));
501 break Err(e);
502 }
503 }
504
505 let delay = self.backoff.next_duration();
506
507 log::warn!(
508 "Connection lost (attempt {}): {}. Reconnecting in {}s...",
509 attempt,
510 e,
511 delay.as_secs()
512 );
513
514 let sleep = tokio::time::sleep(delay);
515 tokio::pin!(sleep);
516
517 loop {
518 tokio::select! {
519 () = &mut sleep => break,
520 cmd = self.cmd_rx.recv() => {
521 match cmd {
522 Some(HandlerCommand::Close) => {
523 log::debug!("Close received during backoff");
524 return Ok(());
525 }
526 None => {
527 log::debug!("Command channel closed during backoff");
528 return Ok(());
529 }
530 Some(cmd) => {
531 log::debug!("Buffering command received during backoff: {cmd:?}");
532 self.buffered_commands.push(cmd);
533 }
534 }
535 }
536 }
537 }
538 }
539 }
540 }
541 }
542
543 async fn run_session(&mut self, attempt: usize) -> anyhow::Result<bool> {
552 if attempt > 1 {
553 log::info!("Reconnecting (attempt {attempt})...");
554 }
555
556 let session_start = tokio::time::Instant::now();
557 let clock = get_atomic_clock_realtime();
558 let mut symbol_map = PitSymbolMap::new();
559 let mut instrument_id_map: AHashMap<u32, InstrumentId> = AHashMap::new();
560 let mut instrument_def_price_precision_map: AHashMap<u32, u8> = AHashMap::new();
561 let mut subscription_price_precision_map: AHashMap<u32, u8> = AHashMap::new();
562
563 let mut buffering_start = None;
564 let mut buffered_deltas: AHashMap<InstrumentId, Vec<OrderBookDelta>> = AHashMap::new();
565 let mut initialized_books = HashSet::new();
566 let timeout = Duration::from_secs(5); let gateway_addr = self.gateway_addr.clone();
569 let api_key = self.credential.api_key().to_owned();
570 let dataset = self.dataset.clone();
571
572 let result = tokio::time::timeout(timeout, async move {
573 let base = databento::LiveClient::builder();
574 let base = if let Some(addr) = gateway_addr {
575 base.addr(addr).await?
576 } else {
577 base
578 };
579 base.user_agent_extension(NAUTILUS_USER_AGENT.into())
580 .key(api_key)?
581 .dataset(dataset)
582 .build()
583 .await
584 })
585 .await?;
586
587 let mut client = match result {
588 Ok(client) => {
589 if attempt > 1 {
590 log::info!("Reconnected successfully");
591 } else {
592 log::info!("Connected");
593 }
594 client
595 }
596 Err(e) => {
597 anyhow::bail!("Failed to connect to Databento LSG: {e}");
598 }
599 };
600
601 let mut start_buffered = false;
603
604 if !self.buffered_commands.is_empty() {
605 log::debug!(
606 "Processing {} buffered commands",
607 self.buffered_commands.len()
608 );
609
610 for cmd in self.buffered_commands.drain(..) {
611 match cmd {
612 HandlerCommand::Subscribe(sub) => {
613 if !self.replay && sub.start.is_some() {
614 self.replay = true;
615 }
616 self.subscriptions.push(sub);
617 }
618 HandlerCommand::SetPricePrecision(symbol, precision) => {
619 self.price_precision_overrides.insert(symbol, precision);
620 }
621 HandlerCommand::Start => {
622 start_buffered = true;
623 }
624 HandlerCommand::Close => {
625 log::warn!("Close command was buffered, shutting down");
626 return Ok(false);
627 }
628 }
629 }
630 }
631
632 let mut running = false;
633
634 if !self.subscriptions.is_empty() {
635 log::info!(
636 "Resubscribing to {} subscriptions",
637 self.subscriptions.len()
638 );
639
640 for sub in self.subscriptions.clone() {
641 client.subscribe(sub).await?;
642 }
643 for sub in &mut self.subscriptions {
645 sub.start = None;
646 }
647 client.start().await?;
648 running = true;
649 log::info!("Resubscription complete");
650 } else if start_buffered {
651 log::debug!("Starting session from buffered Start command");
652 buffering_start = if self.replay {
653 Some(clock.get_time_ns())
654 } else {
655 None
656 };
657 client.start().await?;
658 running = true;
659 }
660
661 loop {
662 if self.msg_tx.is_closed() {
663 log::debug!("Message channel was closed: stopping");
664 return Ok(false);
665 }
666
667 if !running {
672 match self.cmd_rx.recv().await {
673 Some(HandlerCommand::Subscribe(sub)) => {
674 log::debug!("Received command: Subscribe");
675
676 if !self.replay && sub.start.is_some() {
677 self.replay = true;
678 }
679 client.subscribe(sub.clone()).await?;
680 let mut sub_for_reconnect = sub;
681 sub_for_reconnect.start = None;
682 self.subscriptions.push(sub_for_reconnect);
683 continue;
684 }
685 Some(HandlerCommand::SetPricePrecision(symbol, precision)) => {
686 log::debug!(
687 "Received command: SetPricePrecision for {symbol} to {precision}"
688 );
689 self.price_precision_overrides.insert(symbol, precision);
690 continue;
691 }
692 Some(HandlerCommand::Start) => {
693 log::debug!("Received command: Start");
694 buffering_start = if self.replay {
695 Some(clock.get_time_ns())
696 } else {
697 None
698 };
699 client.start().await?;
700 running = true;
701 continue;
702 }
703 Some(HandlerCommand::Close) => {
704 self.send_close_msg();
705 return Ok(false);
706 }
707 None => {
708 log::debug!("Command channel disconnected");
709 return Ok(false);
710 }
711 }
712 }
713
714 let record_opt = tokio::select! {
715 cmd = self.cmd_rx.recv() =>
716 match cmd {
717 Some(HandlerCommand::Subscribe(sub)) => {
718 log::debug!("Received command: Subscribe");
719
720 if sub.start.is_some() {
721 self.replay = true;
722 log::error!(
723 "Ignoring `start` on {} subscribe, session already running, Databento drops replay anchors sent after session start",
724 self.dataset,
725 );
726 }
727 client.subscribe(sub.clone()).await?;
728 let mut sub_for_reconnect = sub;
729 sub_for_reconnect.start = None;
730 self.subscriptions.push(sub_for_reconnect);
731 continue;
732 }
733 Some(HandlerCommand::SetPricePrecision(symbol, precision)) => {
734 log::debug!(
735 "Received command: SetPricePrecision for {symbol} to {precision}"
736 );
737 self.price_precision_overrides.insert(symbol, precision);
738 continue;
739 }
740 Some(HandlerCommand::Start) => {
741 log::warn!("Received Start command but session already running");
742 continue;
743 }
744 Some(HandlerCommand::Close) => {
745 self.send_close_msg();
746 client.close().await?;
747 log::debug!("Closed inner client");
748 return Ok(false);
749 }
750 None => {
751 log::debug!("Command channel disconnected");
752 return Ok(false);
753 }
754 },
755 result = client.next_record() => result,
756 };
757
758 let record = match record_opt {
759 Ok(Some(record)) => record,
760 Ok(None) => {
761 if session_start.elapsed() >= self.success_threshold {
762 log::debug!("Session ended after successful run");
763 return Ok(true);
764 }
765 anyhow::bail!("Session ended by gateway");
766 }
767 Err(e) => {
768 if session_start.elapsed() >= self.success_threshold {
769 log::debug!("Connection error after successful run: {e}");
770 return Ok(true);
771 }
772 anyhow::bail!("Connection error: {e}");
773 }
774 };
775
776 let ts_init = clock.get_time_ns();
777
778 if let Some(msg) = record.get::<dbn::ErrorMsg>() {
780 handle_error_msg(msg);
781 } else if let Some(msg) = record.get::<dbn::SystemMsg>() {
782 if let Some(ack) = handle_system_msg(msg, ts_init) {
783 self.send_msg(DatabentoMessage::SubscriptionAck(ack));
784 }
785 } else if let Some(msg) = record.get::<dbn::SymbolMappingMsg>() {
786 instrument_id_map.remove(&msg.hd.instrument_id);
788 instrument_def_price_precision_map.remove(&msg.hd.instrument_id);
789 update_price_precision_map_with_symbol_mapping_msg(
790 msg,
791 &self.price_precision_overrides,
792 &mut subscription_price_precision_map,
793 )?;
794 handle_symbol_mapping_msg(msg, &mut symbol_map, &mut instrument_id_map)?;
795 } else if let Some(msg) = record.get::<dbn::InstrumentDefMsg>() {
796 if self.use_exchange_as_venue {
797 let exchange = msg.exchange()?;
798 if !exchange.is_empty() {
799 update_instrument_id_map_with_exchange(
800 &symbol_map,
801 &self.symbol_venue_map,
802 &mut instrument_id_map,
803 msg.hd.instrument_id,
804 exchange,
805 )?;
806 }
807 }
808 let maybe_data = handle_instrument_def_msg(
809 msg,
810 &record,
811 &symbol_map,
812 &self.publisher_venue_map,
813 &self.symbol_venue_map,
814 &mut instrument_id_map,
815 ts_init,
816 )?;
817
818 if let Some(data) = maybe_data {
819 instrument_def_price_precision_map
820 .insert(msg.hd.instrument_id, data.price_precision());
821 self.send_msg(DatabentoMessage::Instrument(Box::new(data)));
822 }
823 } else if let Some(msg) = record.get::<dbn::StatusMsg>() {
824 let data = handle_status_msg(
825 msg,
826 &record,
827 &symbol_map,
828 &self.publisher_venue_map,
829 &self.symbol_venue_map,
830 &mut instrument_id_map,
831 ts_init,
832 )?;
833 self.send_msg(DatabentoMessage::Status(data));
834 } else if let Some(msg) = record.get::<dbn::ImbalanceMsg>() {
835 let data = handle_imbalance_msg(
836 msg,
837 &record,
838 &symbol_map,
839 &self.publisher_venue_map,
840 &self.symbol_venue_map,
841 &mut instrument_id_map,
842 &instrument_def_price_precision_map,
843 &subscription_price_precision_map,
844 &self.price_precision_overrides,
845 ts_init,
846 )?;
847 self.send_msg(DatabentoMessage::Imbalance(data));
848 } else if let Some(msg) = record.get::<dbn::StatMsg>() {
849 let maybe_data = handle_statistics_msg(
850 msg,
851 &record,
852 &symbol_map,
853 &self.publisher_venue_map,
854 &self.symbol_venue_map,
855 &mut instrument_id_map,
856 &instrument_def_price_precision_map,
857 &subscription_price_precision_map,
858 &self.price_precision_overrides,
859 ts_init,
860 )?;
861
862 if let Some(data) = maybe_data {
863 self.send_msg(DatabentoMessage::Statistics(data));
864 }
865 } else {
866 let res = handle_record(
868 record,
869 &symbol_map,
870 &self.publisher_venue_map,
871 &self.symbol_venue_map,
872 &mut instrument_id_map,
873 &instrument_def_price_precision_map,
874 &subscription_price_precision_map,
875 &self.price_precision_overrides,
876 ts_init,
877 &initialized_books,
878 self.bars_timestamp_on_close,
879 );
880 let (mut data1, data2) = match res {
881 Ok(decoded) => decoded,
882 Err(e) => {
883 log::error!("Error decoding record: {e}");
884 continue;
885 }
886 };
887
888 if let Some(msg) = record.get::<dbn::MboMsg>() {
889 if let Some(Data::Delta(delta)) = &data1 {
890 initialized_books.insert(delta.instrument_id);
891
892 log::trace!(
893 "Buffering delta: {} {buffering_start:?} flags={}",
894 delta.ts_event,
895 msg.flags.raw(),
896 );
897
898 match process_mbo_delta(
899 *delta,
900 msg.flags.raw(),
901 &mut buffering_start,
902 &mut buffered_deltas,
903 ) {
904 Some(deltas) => data1 = Some(Data::Deltas(Box::new(deltas))),
905 None => continue,
906 }
907 } else {
908 let instrument_id = update_instrument_id_map(
915 &record,
916 &symbol_map,
917 &self.publisher_venue_map,
918 &self.symbol_venue_map,
919 &mut instrument_id_map,
920 )?;
921
922 if let Some(deltas) = flush_mbo_event_boundary(
923 instrument_id,
924 msg.ts_recv.into(),
925 msg.flags.raw(),
926 &mut buffering_start,
927 &mut buffered_deltas,
928 ) {
929 self.send_msg(DatabentoMessage::Data(Data::Deltas(Box::new(deltas))));
930 }
931
932 continue;
933 }
934 }
935
936 if let Some(data) = data1 {
937 self.send_msg(DatabentoMessage::Data(data));
938 }
939
940 if let Some(data) = data2 {
941 self.send_msg(DatabentoMessage::Data(data));
942 }
943 }
944 }
945 }
946
947 fn send_msg(&self, msg: DatabentoMessage) {
949 log::trace!("Sending {msg:?}");
950 match self.msg_tx.send(msg) {
951 Ok(()) => {}
952 Err(e) => log::error!("Error sending message: {e}"),
953 }
954 }
955
956 fn send_close_msg(&self) {
957 if let Err(e) = self.msg_tx.send(DatabentoMessage::Close) {
958 log::debug!("Could not send close message: {e}");
959 }
960 }
961}
962
963fn handle_error_msg(msg: &dbn::ErrorMsg) {
965 log::error!("{msg:?}");
966}
967
968fn handle_system_msg(msg: &dbn::SystemMsg, ts_received: UnixNanos) -> Option<SubscriptionAckEvent> {
970 match msg.code() {
971 Ok(dbn::SystemCode::SubscriptionAck) => {
972 let message = msg.msg().unwrap_or("<invalid utf-8>");
973 log::debug!("Subscription acknowledged: {message}");
974
975 let schema = parse_ack_message(message);
976
977 Some(SubscriptionAckEvent {
978 schema,
979 message: message.to_string(),
980 ts_received,
981 })
982 }
983 Ok(dbn::SystemCode::Heartbeat) => {
984 log::trace!("Heartbeat received");
985 None
986 }
987 Ok(dbn::SystemCode::SlowReaderWarning) => {
988 let message = msg.msg().unwrap_or("<invalid utf-8>");
989 log::warn!("Slow reader warning: {message}");
990 None
991 }
992 Ok(dbn::SystemCode::ReplayCompleted) => {
993 let message = msg.msg().unwrap_or("<invalid utf-8>");
994 log::debug!("Replay completed: {message}");
995 None
996 }
997 _ => {
998 log::debug!("{msg:?}");
999 None
1000 }
1001 }
1002}
1003
1004fn parse_ack_message(message: &str) -> String {
1006 message
1008 .strip_circumfix("Subscription request ", " data succeeded")
1009 .and_then(|rest| rest.split_once(" for "))
1010 .map(|(_, schema)| schema.trim().to_string())
1011 .unwrap_or_default()
1012}
1013
1014fn handle_symbol_mapping_msg(
1020 msg: &dbn::SymbolMappingMsg,
1021 symbol_map: &mut PitSymbolMap,
1022 instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1023) -> anyhow::Result<()> {
1024 symbol_map
1025 .on_symbol_mapping(msg)
1026 .map_err(|e| anyhow::anyhow!("on_symbol_mapping failed for {msg:?}: {e}"))?;
1027 instrument_id_map.remove(&msg.header().instrument_id);
1028 Ok(())
1029}
1030
1031fn update_price_precision_map_with_symbol_mapping_msg(
1032 msg: &dbn::SymbolMappingMsg,
1033 price_precision_overrides: &AHashMap<Symbol, u8>,
1034 subscription_price_precision_map: &mut AHashMap<u32, u8>,
1035) -> anyhow::Result<()> {
1036 subscription_price_precision_map.remove(&msg.hd.instrument_id);
1037
1038 if price_precision_overrides.is_empty() {
1039 return Ok(());
1040 }
1041
1042 let stype_in_symbol = msg
1043 .stype_in_symbol()
1044 .map_err(|e| anyhow::anyhow!("Error decoding `stype_in_symbol`: {e}"))?;
1045 let stype_out_symbol = msg
1046 .stype_out_symbol()
1047 .map_err(|e| anyhow::anyhow!("Error decoding `stype_out_symbol`: {e}"))?;
1048
1049 let price_precision = [stype_in_symbol, stype_out_symbol]
1050 .into_iter()
1051 .find_map(|symbol| {
1052 price_precision_overrides
1053 .get(&Symbol::from_str_unchecked(symbol))
1054 .copied()
1055 });
1056
1057 if let Some(price_precision) = price_precision {
1058 subscription_price_precision_map.insert(msg.hd.instrument_id, price_precision);
1059 }
1060
1061 Ok(())
1062}
1063
1064fn update_instrument_id_map_with_exchange(
1066 symbol_map: &PitSymbolMap,
1067 symbol_venue_map: &AtomicMap<Symbol, Venue>,
1068 instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1069 raw_instrument_id: u32,
1070 exchange: &str,
1071) -> anyhow::Result<InstrumentId> {
1072 let raw_symbol = symbol_map.get(raw_instrument_id).ok_or_else(|| {
1073 anyhow::anyhow!("Cannot resolve raw_symbol for instrument_id {raw_instrument_id}")
1074 })?;
1075 let symbol = Symbol::from(raw_symbol.as_str());
1076 let venue = Venue::from_code(exchange)
1077 .map_err(|e| anyhow::anyhow!("Invalid venue code '{exchange}': {e}"))?;
1078 let instrument_id = InstrumentId::new(symbol, venue);
1079 symbol_venue_map.rcu(|m| {
1080 m.entry(symbol).or_insert(venue);
1081 });
1082 instrument_id_map.insert(raw_instrument_id, instrument_id);
1083 Ok(instrument_id)
1084}
1085
1086fn update_instrument_id_map(
1087 record: &dbn::RecordRef,
1088 symbol_map: &PitSymbolMap,
1089 publisher_venue_map: &IndexMap<PublisherId, Venue>,
1090 symbol_venue_map: &AtomicMap<Symbol, Venue>,
1091 instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1092) -> anyhow::Result<InstrumentId> {
1093 let header = record.header();
1094
1095 if let Some(&instrument_id) = instrument_id_map.get(&header.instrument_id) {
1097 return Ok(instrument_id);
1098 }
1099
1100 let raw_symbol = symbol_map.get_for_rec(record).ok_or_else(|| {
1101 anyhow::anyhow!(
1102 "Cannot resolve `raw_symbol` from `symbol_map` for instrument_id {}",
1103 header.instrument_id
1104 )
1105 })?;
1106
1107 let symbol = Symbol::from_str_unchecked(raw_symbol);
1108
1109 let publisher_id = header.publisher_id;
1110 let venue = if let Some(venue) = symbol_venue_map.get_cloned(&symbol) {
1111 venue
1112 } else {
1113 let venue = publisher_venue_map
1114 .get(&publisher_id)
1115 .ok_or_else(|| anyhow::anyhow!("No venue found for `publisher_id` {publisher_id}"))?;
1116 *venue
1117 };
1118 let instrument_id = InstrumentId::new(symbol, venue);
1119
1120 instrument_id_map.insert(header.instrument_id, instrument_id);
1121 Ok(instrument_id)
1122}
1123
1124fn handle_instrument_def_msg(
1130 msg: &dbn::InstrumentDefMsg,
1131 record: &dbn::RecordRef,
1132 symbol_map: &PitSymbolMap,
1133 publisher_venue_map: &IndexMap<PublisherId, Venue>,
1134 symbol_venue_map: &AtomicMap<Symbol, Venue>,
1135 instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1136 ts_init: UnixNanos,
1137) -> anyhow::Result<Option<InstrumentAny>> {
1138 let instrument_id = update_instrument_id_map(
1139 record,
1140 symbol_map,
1141 publisher_venue_map,
1142 symbol_venue_map,
1143 instrument_id_map,
1144 )?;
1145
1146 decode_instrument_def_msg(msg, instrument_id, Some(ts_init), None)
1147}
1148
1149fn handle_status_msg(
1150 msg: &dbn::StatusMsg,
1151 record: &dbn::RecordRef,
1152 symbol_map: &PitSymbolMap,
1153 publisher_venue_map: &IndexMap<PublisherId, Venue>,
1154 symbol_venue_map: &AtomicMap<Symbol, Venue>,
1155 instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1156 ts_init: UnixNanos,
1157) -> anyhow::Result<InstrumentStatus> {
1158 let instrument_id = update_instrument_id_map(
1159 record,
1160 symbol_map,
1161 publisher_venue_map,
1162 symbol_venue_map,
1163 instrument_id_map,
1164 )?;
1165
1166 decode_status_msg(msg, instrument_id, Some(ts_init))
1167}
1168
1169#[expect(clippy::too_many_arguments)]
1170fn handle_imbalance_msg(
1171 msg: &dbn::ImbalanceMsg,
1172 record: &dbn::RecordRef,
1173 symbol_map: &PitSymbolMap,
1174 publisher_venue_map: &IndexMap<PublisherId, Venue>,
1175 symbol_venue_map: &AtomicMap<Symbol, Venue>,
1176 instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1177 instrument_def_price_precision_map: &AHashMap<u32, u8>,
1178 subscription_price_precision_map: &AHashMap<u32, u8>,
1179 price_precision_overrides: &AHashMap<Symbol, u8>,
1180 ts_init: UnixNanos,
1181) -> anyhow::Result<DatabentoImbalance> {
1182 let instrument_id = update_instrument_id_map(
1183 record,
1184 symbol_map,
1185 publisher_venue_map,
1186 symbol_venue_map,
1187 instrument_id_map,
1188 )?;
1189
1190 let price_precision = resolve_price_precision(
1191 msg.hd.instrument_id,
1192 instrument_id,
1193 instrument_def_price_precision_map,
1194 subscription_price_precision_map,
1195 price_precision_overrides,
1196 );
1197
1198 decode_imbalance_msg(msg, instrument_id, price_precision, Some(ts_init))
1199}
1200
1201#[expect(clippy::too_many_arguments)]
1202fn handle_statistics_msg(
1203 msg: &dbn::StatMsg,
1204 record: &dbn::RecordRef,
1205 symbol_map: &PitSymbolMap,
1206 publisher_venue_map: &IndexMap<PublisherId, Venue>,
1207 symbol_venue_map: &AtomicMap<Symbol, Venue>,
1208 instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1209 instrument_def_price_precision_map: &AHashMap<u32, u8>,
1210 subscription_price_precision_map: &AHashMap<u32, u8>,
1211 price_precision_overrides: &AHashMap<Symbol, u8>,
1212 ts_init: UnixNanos,
1213) -> anyhow::Result<Option<DatabentoStatistics>> {
1214 if !is_supported_stat_type(msg.stat_type) {
1216 log::warn!("Skipping unsupported `stat_type` {}", msg.stat_type);
1217 return Ok(None);
1218 }
1219
1220 let instrument_id = update_instrument_id_map(
1221 record,
1222 symbol_map,
1223 publisher_venue_map,
1224 symbol_venue_map,
1225 instrument_id_map,
1226 )?;
1227
1228 let price_precision = resolve_price_precision(
1229 msg.hd.instrument_id,
1230 instrument_id,
1231 instrument_def_price_precision_map,
1232 subscription_price_precision_map,
1233 price_precision_overrides,
1234 );
1235
1236 decode_statistics_msg(msg, instrument_id, price_precision, Some(ts_init))
1237}
1238
1239#[expect(clippy::too_many_arguments)]
1240fn handle_record(
1241 record: dbn::RecordRef,
1242 symbol_map: &PitSymbolMap,
1243 publisher_venue_map: &IndexMap<PublisherId, Venue>,
1244 symbol_venue_map: &AtomicMap<Symbol, Venue>,
1245 instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1246 instrument_def_price_precision_map: &AHashMap<u32, u8>,
1247 subscription_price_precision_map: &AHashMap<u32, u8>,
1248 price_precision_overrides: &AHashMap<Symbol, u8>,
1249 ts_init: UnixNanos,
1250 initialized_books: &HashSet<InstrumentId>,
1251 bars_timestamp_on_close: bool,
1252) -> anyhow::Result<(Option<Data>, Option<Data>)> {
1253 let instrument_id = update_instrument_id_map(
1254 &record,
1255 symbol_map,
1256 publisher_venue_map,
1257 symbol_venue_map,
1258 instrument_id_map,
1259 )?;
1260
1261 let price_precision = resolve_price_precision(
1262 record.header().instrument_id,
1263 instrument_id,
1264 instrument_def_price_precision_map,
1265 subscription_price_precision_map,
1266 price_precision_overrides,
1267 );
1268
1269 let include_trades = if record.get::<dbn::Mbp1Msg>().is_some()
1272 || record.get::<dbn::TbboMsg>().is_some()
1273 || record.get::<dbn::Cmbp1Msg>().is_some()
1274 {
1275 true } else {
1277 initialized_books.contains(&instrument_id) };
1279
1280 decode_record(
1281 &record,
1282 instrument_id,
1283 price_precision,
1284 Some(ts_init),
1285 include_trades,
1286 bars_timestamp_on_close,
1287 )
1288}
1289
1290fn resolve_price_precision(
1291 record_instrument_id: u32,
1292 instrument_id: InstrumentId,
1293 instrument_def_price_precision_map: &AHashMap<u32, u8>,
1294 subscription_price_precision_map: &AHashMap<u32, u8>,
1295 price_precision_overrides: &AHashMap<Symbol, u8>,
1296) -> u8 {
1297 instrument_def_price_precision_map
1298 .get(&record_instrument_id)
1299 .copied()
1300 .or_else(|| {
1301 subscription_price_precision_map
1302 .get(&record_instrument_id)
1303 .copied()
1304 })
1305 .or_else(|| {
1306 price_precision_overrides
1307 .get(&instrument_id.symbol)
1308 .copied()
1309 })
1310 .unwrap_or(Currency::USD().precision)
1311}
1312
1313fn process_mbo_delta(
1318 delta: OrderBookDelta,
1319 flags: u8,
1320 buffering_start: &mut Option<UnixNanos>,
1321 buffered_deltas: &mut AHashMap<InstrumentId, Vec<OrderBookDelta>>,
1322) -> Option<OrderBookDeltas> {
1323 let is_last = RecordFlag::F_LAST.matches(flags);
1324 let is_snapshot = RecordFlag::F_SNAPSHOT.matches(flags);
1325
1326 if is_last
1328 && !is_snapshot
1329 && buffering_start.is_none()
1330 && !buffered_deltas.contains_key(&delta.instrument_id)
1331 {
1332 let deltas = OrderBookDeltas::new(delta.instrument_id, vec![delta]);
1333 return Some(deltas);
1334 }
1335
1336 let buffer = buffered_deltas.entry(delta.instrument_id).or_default();
1337 buffer.push(delta);
1338
1339 if is_snapshot {
1340 return None;
1341 }
1342
1343 flush_mbo_event_boundary(
1344 delta.instrument_id,
1345 delta.ts_event,
1346 flags,
1347 buffering_start,
1348 buffered_deltas,
1349 )
1350}
1351
1352fn flush_mbo_event_boundary(
1360 instrument_id: InstrumentId,
1361 ts_event: UnixNanos,
1362 flags: u8,
1363 buffering_start: &mut Option<UnixNanos>,
1364 buffered_deltas: &mut AHashMap<InstrumentId, Vec<OrderBookDelta>>,
1365) -> Option<OrderBookDeltas> {
1366 if !RecordFlag::F_LAST.matches(flags) || RecordFlag::F_SNAPSHOT.matches(flags) {
1367 return None;
1368 }
1369
1370 let buffer = buffered_deltas.get(&instrument_id)?;
1371
1372 if buffer.is_empty() {
1373 return None;
1374 }
1375
1376 if let Some(start_ns) = *buffering_start {
1377 if ts_event <= start_ns {
1378 return None;
1379 }
1380 *buffering_start = None;
1381 }
1382
1383 let buffer = buffered_deltas.remove(&instrument_id)?;
1384 let deltas = OrderBookDeltas::new(instrument_id, buffer);
1385 Some(deltas)
1386}
1387
1388#[cfg(test)]
1389mod tests {
1390 use std::path::PathBuf;
1391
1392 use databento::live::Subscription;
1393 use indexmap::IndexMap;
1394 use rstest::*;
1395 use time::macros::datetime;
1396
1397 use super::*;
1398
1399 fn stub_delta(instrument_id: InstrumentId, ts: u64) -> OrderBookDelta {
1400 use nautilus_model::{
1401 data::BookOrder,
1402 enums::{BookAction, OrderSide},
1403 types::{Price, Quantity},
1404 };
1405
1406 OrderBookDelta::new(
1407 instrument_id,
1408 BookAction::Delete,
1409 BookOrder::new(
1410 OrderSide::Sell,
1411 Price::from("100.00"),
1412 Quantity::from("5"),
1413 42,
1414 ),
1415 0, 1,
1417 ts.into(),
1418 ts.into(),
1419 )
1420 }
1421
1422 #[rstest]
1423 fn test_boundary_flag_on_recordless_message_flushes_buffered_event() {
1424 let instrument_id = InstrumentId::from("TEST.GLBX");
1428 let mut buffering_start = None;
1429 let mut buffered = AHashMap::new();
1430
1431 let buffered_result = process_mbo_delta(
1432 stub_delta(instrument_id, 1),
1433 0,
1434 &mut buffering_start,
1435 &mut buffered,
1436 );
1437 assert!(buffered_result.is_none());
1438
1439 let flushed = flush_mbo_event_boundary(
1440 instrument_id,
1441 2.into(),
1442 RecordFlag::F_LAST as u8,
1443 &mut buffering_start,
1444 &mut buffered,
1445 );
1446 assert!(flushed.is_some());
1447 assert!(buffered.is_empty());
1448 }
1449
1450 #[rstest]
1451 fn test_boundary_flag_with_empty_buffer_is_noop() {
1452 let instrument_id = InstrumentId::from("TEST.GLBX");
1453 let start = UnixNanos::from(1);
1454 let mut buffering_start = Some(start);
1455 let mut buffered: AHashMap<InstrumentId, Vec<OrderBookDelta>> = AHashMap::new();
1456 buffered.insert(instrument_id, Vec::new());
1457
1458 let flushed = flush_mbo_event_boundary(
1459 instrument_id,
1460 2.into(),
1461 RecordFlag::F_LAST as u8,
1462 &mut buffering_start,
1463 &mut buffered,
1464 );
1465
1466 assert!(flushed.is_none());
1467 assert_eq!(buffering_start, Some(start));
1468 }
1469
1470 #[rstest]
1471 fn test_empty_boundary_preserves_replay_gate_for_buffered_instrument() {
1472 let empty_id = InstrumentId::from("EMPTY.GLBX");
1473 let buffered_id = InstrumentId::from("BUFFERED.GLBX");
1474 let start = UnixNanos::from(10);
1475 let mut buffering_start = Some(start);
1476 let mut buffered = AHashMap::new();
1477 let _ = process_mbo_delta(
1478 stub_delta(buffered_id, 1),
1479 0,
1480 &mut buffering_start,
1481 &mut buffered,
1482 );
1483
1484 let empty_boundary = flush_mbo_event_boundary(
1485 empty_id,
1486 11.into(),
1487 RecordFlag::F_LAST as u8,
1488 &mut buffering_start,
1489 &mut buffered,
1490 );
1491 let buffered_boundary = flush_mbo_event_boundary(
1492 buffered_id,
1493 5.into(),
1494 RecordFlag::F_LAST as u8,
1495 &mut buffering_start,
1496 &mut buffered,
1497 );
1498
1499 assert!(empty_boundary.is_none());
1500 assert!(buffered_boundary.is_none());
1501 assert_eq!(buffering_start, Some(start));
1502 assert!(buffered.contains_key(&buffered_id));
1503 }
1504
1505 #[rstest]
1506 fn test_boundary_flag_respects_replay_buffering_gate() {
1507 let instrument_id = InstrumentId::from("TEST.GLBX");
1508 let mut buffering_start = Some(UnixNanos::from(10));
1509 let mut buffered = AHashMap::new();
1510 let _ = process_mbo_delta(
1511 stub_delta(instrument_id, 1),
1512 0,
1513 &mut buffering_start,
1514 &mut buffered,
1515 );
1516
1517 let flushed = flush_mbo_event_boundary(
1519 instrument_id,
1520 5.into(),
1521 RecordFlag::F_LAST as u8,
1522 &mut buffering_start,
1523 &mut buffered,
1524 );
1525 assert!(flushed.is_none());
1526 assert!(!buffered.is_empty());
1527
1528 let flushed = flush_mbo_event_boundary(
1530 instrument_id,
1531 11.into(),
1532 RecordFlag::F_LAST as u8,
1533 &mut buffering_start,
1534 &mut buffered,
1535 );
1536 assert!(flushed.is_some());
1537 assert!(buffering_start.is_none());
1538 }
1539
1540 #[rstest]
1541 fn test_boundary_snapshot_flag_never_flushes() {
1542 let instrument_id = InstrumentId::from("TEST.GLBX");
1543 let mut buffering_start = None;
1544 let mut buffered = AHashMap::new();
1545 let _ = process_mbo_delta(
1546 stub_delta(instrument_id, 1),
1547 0,
1548 &mut buffering_start,
1549 &mut buffered,
1550 );
1551
1552 let flags = RecordFlag::F_LAST as u8 | RecordFlag::F_SNAPSHOT as u8;
1553 let flushed = flush_mbo_event_boundary(
1554 instrument_id,
1555 2.into(),
1556 flags,
1557 &mut buffering_start,
1558 &mut buffered,
1559 );
1560 assert!(flushed.is_none());
1561 }
1562
1563 fn create_test_handler(reconnect_timeout_mins: Option<u64>) -> DatabentoFeedHandler {
1564 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1565 let (msg_tx, _msg_rx) = tokio::sync::mpsc::unbounded_channel();
1566
1567 DatabentoFeedHandler::new(
1568 Credential::new("test_key"),
1569 "GLBX.MDP3".to_string(),
1570 cmd_rx,
1571 msg_tx,
1572 IndexMap::new(),
1573 Arc::new(AtomicMap::new()),
1574 false,
1575 false,
1576 reconnect_timeout_mins,
1577 )
1578 }
1579
1580 fn create_test_client() -> DatabentoLiveClient {
1581 DatabentoLiveClient::new(
1582 "test-api-key".to_string(),
1583 "GLBX.MDP3".to_string(),
1584 PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("publishers.json"),
1585 true,
1586 None,
1587 None,
1588 )
1589 .unwrap()
1590 }
1591
1592 #[rstest]
1593 #[case(Some(10))]
1594 #[case(None)]
1595 fn test_backoff_initialization(#[case] reconnect_timeout_mins: Option<u64>) {
1596 let handler = create_test_handler(reconnect_timeout_mins);
1597
1598 assert_eq!(handler.reconnect_timeout_mins, reconnect_timeout_mins);
1599 assert!(handler.subscriptions.is_empty());
1600 assert!(handler.buffered_commands.is_empty());
1601 }
1602
1603 #[rstest]
1604 fn test_subscription_with_and_without_start() {
1605 let start_time = datetime!(2024-01-01 00:00:00 UTC);
1606 let sub_with_start = Subscription::builder()
1607 .symbols("ES.FUT")
1608 .schema(databento::dbn::Schema::Mbp1)
1609 .start(start_time)
1610 .build();
1611
1612 let mut sub_without_start = sub_with_start.clone();
1613 sub_without_start.start = None;
1614
1615 assert!(sub_with_start.start.is_some());
1616 assert!(sub_without_start.start.is_none());
1617 assert_eq!(sub_with_start.schema, sub_without_start.schema);
1618 assert_eq!(sub_with_start.symbols, sub_without_start.symbols);
1619 }
1620
1621 #[rstest]
1622 fn test_handler_initialization_state() {
1623 let handler = create_test_handler(Some(10));
1624
1625 assert!(!handler.replay);
1626 assert_eq!(handler.dataset, "GLBX.MDP3");
1627 assert_eq!(handler.credential.api_key(), "test_key");
1628 assert!(handler.subscriptions.is_empty());
1629 assert!(handler.buffered_commands.is_empty());
1630 }
1631
1632 #[rstest]
1633 fn test_handler_with_no_timeout() {
1634 let handler = create_test_handler(None);
1635
1636 assert_eq!(handler.reconnect_timeout_mins, None);
1637 assert!(!handler.replay);
1638 }
1639
1640 #[rstest]
1641 fn test_handler_with_zero_timeout() {
1642 let handler = create_test_handler(Some(0));
1643
1644 assert_eq!(handler.reconnect_timeout_mins, Some(0));
1645 assert!(!handler.replay);
1646 }
1647
1648 #[rstest]
1649 fn test_subscribe_uses_explicit_parent_stype() {
1650 let mut client = create_test_client();
1651
1652 client
1653 .subscribe(
1654 "definition".to_string(),
1655 vec![InstrumentId::from("ES.FUT.GLBX")],
1656 None,
1657 None,
1658 None,
1659 Some("parent".to_string()),
1660 )
1661 .unwrap();
1662
1663 let command = client.cmd_rx.as_mut().unwrap().try_recv().unwrap();
1664 match command {
1665 HandlerCommand::Subscribe(sub) => {
1666 assert_eq!(sub.schema, dbn::Schema::Definition);
1667 assert_eq!(sub.stype_in, dbn::SType::Parent);
1668 assert_eq!(sub.symbols.to_api_string(), "ES.FUT");
1669 }
1670 other => panic!("expected HandlerCommand::Subscribe, was {other:?}"),
1671 }
1672 }
1673
1674 #[rstest]
1675 fn test_subscribe_rejects_invalid_stype() {
1676 let mut client = create_test_client();
1677
1678 let err = client
1679 .subscribe(
1680 "definition".to_string(),
1681 vec![InstrumentId::from("ES.FUT.GLBX")],
1682 None,
1683 None,
1684 None,
1685 Some("not-a-stype".to_string()),
1686 )
1687 .unwrap_err();
1688
1689 assert!(err.to_string().contains("not-a-stype"));
1690 assert!(!is_command_send_error(&err));
1691 assert!(matches!(
1692 client.cmd_rx.as_mut().unwrap().try_recv(),
1693 Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1694 ));
1695 }
1696
1697 #[rstest]
1698 fn test_subscribe_classifies_command_send_errors() {
1699 let mut client = create_test_client();
1700 client.cmd_rx = None;
1701
1702 let err = client
1703 .subscribe(
1704 "definition".to_string(),
1705 vec![InstrumentId::from("ES.FUT.GLBX")],
1706 None,
1707 None,
1708 None,
1709 Some("parent".to_string()),
1710 )
1711 .unwrap_err();
1712
1713 assert!(is_command_send_error(&err));
1714 }
1715
1716 #[rstest]
1717 fn test_close_after_handler_exit_marks_closed() {
1718 let mut client = create_test_client();
1719 let (handler, _msg_rx) = client.start().unwrap();
1720 drop(handler);
1721
1722 client.close().unwrap();
1723
1724 assert!(!client.is_running());
1725 assert!(client.is_closed());
1726 }
1727
1728 fn test_delta(instrument_id: InstrumentId, ts_event: u64) -> OrderBookDelta {
1729 OrderBookDelta::clear(instrument_id, 0, ts_event.into(), 0.into())
1730 }
1731
1732 #[rstest]
1733 fn test_mbo_delta_without_f_last_buffers() {
1734 let instrument_id = InstrumentId::from("ESM4.GLBX");
1735 let delta = test_delta(instrument_id, 1_000_000_000);
1736 let mut buffering_start = None;
1737 let mut buffered = AHashMap::new();
1738
1739 let result = process_mbo_delta(delta, 0, &mut buffering_start, &mut buffered);
1740
1741 assert!(result.is_none());
1742 assert_eq!(buffered[&instrument_id].len(), 1);
1743 }
1744
1745 #[rstest]
1746 fn test_mbo_single_f_last_emits_without_buffering() {
1747 let instrument_id = InstrumentId::from("ESM4.GLBX");
1748 let ts_event = 1_000_000_000;
1749 let mut delta = test_delta(instrument_id, ts_event);
1750 delta.flags = 128;
1751 delta.sequence = 42;
1752 let mut buffering_start = None;
1753 let mut buffered = AHashMap::new();
1754
1755 let result = process_mbo_delta(delta, delta.flags, &mut buffering_start, &mut buffered);
1756
1757 let emitted = result.expect("single F_LAST delta should emit");
1758 assert_eq!(emitted.instrument_id, instrument_id);
1759 assert_eq!(emitted.deltas.len(), 1);
1760 assert_eq!(emitted.flags, 128);
1761 assert_eq!(emitted.sequence, 42);
1762 assert_eq!(emitted.ts_event, UnixNanos::from(ts_event));
1763 assert_eq!(emitted.deltas[0].instrument_id, instrument_id);
1764 assert_eq!(emitted.deltas[0].flags, 128);
1765 assert_eq!(emitted.deltas[0].sequence, 42);
1766 assert_eq!(emitted.deltas[0].ts_event, UnixNanos::from(ts_event));
1767 assert!(buffering_start.is_none());
1768 assert!(buffered.is_empty());
1769 }
1770
1771 #[rstest]
1772 fn test_mbo_delta_with_f_last_emits() {
1773 let instrument_id = InstrumentId::from("ESM4.GLBX");
1774 let mut buffering_start = None;
1775 let mut buffered = AHashMap::new();
1776
1777 let _ = process_mbo_delta(
1778 test_delta(instrument_id, 1_000_000_000),
1779 0,
1780 &mut buffering_start,
1781 &mut buffered,
1782 );
1783
1784 let result = process_mbo_delta(
1785 test_delta(instrument_id, 2_000_000_000),
1786 128, &mut buffering_start,
1788 &mut buffered,
1789 );
1790
1791 assert!(result.is_some());
1792 assert_eq!(result.unwrap().deltas.len(), 2);
1793 assert!(buffered.is_empty());
1794 }
1795
1796 #[rstest]
1797 fn test_mbo_snapshot_with_f_last_buffers() {
1798 let instrument_id = InstrumentId::from("ESM4.GLBX");
1799 let mut buffering_start = None;
1800 let mut buffered = AHashMap::new();
1801
1802 let result = process_mbo_delta(
1803 test_delta(instrument_id, 1_000_000_000),
1804 128 | 32, &mut buffering_start,
1806 &mut buffered,
1807 );
1808
1809 assert!(result.is_none());
1810 assert_eq!(buffered[&instrument_id].len(), 1);
1811 }
1812
1813 #[rstest]
1814 fn test_mbo_replay_buffers_until_past_start() {
1815 let instrument_id = InstrumentId::from("ESM4.GLBX");
1816 let start_ns = 5_000_000_000u64;
1817 let mut buffering_start = Some(start_ns.into());
1818 let mut buffered = AHashMap::new();
1819
1820 let result = process_mbo_delta(
1821 test_delta(instrument_id, 4_000_000_000),
1822 128, &mut buffering_start,
1824 &mut buffered,
1825 );
1826 assert!(result.is_none());
1827
1828 let result = process_mbo_delta(
1829 test_delta(instrument_id, 5_000_000_000),
1830 128,
1831 &mut buffering_start,
1832 &mut buffered,
1833 );
1834 assert!(result.is_none());
1835
1836 let result = process_mbo_delta(
1838 test_delta(instrument_id, 6_000_000_000),
1839 128,
1840 &mut buffering_start,
1841 &mut buffered,
1842 );
1843 assert!(result.is_some());
1844 assert!(buffering_start.is_none());
1845 }
1846
1847 #[rstest]
1848 fn test_mbo_multiple_deltas_accumulated() {
1849 let instrument_id = InstrumentId::from("ESM4.GLBX");
1850 let mut buffering_start = None;
1851 let mut buffered = AHashMap::new();
1852
1853 for i in 0..5 {
1854 process_mbo_delta(
1855 test_delta(instrument_id, 1_000_000_000 + i),
1856 0,
1857 &mut buffering_start,
1858 &mut buffered,
1859 );
1860 }
1861
1862 let result = process_mbo_delta(
1863 test_delta(instrument_id, 2_000_000_000),
1864 128,
1865 &mut buffering_start,
1866 &mut buffered,
1867 );
1868
1869 assert!(result.is_some());
1870 assert_eq!(result.unwrap().deltas.len(), 6);
1871 }
1872
1873 #[rstest]
1874 fn test_mbo_multi_instrument_isolation() {
1875 let id_a = InstrumentId::from("ESM4.GLBX");
1876 let id_b = InstrumentId::from("NQM4.GLBX");
1877 let mut buffering_start = None;
1878 let mut buffered = AHashMap::new();
1879
1880 process_mbo_delta(
1881 test_delta(id_a, 1_000_000_000),
1882 0,
1883 &mut buffering_start,
1884 &mut buffered,
1885 );
1886 process_mbo_delta(
1887 test_delta(id_b, 1_000_000_000),
1888 0,
1889 &mut buffering_start,
1890 &mut buffered,
1891 );
1892
1893 let result = process_mbo_delta(
1895 test_delta(id_a, 2_000_000_000),
1896 128,
1897 &mut buffering_start,
1898 &mut buffered,
1899 );
1900
1901 assert!(result.is_some());
1902 assert_eq!(result.unwrap().instrument_id, id_a);
1903 assert!(buffered.contains_key(&id_b));
1904 assert!(!buffered.contains_key(&id_a));
1905 }
1906
1907 mod property_tests {
1908 use proptest::prelude::*;
1909 use rstest::rstest;
1910
1911 use super::*;
1912
1913 proptest! {
1914 #[rstest]
1915 fn mbo_buffering_conserves_deltas(
1916 num_non_last in 0usize..=20,
1917 ) {
1918 let instrument_id = InstrumentId::from("ESM4.GLBX");
1919 let mut buffering_start = None;
1920 let mut buffered = AHashMap::new();
1921 let total = num_non_last + 1;
1922
1923 for i in 0..num_non_last {
1924 let result = process_mbo_delta(
1925 test_delta(instrument_id, 1_000_000_000 + i as u64),
1926 0, &mut buffering_start,
1928 &mut buffered,
1929 );
1930 prop_assert!(result.is_none());
1931 }
1932
1933 let result = process_mbo_delta(
1934 test_delta(instrument_id, 2_000_000_000),
1935 128, &mut buffering_start,
1937 &mut buffered,
1938 );
1939
1940 prop_assert!(result.is_some());
1941 let emitted = result.unwrap();
1942 prop_assert_eq!(emitted.deltas.len(), total);
1943 prop_assert!(buffered.is_empty());
1944 }
1945
1946 #[rstest]
1947 fn mbo_snapshots_never_emit(
1948 num_snapshots in 1usize..=20,
1949 ) {
1950 let instrument_id = InstrumentId::from("ESM4.GLBX");
1951 let mut buffering_start = None;
1952 let mut buffered = AHashMap::new();
1953
1954 for i in 0..num_snapshots {
1955 let result = process_mbo_delta(
1956 test_delta(instrument_id, 1_000_000_000 + i as u64),
1957 128 | 32, &mut buffering_start,
1959 &mut buffered,
1960 );
1961 prop_assert!(result.is_none());
1962 }
1963
1964 prop_assert_eq!(buffered[&instrument_id].len(), num_snapshots);
1965 }
1966
1967 #[rstest]
1968 fn mbo_replay_delays_emission(
1969 start_offset in 1u64..=100,
1970 num_before in 1usize..=10,
1971 ) {
1972 let instrument_id = InstrumentId::from("ESM4.GLBX");
1973 let start_ns = 1_000_000_000u64 * start_offset;
1974 let mut buffering_start = Some(start_ns.into());
1975 let mut buffered = AHashMap::new();
1976
1977 for i in 0..num_before {
1978 let ts = start_ns - (num_before as u64 - i as u64);
1979 let result = process_mbo_delta(
1980 test_delta(instrument_id, ts),
1981 128, &mut buffering_start,
1983 &mut buffered,
1984 );
1985 prop_assert!(result.is_none());
1986 }
1987
1988 let result = process_mbo_delta(
1989 test_delta(instrument_id, start_ns + 1),
1990 128,
1991 &mut buffering_start,
1992 &mut buffered,
1993 );
1994
1995 prop_assert!(result.is_some());
1996 prop_assert!(buffering_start.is_none());
1997 let total = num_before + 1;
1998 prop_assert_eq!(result.unwrap().deltas.len(), total);
1999 }
2000 }
2001 }
2002}