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, 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            // Anchors remain only when the first connection failed, so this session replays
642            let replay = self.subscriptions.iter().any(|sub| sub.start.is_some());
643
644            for sub in self.subscriptions.clone() {
645                client.subscribe(sub).await?;
646            }
647            // Strip start timestamps after successful subscription to avoid replaying history on future reconnects
648            for sub in &mut self.subscriptions {
649                sub.start = None;
650            }
651
652            buffering_start = replay.then(|| clock.get_time_ns());
653            client.start().await?;
654            running = true;
655            log::info!("Resubscription complete");
656        } else if start_buffered {
657            log::debug!("Starting session from buffered Start command");
658            buffering_start = if self.replay {
659                Some(clock.get_time_ns())
660            } else {
661                None
662            };
663            client.start().await?;
664            running = true;
665        }
666
667        loop {
668            if self.msg_tx.is_closed() {
669                log::debug!("Message channel was closed: stopping");
670                return Ok(false);
671            }
672
673            // Wait for either a command or a record. When the session has not
674            // started yet (`!running`), only commands are awaited. Once running,
675            // `next_record` is cancel-safe so `tokio::select!` can safely
676            // race both futures.
677            if !running {
678                match self.cmd_rx.recv().await {
679                    Some(HandlerCommand::Subscribe(sub)) => {
680                        log::debug!("Received command: Subscribe");
681
682                        if !self.replay && sub.start.is_some() {
683                            self.replay = true;
684                        }
685                        client.subscribe(sub.clone()).await?;
686                        let mut sub_for_reconnect = sub;
687                        sub_for_reconnect.start = None;
688                        self.subscriptions.push(sub_for_reconnect);
689                        continue;
690                    }
691                    Some(HandlerCommand::SetPricePrecision(symbol, precision)) => {
692                        log::debug!(
693                            "Received command: SetPricePrecision for {symbol} to {precision}"
694                        );
695                        self.price_precision_overrides.insert(symbol, precision);
696                        continue;
697                    }
698                    Some(HandlerCommand::Start) => {
699                        log::debug!("Received command: Start");
700                        buffering_start = if self.replay {
701                            Some(clock.get_time_ns())
702                        } else {
703                            None
704                        };
705                        client.start().await?;
706                        running = true;
707                        continue;
708                    }
709                    Some(HandlerCommand::Close) => {
710                        self.send_close_msg();
711                        return Ok(false);
712                    }
713                    None => {
714                        log::debug!("Command channel disconnected");
715                        return Ok(false);
716                    }
717                }
718            }
719
720            let record_opt = tokio::select! {
721                cmd = self.cmd_rx.recv() =>
722                match cmd {
723                    Some(HandlerCommand::Subscribe(sub)) => {
724                        log::debug!("Received command: Subscribe");
725
726                        if sub.start.is_some() {
727                            self.replay = true;
728                            log::error!(
729                                "Ignoring `start` on {} subscribe, session already running, Databento drops replay anchors sent after session start",
730                                self.dataset,
731                            );
732                        }
733                        client.subscribe(sub.clone()).await?;
734                        let mut sub_for_reconnect = sub;
735                        sub_for_reconnect.start = None;
736                        self.subscriptions.push(sub_for_reconnect);
737                        continue;
738                    }
739                    Some(HandlerCommand::SetPricePrecision(symbol, precision)) => {
740                        log::debug!(
741                            "Received command: SetPricePrecision for {symbol} to {precision}"
742                        );
743                        self.price_precision_overrides.insert(symbol, precision);
744                        continue;
745                    }
746                    Some(HandlerCommand::Start) => {
747                        log::warn!("Received Start command but session already running");
748                        continue;
749                    }
750                    Some(HandlerCommand::Close) => {
751                        self.send_close_msg();
752                        client.close().await?;
753                        log::debug!("Closed inner client");
754                        return Ok(false);
755                    }
756                    None => {
757                        log::debug!("Command channel disconnected");
758                        return Ok(false);
759                    }
760                },
761                result = client.next_record() => result,
762            };
763
764            let record = match record_opt {
765                Ok(Some(record)) => record,
766                Ok(None) => {
767                    if session_start.elapsed() >= self.success_threshold {
768                        log::debug!("Session ended after successful run");
769                        return Ok(true);
770                    }
771                    anyhow::bail!("Session ended by gateway");
772                }
773                Err(e) => {
774                    if session_start.elapsed() >= self.success_threshold {
775                        log::debug!("Connection error after successful run: {e}");
776                        return Ok(true);
777                    }
778                    anyhow::bail!("Connection error: {e}");
779                }
780            };
781
782            let ts_init = clock.get_time_ns();
783
784            // Decode record
785            if let Some(msg) = record.get::<dbn::ErrorMsg>() {
786                handle_error_msg(msg);
787            } else if let Some(msg) = record.get::<dbn::SystemMsg>() {
788                if let Some(ack) = handle_system_msg(msg, ts_init) {
789                    self.send_msg(DatabentoMessage::SubscriptionAck(ack));
790                }
791            } else if let Some(msg) = record.get::<dbn::SymbolMappingMsg>() {
792                // Remove instrument ID index as the raw symbol may have changed
793                instrument_id_map.remove(&msg.hd.instrument_id);
794                instrument_def_price_precision_map.remove(&msg.hd.instrument_id);
795                update_price_precision_map_with_symbol_mapping_msg(
796                    msg,
797                    &self.price_precision_overrides,
798                    &mut subscription_price_precision_map,
799                )?;
800                handle_symbol_mapping_msg(msg, &mut symbol_map, &mut instrument_id_map)?;
801            } else if let Some(msg) = record.get::<dbn::InstrumentDefMsg>() {
802                if self.use_exchange_as_venue {
803                    let exchange = msg.exchange()?;
804                    if !exchange.is_empty() {
805                        update_instrument_id_map_with_exchange(
806                            &symbol_map,
807                            &self.symbol_venue_map,
808                            &mut instrument_id_map,
809                            msg.hd.instrument_id,
810                            exchange,
811                        )?;
812                    }
813                }
814                let maybe_data = handle_instrument_def_msg(
815                    msg,
816                    &record,
817                    &symbol_map,
818                    &self.publisher_venue_map,
819                    &self.symbol_venue_map,
820                    &mut instrument_id_map,
821                    ts_init,
822                )?;
823
824                if let Some(data) = maybe_data {
825                    instrument_def_price_precision_map
826                        .insert(msg.hd.instrument_id, data.price_precision());
827                    self.send_msg(DatabentoMessage::Instrument(Box::new(data)));
828                }
829            } else if let Some(msg) = record.get::<dbn::StatusMsg>() {
830                let data = handle_status_msg(
831                    msg,
832                    &record,
833                    &symbol_map,
834                    &self.publisher_venue_map,
835                    &self.symbol_venue_map,
836                    &mut instrument_id_map,
837                    ts_init,
838                )?;
839                self.send_msg(DatabentoMessage::Status(data));
840            } else if let Some(msg) = record.get::<dbn::ImbalanceMsg>() {
841                let data = handle_imbalance_msg(
842                    msg,
843                    &record,
844                    &symbol_map,
845                    &self.publisher_venue_map,
846                    &self.symbol_venue_map,
847                    &mut instrument_id_map,
848                    &instrument_def_price_precision_map,
849                    &subscription_price_precision_map,
850                    &self.price_precision_overrides,
851                    ts_init,
852                )?;
853                self.send_msg(DatabentoMessage::Imbalance(data));
854            } else if let Some(msg) = record.get::<dbn::StatMsg>() {
855                let maybe_data = handle_statistics_msg(
856                    msg,
857                    &record,
858                    &symbol_map,
859                    &self.publisher_venue_map,
860                    &self.symbol_venue_map,
861                    &mut instrument_id_map,
862                    &instrument_def_price_precision_map,
863                    &subscription_price_precision_map,
864                    &self.price_precision_overrides,
865                    ts_init,
866                )?;
867
868                if let Some(data) = maybe_data {
869                    self.send_msg(DatabentoMessage::Statistics(data));
870                }
871            } else {
872                // Decode a generic record with possible errors
873                let res = handle_record(
874                    record,
875                    &symbol_map,
876                    &self.publisher_venue_map,
877                    &self.symbol_venue_map,
878                    &mut instrument_id_map,
879                    &instrument_def_price_precision_map,
880                    &subscription_price_precision_map,
881                    &self.price_precision_overrides,
882                    ts_init,
883                    &initialized_books,
884                    self.bars_timestamp_on_close,
885                );
886                let (mut data1, data2) = match res {
887                    Ok(decoded) => decoded,
888                    Err(e) => {
889                        log::error!("Error decoding record: {e}");
890                        continue;
891                    }
892                };
893
894                if let Some(msg) = record.get::<dbn::MboMsg>() {
895                    if let Some(Data::BookDelta(delta)) = &data1 {
896                        initialized_books.insert(delta.instrument_id);
897
898                        log::trace!(
899                            "Buffering delta: {} {buffering_start:?} flags={}",
900                            delta.ts_event,
901                            msg.flags.raw(),
902                        );
903
904                        match process_mbo_delta(
905                            *delta,
906                            msg.flags.raw(),
907                            &mut buffering_start,
908                            &mut buffered_deltas,
909                        ) {
910                            Some(deltas) => data1 = Some(Data::BookDeltas(Box::new(deltas))),
911                            None => continue,
912                        }
913                    } else {
914                        // Records that decode to no delta (`Action::Fill`
915                        // attribution, `Action::None` status, or trades) may
916                        // still carry the match-event boundary: honor the raw
917                        // `F_LAST` flag or a buffered partial event is
918                        // stranded and merged into the next event (e.g. a
919                        // non-terminal 'C' followed by 'N' | F_LAST).
920                        let instrument_id = update_instrument_id_map(
921                            &record,
922                            &symbol_map,
923                            &self.publisher_venue_map,
924                            &self.symbol_venue_map,
925                            &mut instrument_id_map,
926                        )?;
927
928                        if let Some(deltas) = flush_mbo_event_boundary(
929                            instrument_id,
930                            msg.ts_recv.into(),
931                            msg.flags.raw(),
932                            &mut buffering_start,
933                            &mut buffered_deltas,
934                        ) {
935                            self.send_msg(DatabentoMessage::Data(Data::BookDeltas(Box::new(
936                                deltas,
937                            ))));
938                        }
939
940                        // Replayed trades stay out of the live stream
941                        if buffering_start.is_some_and(|start| msg.ts_recv <= start) {
942                            continue;
943                        }
944                    }
945                }
946
947                if let Some(data) = data1 {
948                    self.send_msg(DatabentoMessage::Data(data));
949                }
950
951                if let Some(data) = data2 {
952                    self.send_msg(DatabentoMessage::Data(data));
953                }
954            }
955        }
956    }
957
958    /// Sends a message to the message processing task.
959    fn send_msg(&self, msg: DatabentoMessage) {
960        log::trace!("Sending {msg:?}");
961        match self.msg_tx.send(msg) {
962            Ok(()) => {}
963            Err(e) => log::debug!("Error sending message: {e}"),
964        }
965    }
966
967    fn send_close_msg(&self) {
968        if let Err(e) = self.msg_tx.send(DatabentoMessage::Close) {
969            log::debug!("Could not send close message: {e}");
970        }
971    }
972}
973
974/// Handles Databento error messages by logging them.
975fn handle_error_msg(msg: &dbn::ErrorMsg) {
976    log::error!("{msg:?}");
977}
978
979/// Handles Databento system messages, returning a subscription ack event if applicable.
980fn handle_system_msg(msg: &dbn::SystemMsg, ts_received: UnixNanos) -> Option<SubscriptionAckEvent> {
981    match msg.code() {
982        Ok(dbn::SystemCode::SubscriptionAck) => {
983            let message = msg.msg().unwrap_or("<invalid utf-8>");
984            log::debug!("Subscription acknowledged: {message}");
985
986            let schema = parse_ack_message(message);
987
988            Some(SubscriptionAckEvent {
989                schema,
990                message: message.to_string(),
991                ts_received,
992            })
993        }
994        Ok(dbn::SystemCode::Heartbeat) => {
995            log::trace!("Heartbeat received");
996            None
997        }
998        Ok(dbn::SystemCode::SlowReaderWarning) => {
999            let message = msg.msg().unwrap_or("<invalid utf-8>");
1000            log::warn!("Slow reader warning: {message}");
1001            None
1002        }
1003        Ok(dbn::SystemCode::ReplayCompleted) => {
1004            let message = msg.msg().unwrap_or("<invalid utf-8>");
1005            log::debug!("Replay completed: {message}");
1006            None
1007        }
1008        _ => {
1009            log::debug!("{msg:?}");
1010            None
1011        }
1012    }
1013}
1014
1015/// Parses a subscription ack message to extract the schema.
1016fn parse_ack_message(message: &str) -> String {
1017    // Format: "Subscription request N for <schema> data succeeded"
1018    message
1019        .strip_circumfix("Subscription request ", " data succeeded")
1020        .and_then(|rest| rest.split_once(" for "))
1021        .map(|(_, schema)| schema.trim().to_string())
1022        .unwrap_or_default()
1023}
1024
1025/// Handles symbol mapping messages and updates the instrument ID map.
1026///
1027/// # Errors
1028///
1029/// Returns an error if symbol mapping fails.
1030fn handle_symbol_mapping_msg(
1031    msg: &dbn::SymbolMappingMsg,
1032    symbol_map: &mut PitSymbolMap,
1033    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1034) -> anyhow::Result<()> {
1035    symbol_map
1036        .on_symbol_mapping(msg)
1037        .map_err(|e| anyhow::anyhow!("on_symbol_mapping failed for {msg:?}: {e}"))?;
1038    instrument_id_map.remove(&msg.hd.instrument_id);
1039    Ok(())
1040}
1041
1042fn update_price_precision_map_with_symbol_mapping_msg(
1043    msg: &dbn::SymbolMappingMsg,
1044    price_precision_overrides: &AHashMap<Symbol, u8>,
1045    subscription_price_precision_map: &mut AHashMap<u32, u8>,
1046) -> anyhow::Result<()> {
1047    subscription_price_precision_map.remove(&msg.hd.instrument_id);
1048
1049    if price_precision_overrides.is_empty() {
1050        return Ok(());
1051    }
1052
1053    let stype_in_symbol = msg
1054        .stype_in_symbol()
1055        .map_err(|e| anyhow::anyhow!("Error decoding `stype_in_symbol`: {e}"))?;
1056    let stype_out_symbol = msg
1057        .stype_out_symbol()
1058        .map_err(|e| anyhow::anyhow!("Error decoding `stype_out_symbol`: {e}"))?;
1059
1060    let price_precision = [stype_in_symbol, stype_out_symbol]
1061        .into_iter()
1062        .find_map(|symbol| {
1063            price_precision_overrides
1064                .get(&Symbol::from_str_unchecked(symbol))
1065                .copied()
1066        });
1067
1068    if let Some(price_precision) = price_precision {
1069        subscription_price_precision_map.insert(msg.hd.instrument_id, price_precision);
1070    }
1071
1072    Ok(())
1073}
1074
1075/// Updates the instrument ID map using exchange information from the symbol map.
1076fn update_instrument_id_map_with_exchange(
1077    symbol_map: &PitSymbolMap,
1078    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1079    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1080    raw_instrument_id: u32,
1081    exchange: &str,
1082) -> anyhow::Result<InstrumentId> {
1083    let raw_symbol = symbol_map.get(raw_instrument_id).ok_or_else(|| {
1084        anyhow::anyhow!("Cannot resolve raw_symbol for instrument_id {raw_instrument_id}")
1085    })?;
1086    let symbol = Symbol::from(raw_symbol.as_str());
1087    let venue = Venue::from_code(exchange)
1088        .map_err(|e| anyhow::anyhow!("Invalid venue code '{exchange}': {e}"))?;
1089    let instrument_id = InstrumentId::new(symbol, venue);
1090    symbol_venue_map.rcu(|m| {
1091        m.entry(symbol).or_insert(venue);
1092    });
1093    instrument_id_map.insert(raw_instrument_id, instrument_id);
1094    Ok(instrument_id)
1095}
1096
1097fn update_instrument_id_map(
1098    record: &dbn::RecordRef,
1099    symbol_map: &PitSymbolMap,
1100    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1101    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1102    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1103) -> anyhow::Result<InstrumentId> {
1104    let header = record.header();
1105
1106    // Check if instrument ID is already in the map
1107    if let Some(&instrument_id) = instrument_id_map.get(&header.instrument_id) {
1108        return Ok(instrument_id);
1109    }
1110
1111    let raw_symbol = symbol_map.get_for_rec(record).ok_or_else(|| {
1112        anyhow::anyhow!(
1113            "Cannot resolve `raw_symbol` from `symbol_map` for instrument_id {}",
1114            header.instrument_id
1115        )
1116    })?;
1117
1118    let symbol = Symbol::from_str_unchecked(raw_symbol);
1119
1120    let publisher_id = header.publisher_id;
1121    let venue = if let Some(venue) = symbol_venue_map.get_cloned(&symbol) {
1122        venue
1123    } else {
1124        let venue = publisher_venue_map
1125            .get(&publisher_id)
1126            .ok_or_else(|| anyhow::anyhow!("No venue found for `publisher_id` {publisher_id}"))?;
1127        *venue
1128    };
1129    let instrument_id = InstrumentId::new(symbol, venue);
1130
1131    instrument_id_map.insert(header.instrument_id, instrument_id);
1132    Ok(instrument_id)
1133}
1134
1135/// Handles instrument definition messages and decodes them into Nautilus instruments.
1136///
1137/// # Errors
1138///
1139/// Returns an error if instrument decoding fails.
1140fn handle_instrument_def_msg(
1141    msg: &dbn::InstrumentDefMsg,
1142    record: &dbn::RecordRef,
1143    symbol_map: &PitSymbolMap,
1144    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1145    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1146    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1147    ts_init: UnixNanos,
1148) -> anyhow::Result<Option<InstrumentAny>> {
1149    let instrument_id = update_instrument_id_map(
1150        record,
1151        symbol_map,
1152        publisher_venue_map,
1153        symbol_venue_map,
1154        instrument_id_map,
1155    )?;
1156
1157    decode_instrument_def_msg(msg, instrument_id, Some(ts_init), None)
1158}
1159
1160fn handle_status_msg(
1161    msg: &dbn::StatusMsg,
1162    record: &dbn::RecordRef,
1163    symbol_map: &PitSymbolMap,
1164    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1165    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1166    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1167    ts_init: UnixNanos,
1168) -> anyhow::Result<InstrumentStatus> {
1169    let instrument_id = update_instrument_id_map(
1170        record,
1171        symbol_map,
1172        publisher_venue_map,
1173        symbol_venue_map,
1174        instrument_id_map,
1175    )?;
1176
1177    decode_status_msg(msg, instrument_id, Some(ts_init))
1178}
1179
1180#[expect(clippy::too_many_arguments)]
1181fn handle_imbalance_msg(
1182    msg: &dbn::ImbalanceMsg,
1183    record: &dbn::RecordRef,
1184    symbol_map: &PitSymbolMap,
1185    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1186    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1187    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1188    instrument_def_price_precision_map: &AHashMap<u32, u8>,
1189    subscription_price_precision_map: &AHashMap<u32, u8>,
1190    price_precision_overrides: &AHashMap<Symbol, u8>,
1191    ts_init: UnixNanos,
1192) -> anyhow::Result<DatabentoImbalance> {
1193    let instrument_id = update_instrument_id_map(
1194        record,
1195        symbol_map,
1196        publisher_venue_map,
1197        symbol_venue_map,
1198        instrument_id_map,
1199    )?;
1200
1201    let price_precision = resolve_price_precision(
1202        msg.hd.instrument_id,
1203        instrument_id,
1204        instrument_def_price_precision_map,
1205        subscription_price_precision_map,
1206        price_precision_overrides,
1207    );
1208
1209    decode_imbalance_msg(msg, instrument_id, price_precision, Some(ts_init))
1210}
1211
1212#[expect(clippy::too_many_arguments)]
1213fn handle_statistics_msg(
1214    msg: &dbn::StatMsg,
1215    record: &dbn::RecordRef,
1216    symbol_map: &PitSymbolMap,
1217    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1218    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1219    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1220    instrument_def_price_precision_map: &AHashMap<u32, u8>,
1221    subscription_price_precision_map: &AHashMap<u32, u8>,
1222    price_precision_overrides: &AHashMap<Symbol, u8>,
1223    ts_init: UnixNanos,
1224) -> anyhow::Result<Option<DatabentoStatistics>> {
1225    // Precheck before symbol resolution so unmodeled types skip cleanly
1226    if !is_supported_stat_type(msg.stat_type) {
1227        log::warn!("Skipping unsupported `stat_type` {}", msg.stat_type);
1228        return Ok(None);
1229    }
1230
1231    let instrument_id = update_instrument_id_map(
1232        record,
1233        symbol_map,
1234        publisher_venue_map,
1235        symbol_venue_map,
1236        instrument_id_map,
1237    )?;
1238
1239    let price_precision = resolve_price_precision(
1240        msg.hd.instrument_id,
1241        instrument_id,
1242        instrument_def_price_precision_map,
1243        subscription_price_precision_map,
1244        price_precision_overrides,
1245    );
1246
1247    decode_statistics_msg(msg, instrument_id, price_precision, Some(ts_init))
1248}
1249
1250#[expect(clippy::too_many_arguments)]
1251fn handle_record(
1252    record: dbn::RecordRef,
1253    symbol_map: &PitSymbolMap,
1254    publisher_venue_map: &IndexMap<PublisherId, Venue>,
1255    symbol_venue_map: &AtomicMap<Symbol, Venue>,
1256    instrument_id_map: &mut AHashMap<u32, InstrumentId>,
1257    instrument_def_price_precision_map: &AHashMap<u32, u8>,
1258    subscription_price_precision_map: &AHashMap<u32, u8>,
1259    price_precision_overrides: &AHashMap<Symbol, u8>,
1260    ts_init: UnixNanos,
1261    initialized_books: &HashSet<InstrumentId>,
1262    bars_timestamp_on_close: bool,
1263) -> anyhow::Result<(Option<Data>, Option<Data>)> {
1264    let instrument_id = update_instrument_id_map(
1265        &record,
1266        symbol_map,
1267        publisher_venue_map,
1268        symbol_venue_map,
1269        instrument_id_map,
1270    )?;
1271
1272    let price_precision = resolve_price_precision(
1273        record.header().instrument_id,
1274        instrument_id,
1275        instrument_def_price_precision_map,
1276        subscription_price_precision_map,
1277        price_precision_overrides,
1278    );
1279
1280    // For MBP-1 and quote-based schemas, always include trades since they're integral to the data
1281    // For MBO, only include trades after the book is initialized to maintain consistency
1282    let include_trades = if record.get::<dbn::Mbp1Msg>().is_some()
1283        || record.get::<dbn::TbboMsg>().is_some()
1284        || record.get::<dbn::Cmbp1Msg>().is_some()
1285    {
1286        true // These schemas include trade information directly
1287    } else {
1288        initialized_books.contains(&instrument_id) // MBO requires initialized book
1289    };
1290
1291    decode_record(
1292        &record,
1293        instrument_id,
1294        price_precision,
1295        Some(ts_init),
1296        include_trades,
1297        bars_timestamp_on_close,
1298    )
1299}
1300
1301fn resolve_price_precision(
1302    record_instrument_id: u32,
1303    instrument_id: InstrumentId,
1304    instrument_def_price_precision_map: &AHashMap<u32, u8>,
1305    subscription_price_precision_map: &AHashMap<u32, u8>,
1306    price_precision_overrides: &AHashMap<Symbol, u8>,
1307) -> u8 {
1308    instrument_def_price_precision_map
1309        .get(&record_instrument_id)
1310        .copied()
1311        .or_else(|| {
1312            subscription_price_precision_map
1313                .get(&record_instrument_id)
1314                .copied()
1315        })
1316        .or_else(|| {
1317            price_precision_overrides
1318                .get(&instrument_id.symbol)
1319                .copied()
1320        })
1321        .unwrap_or(Currency::USD().precision)
1322}
1323
1324/// Processes an MBO delta through the buffering state machine.
1325///
1326/// Returns `Some(deltas)` when a complete batch is ready to emit (non-snapshot
1327/// F_LAST with replay buffering complete), or `None` when still accumulating.
1328fn process_mbo_delta(
1329    delta: OrderBookDelta,
1330    flags: u8,
1331    buffering_start: &mut Option<UnixNanos>,
1332    buffered_deltas: &mut AHashMap<InstrumentId, Vec<OrderBookDelta>>,
1333) -> Option<OrderBookDeltas> {
1334    let is_last = RecordFlag::F_LAST.matches(flags);
1335    let is_snapshot = RecordFlag::F_SNAPSHOT.matches(flags);
1336
1337    // Most live MBO events are single non-snapshot deltas, avoid map churn on that path
1338    if is_last
1339        && !is_snapshot
1340        && buffering_start.is_none()
1341        && !buffered_deltas.contains_key(&delta.instrument_id)
1342    {
1343        let deltas = OrderBookDeltas::new(delta.instrument_id, vec![delta]);
1344        return Some(deltas);
1345    }
1346
1347    let buffer = buffered_deltas.entry(delta.instrument_id).or_default();
1348    buffer.push(delta);
1349
1350    if is_snapshot {
1351        return None;
1352    }
1353
1354    flush_mbo_event_boundary(
1355        delta.instrument_id,
1356        delta.ts_event,
1357        flags,
1358        buffering_start,
1359        buffered_deltas,
1360    )
1361}
1362
1363/// Flushes the buffered deltas for `instrument_id` when `flags` marks a
1364/// non-snapshot match-event boundary (`F_LAST`).
1365///
1366/// Databento documents that records which decode to no delta (`Action::Fill`,
1367/// `Action::None`) may carry `F_LAST`, so the boundary must be processed
1368/// independently of the decoded payload or a buffered partial event is
1369/// stranded (follow-up to #4445).
1370fn flush_mbo_event_boundary(
1371    instrument_id: InstrumentId,
1372    ts_event: UnixNanos,
1373    flags: u8,
1374    buffering_start: &mut Option<UnixNanos>,
1375    buffered_deltas: &mut AHashMap<InstrumentId, Vec<OrderBookDelta>>,
1376) -> Option<OrderBookDeltas> {
1377    if !RecordFlag::F_LAST.matches(flags) || RecordFlag::F_SNAPSHOT.matches(flags) {
1378        return None;
1379    }
1380
1381    let buffer = buffered_deltas.get(&instrument_id)?;
1382
1383    if buffer.is_empty() {
1384        return None;
1385    }
1386
1387    if let Some(start_ns) = *buffering_start {
1388        if ts_event <= start_ns {
1389            return None;
1390        }
1391        *buffering_start = None;
1392    }
1393
1394    let buffer = buffered_deltas.remove(&instrument_id)?;
1395    let deltas = OrderBookDeltas::new(instrument_id, buffer);
1396    Some(deltas)
1397}
1398
1399#[cfg(test)]
1400mod tests {
1401    use std::path::PathBuf;
1402
1403    use databento::live::Subscription;
1404    use indexmap::IndexMap;
1405    use rstest::*;
1406    use time::macros::datetime;
1407
1408    use super::*;
1409
1410    fn stub_delta(instrument_id: InstrumentId, ts: u64) -> OrderBookDelta {
1411        use nautilus_model::{
1412            data::BookOrder,
1413            enums::{BookAction, OrderSide},
1414            types::{Price, Quantity},
1415        };
1416
1417        OrderBookDelta::new(
1418            instrument_id,
1419            BookAction::Delete,
1420            BookOrder::new(
1421                OrderSide::Sell,
1422                Price::from("100.00"),
1423                Quantity::from("5"),
1424                42,
1425            ),
1426            0, // non-terminal: no F_LAST
1427            1,
1428            ts.into(),
1429            ts.into(),
1430        )
1431    }
1432
1433    #[rstest]
1434    fn test_boundary_flag_on_recordless_message_flushes_buffered_event() {
1435        // A non-terminal 'C' delta buffers; the event terminates on an 'N'
1436        // record carrying F_LAST which decodes to no delta - the raw flag
1437        // must still flush the buffer (follow-up to #4445).
1438        let instrument_id = InstrumentId::from("TEST.GLBX");
1439        let mut buffering_start = None;
1440        let mut buffered = AHashMap::new();
1441
1442        let buffered_result = process_mbo_delta(
1443            stub_delta(instrument_id, 1),
1444            0,
1445            &mut buffering_start,
1446            &mut buffered,
1447        );
1448        assert!(buffered_result.is_none());
1449
1450        let flushed = flush_mbo_event_boundary(
1451            instrument_id,
1452            2.into(),
1453            RecordFlag::F_LAST as u8,
1454            &mut buffering_start,
1455            &mut buffered,
1456        );
1457        assert!(flushed.is_some());
1458        assert!(buffered.is_empty());
1459    }
1460
1461    #[rstest]
1462    fn test_boundary_flag_with_empty_buffer_is_noop() {
1463        let instrument_id = InstrumentId::from("TEST.GLBX");
1464        let start = UnixNanos::from(1);
1465        let mut buffering_start = Some(start);
1466        let mut buffered: AHashMap<InstrumentId, Vec<OrderBookDelta>> = AHashMap::new();
1467        buffered.insert(instrument_id, Vec::new());
1468
1469        let flushed = flush_mbo_event_boundary(
1470            instrument_id,
1471            2.into(),
1472            RecordFlag::F_LAST as u8,
1473            &mut buffering_start,
1474            &mut buffered,
1475        );
1476
1477        assert!(flushed.is_none());
1478        assert_eq!(buffering_start, Some(start));
1479    }
1480
1481    #[rstest]
1482    fn test_empty_boundary_preserves_replay_gate_for_buffered_instrument() {
1483        let empty_id = InstrumentId::from("EMPTY.GLBX");
1484        let buffered_id = InstrumentId::from("BUFFERED.GLBX");
1485        let start = UnixNanos::from(10);
1486        let mut buffering_start = Some(start);
1487        let mut buffered = AHashMap::new();
1488        let _ = process_mbo_delta(
1489            stub_delta(buffered_id, 1),
1490            0,
1491            &mut buffering_start,
1492            &mut buffered,
1493        );
1494
1495        let empty_boundary = flush_mbo_event_boundary(
1496            empty_id,
1497            11.into(),
1498            RecordFlag::F_LAST as u8,
1499            &mut buffering_start,
1500            &mut buffered,
1501        );
1502        let buffered_boundary = flush_mbo_event_boundary(
1503            buffered_id,
1504            5.into(),
1505            RecordFlag::F_LAST as u8,
1506            &mut buffering_start,
1507            &mut buffered,
1508        );
1509
1510        assert!(empty_boundary.is_none());
1511        assert!(buffered_boundary.is_none());
1512        assert_eq!(buffering_start, Some(start));
1513        assert!(buffered.contains_key(&buffered_id));
1514    }
1515
1516    #[rstest]
1517    fn test_boundary_flag_respects_replay_buffering_gate() {
1518        let instrument_id = InstrumentId::from("TEST.GLBX");
1519        let mut buffering_start = Some(UnixNanos::from(10));
1520        let mut buffered = AHashMap::new();
1521        let _ = process_mbo_delta(
1522            stub_delta(instrument_id, 1),
1523            0,
1524            &mut buffering_start,
1525            &mut buffered,
1526        );
1527
1528        // Boundary inside the replay gate: keep buffering
1529        let flushed = flush_mbo_event_boundary(
1530            instrument_id,
1531            5.into(),
1532            RecordFlag::F_LAST as u8,
1533            &mut buffering_start,
1534            &mut buffered,
1535        );
1536        assert!(flushed.is_none());
1537        assert!(!buffered.is_empty());
1538
1539        // Boundary past the gate: flush and clear the gate
1540        let flushed = flush_mbo_event_boundary(
1541            instrument_id,
1542            11.into(),
1543            RecordFlag::F_LAST as u8,
1544            &mut buffering_start,
1545            &mut buffered,
1546        );
1547        assert!(flushed.is_some());
1548        assert!(buffering_start.is_none());
1549    }
1550
1551    #[rstest]
1552    fn test_boundary_snapshot_flag_never_flushes() {
1553        let instrument_id = InstrumentId::from("TEST.GLBX");
1554        let mut buffering_start = None;
1555        let mut buffered = AHashMap::new();
1556        let _ = process_mbo_delta(
1557            stub_delta(instrument_id, 1),
1558            0,
1559            &mut buffering_start,
1560            &mut buffered,
1561        );
1562
1563        let flags = RecordFlag::F_LAST as u8 | RecordFlag::F_SNAPSHOT as u8;
1564        let flushed = flush_mbo_event_boundary(
1565            instrument_id,
1566            2.into(),
1567            flags,
1568            &mut buffering_start,
1569            &mut buffered,
1570        );
1571        assert!(flushed.is_none());
1572    }
1573
1574    fn create_test_handler(reconnect_timeout_mins: Option<u64>) -> DatabentoFeedHandler {
1575        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1576        let (msg_tx, _msg_rx) = tokio::sync::mpsc::unbounded_channel();
1577
1578        DatabentoFeedHandler::new(
1579            Credential::new("test_key"),
1580            "GLBX.MDP3".to_string(),
1581            cmd_rx,
1582            msg_tx,
1583            IndexMap::new(),
1584            Arc::new(AtomicMap::new()),
1585            false,
1586            false,
1587            reconnect_timeout_mins,
1588        )
1589    }
1590
1591    fn create_test_client() -> DatabentoLiveClient {
1592        DatabentoLiveClient::new(
1593            "test-api-key".to_string(),
1594            "GLBX.MDP3".to_string(),
1595            PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("publishers.json"),
1596            true,
1597            None,
1598            None,
1599        )
1600        .unwrap()
1601    }
1602
1603    #[rstest]
1604    #[case(Some(10))]
1605    #[case(None)]
1606    fn test_backoff_initialization(#[case] reconnect_timeout_mins: Option<u64>) {
1607        let handler = create_test_handler(reconnect_timeout_mins);
1608
1609        assert_eq!(handler.reconnect_timeout_mins, reconnect_timeout_mins);
1610        assert!(handler.subscriptions.is_empty());
1611        assert!(handler.buffered_commands.is_empty());
1612    }
1613
1614    #[rstest]
1615    fn test_subscription_with_and_without_start() {
1616        let start_time = datetime!(2024-01-01 00:00:00 UTC);
1617        let sub_with_start = Subscription::builder()
1618            .symbols("ES.FUT")
1619            .schema(databento::dbn::Schema::Mbp1)
1620            .start(start_time)
1621            .build();
1622
1623        let mut sub_without_start = sub_with_start.clone();
1624        sub_without_start.start = None;
1625
1626        assert!(sub_with_start.start.is_some());
1627        assert!(sub_without_start.start.is_none());
1628        assert_eq!(sub_with_start.schema, sub_without_start.schema);
1629        assert_eq!(sub_with_start.symbols, sub_without_start.symbols);
1630    }
1631
1632    #[rstest]
1633    fn test_handler_initialization_state() {
1634        let handler = create_test_handler(Some(10));
1635
1636        assert!(!handler.replay);
1637        assert_eq!(handler.dataset, "GLBX.MDP3");
1638        assert_eq!(handler.credential.api_key(), "test_key");
1639        assert!(handler.subscriptions.is_empty());
1640        assert!(handler.buffered_commands.is_empty());
1641    }
1642
1643    #[rstest]
1644    fn test_handler_with_no_timeout() {
1645        let handler = create_test_handler(None);
1646
1647        assert_eq!(handler.reconnect_timeout_mins, None);
1648        assert!(!handler.replay);
1649    }
1650
1651    #[rstest]
1652    fn test_handler_with_zero_timeout() {
1653        let handler = create_test_handler(Some(0));
1654
1655        assert_eq!(handler.reconnect_timeout_mins, Some(0));
1656        assert!(!handler.replay);
1657    }
1658
1659    #[rstest]
1660    fn test_subscribe_uses_explicit_parent_stype() {
1661        let mut client = create_test_client();
1662
1663        client
1664            .subscribe(
1665                "definition".to_string(),
1666                vec![InstrumentId::from("ES.FUT.GLBX")],
1667                None,
1668                None,
1669                None,
1670                Some("parent".to_string()),
1671            )
1672            .unwrap();
1673
1674        let command = client.cmd_rx.as_mut().unwrap().try_recv().unwrap();
1675        match command {
1676            HandlerCommand::Subscribe(sub) => {
1677                assert_eq!(sub.schema, dbn::Schema::Definition);
1678                assert_eq!(sub.stype_in, dbn::SType::Parent);
1679                assert_eq!(sub.symbols.to_api_string(), "ES.FUT");
1680            }
1681            other => panic!("expected HandlerCommand::Subscribe, was {other:?}"),
1682        }
1683    }
1684
1685    #[rstest]
1686    fn test_subscribe_rejects_invalid_stype() {
1687        let mut client = create_test_client();
1688
1689        let err = client
1690            .subscribe(
1691                "definition".to_string(),
1692                vec![InstrumentId::from("ES.FUT.GLBX")],
1693                None,
1694                None,
1695                None,
1696                Some("not-a-stype".to_string()),
1697            )
1698            .unwrap_err();
1699
1700        assert!(err.to_string().contains("not-a-stype"));
1701        assert!(!is_command_send_error(&err));
1702        assert!(matches!(
1703            client.cmd_rx.as_mut().unwrap().try_recv(),
1704            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1705        ));
1706    }
1707
1708    #[rstest]
1709    fn test_subscribe_classifies_command_send_errors() {
1710        let mut client = create_test_client();
1711        client.cmd_rx = None;
1712
1713        let err = client
1714            .subscribe(
1715                "definition".to_string(),
1716                vec![InstrumentId::from("ES.FUT.GLBX")],
1717                None,
1718                None,
1719                None,
1720                Some("parent".to_string()),
1721            )
1722            .unwrap_err();
1723
1724        assert!(is_command_send_error(&err));
1725    }
1726
1727    #[rstest]
1728    fn test_close_after_handler_exit_marks_closed() {
1729        let mut client = create_test_client();
1730        let (handler, _msg_rx) = client.start().unwrap();
1731        drop(handler);
1732
1733        client.close().unwrap();
1734
1735        assert!(!client.is_running());
1736        assert!(client.is_closed());
1737    }
1738
1739    fn test_delta(instrument_id: InstrumentId, ts_event: u64) -> OrderBookDelta {
1740        OrderBookDelta::clear(instrument_id, 0, ts_event.into(), 0.into())
1741    }
1742
1743    #[rstest]
1744    fn test_mbo_delta_without_f_last_buffers() {
1745        let instrument_id = InstrumentId::from("ESM4.GLBX");
1746        let delta = test_delta(instrument_id, 1_000_000_000);
1747        let mut buffering_start = None;
1748        let mut buffered = AHashMap::new();
1749
1750        let result = process_mbo_delta(delta, 0, &mut buffering_start, &mut buffered);
1751
1752        assert!(result.is_none());
1753        assert_eq!(buffered[&instrument_id].len(), 1);
1754    }
1755
1756    #[rstest]
1757    fn test_mbo_single_f_last_emits_without_buffering() {
1758        let instrument_id = InstrumentId::from("ESM4.GLBX");
1759        let ts_event = 1_000_000_000;
1760        let mut delta = test_delta(instrument_id, ts_event);
1761        delta.flags = 128;
1762        delta.sequence = 42;
1763        let mut buffering_start = None;
1764        let mut buffered = AHashMap::new();
1765
1766        let result = process_mbo_delta(delta, delta.flags, &mut buffering_start, &mut buffered);
1767
1768        let emitted = result.expect("single F_LAST delta should emit");
1769        assert_eq!(emitted.instrument_id, instrument_id);
1770        assert_eq!(emitted.deltas.len(), 1);
1771        assert_eq!(emitted.flags, 128);
1772        assert_eq!(emitted.sequence, 42);
1773        assert_eq!(emitted.ts_event, UnixNanos::from(ts_event));
1774        assert_eq!(emitted.deltas[0].instrument_id, instrument_id);
1775        assert_eq!(emitted.deltas[0].flags, 128);
1776        assert_eq!(emitted.deltas[0].sequence, 42);
1777        assert_eq!(emitted.deltas[0].ts_event, UnixNanos::from(ts_event));
1778        assert!(buffering_start.is_none());
1779        assert!(buffered.is_empty());
1780    }
1781
1782    #[rstest]
1783    fn test_mbo_delta_with_f_last_emits() {
1784        let instrument_id = InstrumentId::from("ESM4.GLBX");
1785        let mut buffering_start = None;
1786        let mut buffered = AHashMap::new();
1787
1788        let _ = process_mbo_delta(
1789            test_delta(instrument_id, 1_000_000_000),
1790            0,
1791            &mut buffering_start,
1792            &mut buffered,
1793        );
1794
1795        let result = process_mbo_delta(
1796            test_delta(instrument_id, 2_000_000_000),
1797            128, // F_LAST
1798            &mut buffering_start,
1799            &mut buffered,
1800        );
1801
1802        assert!(result.is_some());
1803        assert_eq!(result.unwrap().deltas.len(), 2);
1804        assert!(buffered.is_empty());
1805    }
1806
1807    #[rstest]
1808    fn test_mbo_snapshot_with_f_last_buffers() {
1809        let instrument_id = InstrumentId::from("ESM4.GLBX");
1810        let mut buffering_start = None;
1811        let mut buffered = AHashMap::new();
1812
1813        let result = process_mbo_delta(
1814            test_delta(instrument_id, 1_000_000_000),
1815            128 | 32, // F_LAST | F_SNAPSHOT
1816            &mut buffering_start,
1817            &mut buffered,
1818        );
1819
1820        assert!(result.is_none());
1821        assert_eq!(buffered[&instrument_id].len(), 1);
1822    }
1823
1824    #[rstest]
1825    fn test_mbo_replay_buffers_until_past_start() {
1826        let instrument_id = InstrumentId::from("ESM4.GLBX");
1827        let start_ns = 5_000_000_000u64;
1828        let mut buffering_start = Some(start_ns.into());
1829        let mut buffered = AHashMap::new();
1830
1831        let result = process_mbo_delta(
1832            test_delta(instrument_id, 4_000_000_000),
1833            128, // F_LAST
1834            &mut buffering_start,
1835            &mut buffered,
1836        );
1837        assert!(result.is_none());
1838
1839        let result = process_mbo_delta(
1840            test_delta(instrument_id, 5_000_000_000),
1841            128,
1842            &mut buffering_start,
1843            &mut buffered,
1844        );
1845        assert!(result.is_none());
1846
1847        // Delta past start: emits and clears buffering_start
1848        let result = process_mbo_delta(
1849            test_delta(instrument_id, 6_000_000_000),
1850            128,
1851            &mut buffering_start,
1852            &mut buffered,
1853        );
1854        assert!(result.is_some());
1855        assert!(buffering_start.is_none());
1856    }
1857
1858    #[rstest]
1859    fn test_mbo_multiple_deltas_accumulated() {
1860        let instrument_id = InstrumentId::from("ESM4.GLBX");
1861        let mut buffering_start = None;
1862        let mut buffered = AHashMap::new();
1863
1864        for i in 0..5 {
1865            process_mbo_delta(
1866                test_delta(instrument_id, 1_000_000_000 + i),
1867                0,
1868                &mut buffering_start,
1869                &mut buffered,
1870            );
1871        }
1872
1873        let result = process_mbo_delta(
1874            test_delta(instrument_id, 2_000_000_000),
1875            128,
1876            &mut buffering_start,
1877            &mut buffered,
1878        );
1879
1880        assert!(result.is_some());
1881        assert_eq!(result.unwrap().deltas.len(), 6);
1882    }
1883
1884    #[rstest]
1885    fn test_mbo_multi_instrument_isolation() {
1886        let id_a = InstrumentId::from("ESM4.GLBX");
1887        let id_b = InstrumentId::from("NQM4.GLBX");
1888        let mut buffering_start = None;
1889        let mut buffered = AHashMap::new();
1890
1891        process_mbo_delta(
1892            test_delta(id_a, 1_000_000_000),
1893            0,
1894            &mut buffering_start,
1895            &mut buffered,
1896        );
1897        process_mbo_delta(
1898            test_delta(id_b, 1_000_000_000),
1899            0,
1900            &mut buffering_start,
1901            &mut buffered,
1902        );
1903
1904        // F_LAST for A: only A's deltas emitted, B remains
1905        let result = process_mbo_delta(
1906            test_delta(id_a, 2_000_000_000),
1907            128,
1908            &mut buffering_start,
1909            &mut buffered,
1910        );
1911
1912        assert!(result.is_some());
1913        assert_eq!(result.unwrap().instrument_id, id_a);
1914        assert!(buffered.contains_key(&id_b));
1915        assert!(!buffered.contains_key(&id_a));
1916    }
1917
1918    mod property_tests {
1919        use proptest::prelude::*;
1920        use rstest::rstest;
1921
1922        use super::*;
1923
1924        proptest! {
1925            #[rstest]
1926            fn mbo_buffering_conserves_deltas(
1927                num_non_last in 0usize..=20,
1928            ) {
1929                let instrument_id = InstrumentId::from("ESM4.GLBX");
1930                let mut buffering_start = None;
1931                let mut buffered = AHashMap::new();
1932                let total = num_non_last + 1;
1933
1934                for i in 0..num_non_last {
1935                    let result = process_mbo_delta(
1936                        test_delta(instrument_id, 1_000_000_000 + i as u64),
1937                        0, // No F_LAST
1938                        &mut buffering_start,
1939                        &mut buffered,
1940                    );
1941                    prop_assert!(result.is_none());
1942                }
1943
1944                let result = process_mbo_delta(
1945                    test_delta(instrument_id, 2_000_000_000),
1946                    128, // F_LAST
1947                    &mut buffering_start,
1948                    &mut buffered,
1949                );
1950
1951                prop_assert!(result.is_some());
1952                let emitted = result.unwrap();
1953                prop_assert_eq!(emitted.deltas.len(), total);
1954                prop_assert!(buffered.is_empty());
1955            }
1956
1957            #[rstest]
1958            fn mbo_snapshots_never_emit(
1959                num_snapshots in 1usize..=20,
1960            ) {
1961                let instrument_id = InstrumentId::from("ESM4.GLBX");
1962                let mut buffering_start = None;
1963                let mut buffered = AHashMap::new();
1964
1965                for i in 0..num_snapshots {
1966                    let result = process_mbo_delta(
1967                        test_delta(instrument_id, 1_000_000_000 + i as u64),
1968                        128 | 32, // F_LAST | F_SNAPSHOT
1969                        &mut buffering_start,
1970                        &mut buffered,
1971                    );
1972                    prop_assert!(result.is_none());
1973                }
1974
1975                prop_assert_eq!(buffered[&instrument_id].len(), num_snapshots);
1976            }
1977
1978            #[rstest]
1979            fn mbo_replay_delays_emission(
1980                start_offset in 1u64..=100,
1981                num_before in 1usize..=10,
1982            ) {
1983                let instrument_id = InstrumentId::from("ESM4.GLBX");
1984                let start_ns = 1_000_000_000u64 * start_offset;
1985                let mut buffering_start = Some(start_ns.into());
1986                let mut buffered = AHashMap::new();
1987
1988                for i in 0..num_before {
1989                    let ts = start_ns - (num_before as u64 - i as u64);
1990                    let result = process_mbo_delta(
1991                        test_delta(instrument_id, ts),
1992                        128, // F_LAST
1993                        &mut buffering_start,
1994                        &mut buffered,
1995                    );
1996                    prop_assert!(result.is_none());
1997                }
1998
1999                let result = process_mbo_delta(
2000                    test_delta(instrument_id, start_ns + 1),
2001                    128,
2002                    &mut buffering_start,
2003                    &mut buffered,
2004                );
2005
2006                prop_assert!(result.is_some());
2007                prop_assert!(buffering_start.is_none());
2008                let total = num_before + 1;
2009                prop_assert_eq!(result.unwrap().deltas.len(), total);
2010            }
2011        }
2012    }
2013}