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.
51pub(crate) const 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#[derive(Debug)]
138pub struct CatalogWorker {
139    sender: SyncSender<CatalogCommand>,
140    handle: Option<JoinHandle<()>>,
141}
142
143#[expect(
144    clippy::missing_errors_doc,
145    reason = "Worker methods forward channel and catalog errors directly"
146)]
147impl CatalogWorker {
148    #[must_use]
149    pub fn start(mut catalog: DataCatalog) -> Self {
150        let (sender, receiver) = mpsc::sync_channel(COMMAND_QUEUE_CAPACITY);
151        let handle = thread::spawn(move || run_catalog_worker(&mut catalog, receiver));
152
153        Self {
154            sender,
155            handle: Some(handle),
156        }
157    }
158
159    fn send(&self, command: CatalogCommand) -> anyhow::Result<()> {
160        self.sender.send(command).map_err(|_| {
161            anyhow::anyhow!("Catalog worker thread has stopped, so the command cannot be sent")
162        })
163    }
164
165    fn send_with_reply<T>(
166        &self,
167        command: impl FnOnce(Reply<anyhow::Result<T>>) -> CatalogCommand,
168    ) -> anyhow::Result<T> {
169        let (reply, receiver) = mpsc::channel();
170        self.send(command(reply))?;
171        receiver.recv().map_err(|_| {
172            anyhow::anyhow!("Catalog worker thread stopped before replying to the command")
173        })?
174    }
175
176    pub fn query_last_timestamp(
177        &self,
178        data_type: NautilusDataType,
179        identifier: Option<String>,
180    ) -> anyhow::Result<Option<u64>> {
181        self.send_with_reply(|reply| CatalogCommand::QueryLastTimestamp {
182            data_type,
183            identifier,
184            reply,
185        })
186    }
187
188    pub fn get_missing_intervals(
189        &self,
190        start: UnixNanos,
191        end: UnixNanos,
192        data_type: NautilusDataType,
193        identifier: Option<String>,
194    ) -> anyhow::Result<Vec<(u64, u64)>> {
195        self.send_with_reply(|reply| CatalogCommand::GetMissingIntervals {
196            start,
197            end,
198            data_type,
199            identifier,
200            reply,
201        })
202    }
203
204    pub fn get_missing_intervals_for_identifiers(
205        &self,
206        start: UnixNanos,
207        end: UnixNanos,
208        data_type: NautilusDataType,
209        identifiers: Vec<String>,
210    ) -> anyhow::Result<MissingIntervalsByIdentifier> {
211        self.send_with_reply(|reply| CatalogCommand::GetMissingIntervalsForIdentifiers {
212            start,
213            end,
214            data_type,
215            identifiers,
216            reply,
217        })
218    }
219
220    pub fn get_coverage_intervals_for_identifiers(
221        &self,
222        start: UnixNanos,
223        end: UnixNanos,
224        data_type: NautilusDataType,
225        identifiers: Vec<String>,
226    ) -> anyhow::Result<CoverageIntervalsByIdentifier> {
227        self.send_with_reply(|reply| CatalogCommand::GetCoverageIntervalsForIdentifiers {
228            start,
229            end,
230            data_type,
231            identifiers,
232            reply,
233        })
234    }
235
236    pub fn query_batch(&self, query: CatalogQuery) -> anyhow::Result<DataBatch> {
237        self.send_with_reply(|reply| CatalogCommand::QueryBatch { query, reply })
238    }
239
240    /// Enqueues a batch query and invokes `on_complete` with the result on the
241    /// worker thread, without blocking the caller.
242    pub fn query_batch_async(
243        &self,
244        query: CatalogQuery,
245        on_complete: Box<dyn FnOnce(anyhow::Result<DataBatch>) + Send>,
246    ) -> anyhow::Result<()> {
247        self.send(CatalogCommand::QueryBatchAsync { query, on_complete })
248    }
249
250    pub fn open_session(
251        &self,
252        query: CatalogQuery,
253        chunk_size: Option<usize>,
254    ) -> anyhow::Result<UUID4> {
255        self.send_with_reply(|reply| CatalogCommand::OpenSession {
256            query,
257            chunk_size,
258            reply,
259        })
260    }
261
262    pub fn pull_session(&self, session_id: UUID4) -> anyhow::Result<Option<DataBatch>> {
263        self.send_with_reply(|reply| CatalogCommand::PullSession { session_id, reply })
264    }
265
266    /// Enqueues one session pull and invokes `on_complete` on the catalog worker.
267    pub fn pull_session_async(
268        &self,
269        session_id: UUID4,
270        on_complete: Box<dyn FnOnce(anyhow::Result<Option<DataBatch>>) + Send>,
271    ) -> anyhow::Result<()> {
272        self.send(CatalogCommand::PullSessionAsync {
273            session_id,
274            on_complete,
275        })
276    }
277
278    /// Closes a session, returning whether it was still open.
279    pub fn close_session(&self, session_id: UUID4) -> anyhow::Result<bool> {
280        self.send_with_reply(|reply| CatalogCommand::CloseSession { session_id, reply })
281    }
282
283    pub fn query_instruments(
284        &self,
285        query: CatalogInstrumentQuery,
286    ) -> anyhow::Result<Vec<InstrumentAny>> {
287        self.send_with_reply(|reply| CatalogCommand::QueryInstruments { query, reply })
288    }
289
290    pub fn write(&self, job: CatalogWriteJob) -> anyhow::Result<()> {
291        self.send_with_reply(|reply| CatalogCommand::Write { job, reply })
292    }
293
294    pub fn write_async(&self, job: CatalogWriteJob) -> anyhow::Result<()> {
295        self.send(CatalogCommand::WriteAsync { job })
296    }
297
298    pub fn write_instruments(&self, instruments: Vec<InstrumentAny>) -> anyhow::Result<()> {
299        self.send_with_reply(|reply| CatalogCommand::WriteInstruments { instruments, reply })
300    }
301
302    pub fn write_instruments_async(&self, instruments: Vec<InstrumentAny>) -> anyhow::Result<()> {
303        self.send(CatalogCommand::WriteInstrumentsAsync { instruments })
304    }
305
306    pub fn flush(&self) -> anyhow::Result<()> {
307        self.send_with_reply(|reply| CatalogCommand::Flush { reply })
308    }
309}
310
311impl Drop for CatalogWorker {
312    fn drop(&mut self) {
313        let _ = self.sender.send(CatalogCommand::Shutdown);
314        if let Some(handle) = self.handle.take() {
315            let _ = handle.join();
316        }
317    }
318}
319
320#[expect(
321    clippy::needless_pass_by_value,
322    clippy::too_many_lines,
323    reason = "The worker thread owns the receiver for its whole lifetime"
324)]
325fn run_catalog_worker(catalog: &mut DataCatalog, receiver: Receiver<CatalogCommand>) {
326    let mut async_errors: Vec<anyhow::Error> = Vec::new();
327    let mut sessions: AHashMap<UUID4, DataBatchQueryResult> = AHashMap::new();
328
329    while let Ok(command) = receiver.recv() {
330        match command {
331            CatalogCommand::QueryLastTimestamp {
332                data_type,
333                identifier,
334                reply,
335            } => {
336                let result = catalog.query_last_timestamp(data_type, identifier.as_deref());
337                let _ = reply.send(result);
338            }
339            CatalogCommand::GetMissingIntervals {
340                start,
341                end,
342                data_type,
343                identifier,
344                reply,
345            } => {
346                let result = catalog.get_missing_intervals_for_request(
347                    start,
348                    end,
349                    data_type,
350                    identifier.as_deref(),
351                );
352                let _ = reply.send(result);
353            }
354            CatalogCommand::GetMissingIntervalsForIdentifiers {
355                start,
356                end,
357                data_type,
358                identifiers,
359                reply,
360            } => {
361                let result = catalog.get_missing_intervals_for_identifiers(
362                    start,
363                    end,
364                    data_type,
365                    &identifiers,
366                );
367                let _ = reply.send(result);
368            }
369            CatalogCommand::GetCoverageIntervalsForIdentifiers {
370                start,
371                end,
372                data_type,
373                identifiers,
374                reply,
375            } => {
376                let result = catalog.get_coverage_intervals_for_identifiers(
377                    start,
378                    end,
379                    data_type,
380                    &identifiers,
381                );
382                let _ = reply.send(result);
383            }
384            CatalogCommand::QueryBatch { query, reply } => {
385                let result = catalog.query_batch(&query);
386                let _ = reply.send(result);
387            }
388            CatalogCommand::QueryBatchAsync { query, on_complete } => {
389                let result = catalog.query_batch(&query);
390                on_complete(result);
391            }
392            CatalogCommand::OpenSession {
393                query,
394                chunk_size,
395                reply,
396            } => {
397                let result = if sessions.len() >= MAX_OPEN_SESSIONS {
398                    Err(anyhow::anyhow!(
399                        "Catalog session limit of {MAX_OPEN_SESSIONS} reached; close finished sessions before opening more",
400                    ))
401                } else {
402                    open_query_session(catalog, &query, chunk_size).map(|session| {
403                        let session_id = UUID4::new();
404                        sessions.insert(session_id, session);
405                        session_id
406                    })
407                };
408                let _ = reply.send(result);
409            }
410            CatalogCommand::PullSession { session_id, reply } => {
411                let result = pull_session(&mut sessions, session_id);
412                let _ = reply.send(result);
413            }
414            CatalogCommand::PullSessionAsync {
415                session_id,
416                on_complete,
417            } => {
418                on_complete(pull_session(&mut sessions, session_id));
419            }
420            CatalogCommand::CloseSession { session_id, reply } => {
421                let _ = reply.send(Ok(sessions.remove(&session_id).is_some()));
422            }
423            CatalogCommand::QueryInstruments { query, reply } => {
424                let result = catalog.instruments(&query);
425                let _ = reply.send(result);
426            }
427            CatalogCommand::Write { job, reply } => {
428                let result = catalog.write_data_batch(&job.data, job.start, job.end, job.params);
429                let _ = reply.send(result);
430            }
431            CatalogCommand::WriteAsync { job } => {
432                let len = job.data.len();
433                match catalog.write_data_batch(&job.data, job.start, job.end, job.params) {
434                    Ok(()) => log::info!("Catalog worker wrote {len} data rows"),
435                    Err(e) => {
436                        log::error!("Catalog worker failed to write {len} data rows: {e}");
437                        async_errors.push(e);
438                    }
439                }
440            }
441            CatalogCommand::WriteInstruments { instruments, reply } => {
442                let result = catalog.write_instruments(&instruments);
443                let _ = reply.send(result);
444            }
445            CatalogCommand::WriteInstrumentsAsync { instruments } => {
446                let len = instruments.len();
447                match catalog.write_instruments(&instruments) {
448                    Ok(()) => log::info!("Catalog worker wrote {len} instruments"),
449                    Err(e) => {
450                        log::error!("Catalog worker failed to write {len} instruments: {e}");
451                        async_errors.push(e);
452                    }
453                }
454            }
455            CatalogCommand::Flush { reply } => {
456                let _ = reply.send(drain_async_errors(&mut async_errors));
457            }
458            CatalogCommand::Shutdown => break,
459        }
460    }
461}
462
463fn open_query_session(
464    catalog: &mut DataCatalog,
465    query: &CatalogQuery,
466    chunk_size: Option<usize>,
467) -> anyhow::Result<DataBatchQueryResult> {
468    if let Some(mut query_catalog) = catalog.fork_query_catalog()? {
469        query_catalog.reset_session();
470        return query_catalog.query_batch_session(query, chunk_size);
471    }
472
473    catalog.reset_session();
474    catalog.query_batch_session(query, chunk_size)
475}
476
477fn drain_async_errors(errors: &mut Vec<anyhow::Error>) -> anyhow::Result<()> {
478    let failures = errors.len();
479    let details = errors
480        .iter()
481        .map(|e| format!("{e:#}"))
482        .collect::<Vec<_>>()
483        .join("; ");
484    let Some(first) = errors.drain(..).next() else {
485        return Ok(());
486    };
487
488    if failures == 1 {
489        return Err(first);
490    }
491
492    Err(first.context(format!(
493        "{failures} asynchronous catalog writes failed: {details}"
494    )))
495}
496
497fn pull_session(
498    sessions: &mut AHashMap<UUID4, DataBatchQueryResult>,
499    session_id: UUID4,
500) -> anyhow::Result<Option<DataBatch>> {
501    let result = sessions
502        .get_mut(&session_id)
503        .ok_or_else(|| anyhow::anyhow!("Catalog session {session_id} is not open"))?
504        .next_batch();
505
506    if !matches!(result, Ok(Some(_))) {
507        sessions.remove(&session_id);
508    }
509    result
510}
511
512////////////////////////////////////////////////////////////////////////////////
513// Tests
514////////////////////////////////////////////////////////////////////////////////
515
516#[cfg(test)]
517mod tests {
518    use std::{
519        sync::{Arc, Mutex},
520        time::Duration,
521    };
522
523    use nautilus_model::{
524        data::{Data, NautilusRecordType, QuoteTick},
525        identifiers::InstrumentId,
526        types::{Price, Quantity},
527    };
528    use rstest::rstest;
529
530    use super::*;
531    use crate::catalog::{
532        session::{DataBatchQuery, TypedDataBatchSession},
533        traits::{CatalogMetadata, CatalogReader, CatalogRecordQuery, CatalogWriter, RecordBatch},
534    };
535
536    struct FailingSession;
537
538    impl DataBatchQuery for FailingSession {
539        fn next_batch(&mut self) -> anyhow::Result<Option<DataBatch>> {
540            anyhow::bail!("session decode failure")
541        }
542    }
543
544    fn stub_quote() -> QuoteTick {
545        QuoteTick::new(
546            InstrumentId::from("ETHUSDT-PERP.BINANCE"),
547            Price::from("1987.0"),
548            Price::from("1988.0"),
549            Quantity::from("100"),
550            Quantity::from("100"),
551            UnixNanos::from(1),
552            UnixNanos::from(1),
553        )
554    }
555
556    #[derive(Debug)]
557    struct StubCatalog {
558        fail: bool,
559        query_thread: Arc<Mutex<Option<thread::ThreadId>>>,
560    }
561
562    impl CatalogReader for StubCatalog {
563        fn reset_session(&mut self) {}
564
565        fn instruments(
566            &mut self,
567            _query: &CatalogInstrumentQuery,
568        ) -> anyhow::Result<Vec<InstrumentAny>> {
569            Ok(Vec::new())
570        }
571
572        fn query_batch(&mut self, _query: &CatalogQuery) -> anyhow::Result<DataBatch> {
573            *self.query_thread.lock().unwrap() = Some(thread::current().id());
574
575            if self.fail {
576                anyhow::bail!("stub query failure");
577            }
578            Ok(DataBatch::Quote(vec![stub_quote()].into()))
579        }
580
581        fn query_batch_session(
582            &mut self,
583            _query: &CatalogQuery,
584            chunk_size: Option<usize>,
585        ) -> anyhow::Result<DataBatchQueryResult> {
586            *self.query_thread.lock().unwrap() = Some(thread::current().id());
587
588            if self.fail {
589                anyhow::bail!("stub query failure");
590            }
591            Ok(Box::new(TypedDataBatchSession::from_vec(
592                vec![stub_quote()],
593                chunk_size,
594            )))
595        }
596
597        fn query_metadata(
598            &mut self,
599            _query: &CatalogQuery,
600        ) -> anyhow::Result<Vec<CatalogMetadata>> {
601            Ok(Vec::new())
602        }
603
604        fn get_missing_intervals_for_request(
605            &mut self,
606            _start: UnixNanos,
607            _end: UnixNanos,
608            _data_type: NautilusDataType,
609            _identifier: Option<&str>,
610        ) -> anyhow::Result<Vec<(u64, u64)>> {
611            Ok(Vec::new())
612        }
613
614        fn query_last_timestamp(
615            &mut self,
616            _data_type: NautilusDataType,
617            _identifier: Option<&str>,
618        ) -> anyhow::Result<Option<u64>> {
619            Ok(None)
620        }
621
622        fn query_display_record_batches(
623            &mut self,
624            _query: &CatalogQuery,
625        ) -> anyhow::Result<Vec<RecordBatch>> {
626            Ok(Vec::new())
627        }
628
629        fn query_record_batches(
630            &mut self,
631            _query: &CatalogRecordQuery,
632        ) -> anyhow::Result<Vec<RecordBatch>> {
633            Ok(Vec::new())
634        }
635
636        fn query_record_display_batches(
637            &mut self,
638            _query: &CatalogRecordQuery,
639        ) -> anyhow::Result<Vec<RecordBatch>> {
640            Ok(Vec::new())
641        }
642    }
643
644    impl CatalogWriter for StubCatalog {
645        fn write_instruments(&mut self, _instruments: &[InstrumentAny]) -> anyhow::Result<()> {
646            Ok(())
647        }
648
649        fn write_data(
650            &mut self,
651            _data: &[Data],
652            _start: Option<UnixNanos>,
653            _end: Option<UnixNanos>,
654            params: Option<Params>,
655        ) -> anyhow::Result<()> {
656            if self.fail {
657                let message = params
658                    .as_ref()
659                    .and_then(|params| params.get_str("test_error"))
660                    .unwrap_or("stub write failure");
661                anyhow::bail!(message.to_string());
662            }
663            Ok(())
664        }
665
666        fn write_records(
667            &mut self,
668            _record_type: NautilusRecordType,
669            _batches: &[RecordBatch],
670            _params: Option<Params>,
671        ) -> anyhow::Result<()> {
672            Ok(())
673        }
674
675        fn record_empty_coverage(
676            &mut self,
677            _data_type: NautilusDataType,
678            _identifier: Option<&str>,
679            _start: UnixNanos,
680            _end: UnixNanos,
681        ) -> anyhow::Result<()> {
682            Ok(())
683        }
684    }
685
686    fn stub_query() -> CatalogQuery {
687        CatalogQuery::new(NautilusDataType::QuoteTick)
688    }
689
690    #[rstest]
691    fn test_query_batch_async_invokes_callback_on_worker_thread() {
692        let query_thread = Arc::new(Mutex::new(None));
693        let worker = CatalogWorker::start(Box::new(StubCatalog {
694            fail: false,
695            query_thread: query_thread.clone(),
696        }));
697
698        let (tx, rx) = mpsc::channel();
699        worker
700            .query_batch_async(
701                stub_query(),
702                Box::new(move |result| {
703                    let _ = tx.send((thread::current().id(), result.map(|batch| batch.len())));
704                }),
705            )
706            .unwrap();
707
708        let (callback_thread, result) = rx
709            .recv_timeout(Duration::from_secs(5))
710            .expect("callback was not invoked");
711        assert_eq!(result.unwrap(), 1);
712        assert_ne!(callback_thread, thread::current().id());
713        assert_eq!(Some(callback_thread), *query_thread.lock().unwrap());
714    }
715
716    #[rstest]
717    fn test_query_batch_async_passes_query_error_to_callback() {
718        let worker = CatalogWorker::start(Box::new(StubCatalog {
719            fail: true,
720            query_thread: Arc::new(Mutex::new(None)),
721        }));
722
723        let (tx, rx) = mpsc::channel();
724        worker
725            .query_batch_async(
726                stub_query(),
727                Box::new(move |result| {
728                    let _ = tx.send(result.map(|batch| batch.len()));
729                }),
730            )
731            .unwrap();
732
733        let result = rx
734            .recv_timeout(Duration::from_secs(5))
735            .expect("callback was not invoked");
736        assert_eq!(result.unwrap_err().to_string(), "stub query failure");
737    }
738
739    #[rstest]
740    fn test_worker_owns_and_pulls_catalog_session() {
741        let query_thread = Arc::new(Mutex::new(None));
742        let worker = CatalogWorker::start(Box::new(StubCatalog {
743            fail: false,
744            query_thread: query_thread.clone(),
745        }));
746
747        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
748        let batch = worker.pull_session(session_id).unwrap().unwrap();
749        let complete = worker.pull_session(session_id).unwrap();
750        let closed = worker.pull_session(session_id).unwrap_err();
751
752        assert_eq!(batch.len(), 1);
753        assert!(complete.is_none());
754        assert_eq!(
755            closed.to_string(),
756            format!("Catalog session {session_id} is not open")
757        );
758        assert_ne!(*query_thread.lock().unwrap(), Some(thread::current().id()));
759    }
760
761    #[rstest]
762    fn test_close_session_drops_open_session() {
763        let worker = CatalogWorker::start(Box::new(StubCatalog {
764            fail: false,
765            query_thread: Arc::new(Mutex::new(None)),
766        }));
767        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
768
769        let first_close = worker.close_session(session_id).unwrap();
770        let second_close = worker.close_session(session_id).unwrap();
771
772        assert!(first_close);
773        assert!(!second_close);
774    }
775
776    #[rstest]
777    fn errored_session_is_removed_after_pull() {
778        let session_id = UUID4::new();
779        let mut sessions = AHashMap::new();
780        sessions.insert(session_id, Box::new(FailingSession) as DataBatchQueryResult);
781
782        let first = pull_session(&mut sessions, session_id).unwrap_err();
783        let second = pull_session(&mut sessions, session_id).unwrap_err();
784
785        assert_eq!(first.to_string(), "session decode failure");
786        assert_eq!(
787            second.to_string(),
788            format!("Catalog session {session_id} is not open"),
789        );
790        assert!(sessions.is_empty());
791    }
792
793    #[rstest]
794    fn test_flush_reports_every_failed_async_write_then_clears() {
795        let worker = CatalogWorker::start(Box::new(StubCatalog {
796            fail: true,
797            query_thread: Arc::new(Mutex::new(None)),
798        }));
799
800        for message in [
801            "first write failed",
802            "second write failed",
803            "third write failed",
804        ] {
805            let mut params = Params::new();
806            params.insert("test_error".to_string(), message.into());
807            worker
808                .write_async(CatalogWriteJob {
809                    data: DataBatch::Quote(vec![stub_quote()].into()),
810                    start: None,
811                    end: None,
812                    params: Some(params),
813                })
814                .unwrap();
815        }
816
817        let error = worker.flush().unwrap_err();
818
819        assert_eq!(
820            error.to_string(),
821            "3 asynchronous catalog writes failed: first write failed; second write failed; \
822             third write failed",
823        );
824        assert_eq!(error.root_cause().to_string(), "first write failed");
825        assert!(worker.flush().is_ok(), "flush must clear reported failures");
826    }
827
828    #[rstest]
829    fn test_open_session_rejects_beyond_the_session_limit() {
830        let worker = CatalogWorker::start(Box::new(StubCatalog {
831            fail: false,
832            query_thread: Arc::new(Mutex::new(None)),
833        }));
834        let session_ids = (0..MAX_OPEN_SESSIONS)
835            .map(|_| worker.open_session(stub_query(), Some(1)).unwrap())
836            .collect::<Vec<_>>();
837
838        let error = worker.open_session(stub_query(), Some(1)).unwrap_err();
839        assert_eq!(
840            error.to_string(),
841            format!(
842                "Catalog session limit of {MAX_OPEN_SESSIONS} reached; close finished sessions before opening more"
843            ),
844        );
845
846        assert!(worker.close_session(session_ids[0]).unwrap());
847        assert!(
848            worker.open_session(stub_query(), Some(1)).is_ok(),
849            "closing a session must free a slot",
850        );
851    }
852
853    #[rstest]
854    fn test_pull_session_async_invokes_callback_on_worker_thread() {
855        let worker = CatalogWorker::start(Box::new(StubCatalog {
856            fail: false,
857            query_thread: Arc::new(Mutex::new(None)),
858        }));
859        let session_id = worker.open_session(stub_query(), Some(1)).unwrap();
860        let (tx, rx) = mpsc::channel();
861
862        worker
863            .pull_session_async(
864                session_id,
865                Box::new(move |result| {
866                    let _ = tx.send((
867                        thread::current().id(),
868                        result.map(|batch| batch.unwrap().len()),
869                    ));
870                }),
871            )
872            .unwrap();
873        let (callback_thread, result) = rx.recv_timeout(Duration::from_secs(5)).unwrap();
874
875        assert_eq!(result.unwrap(), 1);
876        assert_ne!(callback_thread, thread::current().id());
877    }
878}
nautilus_persistence/catalog/worker.rs

nautilus_persistence/catalog/
worker.rs
