Skip to main content

libdd_shared_runtime/shared_runtime/
mod.rs

1// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4//! [`SharedRuntime`] trait and its implementations: [`ForkSafeRuntime`], [`BasicRuntime`],
5//! and [`LocalRuntime`].
6
7pub(crate) mod pausable_worker;
8
9#[cfg(not(target_arch = "wasm32"))]
10mod basic;
11#[cfg(not(target_arch = "wasm32"))]
12mod fork_safe;
13#[cfg(target_arch = "wasm32")]
14mod local;
15
16#[cfg(not(target_arch = "wasm32"))]
17pub use basic::BasicRuntime;
18#[cfg(not(target_arch = "wasm32"))]
19pub use fork_safe::ForkSafeRuntime;
20#[cfg(target_arch = "wasm32")]
21pub use local::LocalRuntime;
22
23use crate::worker::Worker;
24use libdd_capabilities::MaybeSend;
25use libdd_common::MutexExt;
26use pausable_worker::{PausableWorker, PausableWorkerError};
27use std::sync::{Arc, Mutex};
28use std::{fmt, io};
29
30/// A worker registered on a [`SharedRuntime`].
31pub(crate) type BoxedWorker = Box<dyn Worker + Sync>;
32
33#[derive(Debug)]
34pub(crate) struct WorkerEntry {
35    pub(crate) id: u64,
36    pub(crate) restart_on_fork: bool,
37    pub(crate) worker: PausableWorker<BoxedWorker>,
38}
39
40/// Common interface for all [`SharedRuntime`] implementations.
41///
42/// # Choosing an implementation
43///
44/// | Situation | Runtime |
45/// |-----------|---------|
46/// | Native, host process may call `fork(2)` | [`ForkSafeRuntime`] |
47/// | Native, caller owns a tokio runtime and wants to share it | [`BasicRuntime`] |
48/// | Wasm / single-threaded JS event loop | [`LocalRuntime`] |
49///
50/// Sync entry points (e.g. a blocking `build` or `send`) additionally require
51/// `R: `[`BlockingRuntime`], which only the native implementations satisfy.
52pub trait SharedRuntime {
53    /// Creates a new instance of the runtime with default configuration.
54    ///
55    /// Used as a fallback by callers (e.g. [`crate::shared_runtime`] consumers) that want to
56    /// auto-construct a runtime when one was not supplied; concrete runtimes additionally
57    /// expose richer inherent constructors (e.g. `with_worker_threads`, `from_handle`).
58    fn new() -> Result<Self, SharedRuntimeError>
59    where
60        Self: Sized;
61
62    /// Spawns a worker. `restart_on_fork = true` causes `ForkSafeRuntime::after_fork_child`
63    /// to reset and restart it; `false` drops it without calling shutdown. [`BasicRuntime`]
64    /// and [`LocalRuntime`] ignore this flag — they do not implement a fork protocol.
65    fn spawn_worker<T: Worker + Sync + 'static>(
66        &self,
67        worker: T,
68        restart_on_fork: bool,
69    ) -> Result<WorkerHandle, SharedRuntimeError>;
70
71    /// Shuts down all tracked workers. The runtime itself is not torn down — call
72    /// [`ForkSafeRuntime::shutdown`] (native only) to also drop the tokio runtime.
73    fn shutdown_async(&self) -> impl std::future::Future<Output = ()> + MaybeSend + '_
74    where
75        Self: Sync;
76}
77
78/// Extension of [`SharedRuntime`] for runtimes that can block the current thread on a future.
79#[cfg(not(target_arch = "wasm32"))]
80pub trait BlockingRuntime: SharedRuntime {
81    /// Drives `f` to completion, blocking the current thread.
82    ///
83    /// Returns an [`io::Error`] if the executor cannot be accessed or constructed.
84    fn block_on<F: std::future::Future>(&self, f: F) -> Result<F::Output, io::Error>;
85}
86
87/// Handle to a worker registered on a [`SharedRuntime`].
88///
89/// This handle can be used to stop the worker.
90///
91/// # Warning
92/// If every clone of this handle is dropped without calling [`WorkerHandle::stop`], the worker
93/// remains registered on the [`SharedRuntime`] and can only be torn down by shutting the
94/// runtime down. Workers are expected to detect that their input channel has been closed and
95/// park themselves to avoid spinning, but they will not be freed until the runtime stops.
96#[must_use = "dropping a WorkerHandle without calling stop() leaks the worker until the SharedRuntime is shut down"]
97#[derive(Clone, Debug)]
98pub struct WorkerHandle {
99    pub(crate) worker_id: u64,
100    pub(crate) workers: Arc<Mutex<Vec<WorkerEntry>>>,
101}
102
103#[derive(Debug)]
104pub enum WorkerHandleError {
105    AlreadyStopped,
106    WorkerError(PausableWorkerError),
107}
108
109impl fmt::Display for WorkerHandleError {
110    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
111        match self {
112            Self::AlreadyStopped => {
113                write!(f, "Worker has already been stopped")
114            }
115            Self::WorkerError(err) => write!(f, "Worker error: {}", err),
116        }
117    }
118}
119
120impl std::error::Error for WorkerHandleError {}
121
122impl From<PausableWorkerError> for WorkerHandleError {
123    fn from(err: PausableWorkerError) -> Self {
124        Self::WorkerError(err)
125    }
126}
127
128impl WorkerHandle {
129    /// Stop the worker and execute the shutdown logic.
130    ///
131    /// # Errors
132    /// Returns an error if the worker has already been stopped.
133    ///
134    /// # Cancel safety
135    /// This function is *NOT* cancel safe and shouldn't be called in [Worker::trigger].
136    /// If cancelled, the stopped worker can end up in an invalid state if a fork occurs while
137    /// stopping.
138    pub async fn stop(self) -> Result<(), WorkerHandleError> {
139        let mut worker = {
140            let mut workers_lock = self.workers.lock_or_panic();
141            let Some(position) = workers_lock
142                .iter()
143                .position(|entry| entry.id == self.worker_id)
144            else {
145                return Err(WorkerHandleError::AlreadyStopped);
146            };
147            let WorkerEntry { worker, .. } = workers_lock.swap_remove(position);
148            worker
149        };
150        worker.pause().await?;
151        worker.shutdown().await;
152        Ok(())
153    }
154}
155
156/// Errors that can occur when using a `SharedRuntime` implementation.
157#[derive(Debug)]
158pub enum SharedRuntimeError {
159    /// The runtime is not available or in an invalid state.
160    RuntimeUnavailable,
161    /// Failed to acquire a lock on internal state.
162    LockFailed(String),
163    /// A worker operation failed.
164    WorkerError(PausableWorkerError),
165    /// Failed to create the tokio runtime.
166    RuntimeCreation(io::Error),
167    /// Shutdown timed out.
168    ShutdownTimedOut(std::time::Duration),
169}
170
171impl fmt::Display for SharedRuntimeError {
172    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
173        match self {
174            Self::RuntimeUnavailable => {
175                write!(f, "Runtime is not available or in an invalid state")
176            }
177            Self::LockFailed(msg) => write!(f, "Failed to acquire lock: {}", msg),
178            Self::WorkerError(err) => write!(f, "Worker error: {}", err),
179            Self::RuntimeCreation(err) => {
180                write!(f, "Failed to create runtime: {}", err)
181            }
182            Self::ShutdownTimedOut(duration) => {
183                write!(f, "Shutdown timed out after {:?}", duration)
184            }
185        }
186    }
187}
188
189impl std::error::Error for SharedRuntimeError {}
190
191impl From<PausableWorkerError> for SharedRuntimeError {
192    fn from(err: PausableWorkerError) -> Self {
193        SharedRuntimeError::WorkerError(err)
194    }
195}
196
197impl From<io::Error> for SharedRuntimeError {
198    fn from(err: io::Error) -> Self {
199        SharedRuntimeError::RuntimeCreation(err)
200    }
201}