Skip to main content

nautilus_databento/
live.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Databento live feed handler.
17//!
18//! The feed handler runs a single async task per dataset. It receives
19//! [`HandlerCommand`] messages over an unbounded channel and streams decoded
20//! market data back as [`DatabentoMessage`]s on an unbounded tokio channel.
21//!
22//! The inner loop uses `tokio::select!` to concurrently await the next record
23//! from the Databento gateway and the next command from the engine, giving
24//! near-zero idle CPU and immediate command responsiveness.
25//!
26//! Heartbeat detection is delegated to the upstream `databento` client, which
27//! returns `Error::HeartbeatTimeout` when no data arrives within
28//! `heartbeat_interval + 5 s` (default 35 s). The handler treats this as a
29//! connection error and enters the reconnection backoff loop.
30
31use 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, OrderBookDeltas_API},
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.core.nautilus_pyo3.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    /// Creates a new [`DatabentoLiveClient`] instance.
126    ///
127    /// # Errors
128    ///
129    /// Returns an error if reading or parsing the publishers file fails.
130    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    /// Subscribes to Databento live data for the requested instruments.
172    ///
173    /// # Errors
174    ///
175    /// Returns an error if symbology, schema, timestamp, or precision inputs are invalid,
176    /// or if the command cannot be sent to the feed handler.
177    #[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    /// Starts the live feed handler and returns its message receiver.
244    ///
245    /// # Errors
246    ///
247    /// Returns an error if the client is already closed, already running, or cannot start.
248    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    /// Closes the live client.
289    ///
290    /// # Errors
291    ///
292    /// Returns an error if the client was never started, is already closed, or cannot send
293    /// the close command to the feed handler.
294    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
346/// Handles a raw TCP data feed from the Databento LSG for a single dataset.
347///
348/// [`HandlerCommand`] messages are received synchronously across a channel,
349/// decoded records are sent asynchronously on a tokio channel as [`DatabentoMessage`]s
350/// back to a message processing task.
351///
352/// # Crash Policy
353///
354/// This handler intentionally avoids applying downstream backpressure to the
355/// live feed. If decoded output cannot be drained, memory pressure is the hard
356/// failure mode instead of arbitrary queue limits or delayed market data.
357pub 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    /// Creates a new [`DatabentoFeedHandler`] instance.
390    ///
391    /// # Panics
392    ///
393    /// Panics if exponential backoff creation fails (should never happen with valid hardcoded parameters).
394    #[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        // Choose max delay based on timeout configuration:
408        // - With timeout: 60s max (quick recovery to reconnect within window)
409        // - Without timeout (None): 600s max (patient recovery, respectful of infrastructure)
410        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    /// Sets a custom gateway address, overriding the default Databento LSG endpoint.
440    #[must_use]
441    pub fn with_gateway_addr(mut self, addr: String) -> Self {
442        self.gateway_addr = Some(addr);
443        self
444    }
445
446    /// Sets the duration a session must run before it counts as successful.
447    ///
448    /// A successful session resets the reconnection backoff cycle.
449    /// Defaults to 60 seconds.
450    #[must_use]
451    pub fn with_success_threshold(mut self, threshold: Duration) -> Self {
452        self.success_threshold = threshold;
453        self
454    }
455
456    /// Runs the feed handler main loop, processing commands and streaming market data.
457    ///
458    /// Establishes a connection to the Databento LSG, subscribes to requested data feeds,
459    /// and continuously processes incoming market data messages until shutdown.
460    ///
461    /// Implements automatic reconnection with exponential backoff (1s to 60s with jitter).
462    /// Each successful session resets the reconnection cycle, giving the next disconnect
463    /// a fresh timeout window. Gives up after `reconnect_timeout_mins` if configured.
464    ///
465    /// # Errors
466    ///
467    /// Returns an error if any client operation or message handling fails.
468    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    /// Runs a single session, handling connection, subscriptions, and data streaming.
544    ///
545    /// Returns `Ok(bool)` where the bool indicates if the session ran successfully
546    /// for a meaningful duration (true) or was intentionally closed (false).
547    ///
548    /// # Errors
549    ///
550    /// Returns an error if connection fails, subscription fails, or data streaming encounters an error.
551    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); // Hardcoded timeout for now
567
568        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        // Process any commands buffered during reconnection backoff
602        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            // Strip start timestamps after successful subscription to avoid replaying history on future reconnects
644            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            // Wait for either a command or a record. When the session has not
668            // started yet (`!running`), only commands are awaited. Once running,
669            // `next_record` is cancel-safe so `tokio::select!` can safely
670            // race both futures.
671            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            // Decode record
779            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                // Remove instrument ID index as the raw symbol may have changed
787                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                // Decode a generic record with possible errors
867                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                    } else {
892                        continue;
893                    }
894
895                    if let Some(Data::Delta(delta)) = &data1 {
896                        log::trace!(
897                            "Buffering delta: {} {buffering_start:?} flags={}",
898                            delta.ts_event,
899                            msg.flags.raw(),
900                        );
901
902                        match process_mbo_delta(
903                            *delta,
904                            msg.flags.raw(),
905                            &mut buffering_start,
906                            &mut buffered_deltas,
907                        )? {
908                            Some(deltas) => data1 = Some(Data::Deltas(deltas)),
909                            None => continue,
910                        }
911                    }
912                }
913
914                if let Some(data) = data1 {
915                    self.send_msg(DatabentoMessage::Data(data));
916                }
917
918                if let Some(data) = data2 {
919                    self.send_msg(DatabentoMessage::Data(data));
920                }
921            }
922        }
923    }
924
925    /// Sends a message to the message processing task.
926    fn send_msg(&self, msg: DatabentoMessage) {
927        log::trace!("Sending {msg:?}");
928        match self.msg_tx.send(msg) {
929            Ok(()) => {}
930            Err(e) => log::error!("Error sending message: {e}"),
931        }
932    }
933
934    fn send_close_msg(&self) {
935        if let Err(e) = self.msg_tx.send(DatabentoMessage::Close) {
936            log::debug!("Could not send close message: {e}");
937        }
938    }
939}
940
941/// Handles Databento error messages by logging them.
942fn handle_error_msg(msg: &dbn::ErrorMsg) {
943    log::error!("{msg:?}");
944}
945
946/// Handles Databento system messages, returning a subscription ack event if applicable.
947fn handle_system_msg(msg: &dbn::SystemMsg, ts_received: UnixNanos) -> Option<SubscriptionAckEvent> {
948    match msg.code() {
949        Ok(dbn::SystemCode::SubscriptionAck) => {
950            let message = msg.msg().unwrap_or("<invalid utf-8>");
951            log::debug!("Subscription acknowledged: {message}");
952
953            let schema = parse_ack_message(message);
954
955            Some(SubscriptionAckEvent {
956                schema,
957                message: message.to_string(),
958                ts_received,
959            })
960        }
961        Ok(dbn::SystemCode::Heartbeat) => {
962            log::trace!("Heartbeat received");
963            None
964        }
965        Ok(dbn::SystemCode::SlowReaderWarning) => {
966            let message = msg.msg().unwrap_or("<invalid utf-8>");
967            log::warn!("Slow reader warning: {message}");
968            None
969        }
970        Ok(dbn::SystemCode::ReplayCompleted) => {
971            let message = msg.msg().unwrap_or("<invalid utf-8>");
972            log::debug!("Replay completed: {message}");
973            None
974        }
975        _ => {
976            log::debug!("{msg:?}");
977            None
978        }
979    }
980}
981
982/// Parses a subscription ack message to extract the schema.
983fn parse_ack_message(message: &str) -> String {
984    // Format: "Subscription request N for <schema> data succeeded"
985    message
986        .strip_prefix("Subscription request ")
987        .and_then(|rest| rest.split_once(" for "))
988        .and_then(|(_, after_num)| after_num.strip_suffix(" data succeeded"))
989        .map(|schema| schema.trim().to_string())
990        .unwrap_or_default()
991}
992
993/// Handles symbol mapping messages and updates the instrument ID map.
994///
995/// # Errors
996///
997/// Returns an error if symbol mapping fails.
998fn handle_symbol_mapping_msg(
999    msg: &dbn::SymbolMappingMsg,
1000    symbol_map: &mut PitSymbolMap,
1001    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1002) -> anyhow::Result<()> {
1003    symbol_map
1004        .on_symbol_mapping(msg)
1005        .map_err(|e| anyhow::anyhow!("on_symbol_mapping failed for {msg:?}: {e}"))?;
1006    instrument_id_map.remove(&msg.header().instrument_id);
1007    Ok(())
1008}
1009
1010fn update_price_precision_map_with_symbol_mapping_msg(
1011    msg: &dbn::SymbolMappingMsg,
1012    price_precision_overrides: &AHashMap<Symbol, u8>,
1013    subscription_price_precision_map: &mut AHashMap<u32, u8>,
1014) -> anyhow::Result<()> {
1015    subscription_price_precision_map.remove(&msg.hd.instrument_id);
1016
1017    if price_precision_overrides.is_empty() {
1018        return Ok(());
1019    }
1020
1021    let stype_in_symbol = msg
1022        .stype_in_symbol()
1023        .map_err(|e| anyhow::anyhow!("Error decoding `stype_in_symbol`: {e}"))?;
1024    let stype_out_symbol = msg
1025        .stype_out_symbol()
1026        .map_err(|e| anyhow::anyhow!("Error decoding `stype_out_symbol`: {e}"))?;
1027
1028    let price_precision = [stype_in_symbol, stype_out_symbol]
1029        .into_iter()
1030        .find_map(|symbol| {
1031            price_precision_overrides
1032                .get(&Symbol::from_str_unchecked(symbol))
1033                .copied()
1034        });
1035
1036    if let Some(price_precision) = price_precision {
1037        subscription_price_precision_map.insert(msg.hd.instrument_id, price_precision);
1038    }
1039
1040    Ok(())
1041}
1042
1043/// Updates the instrument ID map using exchange information from the symbol map.
1044fn update_instrument_id_map_with_exchange(
1045    symbol_map: &PitSymbolMap,
1046    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1047    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1048    raw_instrument_id: u32,
1049    exchange: &str,
1050) -> anyhow::Result<InstrumentId> {
1051    let raw_symbol = symbol_map.get(raw_instrument_id).ok_or_else(|| {
1052        anyhow::anyhow!("Cannot resolve raw_symbol for instrument_id {raw_instrument_id}")
1053    })?;
1054    let symbol = Symbol::from(raw_symbol.as_str());
1055    let venue = Venue::from_code(exchange)
1056        .map_err(|e| anyhow::anyhow!("Invalid venue code '{exchange}': {e}"))?;
1057    let instrument_id = InstrumentId::new(symbol, venue);
1058    symbol_venue_map.rcu(|m| {
1059        m.entry(symbol).or_insert(venue);
1060    });
1061    instrument_id_map.insert(raw_instrument_id, instrument_id);
1062    Ok(instrument_id)
1063}
1064
1065fn update_instrument_id_map(
1066    record: &dbn::RecordRef,
1067    symbol_map: &PitSymbolMap,
1068    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1069    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1070    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1071) -> anyhow::Result<InstrumentId> {
1072    let header = record.header();
1073
1074    // Check if instrument ID is already in the map
1075    if let Some(&instrument_id) = instrument_id_map.get(&header.instrument_id) {
1076        return Ok(instrument_id);
1077    }
1078
1079    let raw_symbol = symbol_map.get_for_rec(record).ok_or_else(|| {
1080        anyhow::anyhow!(
1081            "Cannot resolve `raw_symbol` from `symbol_map` for instrument_id {}",
1082            header.instrument_id
1083        )
1084    })?;
1085
1086    let symbol = Symbol::from_str_unchecked(raw_symbol);
1087
1088    let publisher_id = header.publisher_id;
1089    let venue = if let Some(venue) = symbol_venue_map.get_cloned(&symbol) {
1090        venue
1091    } else {
1092        let venue = publisher_venue_map
1093            .get(&publisher_id)
1094            .ok_or_else(|| anyhow::anyhow!("No venue found for `publisher_id` {publisher_id}"))?;
1095        *venue
1096    };
1097    let instrument_id = InstrumentId::new(symbol, venue);
1098
1099    instrument_id_map.insert(header.instrument_id, instrument_id);
1100    Ok(instrument_id)
1101}
1102
1103/// Handles instrument definition messages and decodes them into Nautilus instruments.
1104///
1105/// # Errors
1106///
1107/// Returns an error if instrument decoding fails.
1108fn handle_instrument_def_msg(
1109    msg: &dbn::InstrumentDefMsg,
1110    record: &dbn::RecordRef,
1111    symbol_map: &PitSymbolMap,
1112    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1113    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1114    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1115    ts_init: UnixNanos,
1116) -> anyhow::Result<Option<InstrumentAny>> {
1117    let instrument_id = update_instrument_id_map(
1118        record,
1119        symbol_map,
1120        publisher_venue_map,
1121        symbol_venue_map,
1122        instrument_id_map,
1123    )?;
1124
1125    decode_instrument_def_msg(msg, instrument_id, Some(ts_init), None)
1126}
1127
1128fn handle_status_msg(
1129    msg: &dbn::StatusMsg,
1130    record: &dbn::RecordRef,
1131    symbol_map: &PitSymbolMap,
1132    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1133    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1134    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1135    ts_init: UnixNanos,
1136) -> anyhow::Result<InstrumentStatus> {
1137    let instrument_id = update_instrument_id_map(
1138        record,
1139        symbol_map,
1140        publisher_venue_map,
1141        symbol_venue_map,
1142        instrument_id_map,
1143    )?;
1144
1145    decode_status_msg(msg, instrument_id, Some(ts_init))
1146}
1147
1148#[expect(clippy::too_many_arguments)]
1149fn handle_imbalance_msg(
1150    msg: &dbn::ImbalanceMsg,
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    instrument_def_price_precision_map: &AHashMap<u32, u8>,
1157    subscription_price_precision_map: &AHashMap<u32, u8>,
1158    price_precision_overrides: &AHashMap<Symbol, u8>,
1159    ts_init: UnixNanos,
1160) -> anyhow::Result<DatabentoImbalance> {
1161    let instrument_id = update_instrument_id_map(
1162        record,
1163        symbol_map,
1164        publisher_venue_map,
1165        symbol_venue_map,
1166        instrument_id_map,
1167    )?;
1168
1169    let price_precision = resolve_price_precision(
1170        msg.hd.instrument_id,
1171        instrument_id,
1172        instrument_def_price_precision_map,
1173        subscription_price_precision_map,
1174        price_precision_overrides,
1175    );
1176
1177    decode_imbalance_msg(msg, instrument_id, price_precision, Some(ts_init))
1178}
1179
1180#[expect(clippy::too_many_arguments)]
1181fn handle_statistics_msg(
1182    msg: &dbn::StatMsg,
1183    record: &dbn::RecordRef,
1184    symbol_map: &PitSymbolMap,
1185    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1186    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1187    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1188    instrument_def_price_precision_map: &AHashMap<u32, u8>,
1189    subscription_price_precision_map: &AHashMap<u32, u8>,
1190    price_precision_overrides: &AHashMap<Symbol, u8>,
1191    ts_init: UnixNanos,
1192) -> anyhow::Result<Option<DatabentoStatistics>> {
1193    // Precheck before symbol resolution so unmodeled types skip cleanly
1194    if !is_supported_stat_type(msg.stat_type) {
1195        log::warn!("Skipping unsupported `stat_type` {}", msg.stat_type);
1196        return Ok(None);
1197    }
1198
1199    let instrument_id = update_instrument_id_map(
1200        record,
1201        symbol_map,
1202        publisher_venue_map,
1203        symbol_venue_map,
1204        instrument_id_map,
1205    )?;
1206
1207    let price_precision = resolve_price_precision(
1208        msg.hd.instrument_id,
1209        instrument_id,
1210        instrument_def_price_precision_map,
1211        subscription_price_precision_map,
1212        price_precision_overrides,
1213    );
1214
1215    decode_statistics_msg(msg, instrument_id, price_precision, Some(ts_init))
1216}
1217
1218#[expect(clippy::too_many_arguments)]
1219fn handle_record(
1220    record: dbn::RecordRef,
1221    symbol_map: &PitSymbolMap,
1222    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1223    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1224    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1225    instrument_def_price_precision_map: &AHashMap<u32, u8>,
1226    subscription_price_precision_map: &AHashMap<u32, u8>,
1227    price_precision_overrides: &AHashMap<Symbol, u8>,
1228    ts_init: UnixNanos,
1229    initialized_books: &HashSet<InstrumentId>,
1230    bars_timestamp_on_close: bool,
1231) -> anyhow::Result<(Option<Data>, Option<Data>)> {
1232    let instrument_id = update_instrument_id_map(
1233        &record,
1234        symbol_map,
1235        publisher_venue_map,
1236        symbol_venue_map,
1237        instrument_id_map,
1238    )?;
1239
1240    let price_precision = resolve_price_precision(
1241        record.header().instrument_id,
1242        instrument_id,
1243        instrument_def_price_precision_map,
1244        subscription_price_precision_map,
1245        price_precision_overrides,
1246    );
1247
1248    // For MBP-1 and quote-based schemas, always include trades since they're integral to the data
1249    // For MBO, only include trades after the book is initialized to maintain consistency
1250    let include_trades = if record.get::<dbn::Mbp1Msg>().is_some()
1251        || record.get::<dbn::TbboMsg>().is_some()
1252        || record.get::<dbn::Cmbp1Msg>().is_some()
1253    {
1254        true // These schemas include trade information directly
1255    } else {
1256        initialized_books.contains(&instrument_id) // MBO requires initialized book
1257    };
1258
1259    decode_record(
1260        &record,
1261        instrument_id,
1262        price_precision,
1263        Some(ts_init),
1264        include_trades,
1265        bars_timestamp_on_close,
1266    )
1267}
1268
1269fn resolve_price_precision(
1270    record_instrument_id: u32,
1271    instrument_id: InstrumentId,
1272    instrument_def_price_precision_map: &AHashMap<u32, u8>,
1273    subscription_price_precision_map: &AHashMap<u32, u8>,
1274    price_precision_overrides: &AHashMap<Symbol, u8>,
1275) -> u8 {
1276    instrument_def_price_precision_map
1277        .get(&record_instrument_id)
1278        .copied()
1279        .or_else(|| {
1280            subscription_price_precision_map
1281                .get(&record_instrument_id)
1282                .copied()
1283        })
1284        .or_else(|| {
1285            price_precision_overrides
1286                .get(&instrument_id.symbol)
1287                .copied()
1288        })
1289        .unwrap_or(Currency::USD().precision)
1290}
1291
1292/// Processes an MBO delta through the buffering state machine.
1293///
1294/// Returns `Some(deltas)` when a complete batch is ready to emit (non-snapshot
1295/// F_LAST with replay buffering complete), or `None` when still accumulating.
1296fn process_mbo_delta(
1297    delta: OrderBookDelta,
1298    flags: u8,
1299    buffering_start: &mut Option<UnixNanos>,
1300    buffered_deltas: &mut AHashMap<InstrumentId, Vec<OrderBookDelta>>,
1301) -> anyhow::Result<Option<OrderBookDeltas_API>> {
1302    let is_last = RecordFlag::F_LAST.matches(flags);
1303    let is_snapshot = RecordFlag::F_SNAPSHOT.matches(flags);
1304
1305    // Most live MBO events are single non-snapshot deltas, avoid map churn on that path
1306    if is_last
1307        && !is_snapshot
1308        && buffering_start.is_none()
1309        && !buffered_deltas.contains_key(&delta.instrument_id)
1310    {
1311        let deltas = OrderBookDeltas::new(delta.instrument_id, vec![delta]);
1312        return Ok(Some(OrderBookDeltas_API::new(deltas)));
1313    }
1314
1315    let buffer = buffered_deltas.entry(delta.instrument_id).or_default();
1316    buffer.push(delta);
1317
1318    if !is_last {
1319        return Ok(None);
1320    }
1321
1322    if is_snapshot {
1323        return Ok(None);
1324    }
1325
1326    if let Some(start_ns) = *buffering_start {
1327        if delta.ts_event <= start_ns {
1328            return Ok(None);
1329        }
1330        *buffering_start = None;
1331    }
1332
1333    let buffer = buffered_deltas
1334        .remove(&delta.instrument_id)
1335        .ok_or_else(|| {
1336            anyhow::anyhow!(
1337                "Internal error: no buffered deltas for instrument {id}",
1338                id = delta.instrument_id
1339            )
1340        })?;
1341    let deltas = OrderBookDeltas::new(delta.instrument_id, buffer);
1342    Ok(Some(OrderBookDeltas_API::new(deltas)))
1343}
1344
1345#[cfg(test)]
1346mod tests {
1347    use std::path::PathBuf;
1348
1349    use databento::live::Subscription;
1350    use indexmap::IndexMap;
1351    use rstest::*;
1352    use time::macros::datetime;
1353
1354    use super::*;
1355
1356    fn create_test_handler(reconnect_timeout_mins: Option<u64>) -> DatabentoFeedHandler {
1357        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1358        let (msg_tx, _msg_rx) = tokio::sync::mpsc::unbounded_channel();
1359
1360        DatabentoFeedHandler::new(
1361            Credential::new("test_key"),
1362            "GLBX.MDP3".to_string(),
1363            cmd_rx,
1364            msg_tx,
1365            IndexMap::new(),
1366            Arc::new(AtomicMap::new()),
1367            false,
1368            false,
1369            reconnect_timeout_mins,
1370        )
1371    }
1372
1373    fn create_test_client() -> DatabentoLiveClient {
1374        DatabentoLiveClient::new(
1375            "test-api-key".to_string(),
1376            "GLBX.MDP3".to_string(),
1377            PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("publishers.json"),
1378            true,
1379            None,
1380            None,
1381        )
1382        .unwrap()
1383    }
1384
1385    #[rstest]
1386    #[case(Some(10))]
1387    #[case(None)]
1388    fn test_backoff_initialization(#[case] reconnect_timeout_mins: Option<u64>) {
1389        let handler = create_test_handler(reconnect_timeout_mins);
1390
1391        assert_eq!(handler.reconnect_timeout_mins, reconnect_timeout_mins);
1392        assert!(handler.subscriptions.is_empty());
1393        assert!(handler.buffered_commands.is_empty());
1394    }
1395
1396    #[rstest]
1397    fn test_subscription_with_and_without_start() {
1398        let start_time = datetime!(2024-01-01 00:00:00 UTC);
1399        let sub_with_start = Subscription::builder()
1400            .symbols("ES.FUT")
1401            .schema(databento::dbn::Schema::Mbp1)
1402            .start(start_time)
1403            .build();
1404
1405        let mut sub_without_start = sub_with_start.clone();
1406        sub_without_start.start = None;
1407
1408        assert!(sub_with_start.start.is_some());
1409        assert!(sub_without_start.start.is_none());
1410        assert_eq!(sub_with_start.schema, sub_without_start.schema);
1411        assert_eq!(sub_with_start.symbols, sub_without_start.symbols);
1412    }
1413
1414    #[rstest]
1415    fn test_handler_initialization_state() {
1416        let handler = create_test_handler(Some(10));
1417
1418        assert!(!handler.replay);
1419        assert_eq!(handler.dataset, "GLBX.MDP3");
1420        assert_eq!(handler.credential.api_key(), "test_key");
1421        assert!(handler.subscriptions.is_empty());
1422        assert!(handler.buffered_commands.is_empty());
1423    }
1424
1425    #[rstest]
1426    fn test_handler_with_no_timeout() {
1427        let handler = create_test_handler(None);
1428
1429        assert_eq!(handler.reconnect_timeout_mins, None);
1430        assert!(!handler.replay);
1431    }
1432
1433    #[rstest]
1434    fn test_handler_with_zero_timeout() {
1435        let handler = create_test_handler(Some(0));
1436
1437        assert_eq!(handler.reconnect_timeout_mins, Some(0));
1438        assert!(!handler.replay);
1439    }
1440
1441    #[rstest]
1442    fn test_subscribe_uses_explicit_parent_stype() {
1443        let mut client = create_test_client();
1444
1445        client
1446            .subscribe(
1447                "definition".to_string(),
1448                vec![InstrumentId::from("ES.FUT.GLBX")],
1449                None,
1450                None,
1451                None,
1452                Some("parent".to_string()),
1453            )
1454            .unwrap();
1455
1456        let command = client.cmd_rx.as_mut().unwrap().try_recv().unwrap();
1457        match command {
1458            HandlerCommand::Subscribe(sub) => {
1459                assert_eq!(sub.schema, dbn::Schema::Definition);
1460                assert_eq!(sub.stype_in, dbn::SType::Parent);
1461                assert_eq!(sub.symbols.to_api_string(), "ES.FUT");
1462            }
1463            other => panic!("expected HandlerCommand::Subscribe, was {other:?}"),
1464        }
1465    }
1466
1467    #[rstest]
1468    fn test_subscribe_rejects_invalid_stype() {
1469        let mut client = create_test_client();
1470
1471        let err = client
1472            .subscribe(
1473                "definition".to_string(),
1474                vec![InstrumentId::from("ES.FUT.GLBX")],
1475                None,
1476                None,
1477                None,
1478                Some("not-a-stype".to_string()),
1479            )
1480            .unwrap_err();
1481
1482        assert!(err.to_string().contains("not-a-stype"));
1483        assert!(!is_command_send_error(&err));
1484        assert!(matches!(
1485            client.cmd_rx.as_mut().unwrap().try_recv(),
1486            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1487        ));
1488    }
1489
1490    #[rstest]
1491    fn test_subscribe_classifies_command_send_errors() {
1492        let mut client = create_test_client();
1493        client.cmd_rx = None;
1494
1495        let err = client
1496            .subscribe(
1497                "definition".to_string(),
1498                vec![InstrumentId::from("ES.FUT.GLBX")],
1499                None,
1500                None,
1501                None,
1502                Some("parent".to_string()),
1503            )
1504            .unwrap_err();
1505
1506        assert!(is_command_send_error(&err));
1507    }
1508
1509    #[rstest]
1510    fn test_close_after_handler_exit_marks_closed() {
1511        let mut client = create_test_client();
1512        let (handler, _msg_rx) = client.start().unwrap();
1513        drop(handler);
1514
1515        client.close().unwrap();
1516
1517        assert!(!client.is_running());
1518        assert!(client.is_closed());
1519    }
1520
1521    fn test_delta(instrument_id: InstrumentId, ts_event: u64) -> OrderBookDelta {
1522        OrderBookDelta::clear(instrument_id, 0, ts_event.into(), 0.into())
1523    }
1524
1525    #[rstest]
1526    fn test_mbo_delta_without_f_last_buffers() {
1527        let instrument_id = InstrumentId::from("ESM4.GLBX");
1528        let delta = test_delta(instrument_id, 1_000_000_000);
1529        let mut buffering_start = None;
1530        let mut buffered = AHashMap::new();
1531
1532        let result = process_mbo_delta(delta, 0, &mut buffering_start, &mut buffered).unwrap();
1533
1534        assert!(result.is_none());
1535        assert_eq!(buffered[&instrument_id].len(), 1);
1536    }
1537
1538    #[rstest]
1539    fn test_mbo_single_f_last_emits_without_buffering() {
1540        let instrument_id = InstrumentId::from("ESM4.GLBX");
1541        let ts_event = 1_000_000_000;
1542        let mut delta = test_delta(instrument_id, ts_event);
1543        delta.flags = 128;
1544        delta.sequence = 42;
1545        let mut buffering_start = None;
1546        let mut buffered = AHashMap::new();
1547
1548        let result =
1549            process_mbo_delta(delta, delta.flags, &mut buffering_start, &mut buffered).unwrap();
1550
1551        let emitted = result.expect("single F_LAST delta should emit");
1552        assert_eq!(emitted.instrument_id, instrument_id);
1553        assert_eq!(emitted.deltas.len(), 1);
1554        assert_eq!(emitted.flags, 128);
1555        assert_eq!(emitted.sequence, 42);
1556        assert_eq!(emitted.ts_event, UnixNanos::from(ts_event));
1557        assert_eq!(emitted.deltas[0].instrument_id, instrument_id);
1558        assert_eq!(emitted.deltas[0].flags, 128);
1559        assert_eq!(emitted.deltas[0].sequence, 42);
1560        assert_eq!(emitted.deltas[0].ts_event, UnixNanos::from(ts_event));
1561        assert!(buffering_start.is_none());
1562        assert!(buffered.is_empty());
1563    }
1564
1565    #[rstest]
1566    fn test_mbo_delta_with_f_last_emits() {
1567        let instrument_id = InstrumentId::from("ESM4.GLBX");
1568        let mut buffering_start = None;
1569        let mut buffered = AHashMap::new();
1570
1571        process_mbo_delta(
1572            test_delta(instrument_id, 1_000_000_000),
1573            0,
1574            &mut buffering_start,
1575            &mut buffered,
1576        )
1577        .unwrap();
1578
1579        let result = process_mbo_delta(
1580            test_delta(instrument_id, 2_000_000_000),
1581            128, // F_LAST
1582            &mut buffering_start,
1583            &mut buffered,
1584        )
1585        .unwrap();
1586
1587        assert!(result.is_some());
1588        assert_eq!(result.unwrap().deltas.len(), 2);
1589        assert!(buffered.is_empty());
1590    }
1591
1592    #[rstest]
1593    fn test_mbo_snapshot_with_f_last_buffers() {
1594        let instrument_id = InstrumentId::from("ESM4.GLBX");
1595        let mut buffering_start = None;
1596        let mut buffered = AHashMap::new();
1597
1598        let result = process_mbo_delta(
1599            test_delta(instrument_id, 1_000_000_000),
1600            128 | 32, // F_LAST | F_SNAPSHOT
1601            &mut buffering_start,
1602            &mut buffered,
1603        )
1604        .unwrap();
1605
1606        assert!(result.is_none());
1607        assert_eq!(buffered[&instrument_id].len(), 1);
1608    }
1609
1610    #[rstest]
1611    fn test_mbo_replay_buffers_until_past_start() {
1612        let instrument_id = InstrumentId::from("ESM4.GLBX");
1613        let start_ns = 5_000_000_000u64;
1614        let mut buffering_start = Some(start_ns.into());
1615        let mut buffered = AHashMap::new();
1616
1617        let result = process_mbo_delta(
1618            test_delta(instrument_id, 4_000_000_000),
1619            128, // F_LAST
1620            &mut buffering_start,
1621            &mut buffered,
1622        )
1623        .unwrap();
1624        assert!(result.is_none());
1625
1626        let result = process_mbo_delta(
1627            test_delta(instrument_id, 5_000_000_000),
1628            128,
1629            &mut buffering_start,
1630            &mut buffered,
1631        )
1632        .unwrap();
1633        assert!(result.is_none());
1634
1635        // Delta past start: emits and clears buffering_start
1636        let result = process_mbo_delta(
1637            test_delta(instrument_id, 6_000_000_000),
1638            128,
1639            &mut buffering_start,
1640            &mut buffered,
1641        )
1642        .unwrap();
1643        assert!(result.is_some());
1644        assert!(buffering_start.is_none());
1645    }
1646
1647    #[rstest]
1648    fn test_mbo_multiple_deltas_accumulated() {
1649        let instrument_id = InstrumentId::from("ESM4.GLBX");
1650        let mut buffering_start = None;
1651        let mut buffered = AHashMap::new();
1652
1653        for i in 0..5 {
1654            process_mbo_delta(
1655                test_delta(instrument_id, 1_000_000_000 + i),
1656                0,
1657                &mut buffering_start,
1658                &mut buffered,
1659            )
1660            .unwrap();
1661        }
1662
1663        let result = process_mbo_delta(
1664            test_delta(instrument_id, 2_000_000_000),
1665            128,
1666            &mut buffering_start,
1667            &mut buffered,
1668        )
1669        .unwrap();
1670
1671        assert!(result.is_some());
1672        assert_eq!(result.unwrap().deltas.len(), 6);
1673    }
1674
1675    #[rstest]
1676    fn test_mbo_multi_instrument_isolation() {
1677        let id_a = InstrumentId::from("ESM4.GLBX");
1678        let id_b = InstrumentId::from("NQM4.GLBX");
1679        let mut buffering_start = None;
1680        let mut buffered = AHashMap::new();
1681
1682        process_mbo_delta(
1683            test_delta(id_a, 1_000_000_000),
1684            0,
1685            &mut buffering_start,
1686            &mut buffered,
1687        )
1688        .unwrap();
1689        process_mbo_delta(
1690            test_delta(id_b, 1_000_000_000),
1691            0,
1692            &mut buffering_start,
1693            &mut buffered,
1694        )
1695        .unwrap();
1696
1697        // F_LAST for A: only A's deltas emitted, B remains
1698        let result = process_mbo_delta(
1699            test_delta(id_a, 2_000_000_000),
1700            128,
1701            &mut buffering_start,
1702            &mut buffered,
1703        )
1704        .unwrap();
1705
1706        assert!(result.is_some());
1707        assert_eq!(result.unwrap().instrument_id, id_a);
1708        assert!(buffered.contains_key(&id_b));
1709        assert!(!buffered.contains_key(&id_a));
1710    }
1711
1712    mod property_tests {
1713        use proptest::prelude::*;
1714        use rstest::rstest;
1715
1716        use super::*;
1717
1718        proptest! {
1719            #[rstest]
1720            fn mbo_buffering_conserves_deltas(
1721                num_non_last in 0usize..=20,
1722            ) {
1723                let instrument_id = InstrumentId::from("ESM4.GLBX");
1724                let mut buffering_start = None;
1725                let mut buffered = AHashMap::new();
1726                let total = num_non_last + 1;
1727
1728                for i in 0..num_non_last {
1729                    let result = process_mbo_delta(
1730                        test_delta(instrument_id, 1_000_000_000 + i as u64),
1731                        0, // No F_LAST
1732                        &mut buffering_start,
1733                        &mut buffered,
1734                    ).unwrap();
1735                    prop_assert!(result.is_none());
1736                }
1737
1738                let result = process_mbo_delta(
1739                    test_delta(instrument_id, 2_000_000_000),
1740                    128, // F_LAST
1741                    &mut buffering_start,
1742                    &mut buffered,
1743                ).unwrap();
1744
1745                prop_assert!(result.is_some());
1746                let emitted = result.unwrap();
1747                prop_assert_eq!(emitted.deltas.len(), total);
1748                prop_assert!(buffered.is_empty());
1749            }
1750
1751            #[rstest]
1752            fn mbo_snapshots_never_emit(
1753                num_snapshots in 1usize..=20,
1754            ) {
1755                let instrument_id = InstrumentId::from("ESM4.GLBX");
1756                let mut buffering_start = None;
1757                let mut buffered = AHashMap::new();
1758
1759                for i in 0..num_snapshots {
1760                    let result = process_mbo_delta(
1761                        test_delta(instrument_id, 1_000_000_000 + i as u64),
1762                        128 | 32, // F_LAST | F_SNAPSHOT
1763                        &mut buffering_start,
1764                        &mut buffered,
1765                    ).unwrap();
1766                    prop_assert!(result.is_none());
1767                }
1768
1769                prop_assert_eq!(buffered[&instrument_id].len(), num_snapshots);
1770            }
1771
1772            #[rstest]
1773            fn mbo_replay_delays_emission(
1774                start_offset in 1u64..=100,
1775                num_before in 1usize..=10,
1776            ) {
1777                let instrument_id = InstrumentId::from("ESM4.GLBX");
1778                let start_ns = 1_000_000_000u64 * start_offset;
1779                let mut buffering_start = Some(start_ns.into());
1780                let mut buffered = AHashMap::new();
1781
1782                for i in 0..num_before {
1783                    let ts = start_ns - (num_before as u64 - i as u64);
1784                    let result = process_mbo_delta(
1785                        test_delta(instrument_id, ts),
1786                        128, // F_LAST
1787                        &mut buffering_start,
1788                        &mut buffered,
1789                    ).unwrap();
1790                    prop_assert!(result.is_none());
1791                }
1792
1793                let result = process_mbo_delta(
1794                    test_delta(instrument_id, start_ns + 1),
1795                    128,
1796                    &mut buffering_start,
1797                    &mut buffered,
1798                ).unwrap();
1799
1800                prop_assert!(result.is_some());
1801                prop_assert!(buffering_start.is_none());
1802                let total = num_before + 1;
1803                prop_assert_eq!(result.unwrap().deltas.len(), total);
1804            }
1805        }
1806    }
1807}