nautilus_persistence/catalog/
worker.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//! Single-owner catalog worker for serialized query/write access.
17
18use std::{
19    sync::mpsc::{self, Receiver, Sender, SyncSender},
20    thread::{self, JoinHandle},
21};
22
23use ahash::AHashMap;
24use nautilus_core::{Params, UUID4, UnixNanos};
25use nautilus_model::{
26    data::{DataBatch, NautilusDataType},
27    instruments::InstrumentAny,
28};
29
30use super::{
31    traits::DataCatalog,
32    types::{CatalogInstrumentQuery, CatalogQuery},
33};
34use crate::{catalog::session::DataBatchQueryResult, common::coverage::CoverageIntervals};
35
36type Reply<T> = Sender<T>;
37type MissingIntervalsByIdentifier = AHashMap<String, Vec<(u64, u64)>>;
38type CoverageIntervalsByIdentifier = AHashMap<String, CoverageIntervals>;
39
40/// Commands buffered before senders block.
41///
42/// The queue is bounded so a producer that outruns the catalog backend applies backpressure
43/// instead of growing without limit. Both `write_async` and `query_batch_async` block once it is
44/// full.
45const COMMAND_QUEUE_CAPACITY: usize = 10;
46
47/// Query sessions the worker keeps open at once.
48///
49/// Each open session holds its backend query state until it is drained or closed, so the count is
50/// bounded to turn a caller that never closes sessions into an error instead of an unbounded leak.
51const MAX_OPEN_SESSIONS: usize = 64;
52
53#[derive(Debug)]
54pub struct CatalogWriteJob {
55    pub data: DataBatch,
56    pub start: Option<UnixNanos>,
57    pub end: Option<UnixNanos>,
58    pub params: Option<Params>,
59}
60
61pub enum CatalogCommand {
62    QueryLastTimestamp {
63        data_type: NautilusDataType,
64        identifier: Option<String>,
65        reply: Reply<anyhow::Result<Option<u64>>>,
66    },
67    GetMissingIntervals {
68        start: UnixNanos,
69        end: UnixNanos,
70        data_type: NautilusDataType,
71        identifier: Option<String>,
72        reply: Reply<anyhow::Result<Vec<(u64, u64)>>>,
73    },
74    GetMissingIntervalsForIdentifiers {
75        start: UnixNanos,
76        end: UnixNanos,
77        data_type: NautilusDataType,
78        identifiers: Vec<String>,
79        reply: Reply<anyhow::Result<MissingIntervalsByIdentifier>>,
80    },
81    GetCoverageIntervalsForIdentifiers {
82        start: UnixNanos,
83        end: UnixNanos,
84        data_type: NautilusDataType,
85        identifiers: Vec<String>,
86        reply: Reply<anyhow::Result<CoverageIntervalsByIdentifier>>,
87    },
88    QueryBatch {
89        query: CatalogQuery,
90        reply: Reply<anyhow::Result<DataBatch>>,
91    },
92    QueryBatchAsync {
93        query: CatalogQuery,
94        on_complete: Box<dyn FnOnce(anyhow::Result<DataBatch>) + Send>,
95    },
96    OpenSession {
97        query: CatalogQuery,
98        chunk_size: Option<usize>,
99        reply: Reply<anyhow::Result<UUID4>>,
100    },
101    PullSession {
102        session_id: UUID4,
103        reply: Reply<anyhow::Result<Option<DataBatch>>>,
104    },
105    PullSessionAsync {
106        session_id: UUID4,
107        on_complete: Box<dyn FnOnce(anyhow::Result<Option<DataBatch>>) + Send>,
108    },
109    CloseSession {
110        session_id: UUID4,
111        reply: Reply<anyhow::Result<bool>>,
112    },
113    QueryInstruments {
114        query: CatalogInstrumentQuery,
115        reply: Reply<anyhow::Result<Vec<InstrumentAny>>>,
116    },
117    Write {
118        job: CatalogWriteJob,
119        reply: Reply<anyhow::Result<()>>,
120    },
121    WriteAsync {
122        job: CatalogWriteJob,
123    },
124    WriteInstruments {
125        instruments: Vec<InstrumentAny>,
126        reply: Reply<anyhow::Result<()>>,
127    },
128    WriteInstrumentsAsync {
129        instruments: Vec<InstrumentAny>,
130    },
131    Flush {
132        reply: Reply<anyhow::Result<()>>,
133    },
134    Shutdown,
135}
136
137/// Serializes catalog work on a bounded queue owned by one worker thread.
138///
139/// Asynchronous queries deliver their result through callbacks, but submitting work can block when
140/// the queue is full. Close and shutdown wait behind in-flight backend operations; dropping the
141/// worker joins its thread. Caller-side timeouts do not cancel backend queries.
142#[derive(Debug)]
143pub struct CatalogWorker {
144    sender: SyncSender<CatalogCommand>,
145    handle: Option<JoinHandle<()>>,
146}
147
148#[expect(
149    clippy::missing_errors_doc,
150    reason = "Worker methods forward channel and catalog errors directly"
151)]
152impl CatalogWorker {
153    #[must_use]
154    pub fn start(mut catalog: DataCatalog) -> Self {
155        let (sender, receiver) = mpsc::sync_channel(COMMAND_QUEUE_CAPACITY);
156        let handle = thread::spawn(move || run_catalog_worker(&mut catalog, receiver));
157
158        Self {
159            sender,
160            handle: Some(handle),
161        }
162    }
163
164    fn send(&self, command: CatalogCommand) -> anyhow::Result<()> {
165        self.sender.send(command).map_err(|_| {
166            anyhow::anyhow!("Catalog worker thread has stopped, so the command cannot be sent")
167        })
168    }
169
170    fn send_with_reply<T>(
171        &self,
172        command: impl FnOnce(Reply<anyhow::Result<T>>) -> CatalogCommand,
173    ) -> anyhow::Result<T> {
174        let (reply, receiver) = mpsc::channel();
175        self.send(command(reply))?;
176        receiver.recv().map_err(|_| {
177            anyhow::anyhow!("Catalog worker thread stopped before replying to the command")
178        })?
179    }
180
181    pub fn query_last_timestamp(
182        &self,
183        data_type: NautilusDataType,
184        identifier: Option<String>,
185    ) -> anyhow::Result<Option<u64>> {
186        self.send_with_reply(|reply| CatalogCommand::QueryLastTimestamp {
187            data_type,
188            identifier,
189            reply,
190        })
191    }
192
193    pub fn get_missing_intervals(
194        &self,
195        start: UnixNanos,
196        end: UnixNanos,
197        data_type: NautilusDataType,
198        identifier: Option<String>,
199    ) -> anyhow::Result<Vec<(u64, u64)>> {
200        self.send_with_reply(|reply| CatalogCommand::GetMissingIntervals {
201            start,
202            end,
203            data_type,
204            identifier,
205            reply,
206        })
207    }
208
209    pub fn get_missing_intervals_for_identifiers(
210        &self,
211        start: UnixNanos,
212        end: UnixNanos,
213        data_type: NautilusDataType,
214        identifiers: Vec<String>,
215    ) -> anyhow::Result<MissingIntervalsByIdentifier> {
216        self.send_with_reply(|reply| CatalogCommand::GetMissingIntervalsForIdentifiers {
217            start,
218            end,
219            data_type,
220            identifiers,
221            reply,
222        })
223    }
224
225    pub fn get_coverage_intervals_for_identifiers(
226        &self,
227        start: UnixNanos,
228        end: UnixNanos,
229        data_type: NautilusDataType,
230        identifiers: Vec<String>,
231    ) -> anyhow::Result<CoverageIntervalsByIdentifier> {
232        self.send_with_reply(|reply| CatalogCommand::GetCoverageIntervalsForIdentifiers {
233            start,
234            end,
235            data_type,
236            identifiers,
237            reply,
238        })
239    }
240
241    pub fn query_batch(&self, query: CatalogQuery) -> anyhow::Result<DataBatch> {
242        self.send_with_reply(|reply| CatalogCommand::QueryBatch { query, reply })
243    }
244
245    /// Enqueues a batch query and invokes `on_complete` with the result on the
246    /// worker thread, without blocking the caller.
247    pub fn query_batch_async(
248        &self,
249        query: CatalogQuery,
250        on_complete: Box<dyn FnOnce(anyhow::Result<DataBatch>) + Send>,
251    ) -> anyhow::Result<()> {
252        self.send(CatalogCommand::QueryBatchAsync { query, on_complete })
253    }
254
255    pub fn open_session(
256        &self,
257        query: CatalogQuery,
258        chunk_size: Option<usize>,
259    ) -> anyhow::Result<UUID4> {
260        self.send_with_reply(|reply| CatalogCommand::OpenSession {
261            query,
262            chunk_size,
263            reply,
264        })
265    }
266
267    pub fn pull_session(&self, session_id: UUID4) -> anyhow::Result<Option<DataBatch>> {
268        self.send_with_reply(|reply| CatalogCommand::PullSession { session_id, reply })
269    }
270
271    /// Enqueues one session pull and invokes `on_complete` on the catalog worker.
272    pub fn pull_session_async(
273        &self,
274        session_id: UUID4,
275        on_complete: Box<dyn FnOnce(anyhow::Result<Option<DataBatch>>) + Send>,
276    ) -> anyhow::Result<()> {
277        self.send(CatalogCommand::PullSessionAsync {
278            session_id,
279            on_complete,
280        })
281    }
282
283    /// Closes a session, returning whether it was still open.
284    pub fn close_session(&self, session_id: UUID4) -> anyhow::Result<bool> {
285        self.send_with_reply(|reply| CatalogCommand::CloseSession { session_id, reply })
286    }
287
288    pub fn query_instruments(
289        &self,
290        query: CatalogInstrumentQuery,
291    ) -> anyhow::Result<Vec<InstrumentAny>> {
292        self.send_with_reply(|reply| CatalogCommand::QueryInstruments { query, reply })
293    }
294
295    pub fn write(&self, job: CatalogWriteJob) -> anyhow::Result<()> {
296        self.send_with_reply(|reply| CatalogCommand::Write { job, reply })
297    }
298
299    pub fn write_async(&self, job: CatalogWriteJob) -> anyhow::Result<()> {
300        self.send(CatalogCommand::WriteAsync { job })
301    }
302
303    pub fn write_instruments(&self, instruments: Vec<InstrumentAny>) -> anyhow::Result<()> {
304        self.send_with_reply(|reply| CatalogCommand::WriteInstruments { instruments, reply })
305    }
306
307    pub fn write_instruments_async(&self, instruments: Vec<InstrumentAny>) -> anyhow::Result<()> {
308        self.send(CatalogCommand::WriteInstrumentsAsync { instruments })
309    }
310
311    pub fn flush(&self) -> anyhow::Result<()> {
312        self.send_with_reply(|reply| CatalogCommand::Flush { reply })
313    }
314}
315
316impl Drop for CatalogWorker {
317    fn drop(&mut self) {
318        let _ = self.sender.send(CatalogCommand::Shutdown);
319
320        if let Some(handle) = self.handle.take() {
321            let _ = handle.join();
322        }
323    }
324}
325
326#[expect(
327    clippy::needless_pass_by_value,
328    clippy::too_many_lines,
329    reason = "The worker thread owns the receiver for its whole lifetime"
330)]
331fn run_catalog_worker(catalog: &mut DataCatalog, receiver: Receiver<CatalogCommand>) {
332    let mut async_errors: Vec<anyhow::Error> = Vec::new();
333    let mut sessions: AHashMap<UUID4, DataBatchQueryResult> = AHashMap::new();
334
335    while let Ok(command) = receiver.recv() {
336        match command {
337            CatalogCommand::QueryLastTimestamp {
338                data_type,
339                identifier,
340                reply,
341            } => {
342                let result = catalog.query_last_timestamp(data_type, identifier.as_deref());
343                let _ = reply.send(result);
344            }
345            CatalogCommand::GetMissingIntervals {
346                start,
347                end,
348                data_type,
349                identifier,
350                reply,
351            } => {
352                let result = catalog.get_missing_intervals_for_request(
353                    start,
354                    end,
355                    data_type,
356                    identifier.as_deref(),
357                );
358                let _ = reply.send(result);
359            }
360            CatalogCommand::GetMissingIntervalsForIdentifiers {
361                start,
362                end,
363                data_type,
364                identifiers,
365                reply,
366            } => {
367                let result = catalog.get_missing_intervals_for_identifiers(
368                    start,
369                    end,
370                    data_type,
371                    &identifiers,
372                );
373                let _ = reply.send(result);
374            }
375            CatalogCommand::GetCoverageIntervalsForIdentifiers {
376                start,
377                end,
378                data_type,
379                identifiers,
380                reply,
381            } => {
382                let result = catalog.get_coverage_intervals_for_identifiers(
383                    start,
384                    end,
385                    data_type,
386                    &identifiers,
387                );
388                let _ = reply.send(result);
389            }
390            CatalogCommand::QueryBatch { query, reply } => {
391                let result = catalog.query_batch(&query);
392                let _ = reply.send(result);
393            }
394            CatalogCommand::QueryBatchAsync { query, on_complete } => {
395                let result = catalog.query_batch(&query);
396                on_complete(result);
397            }
398            CatalogCommand::OpenSession {
399                query,
400                chunk_size,
401                reply,
402            } => {
403                let result = if sessions.len() >= MAX_OPEN_SESSIONS {
404                    Err(anyhow::anyhow!(
405                        "Catalog session limit of {MAX_OPEN_SESSIONS} reached; close finished sessions before opening more",
406                    ))
407                } else {
408                    open_query_session(catalog, &query, chunk_size).map(|session| {
409                        let session_id = UUID4::new();
410                        sessions.insert(session_id, session);
411                        session_id
412                    })
413                };
414
415                let _ = reply.send(result);
416            }
417            CatalogCommand::PullSession { session_id, reply } => {
418                let result = pull_session(&mut sessions, session_id);
419                let _ = reply.send(result);
420            }
421            CatalogCommand::PullSessionAsync {
422                session_id,
423                on_complete,
424            } => {
425                on_complete(pull_session(&mut sessions, session_id));
426            }
427            CatalogCommand::CloseSession { session_id, reply } => {
428                let _ = reply.send(Ok(sessions.remove(&session_id).is_some()));
429            }
430            CatalogCommand::QueryInstruments { query, reply } => {
431                let result = catalog.instruments(&query);
432                let _ = reply.send(result);
433            }
434            CatalogCommand::Write { job, reply } => {
435                let result = catalog.write_data_batch(&job.data, job.start, job.end, job.params);
436                let _ = reply.send(result);
437            }
438            CatalogCommand::WriteAsync { job } => {
439                let len = job.data.len();
440                match catalog.write_data_batch(&job.data, job.start, job.end, job.params) {
441                    Ok(()) => log::info!("Catalog worker wrote {len} data rows"),
442                    Err(e) => {
443                        log::error!("Catalog worker failed to write {len} data rows: {e}");
444                        async_errors.push(e);
445                    }
446                }
447            }
448            CatalogCommand::WriteInstruments { instruments, reply } => {
449                let result = catalog.write_instruments(&instruments);
450                let _ = reply.send(result);
451            }
452            CatalogCommand::WriteInstrumentsAsync { instruments } => {
453                let len = instruments.len();
454                match catalog.write_instruments(&instruments) {
455                    Ok(()) => log::info!("Catalog worker wrote {len} instruments"),
456                    Err(e) => {
457                        log::error!("Catalog worker failed to write {len} instruments: {e}");
458                        async_errors.push(e);
459                    }
460                }
461            }
462            CatalogCommand::Flush { reply } => {
463                let _ = reply.send(drain_async_errors(&mut async_errors));
464            }
465            CatalogCommand::Shutdown => break,
466        }
467    }
468}
469
470fn open_query_session(
471    catalog: &mut DataCatalog,
472    query: &CatalogQuery,
473    chunk_size: Option<usize>,
474) -> anyhow::Result<DataBatchQueryResult> {
475    if let Some(mut query_catalog) = catalog.fork_query_catalog()? {
476        query_catalog.reset_session();
477        return query_catalog.query_batch_session(query, chunk_size);
478    }
479
480    catalog.reset_session();
481    catalog.query_batch_session(query, chunk_size)
482}
483
484fn drain_async_errors(errors: &mut Vec<anyhow::Error>) -> anyhow::Result<()> {
485    let failures = errors.len();
486    let details = errors
487        .iter()
488        .map(|e| format!("{e:#}"))
489        .collect::<Vec<_>>()
490        .join("; ");
491
492    let Some(first) = errors.drain(..).next() else {
493        return Ok(());
494    };
495
496    if failures == 1 {
497        return Err(first);
498    }
499
500    Err(first.context(format!(
501        "{failures} asynchronous catalog writes failed: {details}"
502    )))
503}
504
505fn pull_session(
506    sessions: &mut AHashMap<UUID4, DataBatchQueryResult>,
507    session_id: UUID4,
508) -> anyhow::Result<Option<DataBatch>> {
509    let result = sessions
510        .get_mut(&session_id)
511        .ok_or_else(|| anyhow::anyhow!("Catalog session {session_id} is not open"))?
512        .next_batch();
513
514    if !matches!(result, Ok(Some(_))) {
515        sessions.remove(&session_id);
516    }
517
518    result
519}
520
521#[cfg(test)]
522mod tests {
523    use std::{
524        sync::{Arc, Mutex},
525        time::{Duration, Instant},
526    };
527
528    use nautilus_core::ClosedInterval;
529    use nautilus_model::{
530        data::{Data, NautilusRecordType, QuoteTick},
531        identifiers::InstrumentId,
532        instruments::stubs::audusd_sim,
533        types::{Price, Quantity},
534    };
535    use rstest::rstest;
536
537    use super::*;
538    use crate::catalog::{
539        session::{DataBatchQuery, TypedDataBatchSession},
540        traits::{CatalogReader, CatalogWriter, RecordBatch},
541    };
542
543    struct FailingSession;
544
545    impl DataBatchQuery for FailingSession {
546        fn next_batch(&mut self) -> anyhow::Result<Option<DataBatch>> {
547            anyhow::bail!("session decode failure")
548        }
549    }
550
551    fn stub_quote() -> QuoteTick {
552        QuoteTick::new(
553            InstrumentId::from("ETHUSDT-PERP.BINANCE"),
554            Price::from("1987.0"),
555            Price::from("1988.0"),
556            Quantity::from("100"),
557            Quantity::from("100"),
558            UnixNanos::from(1),
559            UnixNanos::from(1),
560        )
561    }
562
563    #[derive(Debug)]
564    struct StubCatalog {
565        fail: bool,
566        query_thread: Arc<Mutex<Option<thread::ThreadId>>>,
567    }
568
569    impl CatalogReader for StubCatalog {
570        fn reset_session(&mut self) {}
571
572        fn instruments(
573            &mut self,
574            _query: &CatalogInstrumentQuery,
575        ) -> anyhow::Result<Vec<InstrumentAny>> {
576            Ok(Vec::new())
577        }
578
579        fn query_batch(&mut self, _query: &CatalogQuery) -> anyhow::Result<DataBatch> {
580            *self.query_thread.lock().unwrap() = Some(thread::current().id());
581
582            if self.fail {
583                anyhow::bail!("stub query failure");
584            }
585
586            Ok(DataBatch::Quote(vec![stub_quote()].into()))
587        }
588
589        fn query_batch_session(
590            &mut self,
591            _query: &CatalogQuery,
592            chunk_size: Option<usize>,
593        ) -> anyhow::Result<DataBatchQueryResult> {
594            *self.query_thread.lock().unwrap() = Some(thread::current().id());
595
596            if self.fail {
597                anyhow::bail!("stub query failure");
598            }
599
600            Ok(Box::new(TypedDataBatchSession::from_vec(
601                vec![stub_quote()],
602                chunk_size,
603            )))
604        }
605    }
606
607    impl CatalogWriter for StubCatalog {
608        fn write_instruments(&mut self, _instruments: &[InstrumentAny]) -> anyhow::Result<()> {
609            if self.fail {
610                anyhow::bail!("stub instrument write failure");
611            }
612
613            Ok(())
614        }
615
616        fn write_data(
617            &mut self,
618            _data: &[Data],
619            _start: Option<UnixNanos>,
620            _end: Option<UnixNanos>,
621            params: Option<Params>,
622        ) -> anyhow::Result<()> {
623            if self.fail {
624                let message = params
625                    .as_ref()
626                    .and_then(|params| params.get_str("test_error"))
627                    .unwrap_or("stub write failure");
628                anyhow::bail!(message.to_string());
629            }
630
631            Ok(())
632        }
633
634        fn write_records(
635            &mut self,
636            _record_type: NautilusRecordType,
637            _batches: &[RecordBatch],
638            _params: Option<Params>,
639        ) -> anyhow::Result<()> {
640            Ok(())
641        }
642
643        fn record_empty_coverage(
644            &mut self,
645            _data_type: NautilusDataType,
646            _identifier: Option<&str>,
647            _start: UnixNanos,
648            _end: UnixNanos,
649        ) -> anyhow::Result<()> {
650            Ok(())
651        }
652    }
653
654    fn stub_query() -> CatalogQuery {
655        CatalogQuery::new(NautilusDataType::QuoteTick)
656    }
657
658    #[rstest]
659    fn test_query_batch_async_invokes_callback_on_worker_thread() {
660        let query_thread = Arc::new(Mutex::new(None));
661
662        let worker = CatalogWorker::start(Box::new(StubCatalog {
663            fail: false,
664            query_thread: query_thread.clone(),
665        }));
666
667        let (tx, rx) = mpsc::channel();
668        worker
669            .query_batch_async(
670                stub_query(),
671                Box::new(move |result| {
672                    let _ = tx.send((thread::current().id(), result.map(|batch| batch.len())));
673                }),
674            )
675            .unwrap();
676
677        let (callback_thread, result) = rx
678            .recv_timeout(Duration::from_secs(5))
679            .expect("callback was not invoked");
680        assert_eq!(result.unwrap(), 1);
681        assert_ne!(callback_thread, thread::current().id());
682        assert_eq!(Some(callback_thread), *query_thread.lock().unwrap());
683    }
684
685    #[rstest]
686    fn test_query_batch_async_passes_query_error_to_callback() {
687        let worker = CatalogWorker::start(Box::new(StubCatalog {
688            fail: true,
689            query_thread: Arc::new(Mutex::new(None)),
690        }));
691
692        let (tx, rx) = mpsc::channel();
693        worker
694            .query_batch_async(
695                stub_query(),
696                Box::new(move |result| {
697                    let _ = tx.send(result.map(|batch| batch.len()));
698                }),
699            )
700            .unwrap();
701
702        let result = rx
703            .recv_timeout(Duration::from_secs(5))
704            .expect("callback was not invoked");
705        assert_eq!(result.unwrap_err().to_string(), "stub query failure");
706    }
707
708    #[rstest]
709    fn test_worker_owns_and_pulls_catalog_session() {
710        let query_thread = Arc::new(Mutex::new(None));
711
712        let worker = CatalogWorker::start(Box::new(StubCatalog {
713            fail: false,
714            query_thread: query_thread.clone(),
715        }));
716
717        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
718        let batch = worker.pull_session(session_id).unwrap().unwrap();
719        let complete = worker.pull_session(session_id).unwrap();
720        let closed = worker.pull_session(session_id).unwrap_err();
721
722        assert_eq!(batch.len(), 1);
723        assert!(complete.is_none());
724        assert_eq!(
725            closed.to_string(),
726            format!("Catalog session {session_id} is not open")
727        );
728        assert_ne!(*query_thread.lock().unwrap(), Some(thread::current().id()));
729    }
730
731    #[rstest]
732    fn test_close_session_drops_open_session() {
733        let worker = CatalogWorker::start(Box::new(StubCatalog {
734            fail: false,
735            query_thread: Arc::new(Mutex::new(None)),
736        }));
737
738        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
739
740        let first_close = worker.close_session(session_id).unwrap();
741        let second_close = worker.close_session(session_id).unwrap();
742
743        assert!(first_close);
744        assert!(!second_close);
745    }
746
747    #[rstest]
748    fn errored_session_is_removed_after_pull() {
749        let session_id = UUID4::new();
750        let mut sessions = AHashMap::new();
751        sessions.insert(session_id, Box::new(FailingSession) as DataBatchQueryResult);
752
753        let first = pull_session(&mut sessions, session_id).unwrap_err();
754        let second = pull_session(&mut sessions, session_id).unwrap_err();
755
756        assert_eq!(first.to_string(), "session decode failure");
757        assert_eq!(
758            second.to_string(),
759            format!("Catalog session {session_id} is not open"),
760        );
761        assert!(sessions.is_empty());
762    }
763
764    #[rstest]
765    fn test_flush_reports_every_failed_async_write_then_clears() {
766        let worker = CatalogWorker::start(Box::new(StubCatalog {
767            fail: true,
768            query_thread: Arc::new(Mutex::new(None)),
769        }));
770
771        for message in [
772            "first write failed",
773            "second write failed",
774            "third write failed",
775        ] {
776            let mut params = Params::new();
777            params.insert("test_error".to_string(), message.into());
778            worker
779                .write_async(CatalogWriteJob {
780                    data: DataBatch::Quote(vec![stub_quote()].into()),
781                    start: None,
782                    end: None,
783                    params: Some(params),
784                })
785                .unwrap();
786        }
787
788        let error = worker.flush().unwrap_err();
789
790        assert_eq!(
791            error.to_string(),
792            "3 asynchronous catalog writes failed: first write failed; second write failed; \
793             third write failed",
794        );
795        assert_eq!(error.root_cause().to_string(), "first write failed");
796        assert!(worker.flush().is_ok(), "flush must clear reported failures");
797    }
798
799    #[rstest]
800    fn test_open_session_rejects_beyond_the_session_limit() {
801        let worker = CatalogWorker::start(Box::new(StubCatalog {
802            fail: false,
803            query_thread: Arc::new(Mutex::new(None)),
804        }));
805
806        let session_ids = (0..MAX_OPEN_SESSIONS)
807            .map(|_| worker.open_session(stub_query(), Some(1)).unwrap())
808            .collect::<Vec<_>>();
809
810        let error = worker.open_session(stub_query(), Some(1)).unwrap_err();
811        assert_eq!(
812            error.to_string(),
813            format!(
814                "Catalog session limit of {MAX_OPEN_SESSIONS} reached; close finished sessions before opening more"
815            ),
816        );
817
818        assert!(worker.close_session(session_ids[0]).unwrap());
819        assert!(
820            worker.open_session(stub_query(), Some(1)).is_ok(),
821            "closing a session must free a slot",
822        );
823    }
824
825    #[rstest]
826    fn test_pull_session_async_invokes_callback_on_worker_thread() {
827        let worker = CatalogWorker::start(Box::new(StubCatalog {
828            fail: false,
829            query_thread: Arc::new(Mutex::new(None)),
830        }));
831
832        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
833        let (tx, rx) = mpsc::channel();
834
835        worker
836            .pull_session_async(
837                session_id,
838                Box::new(move |result| {
839                    let _ = tx.send((
840                        thread::current().id(),
841                        result.map(|batch| batch.unwrap().len()),
842                    ));
843                }),
844            )
845            .unwrap();
846
847        let (callback_thread, result) = rx.recv_timeout(Duration::from_secs(5)).unwrap();
848
849        assert_eq!(result.unwrap(), 1);
850        assert_ne!(callback_thread, thread::current().id());
851    }
852
853    #[rstest]
854    fn test_flush_returns_single_async_failure_without_context() {
855        let worker = CatalogWorker::start(Box::new(StubCatalog {
856            fail: true,
857            query_thread: Arc::new(Mutex::new(None)),
858        }));
859
860        worker
861            .write_instruments_async(vec![InstrumentAny::CurrencyPair(audusd_sim())])
862            .unwrap();
863        let error = worker.flush().unwrap_err();
864
865        assert_eq!(error.to_string(), "stub instrument write failure");
866        assert_eq!(error.chain().count(), 1);
867    }
868
869    #[derive(Debug, Default)]
870    struct RecordingCatalog {
871        label: &'static str,
872        calls: Arc<Mutex<Vec<String>>>,
873        write_gate: Option<Receiver<()>>,
874        fork: bool,
875        panic_on_query: bool,
876    }
877
878    impl RecordingCatalog {
879        fn record(&self, call: &str) {
880            self.calls
881                .lock()
882                .unwrap()
883                .push(format!("{}{call}", self.label));
884        }
885    }
886
887    impl CatalogReader for RecordingCatalog {
888        fn fork_query_catalog(&self) -> anyhow::Result<Option<DataCatalog>> {
889            if !self.fork {
890                return Ok(None);
891            }
892
893            Ok(Some(Box::new(Self {
894                label: "fork ",
895                calls: self.calls.clone(),
896                ..Self::default()
897            })))
898        }
899
900        fn reset_session(&mut self) {
901            self.record("reset_session");
902        }
903
904        fn instruments(
905            &mut self,
906            query: &CatalogInstrumentQuery,
907        ) -> anyhow::Result<Vec<InstrumentAny>> {
908            self.record(&format!("instruments {:?}", query.instrument_ids));
909            Ok(vec![InstrumentAny::CurrencyPair(audusd_sim())])
910        }
911
912        #[expect(
913            clippy::panic_in_result_fn,
914            reason = "the stub simulates a catalog backend that panics mid-query"
915        )]
916        fn query_batch(&mut self, _query: &CatalogQuery) -> anyhow::Result<DataBatch> {
917            assert!(!self.panic_on_query, "stub query panic");
918
919            Ok(DataBatch::Quote(vec![stub_quote()].into()))
920        }
921
922        fn query_batch_session(
923            &mut self,
924            _query: &CatalogQuery,
925            chunk_size: Option<usize>,
926        ) -> anyhow::Result<DataBatchQueryResult> {
927            self.record(&format!("query_batch_session {chunk_size:?}"));
928            Ok(Box::new(TypedDataBatchSession::from_vec(
929                vec![stub_quote()],
930                chunk_size,
931            )))
932        }
933
934        fn get_missing_intervals_for_request(
935            &mut self,
936            start: UnixNanos,
937            end: UnixNanos,
938            _data_type: NautilusDataType,
939            identifier: Option<&str>,
940        ) -> anyhow::Result<Vec<(u64, u64)>> {
941            if identifier == Some("covered") {
942                return Ok(Vec::new());
943            }
944
945            Ok(vec![(start.as_u64(), end.as_u64())])
946        }
947
948        fn query_last_timestamp(
949            &mut self,
950            data_type: NautilusDataType,
951            identifier: Option<&str>,
952        ) -> anyhow::Result<Option<u64>> {
953            self.record(&format!("query_last_timestamp {data_type} {identifier:?}"));
954            Ok(Some(42))
955        }
956    }
957
958    impl CatalogWriter for RecordingCatalog {
959        fn write_instruments(&mut self, instruments: &[InstrumentAny]) -> anyhow::Result<()> {
960            self.record(&format!("write_instruments {}", instruments.len()));
961            Ok(())
962        }
963
964        fn write_data(
965            &mut self,
966            data: &[Data],
967            start: Option<UnixNanos>,
968            end: Option<UnixNanos>,
969            _params: Option<Params>,
970        ) -> anyhow::Result<()> {
971            if let Some(gate) = &self.write_gate {
972                gate.recv()?;
973            }
974
975            self.record(&format!(
976                "write_data {} {:?} {:?}",
977                data.len(),
978                start.map(|ts| ts.as_u64()),
979                end.map(|ts| ts.as_u64()),
980            ));
981            Ok(())
982        }
983
984        fn write_records(
985            &mut self,
986            _record_type: NautilusRecordType,
987            _batches: &[RecordBatch],
988            _params: Option<Params>,
989        ) -> anyhow::Result<()> {
990            Ok(())
991        }
992
993        fn record_empty_coverage(
994            &mut self,
995            _data_type: NautilusDataType,
996            _identifier: Option<&str>,
997            _start: UnixNanos,
998            _end: UnixNanos,
999        ) -> anyhow::Result<()> {
1000            Ok(())
1001        }
1002    }
1003
1004    fn quote_write_job() -> CatalogWriteJob {
1005        CatalogWriteJob {
1006            data: DataBatch::Quote(vec![stub_quote()].into()),
1007            start: Some(UnixNanos::from(1)),
1008            end: Some(UnixNanos::from(2)),
1009            params: None,
1010        }
1011    }
1012
1013    #[rstest]
1014    fn test_worker_forwards_queries_and_writes_to_catalog() {
1015        let calls = Arc::new(Mutex::new(Vec::new()));
1016
1017        let worker = CatalogWorker::start(Box::new(RecordingCatalog {
1018            calls: calls.clone(),
1019            ..RecordingCatalog::default()
1020        }));
1021
1022        let identifiers = vec!["missing".to_string(), "covered".to_string()];
1023
1024        let last_timestamp = worker
1025            .query_last_timestamp(NautilusDataType::QuoteTick, Some("missing".to_string()))
1026            .unwrap();
1027        let missing = worker
1028            .get_missing_intervals(
1029                UnixNanos::from(3),
1030                UnixNanos::from(9),
1031                NautilusDataType::QuoteTick,
1032                Some("missing".to_string()),
1033            )
1034            .unwrap();
1035        let missing_by_identifier = worker
1036            .get_missing_intervals_for_identifiers(
1037                UnixNanos::from(1),
1038                UnixNanos::from(10),
1039                NautilusDataType::QuoteTick,
1040                identifiers.clone(),
1041            )
1042            .unwrap();
1043        let coverage_by_identifier = worker
1044            .get_coverage_intervals_for_identifiers(
1045                UnixNanos::from(1),
1046                UnixNanos::from(10),
1047                NautilusDataType::QuoteTick,
1048                identifiers,
1049            )
1050            .unwrap();
1051        let batch = worker.query_batch(stub_query()).unwrap();
1052        let instruments = worker
1053            .query_instruments(
1054                CatalogInstrumentQuery::new()
1055                    .with_instrument_ids(Some(vec!["AUD/USD.SIM".to_string()])),
1056            )
1057            .unwrap();
1058        worker.write(quote_write_job()).unwrap();
1059        worker
1060            .write_instruments(vec![InstrumentAny::CurrencyPair(audusd_sim())])
1061            .unwrap();
1062        worker
1063            .write_instruments_async(vec![
1064                InstrumentAny::CurrencyPair(audusd_sim()),
1065                InstrumentAny::CurrencyPair(audusd_sim()),
1066            ])
1067            .unwrap();
1068        worker.flush().unwrap();
1069
1070        assert_eq!(last_timestamp, Some(42));
1071        assert_eq!(missing, vec![(3, 9)]);
1072        assert_eq!(
1073            missing_by_identifier,
1074            AHashMap::from_iter([
1075                ("missing".to_string(), vec![(1, 10)]),
1076                ("covered".to_string(), Vec::new()),
1077            ]),
1078        );
1079        assert_eq!(
1080            coverage_by_identifier,
1081            AHashMap::from_iter([
1082                ("missing".to_string(), CoverageIntervals::default()),
1083                (
1084                    "covered".to_string(),
1085                    CoverageIntervals {
1086                        data: vec![ClosedInterval::new(1, 10).unwrap()],
1087                        empty: Vec::new(),
1088                    },
1089                ),
1090            ]),
1091        );
1092        assert_eq!(batch.len(), 1);
1093        assert_eq!(instruments, vec![InstrumentAny::CurrencyPair(audusd_sim())]);
1094        assert_eq!(
1095            *calls.lock().unwrap(),
1096            vec![
1097                "query_last_timestamp QuoteTick Some(\"missing\")",
1098                "instruments Some([\"AUD/USD.SIM\"])",
1099                "write_data 1 Some(1) Some(2)",
1100                "write_instruments 1",
1101                "write_instruments 2",
1102            ],
1103        );
1104    }
1105
1106    #[rstest]
1107    fn test_drop_waits_for_queued_async_writes() {
1108        let calls = Arc::new(Mutex::new(Vec::new()));
1109        let (release, gate) = mpsc::channel();
1110
1111        let worker = CatalogWorker::start(Box::new(RecordingCatalog {
1112            calls: calls.clone(),
1113            write_gate: Some(gate),
1114            ..RecordingCatalog::default()
1115        }));
1116
1117        worker.write_async(quote_write_job()).unwrap();
1118
1119        let releaser = thread::spawn(move || {
1120            thread::sleep(Duration::from_millis(50));
1121            release.send(()).unwrap();
1122        });
1123
1124        drop(worker);
1125
1126        assert_eq!(*calls.lock().unwrap(), vec!["write_data 1 Some(1) Some(2)"]);
1127        releaser.join().unwrap();
1128    }
1129
1130    #[rstest]
1131    #[case::shared_catalog(false, &["reset_session", "query_batch_session Some(1)"])]
1132    #[case::forked_catalog(true, &["fork reset_session", "fork query_batch_session Some(1)"])]
1133    fn test_open_session_queries_forked_catalog_when_available(
1134        #[case] fork: bool,
1135        #[case] expected_calls: &[&str],
1136    ) {
1137        let calls = Arc::new(Mutex::new(Vec::new()));
1138
1139        let worker = CatalogWorker::start(Box::new(RecordingCatalog {
1140            calls: calls.clone(),
1141            fork,
1142            ..RecordingCatalog::default()
1143        }));
1144
1145        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
1146        let batch = worker.pull_session(session_id).unwrap().unwrap();
1147
1148        assert_eq!(batch.len(), 1);
1149        assert_eq!(*calls.lock().unwrap(), expected_calls);
1150    }
1151
1152    #[rstest]
1153    fn test_worker_reports_stopped_thread_after_catalog_panic() {
1154        let worker = CatalogWorker::start(Box::new(RecordingCatalog {
1155            panic_on_query: true,
1156            ..RecordingCatalog::default()
1157        }));
1158
1159        let reply_error = worker.query_batch(stub_query()).unwrap_err();
1160        let handle = worker.handle.as_ref().unwrap();
1161        let deadline = Instant::now() + Duration::from_secs(5);
1162
1163        while !handle.is_finished() {
1164            assert!(Instant::now() < deadline, "worker thread did not stop");
1165            thread::sleep(Duration::from_millis(1));
1166        }
1167
1168        let send_error = worker.flush().unwrap_err();
1169
1170        assert_eq!(
1171            reply_error.to_string(),
1172            "Catalog worker thread stopped before replying to the command",
1173        );
1174        assert_eq!(
1175            send_error.to_string(),
1176            "Catalog worker thread has stopped, so the command cannot be sent",
1177        );
1178    }
1179
1180    #[rstest]
1181    #[case(false)]
1182    #[case(true)]
1183    fn session_close_and_shutdown_wait_for_inflight_backend(#[case] shutdown: bool) {
1184        let query_thread = Arc::new(Mutex::new(None));
1185
1186        let worker = CatalogWorker::start(Box::new(StubCatalog {
1187            fail: false,
1188            query_thread: Arc::clone(&query_thread),
1189        }));
1190
1191        let session = worker.open_session(stub_query(), Some(1)).unwrap();
1192        let backend_gate = query_thread.lock().unwrap();
1193        let (query_tx, query_rx) = mpsc::channel();
1194        worker
1195            .query_batch_async(
1196                stub_query(),
1197                Box::new(move |result| {
1198                    query_tx.send(result.map(|batch| batch.len())).unwrap();
1199                }),
1200            )
1201            .unwrap();
1202
1203        let (started_tx, started_rx) = mpsc::channel();
1204        let (closed_tx, closed_rx) = mpsc::channel();
1205
1206        // Close reports before the worker drops, so the drop's join cannot mask a close that
1207        // returns early
1208        let control = thread::spawn(move || {
1209            started_tx.send(()).unwrap();
1210
1211            if shutdown {
1212                drop(worker);
1213                closed_tx.send(None).unwrap();
1214            } else {
1215                closed_tx
1216                    .send(Some(worker.close_session(session).unwrap()))
1217                    .unwrap();
1218            }
1219        });
1220
1221        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
1222        let blocked = closed_rx.recv_timeout(Duration::from_millis(25));
1223        drop(backend_gate);
1224        let rows = query_rx
1225            .recv_timeout(Duration::from_secs(5))
1226            .unwrap()
1227            .unwrap();
1228        let closed = closed_rx.recv_timeout(Duration::from_secs(5)).unwrap();
1229        control.join().unwrap();
1230        assert_eq!(blocked, Err(mpsc::RecvTimeoutError::Timeout));
1231        assert_eq!(rows, 1);
1232        assert_eq!(closed, (!shutdown).then_some(true));
1233    }
1234}
nautilus_persistence/catalog/worker.rs

nautilus_persistence/catalog/
worker.rs
