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