a3s_vec/collection/
async_api.rs

1//! Tokio entry points for running synchronous collection queries off-runtime.
2
3use super::Collection;
4use crate::doc::Doc;
5use crate::error::{Error, Result};
6use crate::multi_query::MultiQuery;
7use crate::query::{GroupBySearchQuery, SearchQuery};
8use std::collections::HashMap;
9
10impl Collection {
11    /// Executes a query on Tokio's blocking pool.
12    ///
13    /// This keeps positioned `DiskANN` reads and exact refinement off async
14    /// runtime worker threads. The query still uses the same synchronous
15    /// snapshot, planner, fallback, scoring, and telemetry path as [`Self::query`].
16    /// Once scheduled, dropping the returned future does not cancel the
17    /// blocking query.
18    pub async fn query_async(&self, query: &SearchQuery) -> Result<Vec<Doc>> {
19        let collection = self.clone();
20        let query = query.clone();
21        run_blocking("query", move || collection.query(&query)).await
22    }
23
24    /// Executes every multi-query branch on Tokio's blocking pool.
25    ///
26    /// The full fusion operation runs in one blocking task so all branches use
27    /// the same captured collection snapshot as [`Self::multi_query`].
28    /// Once scheduled, dropping the returned future does not cancel the task.
29    pub async fn multi_query_async(&self, query: &MultiQuery) -> Result<Vec<Doc>> {
30        let collection = self.clone();
31        let query = query.clone();
32        run_blocking("multi-query", move || collection.multi_query(&query)).await
33    }
34
35    /// Executes a grouped vector query on Tokio's blocking pool.
36    ///
37    /// Once scheduled, dropping the returned future does not cancel the task.
38    pub async fn group_by_async(
39        &self,
40        query: &GroupBySearchQuery,
41    ) -> Result<HashMap<String, Vec<Doc>>> {
42        let collection = self.clone();
43        let query = query.clone();
44        run_blocking("group-by query", move || collection.group_by(&query)).await
45    }
46}
47
48async fn run_blocking<T, F>(operation: &'static str, task: F) -> Result<T>
49where
50    T: Send + 'static,
51    F: FnOnce() -> Result<T> + Send + 'static,
52{
53    let runtime = tokio::runtime::Handle::try_current().map_err(|_| {
54        Error::failed_precondition(format!(
55            "async {operation} requires an active Tokio runtime"
56        ))
57    })?;
58    runtime
59        .spawn_blocking(task)
60        .await
61        .map_err(|error| Error::internal(format!("async {operation} task failed: {error}")))?
62}
63
64#[cfg(test)]
65mod tests {
66    use super::run_blocking;
67    use crate::ErrorCode;
68
69    #[test]
70    fn missing_tokio_runtime_is_a_typed_error() {
71        let future = run_blocking("test operation", || Ok::<_, crate::Error>(()));
72        let error = poll_once_without_runtime(future).expect_err("runtime must be required");
73        assert_eq!(error.code, ErrorCode::FailedPrecondition);
74        assert_eq!(
75            error.message,
76            "async test operation requires an active Tokio runtime"
77        );
78    }
79
80    #[test]
81    fn query_work_leaves_the_runtime_worker_thread() {
82        let caller = std::thread::current().id();
83        let runtime = tokio::runtime::Builder::new_current_thread()
84            .build()
85            .expect("Tokio runtime must build");
86        let worker = runtime
87            .block_on(run_blocking("test operation", || {
88                Ok::<_, crate::Error>(std::thread::current().id())
89            }))
90            .expect("blocking task must succeed");
91        assert_ne!(worker, caller);
92    }
93
94    fn poll_once_without_runtime<F>(future: F) -> F::Output
95    where
96        F: std::future::Future,
97    {
98        use std::sync::Arc;
99        use std::task::{Context, Poll, Wake, Waker};
100
101        struct NoopWake;
102
103        impl Wake for NoopWake {
104            fn wake(self: Arc<Self>) {}
105        }
106
107        let waker = Waker::from(Arc::new(NoopWake));
108        let mut context = Context::from_waker(&waker);
109        let mut future = std::pin::pin!(future);
110        match future.as_mut().poll(&mut context) {
111            Poll::Ready(output) => output,
112            Poll::Pending => panic!("future unexpectedly waited without a Tokio runtime"),
113        }
114    }
115}
a3s_vec/collection/async_api.rs

a3s_vec/collection/
async_api.rs
