diff --git a/Cargo.lock b/Cargo.lock index bb5abb64fdd..c63466befea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8028,6 +8028,7 @@ dependencies = [ "spacetimedb-fs-utils", "spacetimedb-paths", "spacetimedb-primitives", + "spacetimedb-runtime", "spacetimedb-sats", "tempfile", "thiserror 1.0.69", diff --git a/crates/commitlog/Cargo.toml b/crates/commitlog/Cargo.toml index 0c30960248f..62f534a3977 100644 --- a/crates/commitlog/Cargo.toml +++ b/crates/commitlog/Cargo.toml @@ -41,6 +41,7 @@ zstd-framed.workspace = true # For the 'test' feature env_logger = { workspace = true, optional = true } pretty_assertions.workspace = true +spacetimedb-runtime.workspace = true [dev-dependencies] # Enable streaming in tests diff --git a/crates/commitlog/src/stream/reader.rs b/crates/commitlog/src/stream/reader.rs index 48984c4bf02..c433f041686 100644 --- a/crates/commitlog/src/stream/reader.rs +++ b/crates/commitlog/src/stream/reader.rs @@ -4,10 +4,8 @@ use async_stream::try_stream; use bytes::{Buf as _, Bytes}; use futures::Stream; use log::{trace, warn}; -use tokio::{ - io::{self, AsyncBufRead, AsyncReadExt as _, AsyncSeek, AsyncSeekExt as _}, - task::spawn_blocking, -}; +use spacetimedb_runtime::spawn_blocking; +use tokio::io::{self, AsyncBufRead, AsyncReadExt as _, AsyncSeek, AsyncSeekExt as _}; use tokio_util::io::SyncIoBridge; use crate::{ @@ -107,8 +105,7 @@ fn read_segment( } segment.into_inner() }) - .await - .unwrap(); + .await; } let checksum_len = CHECKSUM_LEN[segment_header.checksum_algorithm as usize]; diff --git a/crates/commitlog/src/stream/writer.rs b/crates/commitlog/src/stream/writer.rs index b99622b83f4..fdea0a6d850 100644 --- a/crates/commitlog/src/stream/writer.rs +++ b/crates/commitlog/src/stream/writer.rs @@ -5,10 +5,8 @@ use std::{ use futures::TryFutureExt; use log::{debug, error, info, trace, warn}; -use tokio::{ - io::{AsyncBufRead, AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt}, - task::spawn_blocking, -}; +use spacetimedb_runtime::spawn_blocking; +use tokio::io::{AsyncBufRead, AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt}; use crate::{ commit, error, @@ -217,7 +215,6 @@ where move || create_segment(repo, last_written_tx_range, commitlog_options, header) }) .await - .unwrap() .map(|(segment, index)| (segment.into_async_writer(), index))?; stream.consume(segment::Header::LEN as _); @@ -425,8 +422,7 @@ impl CurrentSegment { .ok(); index }) - .await - .unwrap(); + .await; self.offset_index = Some(index); } diff --git a/crates/runtime-core/src/sim/executor/mod.rs b/crates/runtime-core/src/sim/executor/mod.rs index e8f4ce116d1..a30897179b2 100644 --- a/crates/runtime-core/src/sim/executor/mod.rs +++ b/crates/runtime-core/src/sim/executor/mod.rs @@ -175,22 +175,23 @@ impl Node { self.handle.restart_node(self.id); } - /// Spawn a `Send` future onto this simulated node. + /// Spawn a future onto this simulated node. pub fn spawn(&self, future: F) -> JoinHandle where - F: Future + Send + 'static, - F::Output: Send + 'static, + F: Future + 'static, + F::Output: 'static, { - self.handle.spawn_on(self.id, future) + self.handle.executor.assert_main_or_node(self.id); + self.handle.executor.spawn_on(self.id, future) } - /// Spawn a non-`Send` future onto this simulated node. + /// Spawn a future onto this simulated node. pub fn spawn_local(&self, future: F) -> JoinHandle where F: Future + 'static, F::Output: 'static, { - self.handle.spawn_local_on(self.id, future) + self.spawn(future) } } @@ -257,13 +258,13 @@ impl Runtime { self.handle().resume(node); } - /// Spawn a `Send` future onto a specific simulated node. - pub fn spawn_on(&self, node: NodeId, future: F) -> JoinHandle + /// Spawn a future onto the currently running node, or `MAIN` outside node work. + pub fn spawn(&self, future: F) -> JoinHandle where - F: Future + Send + 'static, - F::Output: Send + 'static, + F: Future + 'static, + F::Output: 'static, { - self.handle().spawn_on(node, future) + self.executor.spawn(future) } pub fn enable_buggify(&self) { @@ -369,24 +370,13 @@ impl Handle { self.executor.restart_node(node); } - /// Spawn a `Send` future onto a specific simulated node. - pub fn spawn_on(&self, node: NodeId, future: F) -> JoinHandle - where - F: Future + Send + 'static, - F::Output: Send + 'static, - { - self.executor.spawn_on(node, future) - } - - /// Spawn a non-`Send` future onto a specific simulated node. - /// - /// This is only valid because the simulation executor is single-threaded. - pub fn spawn_local_on(&self, node: NodeId, future: F) -> JoinHandle + /// Spawn a future onto the currently running node, or `MAIN` outside node work. + pub fn spawn(&self, future: F) -> JoinHandle where F: Future + 'static, F::Output: 'static, { - self.executor.spawn_local_on(node, future) + self.executor.spawn(future) } /// Return the current virtual time for this runtime. @@ -455,6 +445,7 @@ impl Handle { struct Executor { queue: Receiver, sender: Sender, + current_task: Mutex>, nodes: spin::Mutex>>, node_faults: NodeFaultOptions, next_node: AtomicU64, @@ -483,6 +474,7 @@ impl Executor { Self { queue: queue.receiver(), sender: queue.sender(), + current_task: Mutex::new(None), nodes: spin::Mutex::new(nodes), node_faults: config.node_faults, next_node: AtomicU64::new(1), @@ -585,28 +577,17 @@ impl Executor { } } - /// Spawn a `Send` task and enqueue its runnable on the shared runtime queue. - fn spawn_on(&self, node: NodeId, future: F) -> JoinHandle + /// Spawn a task onto the node whose task is currently being polled. + fn spawn(&self, future: F) -> JoinHandle where - F: Future + Send + 'static, - F::Output: Send + 'static, + F: Future + 'static, + F::Output: 'static, { - let abort = AbortHandle::new(); - let abortable = Abortable::new(future, abort.clone()); - let sender = self.sender.clone(); - let (runnable, task) = async_task::Builder::new() - .metadata(self.task_meta(node)) - .spawn(move |_| abortable, move |runnable| sender.send(runnable)); - runnable.schedule(); - - JoinHandle { - task: task.fallible(), - abort, - } + self.spawn_on(self.current_node(), future) } - /// Spawn a non-`Send` task on the single-threaded runtime. - fn spawn_local_on(&self, node: NodeId, future: F) -> JoinHandle + /// Spawn a task and enqueue its runnable on the shared runtime queue. + fn spawn_on(&self, node: NodeId, future: F) -> JoinHandle where F: Future + 'static, F::Output: 'static, @@ -778,6 +759,7 @@ impl Executor { state.paused_queue.lock().push(runnable); continue; } + let _current_task = self.enter_current_task(meta); runnable.run(); // Advance virtual time by 100ns-1us per task poll to model execution cost. // Using the runtime RNG keeps overhead deterministic by seed. @@ -805,11 +787,48 @@ impl Executor { TaskMeta::new(node, state.generation()) } + fn current_node(&self) -> NodeId { + self.current_task + .lock() + .as_ref() + .map(|meta| meta.node) + .unwrap_or(NodeId::MAIN) + } + + fn assert_main_or_node(&self, node: NodeId) { + let caller = self.current_node(); + assert!( + caller == NodeId::MAIN || caller == node, + "node {caller} cannot spawn task on node {node}" + ); + } + + fn enter_current_task(&self, meta: TaskMeta) -> CurrentTaskGuard<'_> { + let mut current = self.current_task.lock(); + // The executor must not poll another runnable while one task's node + // context is installed; otherwise ambient spawn would inherit the + // wrong node/generation after reentrant scheduling. + assert!(current.is_none(), "nested simulated task polling"); + *current = Some(meta); + CurrentTaskGuard { executor: self } + } + fn node_state(&self, node: NodeId) -> Arc { self.node_record(node).state.clone() } } +struct CurrentTaskGuard<'a> { + executor: &'a Executor, +} + +impl Drop for CurrentTaskGuard<'_> { + fn drop(&mut self) { + let current = self.executor.current_task.lock().take(); + assert!(current.is_some(), "current simulated task guard dropped without task"); + } +} + fn poll_finished_task(task: &mut async_task::Task) -> Option { if !task.is_finished() { return None; @@ -1036,6 +1055,20 @@ mod tests { assert_eq!(value, 11); } + #[test] + #[should_panic(expected = "cannot spawn task on node")] + fn node_cannot_spawn_task_on_another_node() { + let mut runtime = Runtime::new(3); + let node_a = runtime.create_node().name("a").build(); + let node_b = runtime.create_node().name("b").build(); + + let task = node_a.spawn(async move { + let _child = node_b.spawn(async {}); + }); + + runtime.block_on(task).expect("parent task should panic first"); + } + #[test] fn runtime_config_sets_seed() { let runtime = Runtime::with_config(RuntimeConfig::new(77)); diff --git a/crates/runtime-core/src/sim/time/mod.rs b/crates/runtime-core/src/sim/time/mod.rs index 7af1ab3bb70..124508281d3 100644 --- a/crates/runtime-core/src/sim/time/mod.rs +++ b/crates/runtime-core/src/sim/time/mod.rs @@ -268,14 +268,14 @@ mod tests { async move { let slow_order = Arc::clone(&order); let slow_handle = handle.clone(); - let slow = handle.spawn_on(sim::NodeId::MAIN, async move { + let slow = handle.spawn(async move { slow_handle.sleep(Duration::from_millis(10)).await; slow_order.lock().push(10); }); let fast_order = Arc::clone(&order); let fast_handle = handle.clone(); - let fast = handle.spawn_on(sim::NodeId::MAIN, async move { + let fast = handle.spawn(async move { fast_handle.sleep(Duration::from_millis(3)).await; fast_order.lock().push(3); }); diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs index 8f876aca6e3..da0676f3007 100644 --- a/crates/runtime/src/lib.rs +++ b/crates/runtime/src/lib.rs @@ -53,6 +53,16 @@ pub enum Handle { Simulation(sim::Handle), } +impl fmt::Debug for Handle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Tokio(_) => f.write_str("Handle::Tokio"), + #[cfg(feature = "simulation")] + Self::Simulation(_) => f.write_str("Handle::Simulation"), + } + } +} + pub struct JoinHandle { inner: JoinHandleInner, } @@ -134,6 +144,16 @@ impl fmt::Display for JoinError { impl std::error::Error for JoinError {} +impl JoinError { + pub fn is_panic(&self) -> bool { + match &self.inner { + JoinErrorInner::Tokio(error) => error.is_panic(), + #[cfg(feature = "simulation")] + JoinErrorInner::Simulation(_) => false, + } + } +} + impl JoinHandleInner { fn abort_handle(&self) -> AbortHandle { match self { @@ -208,6 +228,28 @@ impl Unpin for JoinHandle {} #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct RuntimeTimeout; +#[must_use = "runtime context exits immediately unless the guard is held"] +pub struct EnterGuard<'a> { + _inner: EnterGuardInner<'a>, +} + +#[allow(dead_code)] +enum EnterGuardInner<'a> { + Tokio(tokio::runtime::EnterGuard<'a>), + #[cfg(feature = "simulation")] + Simulation(sim_std::EnterGuard), +} + +impl Drop for EnterGuard<'_> { + fn drop(&mut self) { + match &self._inner { + EnterGuardInner::Tokio(_) => {} + #[cfg(feature = "simulation")] + EnterGuardInner::Simulation(_) => {} + } + } +} + impl fmt::Display for RuntimeTimeout { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str("runtime operation timed out") @@ -216,6 +258,36 @@ impl fmt::Display for RuntimeTimeout { impl std::error::Error for RuntimeTimeout {} +/// Spawn a task on the current Runtime +pub fn spawn(future: impl Future + Send + 'static) -> JoinHandle { + Handle::current().spawn(future) +} + +/// Run blocking work on the current runtime. +/// +/// Tokio runs `f` on its blocking thread pool. The simulation backend runs `f` +/// as a normal simulated task on the single executor thread, so it preserves +/// deterministic scheduling but does not provide blocking-pool parallelism. +pub async fn spawn_blocking(f: F) -> R +where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, +{ + Handle::current().spawn_blocking(f).await +} + +/// Spawn a task on the current Tokio runtime, bypassing simulation enter state. +pub fn tokio_spawn(future: impl Future + Send + 'static) -> JoinHandle { + JoinHandle { + inner: JoinHandleInner::Tokio(tokio::spawn(future)), + } +} + +/// Sleep on the currently active Runtime +pub async fn sleep(duration: Duration) { + Handle::current().sleep(duration).await +} + impl Handle { pub fn tokio(handle: TokioHandle) -> Self { Self::Tokio(handle) @@ -224,6 +296,16 @@ impl Handle { pub fn tokio_current() -> Self { Self::tokio(TokioHandle::current()) } + + /// Return the runtime handle active on this OS thread. + pub fn current() -> Handle { + #[cfg(feature = "simulation")] + if let Some(handle) = sim_std::try_current_handle() { + return handle; + } + + Handle::tokio_current() + } } #[cfg(feature = "simulation")] @@ -234,6 +316,18 @@ impl Handle { } impl Handle { + pub fn enter(&self) -> EnterGuard<'_> { + match self { + Self::Tokio(handle) => EnterGuard { + _inner: EnterGuardInner::Tokio(handle.enter()), + }, + #[cfg(feature = "simulation")] + Self::Simulation(_) => EnterGuard { + _inner: EnterGuardInner::Simulation(sim_std::enter(self.clone())), + }, + } + } + pub fn spawn(&self, future: impl Future + Send + 'static) -> JoinHandle { match self { Self::Tokio(handle) => JoinHandle { @@ -241,11 +335,19 @@ impl Handle { }, #[cfg(feature = "simulation")] Self::Simulation(handle) => JoinHandle { - inner: JoinHandleInner::Simulation(handle.spawn_on(sim::NodeId::MAIN, future)), + inner: JoinHandleInner::Simulation(handle.spawn(future)), }, } } + pub fn block_on(&self, future: F) -> F::Output { + match self { + Self::Tokio(handle) => handle.block_on(future), + #[cfg(feature = "simulation")] + Self::Simulation(handle) => handle.block_on(future), + } + } + pub async fn spawn_blocking(&self, f: F) -> R where F: FnOnce() -> R + Send + 'static, @@ -265,7 +367,7 @@ impl Handle { // the simulation backend. #[cfg(feature = "simulation")] Self::Simulation(handle) => handle - .spawn_on(sim::NodeId::MAIN, async move { f() }) + .spawn(async move { f() }) .await .expect("simulation spawn_blocking task should not be cancelled"), } @@ -329,6 +431,73 @@ mod tests { assert!(flag.load(Ordering::Acquire)); } + #[cfg(feature = "simulation")] + #[test] + fn ambient_runtime_uses_entered_simulation_handle() { + use crate::sim::Runtime; + let mut rt = Runtime::new(9); + let handle = Handle::simulation(rt.handle()); + let _entered = handle.enter(); + + let output = rt.block_on(async { + let task = crate::spawn(async { + crate::sleep(std::time::Duration::from_millis(1)).await; + 7 + }); + task.await.expect("ambient simulation task should complete") + }); + + assert_eq!(output, 7); + assert!(rt.elapsed() >= std::time::Duration::from_millis(1)); + } + + #[cfg(feature = "simulation")] + #[test] + fn simulation_enter_is_reentrant_for_same_handle() { + use crate::sim::Runtime; + let rt = Runtime::new(10); + let handle = Handle::simulation(rt.handle()); + let outer = handle.enter(); + + { + let _inner = handle.enter(); + assert!(matches!(crate::current(), Handle::Simulation(_))); + } + + assert!(matches!(crate::current(), Handle::Simulation(_))); + drop(outer); + } + + #[cfg(feature = "simulation")] + #[test] + fn ambient_spawn_inside_node_inherits_node_pause() { + use crate::sim::Runtime; + let mut rt = Runtime::new(11); + let handle = Handle::simulation(rt.handle()); + let _entered = handle.enter(); + let node = rt.create_node().name("worker").build(); + let child_ran = Arc::new(AtomicBool::new(false)); + + rt.block_on(async { + let flag = Arc::clone(&child_ran); + let node_for_parent = node.clone(); + let parent = node.spawn(async move { + let child = crate::spawn(async move { + flag.store(true, Ordering::Release); + }); + drop(child); + node_for_parent.pause(); + }); + + parent.await.expect("parent task should complete"); + assert!(!child_ran.load(Ordering::Acquire)); + + node.resume(); + handle.sleep(std::time::Duration::from_millis(1)).await; + assert!(child_ran.load(Ordering::Acquire)); + }); + } + #[cfg(feature = "simulation")] #[test] fn abort_cancels_task_in_simulation() { diff --git a/crates/runtime/src/sim_std.rs b/crates/runtime/src/sim_std.rs index 8153482519e..a92a2fedef8 100644 --- a/crates/runtime/src/sim_std.rs +++ b/crates/runtime/src/sim_std.rs @@ -7,11 +7,11 @@ #![allow(clippy::disallowed_macros)] -use core::{cell::Cell, future::Future}; +use core::{cell::RefCell, future::Future, marker::PhantomData}; use std::boxed::Box; use std::sync::OnceLock; -use crate::sim; +use crate::{sim, Handle}; // Public entry points. @@ -21,7 +21,7 @@ use crate::sim; /// tests that execute inside a hosted process. While the future runs, this /// marks the thread as inside simulation so OS thread spawns can be rejected. pub fn block_on(runtime: &mut sim::Runtime, future: F) -> F::Output { - let _guard = enter(); + let _guard = enter(Handle::simulation(runtime.handle())); runtime.block_on(future) } @@ -70,31 +70,71 @@ fn panic_with_seed(seed: u64, payload: Box) -> ! { // Ambient hosted state used only while sim_std is driving a simulation runtime. // -// The simulator itself stays explicit-handle based. This flag only marks the -// current OS thread as simulation-owned so host thread creation and randomness -// hooks can reject accidental escapes from deterministic simulation. +// The simulator itself stays explicit-handle based. This thread-local slot lets +// production-shaped runtime APIs resolve to the simulated handle while hosted +// DST code runs, and lets OS hooks reject escapes from deterministic execution. thread_local! { - static IN_SIMULATION: Cell = const { Cell::new(false) }; + static SIM_RUNTIME: RefCell> = const { RefCell::new(None) }; } #[must_use = "simulation runtime context exits immediately unless the guard is held"] -struct EnterGuard; +pub struct EnterGuard { + // `active` means this guard installed the thread-local runtime and is + // responsible for clearing it on drop. Reentrant enters for the same + // simulation runtime return an inactive guard so sync helper functions can + // call `runtime.enter()` inside an already-entered DST run without ending + // the outer context when the helper returns. + active: bool, + _not_send: PhantomData>, +} + +/// Enter a simulated runtime on the current OS thread. +/// +/// This is hosted glue for DST and tests. It intentionally lives outside +/// runtime-core because `thread_local!` and process hooks are std-only. +pub(crate) fn enter(handle: Handle) -> EnterGuard { + assert!( + matches!(&handle, Handle::Simulation(_)), + "sim_std::enter requires a simulation runtime handle" + ); + let active = SIM_RUNTIME.with(|current| { + let mut current = current.borrow_mut(); + if current.is_some() { + return false; + } -fn enter() -> EnterGuard { - IN_SIMULATION.with(|current| { - assert!(!current.replace(true), "nested hosted simulation block_on"); + *current = Some(handle); + true }); - EnterGuard + EnterGuard { + active, + _not_send: PhantomData, + } +} + +/// Return the simulated runtime currently entered on this OS thread, if any. +pub fn try_current_handle() -> Option { + SIM_RUNTIME.with(|current| current.borrow().clone()) +} + +/// Return the simulated runtime currently entered on this OS thread. +pub fn current_handle() -> Handle { + try_current_handle().expect("simulation runtime API used outside runtime.enter") } fn in_simulation() -> bool { - IN_SIMULATION.with(Cell::get) + try_current_handle().is_some() } impl Drop for EnterGuard { fn drop(&mut self) { - IN_SIMULATION.with(|current| { - assert!(current.replace(false), "simulation context guard dropped without enter"); + if !self.active { + return; + } + + SIM_RUNTIME.with(|current| { + let old = current.borrow_mut().take(); + assert!(old.is_some(), "simulation context guard dropped without enter"); }); } }