cognee_core/
runtime.rs

1use tokio::runtime::{Builder, EnterGuard, Handle, Runtime};
2
3use crate::error::CoreError;
4
5/// A wrapper around a Tokio [`Runtime`].
6///
7/// Provides convenient constructors for common configurations and exposes the
8/// most-used entry points (`block_on`, `spawn`, `handle`, `enter`).
9pub struct AsyncRuntime {
10    inner: Runtime,
11}
12
13impl AsyncRuntime {
14    /// Create a multi-threaded runtime with Tokio's defaults (one worker thread
15    /// per logical CPU, all feature flags enabled).
16    pub fn new() -> Result<Self, CoreError> {
17        let rt = Runtime::new().map_err(|e| CoreError::Runtime(e.to_string()))?;
18        Ok(Self { inner: rt })
19    }
20
21    /// Create a multi-threaded runtime with an explicit worker-thread count.
22    pub fn multi_thread(num_workers: usize) -> Result<Self, CoreError> {
23        let rt = Builder::new_multi_thread()
24            .worker_threads(num_workers)
25            .enable_all()
26            .build()
27            .map_err(|e| CoreError::Runtime(e.to_string()))?;
28        Ok(Self { inner: rt })
29    }
30
31    /// Create a single-threaded (current-thread) runtime.
32    pub fn current_thread() -> Result<Self, CoreError> {
33        let rt = Builder::new_current_thread()
34            .enable_all()
35            .build()
36            .map_err(|e| CoreError::Runtime(e.to_string()))?;
37        Ok(Self { inner: rt })
38    }
39
40    /// Return a cloneable [`Handle`] to this runtime.
41    pub fn handle(&self) -> Handle {
42        self.inner.handle().clone()
43    }
44
45    /// Block the current thread until `future` completes and return its output.
46    pub fn block_on<F: std::future::Future>(&self, future: F) -> F::Output {
47        self.inner.block_on(future)
48    }
49
50    /// Spawn a future onto the runtime without waiting for its result.
51    pub fn spawn<F>(&self, future: F) -> tokio::task::JoinHandle<F::Output>
52    where
53        F: std::future::Future + Send + 'static,
54        F::Output: Send + 'static,
55    {
56        self.inner.spawn(future)
57    }
58
59    /// Enter the runtime context so that async code can call `Handle::current()`.
60    pub fn enter(&self) -> EnterGuard<'_> {
61        self.inner.enter()
62    }
63}
cognee_core/runtime.rs

cognee_core/
runtime.rs
