handle single-threaded rt

This commit is contained in:
Aiden McClelland
2024-06-24 17:16:52 -06:00
parent 5a24903031
commit f608480034
3 changed files with 25 additions and 17 deletions

View File

@@ -1,7 +1,9 @@
use tokio::runtime::Handle; use std::sync::Arc;
use tokio::runtime::Runtime;
pub trait Context: Send + Sync + 'static { pub trait Context: Send + Sync + 'static {
fn runtime(&self) -> Handle { fn runtime(&self) -> Option<Arc<Runtime>> {
Handle::current() None
} }
} }

View File

@@ -213,10 +213,11 @@ pub trait HandlerFor<Context: crate::Context>:
&self, &self,
handle_args: HandlerArgsFor<Context, Self>, handle_args: HandlerArgsFor<Context, Self>,
) -> Result<Self::Ok, Self::Err> { ) -> Result<Self::Ok, Self::Err> {
handle_args if let Some(rt) = handle_args.context.runtime() {
.context rt.block_on(self.handle_async(handle_args))
.runtime() } else {
.block_on(self.handle_async(handle_args)) tokio::runtime::Handle::current().block_on(self.handle_async(handle_args))
}
} }
fn handle_async( fn handle_async(
&self, &self,
@@ -234,11 +235,13 @@ pub trait HandlerFor<Context: crate::Context>:
) -> impl Future<Output = Result<Self::Ok, Self::Err>> + Send + 'a { ) -> impl Future<Output = Result<Self::Ok, Self::Err>> + Send + 'a {
async move { async move {
let s = self.clone(); let s = self.clone();
handle_args if let Some(rt) = handle_args.context.runtime() {
.context rt.spawn_blocking(move || s.handle_sync(handle_args)).await
.runtime() } else {
tokio::runtime::Handle::current()
.spawn_blocking(move || s.handle_sync(handle_args)) .spawn_blocking(move || s.handle_sync(handle_args))
.await .await
}
.unwrap() .unwrap()
} }
} }

View File

@@ -48,16 +48,19 @@ impl CliConfig {
struct CliContextSeed { struct CliContextSeed {
host: PathBuf, host: PathBuf,
rt: OnceCell<Runtime>, rt: OnceCell<Arc<Runtime>>,
} }
#[derive(Clone)] #[derive(Clone)]
struct CliContext(Arc<CliContextSeed>); struct CliContext(Arc<CliContextSeed>);
impl Context for CliContext { impl Context for CliContext {
fn runtime(&self) -> Handle { fn runtime(&self) -> Option<Arc<Runtime>> {
if self.0.rt.get().is_none() { if self.0.rt.get().is_none() {
self.0.rt.set(Runtime::new().unwrap()).unwrap(); let rt = Arc::new(Runtime::new().unwrap());
self.0.rt.set(rt.clone()).unwrap_or_default();
Some(rt)
} else {
self.0.rt.get().cloned()
} }
self.0.rt.get().unwrap().handle().clone()
} }
} }