reifydb_runtime/context/
spawner.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
5use std::panic;
6use std::{
7	fmt,
8	future::Future,
9	pin::Pin,
10	task::{Context, Poll},
11};
12
13#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
14use tokio::{runtime::Handle, task::JoinHandle};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum TaskError {
18	Shutdown,
19}
20
21#[derive(Clone, Debug)]
22pub enum Spawner {
23	#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
24	Host(Handle),
25	Inline,
26}
27
28impl Spawner {
29	#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
30	pub fn host(handle: Handle) -> Self {
31		Spawner::Host(handle)
32	}
33
34	pub fn inline() -> Self {
35		Spawner::Inline
36	}
37
38	pub fn spawn<F>(&self, future: F) -> TaskHandle<F::Output>
39	where
40		F: Future + Send + 'static,
41		F::Output: Send + 'static,
42	{
43		match self {
44			#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
45			Spawner::Host(handle) => TaskHandle::Host(HostTask {
46				inner: handle.spawn(future),
47			}),
48			Spawner::Inline => TaskHandle::Inline(InlineTask {
49				future: Some(Box::pin(future)),
50			}),
51		}
52	}
53}
54
55pub enum TaskHandle<T> {
56	#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
57	Host(HostTask<T>),
58	Inline(InlineTask<T>),
59}
60
61impl<T> Future for TaskHandle<T> {
62	type Output = Result<T, TaskError>;
63
64	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
65		match self.get_mut() {
66			#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
67			TaskHandle::Host(task) => task.poll(cx),
68			TaskHandle::Inline(task) => task.poll(cx),
69		}
70	}
71}
72
73impl<T> fmt::Debug for TaskHandle<T> {
74	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75		match self {
76			#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
77			TaskHandle::Host(_) => f.write_str("TaskHandle::Host"),
78			TaskHandle::Inline(_) => f.write_str("TaskHandle::Inline"),
79		}
80	}
81}
82
83#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
84pub struct HostTask<T> {
85	inner: JoinHandle<T>,
86}
87
88#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
89impl<T> HostTask<T> {
90	fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TaskError>> {
91		match Pin::new(&mut self.inner).poll(cx) {
92			Poll::Pending => Poll::Pending,
93			Poll::Ready(Ok(value)) => Poll::Ready(Ok(value)),
94			Poll::Ready(Err(error)) if error.is_panic() => panic::resume_unwind(error.into_panic()),
95			Poll::Ready(Err(_)) => Poll::Ready(Err(TaskError::Shutdown)),
96		}
97	}
98}
99
100#[cfg(all(reifydb_target = "host", not(reifydb_dst)))]
101impl<T> Drop for HostTask<T> {
102	fn drop(&mut self) {
103		self.inner.abort();
104	}
105}
106
107pub struct InlineTask<T> {
108	future: Option<Pin<Box<dyn Future<Output = T> + Send>>>,
109}
110
111impl<T> InlineTask<T> {
112	fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TaskError>> {
113		let Some(future) = self.future.as_mut() else {
114			return Poll::Ready(Err(TaskError::Shutdown));
115		};
116		match future.as_mut().poll(cx) {
117			Poll::Pending => Poll::Pending,
118			Poll::Ready(value) => {
119				self.future = None;
120				Poll::Ready(Ok(value))
121			}
122		}
123	}
124}
reifydb_runtime/context/spawner.rs

reifydb_runtime/context/
spawner.rs
