From 1dd2349d52aeea038d77fb0816f781c6b714fe77 Mon Sep 17 00:00:00 2001 From: Wondertan Date: Thu, 10 Sep 2026 09:40:30 +0200 Subject: [PATCH] fix(common): keep executor alive while contexts exist Contexts can outlive the Executor which created them. Retain a shared shutdown guard so dropping an owner does not cancel active context tasks. Assisted-by: GPT-5 Signed-off-by: Wondertan --- crates/common/src/context.rs | 10 ++-- crates/common/src/executor.rs | 108 ++++++++++++++++++++++++---------- 2 files changed, 82 insertions(+), 36 deletions(-) diff --git a/crates/common/src/context.rs b/crates/common/src/context.rs index a666a59e..4cc27e49 100644 --- a/crates/common/src/context.rs +++ b/crates/common/src/context.rs @@ -19,7 +19,7 @@ pub use test::{ test_mt_context, test_mt_context_with_spawn, test_st_context, }; -use crate::{ContextId, executor::Inner, io::Io, mux::Mux}; +use crate::{ContextId, executor::ExecutorHandle, io::Io, mux::Mux}; /// A task execution context. /// @@ -42,7 +42,7 @@ enum Mode { Single, Multi { mux: Arc, - executor: Option>, + executor: Option, }, } @@ -122,7 +122,7 @@ impl Context { id: ContextId, io: Io, mux: Arc, - executor: Arc, + executor: ExecutorHandle, ) -> Self { Self { id, @@ -342,7 +342,7 @@ impl Context { Ok(future::try_join4(task_a, task_b, task_c, task_d).await) } - fn executor(&self) -> Option<&Arc> { + fn executor(&self) -> Option<&ExecutorHandle> { if let Mode::Multi { executor, .. } = &self.mode { executor.as_ref() } else { @@ -354,7 +354,7 @@ impl Context { /// Spawns `fut` on `executor` if one is provided, otherwise yields the future /// as-is. The output type is identical either way. fn run( - executor: Option<&Arc>, + executor: Option<&ExecutorHandle>, fut: F, ) -> impl std::future::Future + Send where diff --git a/crates/common/src/executor.rs b/crates/common/src/executor.rs index 890a52f0..639f025e 100644 --- a/crates/common/src/executor.rs +++ b/crates/common/src/executor.rs @@ -19,9 +19,27 @@ use crate::{Context, ContextId, mux::Mux}; /// A work-stealing async executor. #[derive(Debug)] pub struct Executor { + handle: ExecutorHandle, +} + +/// Keeps the executor alive while contexts can still spawn tasks. +#[derive(Clone, Debug)] +pub(crate) struct ExecutorHandle { + inner: Arc, + _shutdown: Arc, +} + +#[derive(Debug)] +struct ShutdownGuard { inner: Arc, } +impl Drop for ShutdownGuard { + fn drop(&mut self) { + self.inner.shutdown(); + } +} + /// Per-worker parking state. struct WorkerState { unparker: Unparker, @@ -60,6 +78,27 @@ impl std::fmt::Debug for Inner { } } +impl Inner { + fn shutdown(&self) { + self.shutdown.store(true, Ordering::SeqCst); + + // Drain the injector before unparking workers. Any push that races + // with this drain is handled by the shutdown check in the schedule + // callback (see `spawn_on`), which drops the runnable. + loop { + match self.injector.steal() { + Steal::Success(_) => continue, + Steal::Empty => break, + Steal::Retry => continue, + } + } + + for w in self.workers.iter() { + w.unparker.unpark(); + } + } +} + /// A worker spawn callback. /// /// Receives a worker entry-point and dispatches it on a thread (or @@ -171,7 +210,14 @@ impl ExecutorBuilder { .expect("failed to spawn worker thread"); } - Executor { inner } + Executor { + handle: ExecutorHandle { + _shutdown: Arc::new(ShutdownGuard { + inner: inner.clone(), + }), + inner, + }, + } } } @@ -269,27 +315,12 @@ impl Executor { /// a [`Runnable`] cancels its task, so awaiters of `Task` propagate /// cancellation rather than hanging on a worker that has exited. pub fn shutdown(&self) { - self.inner.shutdown.store(true, Ordering::SeqCst); - - // Drain the injector before unparking workers. Any push that races - // with this drain is handled by the shutdown check in the schedule - // callback (see `spawn_on`), which drops the runnable. - loop { - match self.inner.injector.steal() { - Steal::Success(_) => continue, - Steal::Empty => break, - Steal::Retry => continue, - } - } - - for w in self.inner.workers.iter() { - w.unparker.unpark(); - } + self.handle.inner.shutdown(); } /// Returns `true` if the executor has been shut down. pub fn is_shutdown(&self) -> bool { - self.inner.shutdown.load(Ordering::SeqCst) + self.handle.inner.shutdown.load(Ordering::SeqCst) } /// Creates a new context. @@ -297,31 +328,29 @@ impl Executor { /// Each context produced by an executor is given a distinct ID under the /// executor's configured prefix. pub fn new_context(&self) -> Result { - let index = self.inner.next_context.fetch_add(1, Ordering::Relaxed); - let id = self.inner.prefix.child(index); - let io = self.inner.mux.open(id.as_ref())?; + let index = self + .handle + .inner + .next_context + .fetch_add(1, Ordering::Relaxed); + let id = self.handle.inner.prefix.child(index); + let io = self.handle.inner.mux.open(id.as_ref())?; Ok(Context::with_executor( id, io, - self.inner.mux.clone(), - self.inner.clone(), + self.handle.inner.mux.clone(), + self.handle.clone(), )) } } -impl Drop for Executor { - fn drop(&mut self) { - self.shutdown(); - } -} - /// Spawns a future on the given executor inner. -pub(crate) fn spawn_on(inner: &Arc, future: F) -> Task +pub(crate) fn spawn_on(handle: &ExecutorHandle, future: F) -> Task where F: std::future::Future + Send + 'static, F::Output: Send + 'static, { - let inner = Arc::clone(inner); + let inner = Arc::clone(&handle.inner); let schedule = move |runnable: Runnable| { // After shutdown, no worker will run this. Dropping the runnable // cancels the task so the awaiter doesn't hang. SeqCst pairs with @@ -376,6 +405,23 @@ mod tests { executor.shutdown(); } + #[test] + fn test_context_outlives_executor() { + let (mux_a, _mux_b) = test_framed_mux(1024); + let executor = Executor::builder().num_threads(2).build(mux_a); + let mut ctx = executor.new_context().unwrap(); + + drop(executor); + + let output = futures::executor::block_on(ctx.join( + |_ctx| Box::pin(async move { 21 }), + |_ctx| Box::pin(async move { 21 }), + )) + .unwrap(); + + assert_eq!(output, (21, 21)); + } + #[test] fn test_executor_map() { let (mux_a, _mux_b) = test_framed_mux(1024);