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 params = Params::new([("test_error".to_string(), message.into())].into());
777            worker
778                .write_async(CatalogWriteJob {
779                    data: DataBatch::Quote(vec![stub_quote()].into()),
780                    start: None,
781                    end: None,
782                    params: Some(params),
783                })
784                .unwrap();
785        }
786
787        let error = worker.flush().unwrap_err();
788
789        assert_eq!(
790            error.to_string(),
791            "3 asynchronous catalog writes failed: first write failed; second write failed; \
792             third write failed",
793        );
794        assert_eq!(error.root_cause().to_string(), "first write failed");
795        assert!(worker.flush().is_ok(), "flush must clear reported failures");
796    }
797
798    #[rstest]
799    fn test_open_session_rejects_beyond_the_session_limit() {
800        let worker = CatalogWorker::start(Box::new(StubCatalog {
801            fail: false,
802            query_thread: Arc::new(Mutex::new(None)),
803        }));
804
805        let session_ids = (0..MAX_OPEN_SESSIONS)
806            .map(|_| worker.open_session(stub_query(), Some(1)).unwrap())
807            .collect::<Vec<_>>();
808
809        let error = worker.open_session(stub_query(), Some(1)).unwrap_err();
810        assert_eq!(
811            error.to_string(),
812            format!(
813                "Catalog session limit of {MAX_OPEN_SESSIONS} reached; close finished sessions before opening more"
814            ),
815        );
816
817        assert!(worker.close_session(session_ids[0]).unwrap());
818        assert!(
819            worker.open_session(stub_query(), Some(1)).is_ok(),
820            "closing a session must free a slot",
821        );
822    }
823
824    #[rstest]
825    fn test_pull_session_async_invokes_callback_on_worker_thread() {
826        let worker = CatalogWorker::start(Box::new(StubCatalog {
827            fail: false,
828            query_thread: Arc::new(Mutex::new(None)),
829        }));
830
831        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
832        let (tx, rx) = mpsc::channel();
833
834        worker
835            .pull_session_async(
836                session_id,
837                Box::new(move |result| {
838                    let _ = tx.send((
839                        thread::current().id(),
840                        result.map(|batch| batch.unwrap().len()),
841                    ));
842                }),
843            )
844            .unwrap();
845
846        let (callback_thread, result) = rx.recv_timeout(Duration::from_secs(5)).unwrap();
847
848        assert_eq!(result.unwrap(), 1);
849        assert_ne!(callback_thread, thread::current().id());
850    }
851
852    #[rstest]
853    fn test_flush_returns_single_async_failure_without_context() {
854        let worker = CatalogWorker::start(Box::new(StubCatalog {
855            fail: true,
856            query_thread: Arc::new(Mutex::new(None)),
857        }));
858
859        worker
860            .write_instruments_async(vec![InstrumentAny::CurrencyPair(audusd_sim())])
861            .unwrap();
862        let error = worker.flush().unwrap_err();
863
864        assert_eq!(error.to_string(), "stub instrument write failure");
865        assert_eq!(error.chain().count(), 1);
866    }
867
868    #[derive(Debug, Default)]
869    struct RecordingCatalog {
870        label: &'static str,
871        calls: Arc<Mutex<Vec<String>>>,
872        write_gate: Option<Receiver<()>>,
873        fork: bool,
874        panic_on_query: bool,
875    }
876
877    impl RecordingCatalog {
878        fn record(&self, call: &str) {
879            self.calls
880                .lock()
881                .unwrap()
882                .push(format!("{}{call}", self.label));
883        }
884    }
885
886    impl CatalogReader for RecordingCatalog {
887        fn fork_query_catalog(&self) -> anyhow::Result<Option<DataCatalog>> {
888            if !self.fork {
889                return Ok(None);
890            }
891
892            Ok(Some(Box::new(Self {
893                label: "fork ",
894                calls: self.calls.clone(),
895                ..Self::default()
896            })))
897        }
898
899        fn reset_session(&mut self) {
900            self.record("reset_session");
901        }
902
903        fn instruments(
904            &mut self,
905            query: &CatalogInstrumentQuery,
906        ) -> anyhow::Result<Vec<InstrumentAny>> {
907            self.record(&format!("instruments {:?}", query.instrument_ids));
908            Ok(vec![InstrumentAny::CurrencyPair(audusd_sim())])
909        }
910
911        #[expect(
912            clippy::panic_in_result_fn,
913            reason = "the stub simulates a catalog backend that panics mid-query"
914        )]
915        fn query_batch(&mut self, _query: &CatalogQuery) -> anyhow::Result<DataBatch> {
916            assert!(!self.panic_on_query, "stub query panic");
917
918            Ok(DataBatch::Quote(vec![stub_quote()].into()))
919        }
920
921        fn query_batch_session(
922            &mut self,
923            _query: &CatalogQuery,
924            chunk_size: Option<usize>,
925        ) -> anyhow::Result<DataBatchQueryResult> {
926            self.record(&format!("query_batch_session {chunk_size:?}"));
927            Ok(Box::new(TypedDataBatchSession::from_vec(
928                vec![stub_quote()],
929                chunk_size,
930            )))
931        }
932
933        fn get_missing_intervals_for_request(
934            &mut self,
935            start: UnixNanos,
936            end: UnixNanos,
937            _data_type: NautilusDataType,
938            identifier: Option<&str>,
939        ) -> anyhow::Result<Vec<(u64, u64)>> {
940            if identifier == Some("covered") {
941                return Ok(Vec::new());
942            }
943
944            Ok(vec![(start.as_u64(), end.as_u64())])
945        }
946
947        fn query_last_timestamp(
948            &mut self,
949            data_type: NautilusDataType,
950            identifier: Option<&str>,
951        ) -> anyhow::Result<Option<u64>> {
952            self.record(&format!("query_last_timestamp {data_type} {identifier:?}"));
953            Ok(Some(42))
954        }
955    }
956
957    impl CatalogWriter for RecordingCatalog {
958        fn write_instruments(&mut self, instruments: &[InstrumentAny]) -> anyhow::Result<()> {
959            self.record(&format!("write_instruments {}", instruments.len()));
960            Ok(())
961        }
962
963        fn write_data(
964            &mut self,
965            data: &[Data],
966            start: Option<UnixNanos>,
967            end: Option<UnixNanos>,
968            _params: Option<Params>,
969        ) -> anyhow::Result<()> {
970            if let Some(gate) = &self.write_gate {
971                gate.recv()?;
972            }
973
974            self.record(&format!(
975                "write_data {} {:?} {:?}",
976                data.len(),
977                start.map(|ts| ts.as_u64()),
978                end.map(|ts| ts.as_u64()),
979            ));
980            Ok(())
981        }
982
983        fn write_records(
984            &mut self,
985            _record_type: NautilusRecordType,
986            _batches: &[RecordBatch],
987            _params: Option<Params>,
988        ) -> anyhow::Result<()> {
989            Ok(())
990        }
991
992        fn record_empty_coverage(
993            &mut self,
994            _data_type: NautilusDataType,
995            _identifier: Option<&str>,
996            _start: UnixNanos,
997            _end: UnixNanos,
998        ) -> anyhow::Result<()> {
999            Ok(())
1000        }
1001    }
1002
1003    fn quote_write_job() -> CatalogWriteJob {
1004        CatalogWriteJob {
1005            data: DataBatch::Quote(vec![stub_quote()].into()),
1006            start: Some(UnixNanos::from(1)),
1007            end: Some(UnixNanos::from(2)),
1008            params: None,
1009        }
1010    }
1011
1012    #[rstest]
1013    fn test_worker_forwards_queries_and_writes_to_catalog() {
1014        let calls = Arc::new(Mutex::new(Vec::new()));
1015
1016        let worker = CatalogWorker::start(Box::new(RecordingCatalog {
1017            calls: calls.clone(),
1018            ..RecordingCatalog::default()
1019        }));
1020
1021        let identifiers = vec!["missing".to_string(), "covered".to_string()];
1022
1023        let last_timestamp = worker
1024            .query_last_timestamp(NautilusDataType::QuoteTick, Some("missing".to_string()))
1025            .unwrap();
1026        let missing = worker
1027            .get_missing_intervals(
1028                UnixNanos::from(3),
1029                UnixNanos::from(9),
1030                NautilusDataType::QuoteTick,
1031                Some("missing".to_string()),
1032            )
1033            .unwrap();
1034        let missing_by_identifier = worker
1035            .get_missing_intervals_for_identifiers(
1036                UnixNanos::from(1),
1037                UnixNanos::from(10),
1038                NautilusDataType::QuoteTick,
1039                identifiers.clone(),
1040            )
1041            .unwrap();
1042        let coverage_by_identifier = worker
1043            .get_coverage_intervals_for_identifiers(
1044                UnixNanos::from(1),
1045                UnixNanos::from(10),
1046                NautilusDataType::QuoteTick,
1047                identifiers,
1048            )
1049            .unwrap();
1050        let batch = worker.query_batch(stub_query()).unwrap();
1051        let instruments = worker
1052            .query_instruments(
1053                CatalogInstrumentQuery::new()
1054                    .with_instrument_ids(Some(vec!["AUD/USD.SIM".to_string()])),
1055            )
1056            .unwrap();
1057        worker.write(quote_write_job()).unwrap();
1058        worker
1059            .write_instruments(vec![InstrumentAny::CurrencyPair(audusd_sim())])
1060            .unwrap();
1061        worker
1062            .write_instruments_async(vec![
1063                InstrumentAny::CurrencyPair(audusd_sim()),
1064                InstrumentAny::CurrencyPair(audusd_sim()),
1065            ])
1066            .unwrap();
1067        worker.flush().unwrap();
1068
1069        assert_eq!(last_timestamp, Some(42));
1070        assert_eq!(missing, vec![(3, 9)]);
1071        assert_eq!(
1072            missing_by_identifier,
1073            AHashMap::from_iter([
1074                ("missing".to_string(), vec![(1, 10)]),
1075                ("covered".to_string(), Vec::new()),
1076            ]),
1077        );
1078        assert_eq!(
1079            coverage_by_identifier,
1080            AHashMap::from_iter([
1081                ("missing".to_string(), CoverageIntervals::default()),
1082                (
1083                    "covered".to_string(),
1084                    CoverageIntervals {
1085                        data: vec![ClosedInterval::new(1, 10).unwrap()],
1086                        empty: Vec::new(),
1087                    },
1088                ),
1089            ]),
1090        );
1091        assert_eq!(batch.len(), 1);
1092        assert_eq!(instruments, vec![InstrumentAny::CurrencyPair(audusd_sim())]);
1093        assert_eq!(
1094            *calls.lock().unwrap(),
1095            vec![
1096                "query_last_timestamp QuoteTick Some(\"missing\")",
1097                "instruments Some([\"AUD/USD.SIM\"])",
1098                "write_data 1 Some(1) Some(2)",
1099                "write_instruments 1",
1100                "write_instruments 2",
1101            ],
1102        );
1103    }
1104
1105    #[rstest]
1106    fn test_drop_waits_for_queued_async_writes() {
1107        let calls = Arc::new(Mutex::new(Vec::new()));
1108        let (release, gate) = mpsc::channel();
1109
1110        let worker = CatalogWorker::start(Box::new(RecordingCatalog {
1111            calls: calls.clone(),
1112            write_gate: Some(gate),
1113            ..RecordingCatalog::default()
1114        }));
1115
1116        worker.write_async(quote_write_job()).unwrap();
1117
1118        let releaser = thread::spawn(move || {
1119            thread::sleep(Duration::from_millis(50));
1120            release.send(()).unwrap();
1121        });
1122
1123        drop(worker);
1124
1125        assert_eq!(*calls.lock().unwrap(), vec!["write_data 1 Some(1) Some(2)"]);
1126        releaser.join().unwrap();
1127    }
1128
1129    #[rstest]
1130    #[case::shared_catalog(false, &["reset_session", "query_batch_session Some(1)"])]
1131    #[case::forked_catalog(true, &["fork reset_session", "fork query_batch_session Some(1)"])]
1132    fn test_open_session_queries_forked_catalog_when_available(
1133        #[case] fork: bool,
1134        #[case] expected_calls: &[&str],
1135    ) {
1136        let calls = Arc::new(Mutex::new(Vec::new()));
1137
1138        let worker = CatalogWorker::start(Box::new(RecordingCatalog {
1139            calls: calls.clone(),
1140            fork,
1141            ..RecordingCatalog::default()
1142        }));
1143
1144        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
1145        let batch = worker.pull_session(session_id).unwrap().unwrap();
1146
1147        assert_eq!(batch.len(), 1);
1148        assert_eq!(*calls.lock().unwrap(), expected_calls);
1149    }
1150
1151    #[rstest]
1152    fn test_worker_reports_stopped_thread_after_catalog_panic() {
1153        let worker = CatalogWorker::start(Box::new(RecordingCatalog {
1154            panic_on_query: true,
1155            ..RecordingCatalog::default()
1156        }));
1157
1158        let reply_error = worker.query_batch(stub_query()).unwrap_err();
1159        let handle = worker.handle.as_ref().unwrap();
1160        let deadline = Instant::now() + Duration::from_secs(5);
1161
1162        while !handle.is_finished() {
1163            assert!(Instant::now() < deadline, "worker thread did not stop");
1164            thread::sleep(Duration::from_millis(1));
1165        }
1166
1167        let send_error = worker.flush().unwrap_err();
1168
1169        assert_eq!(
1170            reply_error.to_string(),
1171            "Catalog worker thread stopped before replying to the command",
1172        );
1173        assert_eq!(
1174            send_error.to_string(),
1175            "Catalog worker thread has stopped, so the command cannot be sent",
1176        );
1177    }
1178
1179    #[rstest]
1180    #[case(false)]
1181    #[case(true)]
1182    fn session_close_and_shutdown_wait_for_inflight_backend(#[case] shutdown: bool) {
1183        let query_thread = Arc::new(Mutex::new(None));
1184
1185        let worker = CatalogWorker::start(Box::new(StubCatalog {
1186            fail: false,
1187            query_thread: Arc::clone(&query_thread),
1188        }));
1189
1190        let session = worker.open_session(stub_query(), Some(1)).unwrap();
1191        let backend_gate = query_thread.lock().unwrap();
1192        let (query_tx, query_rx) = mpsc::channel();
1193        worker
1194            .query_batch_async(
1195                stub_query(),
1196                Box::new(move |result| {
1197                    query_tx.send(result.map(|batch| batch.len())).unwrap();
1198                }),
1199            )
1200            .unwrap();
1201
1202        let (started_tx, started_rx) = mpsc::channel();
1203        let (closed_tx, closed_rx) = mpsc::channel();
1204
1205        // Close reports before the worker drops, so the drop's join cannot mask a close that
1206        // returns early
1207        let control = thread::spawn(move || {
1208            started_tx.send(()).unwrap();
1209
1210            if shutdown {
1211                drop(worker);
1212                closed_tx.send(None).unwrap();
1213            } else {
1214                closed_tx
1215                    .send(Some(worker.close_session(session).unwrap()))
1216                    .unwrap();
1217            }
1218        });
1219
1220        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
1221        let blocked = closed_rx.recv_timeout(Duration::from_millis(25));
1222        drop(backend_gate);
1223        let rows = query_rx
1224            .recv_timeout(Duration::from_secs(5))
1225            .unwrap()
1226            .unwrap();
1227        let closed = closed_rx.recv_timeout(Duration::from_secs(5)).unwrap();
1228        control.join().unwrap();
1229        assert_eq!(blocked, Err(mpsc::RecvTimeoutError::Timeout));
1230        assert_eq!(rows, 1);
1231        assert_eq!(closed, (!shutdown).then_some(true));
1232    }
1233}
nautilus_persistence/catalog/worker.rs

nautilus_persistence/catalog/
worker.rs
