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},
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    /// 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
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                        // Records that decode to no delta (`Action::Fill`
909                        // attribution, `Action::None` status, or trades) may
910                        // still carry the match-event boundary: honor the raw
911                        // `F_LAST` flag or a buffered partial event is
912                        // stranded and merged into the next event (e.g. a
913                        // non-terminal 'C' followed by 'N' | F_LAST).
914                        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    /// Sends a message to the message processing task.
948    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
963/// Handles Databento error messages by logging them.
964fn handle_error_msg(msg: &dbn::ErrorMsg) {
965    log::error!("{msg:?}");
966}
967
968/// Handles Databento system messages, returning a subscription ack event if applicable.
969fn 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
1004/// Parses a subscription ack message to extract the schema.
1005fn parse_ack_message(message: &str) -> String {
1006    // Format: "Subscription request N for <schema> data succeeded"
1007    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
1014/// Handles symbol mapping messages and updates the instrument ID map.
1015///
1016/// # Errors
1017///
1018/// Returns an error if symbol mapping fails.
1019fn 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
1064/// Updates the instrument ID map using exchange information from the symbol map.
1065fn 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    // Check if instrument ID is already in the map
1096    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
1124/// Handles instrument definition messages and decodes them into Nautilus instruments.
1125///
1126/// # Errors
1127///
1128/// Returns an error if instrument decoding fails.
1129fn 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    // Precheck before symbol resolution so unmodeled types skip cleanly
1215    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    // For MBP-1 and quote-based schemas, always include trades since they're integral to the data
1270    // For MBO, only include trades after the book is initialized to maintain consistency
1271    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 // These schemas include trade information directly
1276    } else {
1277        initialized_books.contains(&instrument_id) // MBO requires initialized book
1278    };
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
1313/// Processes an MBO delta through the buffering state machine.
1314///
1315/// Returns `Some(deltas)` when a complete batch is ready to emit (non-snapshot
1316/// F_LAST with replay buffering complete), or `None` when still accumulating.
1317fn 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    // Most live MBO events are single non-snapshot deltas, avoid map churn on that path
1327    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
1352/// Flushes the buffered deltas for `instrument_id` when `flags` marks a
1353/// non-snapshot match-event boundary (`F_LAST`).
1354///
1355/// Databento documents that records which decode to no delta (`Action::Fill`,
1356/// `Action::None`) may carry `F_LAST`, so the boundary must be processed
1357/// independently of the decoded payload or a buffered partial event is
1358/// stranded (follow-up to #4445).
1359fn 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, // non-terminal: no F_LAST
1416            1,
1417            ts.into(),
1418            ts.into(),
1419        )
1420    }
1421
1422    #[rstest]
1423    fn test_boundary_flag_on_recordless_message_flushes_buffered_event() {
1424        // A non-terminal 'C' delta buffers; the event terminates on an 'N'
1425        // record carrying F_LAST which decodes to no delta - the raw flag
1426        // must still flush the buffer (follow-up to #4445).
1427        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        // Boundary inside the replay gate: keep buffering
1518        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        // Boundary past the gate: flush and clear the gate
1529        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, // F_LAST
1787            &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, // F_LAST | F_SNAPSHOT
1805            &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, // F_LAST
1823            &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        // Delta past start: emits and clears buffering_start
1837        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        // F_LAST for A: only A's deltas emitted, B remains
1894        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, // No F_LAST
1927                        &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, // F_LAST
1936                    &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, // F_LAST | F_SNAPSHOT
1958                        &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, // F_LAST
1982                        &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}