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..92775650404 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, 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 _); @@ -378,7 +375,7 @@ where fn drop(&mut self) { if let Some(current_segment) = self.current_segment.take() { trace!("closing current segment on writer drop"); - tokio::spawn( + spawn( current_segment .close() .inspect_err(|e| warn!("error closing segment on drop: {e}")), @@ -425,8 +422,7 @@ impl CurrentSegment { .ok(); index }) - .await - .unwrap(); + .await; self.offset_index = Some(index); } diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index bc9e4fcc36e..432646888fe 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -7,6 +7,7 @@ use spacetimedb_sats::ArrayValue; use super::model::Model; use crate::schema::{SchemaPlan, TablePlan, Type}; +use crate::traits::InteractionGen; pub type Row = ProductValue; @@ -260,6 +261,14 @@ impl WorkloadGen { } } +impl InteractionGen for WorkloadGen { + type Interaction = Interaction; + + fn next_interaction(&mut self) -> Self::Interaction { + self.next_interaction() + } +} + impl Debug for WorkloadGen { fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> { write!(f, "{:?}", self.stats()) diff --git a/crates/dst/src/lib.rs b/crates/dst/src/lib.rs index 8d12c575e4c..b89ba60cd58 100644 --- a/crates/dst/src/lib.rs +++ b/crates/dst/src/lib.rs @@ -3,4 +3,4 @@ pub mod schema; pub mod sim; pub mod traits; -pub use traits::{Properties, TargetDriver, TestSuite, TestSuiteParts}; +pub use traits::{InteractionGen, Properties, TargetDriver, TestSuite, TestSuiteParts}; diff --git a/crates/dst/src/traits.rs b/crates/dst/src/traits.rs index 2185e0ec918..6f8e96c6dee 100644 --- a/crates/dst/src/traits.rs +++ b/crates/dst/src/traits.rs @@ -16,6 +16,19 @@ pub trait Properties { fn observe(&mut self, interaction: &I, observation: &O) -> Result<(), Error>; } +/// Generates interactions, and can feed observations back to the generator so +/// its internal state stays in sync with the target. +pub trait InteractionGen: std::fmt::Debug { + type Interaction: std::fmt::Debug; + + fn next_interaction(&mut self) -> Self::Interaction; + + /// Feed an observation back to the generator. Defaults to ignoring it. + fn observe(&mut self, _interaction: &Self::Interaction, _observation: &O) -> Result<(), Error> { + Ok(()) + } +} + pub type TestSuiteParts = ( ::Interactions, ::Target, @@ -24,7 +37,10 @@ pub type TestSuiteParts = ( pub trait TestSuite { type Interaction: std::fmt::Debug; - type Interactions: Iterator + std::fmt::Debug; + type Interactions: InteractionGen< + >::Observation, + Interaction = Self::Interaction, + >; type Target: TargetDriver; type Properties: Properties>::Observation>; @@ -39,19 +55,16 @@ pub trait TestSuite { async move { let (mut interactions, mut target, mut properties) = self.build(rng).await?; - let result = async { - for interaction in interactions.by_ref().take(max_interactions) { - let observation = target.execute(&interaction).await?; - properties.observe(&interaction, &observation)?; - } - - Ok(()) + for _ in 0..max_interactions { + let interaction = interactions.next_interaction(); + let observation = target.execute(&interaction).await?; + interactions.observe(&interaction, &observation)?; + properties.observe(&interaction, &observation)?; } - .await; tracing::info!(interaction_counts = ?interactions, "final interaction counts"); - result + Ok(()) } } } diff --git a/crates/runtime-core/src/sim/executor/mod.rs b/crates/runtime-core/src/sim/executor/mod.rs index e8f4ce116d1..9b149c8701a 100644 --- a/crates/runtime-core/src/sim/executor/mod.rs +++ b/crates/runtime-core/src/sim/executor/mod.rs @@ -39,6 +39,7 @@ const READY_TASK_BUDGET: usize = 256; pub struct RuntimeConfig { pub seed: u64, pub network: Option, + pub buggify: bool, pub node_faults: NodeFaultOptions, } @@ -47,6 +48,7 @@ impl RuntimeConfig { Self { seed, network: None, + buggify: false, node_faults: NodeFaultOptions::default(), } } @@ -56,6 +58,11 @@ impl RuntimeConfig { self } + pub fn enable_buggify(mut self) -> Self { + self.buggify = true; + self + } + pub fn with_node_faults(mut self, node_faults: NodeFaultOptions) -> Self { self.node_faults = node_faults; self @@ -131,10 +138,9 @@ impl NodeBuilder { } } -/// Handle to one simulated node in the runtime. +/// Thin wrapper around a node-bound runtime handle. #[derive(Clone)] pub struct Node { - id: NodeId, handle: Handle, config: Arc, } @@ -142,7 +148,17 @@ pub struct Node { impl Node { /// Return the stable identifier for this simulated node. pub fn id(&self) -> NodeId { - self.id + self.handle.node_id() + } + + /// Return a cloneable handle bound to this simulated node. + pub fn handle(&self) -> Handle { + self.handle.clone() + } + + /// Convert this node wrapper into its node-bound handle. + pub fn into_handle(self) -> Handle { + self.handle } /// Return the optional human-readable name for this node. @@ -152,45 +168,45 @@ impl Node { /// Return the simulated network endpoint for this node. pub fn net(&self) -> Option { - self.handle.network().map(|net| net.on_node(self.id)) - } - - /// Crash this node and invalidate all tasks spawned before the crash. - pub fn crash(&self) { - self.handle.crash_node(self.id); + self.handle.net() } /// Pause scheduling for this node. pub fn pause(&self) { - self.handle.pause(self.id); + self.handle.pause(); + } + + /// Crash this node and invalidate all tasks spawned before the crash. + pub fn crash(&self) { + self.handle.crash_node(); } /// Resume scheduling for this node. pub fn resume(&self) { - self.handle.resume(self.id); + self.handle.resume(); } /// Restart this node and invalidate all tasks spawned before the restart. pub fn restart(&self) { - self.handle.restart_node(self.id); + self.handle.restart_node(); } - /// 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.spawn(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.handle.spawn_local(future) } } @@ -221,7 +237,7 @@ impl Runtime { /// While the future runs, spawned tasks share the same deterministic /// scheduler, timer wheel, and runtime RNG. pub fn block_on(&mut self, future: F) -> F::Output { - self.executor.block_on(future) + self.executor.block_on_on(NodeId::MAIN, future) } /// Return the amount of virtual time elapsed in this runtime. @@ -233,6 +249,7 @@ impl Runtime { pub fn handle(&self) -> Handle { Handle { executor: Arc::clone(&self.executor), + node: NodeId::MAIN, } } @@ -249,21 +266,21 @@ impl Runtime { /// Tasks already queued for the node are retained and will run only after /// the node is resumed. pub fn pause(&self, node: NodeId) { - self.handle().pause(node); + self.executor.pause(node); } /// Resume scheduling for a previously paused node. pub fn resume(&self, node: NodeId) { - self.handle().resume(node); + self.executor.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) { @@ -319,14 +336,32 @@ impl Runtime { #[derive(Clone)] pub struct Handle { executor: Arc, + node: NodeId, } impl Handle { + /// Return the stable identifier this handle is bound to. + pub fn node_id(&self) -> NodeId { + self.node + } + + fn with_node(&self, node: NodeId) -> Self { + Self { + executor: Arc::clone(&self.executor), + node, + } + } + /// Return the shared simulated network for this runtime. pub fn network(&self) -> Option { self.executor.net.clone() } + /// Return the simulated network endpoint for this handle's node. + pub fn net(&self) -> Option { + self.network().map(|net| net.on_node(self.node)) + } + /// Create a new simulated node owned by this runtime. pub fn create_node(&self) -> NodeBuilder { NodeBuilder { @@ -343,50 +378,56 @@ impl Handle { let id = self.executor.create_node(config); let config = self.node_config(id); Node { - id, - handle: self.clone(), + handle: self.with_node(id), config, } } /// Pause scheduling for a node. - pub fn pause(&self, node: NodeId) { - self.executor.pause(node); + pub fn pause(&self) { + self.executor.pause(self.node); } /// Crash a node until it is restarted. - pub fn crash_node(&self, node: NodeId) { - self.executor.crash_node(node); + pub fn crash_node(&self) { + self.executor.crash_node(self.node); } /// Resume scheduling for a node and requeue any buffered tasks for it. - pub fn resume(&self, node: NodeId) { - self.executor.resume(node); + pub fn resume(&self) { + self.executor.resume(self.node); } /// Restart a node and invalidate all tasks spawned before the restart. - pub fn restart_node(&self, node: NodeId) { - self.executor.restart_node(node); + pub fn restart_node(&self) { + self.executor.restart_node(self.node); } - /// Spawn a `Send` future onto a specific simulated node. - pub fn spawn_on(&self, node: NodeId, future: F) -> JoinHandle + /// Spawn a future onto this handle's node. + /// + /// The main runtime handle keeps ambient spawn semantics and inherits the + /// node currently being polled. Node-bound handles always target their + /// bound node. + pub fn spawn(&self, future: F) -> JoinHandle where - F: Future + Send + 'static, - F::Output: Send + 'static, + F: Future + 'static, + F::Output: 'static, { - self.executor.spawn_on(node, future) + if self.node == NodeId::MAIN { + self.executor.spawn(future) + } else { + self.executor.assert_main_or_node(self.node); + self.executor.spawn_on(self.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 non-`Send` future onto this handle's node. + pub fn spawn_local(&self, future: F) -> JoinHandle where F: Future + 'static, F::Output: 'static, { - self.executor.spawn_local_on(node, future) + self.spawn(future) } /// Return the current virtual time for this runtime. @@ -419,7 +460,7 @@ impl Handle { } pub fn block_on(&self, future: F) -> F::Output { - self.executor.block_on(future) + self.executor.block_on_on(self.node, future) } pub fn enable_buggify(&self) { @@ -455,6 +496,7 @@ impl Handle { struct Executor { queue: Receiver, sender: Sender, + current_tasks: Mutex>, nodes: spin::Mutex>>, node_faults: NodeFaultOptions, next_node: AtomicU64, @@ -480,9 +522,14 @@ impl Executor { net.register_node(NodeId::MAIN); } + if config.buggify { + rng.enable_buggify(); + } + Self { queue: queue.receiver(), sender: queue.sender(), + current_tasks: Mutex::new(Vec::new()), nodes: spin::Mutex::new(nodes), node_faults: config.node_faults, next_node: AtomicU64::new(1), @@ -585,28 +632,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, @@ -622,7 +658,7 @@ impl Executor { runnable.schedule(); JoinHandle { - task: task.fallible(), + task: Some(task.fallible()), abort, } } @@ -634,11 +670,13 @@ impl Executor { /// simulated fault sources at one captured instant, then advances virtual /// time only when no current-time source can make progress. If neither /// runnable work nor timers remain, the simulation is considered deadlocked. - fn block_on(&self, future: F) -> F::Output { + fn block_on_on(&self, node: NodeId, future: F) -> F::Output { + self.assert_main_or_node(node); + let sender = self.sender.clone(); let (runnable, mut task) = unsafe { async_task::Builder::new() - .metadata(self.task_meta(NodeId::MAIN)) + .metadata(self.task_meta(node)) .spawn_unchecked(move |_| future, move |runnable| sender.send(runnable)) }; runnable.schedule(); @@ -778,6 +816,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. @@ -801,15 +840,54 @@ impl Executor { } fn task_meta(&self, node: NodeId) -> TaskMeta { + if let Some(current) = self.current_task() + && current.node == node + { + return current; + } + let state = self.node_state(node); TaskMeta::new(node, state.generation()) } + fn current_task(&self) -> Option { + self.current_tasks.lock().last().copied() + } + + fn current_node(&self) -> NodeId { + self.current_task().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<'_> { + self.current_tasks.lock().push(meta); + CurrentTaskGuard { executor: self, meta } + } + fn node_state(&self, node: NodeId) -> Arc { self.node_record(node).state.clone() } } +struct CurrentTaskGuard<'a> { + executor: &'a Executor, + meta: TaskMeta, +} + +impl Drop for CurrentTaskGuard<'_> { + fn drop(&mut self) { + let current = self.executor.current_tasks.lock().pop(); + assert_eq!(current, Some(self.meta), "current simulated task stack corrupted"); + } +} + fn poll_finished_task(task: &mut async_task::Task) -> Option { if !task.is_finished() { return None; @@ -980,6 +1058,8 @@ mod tests { use super::*; use crate::sim::RuntimeConfig; + struct Spawned(JoinHandle); + struct DropFlag(Arc); impl Drop for DropFlag { @@ -1036,6 +1116,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)); @@ -1204,6 +1298,94 @@ mod tests { assert_eq!(value, 17); } + #[test] + fn node_bound_handle_block_on_runs_on_bound_node() { + let mut runtime = Runtime::new(14); + let main = runtime.handle(); + let node = runtime.create_node().name("bound").build(); + let node_handle = node.handle(); + let child_ran = Arc::new(AtomicBool::new(false)); + + let child = node_handle + .block_on({ + let main = main.clone(); + let node = node.clone(); + let child_ran = Arc::clone(&child_ran); + async move { + let child = main.spawn(async move { + child_ran.store(true, Ordering::Release); + }); + node.pause(); + Spawned(child) + } + }) + .0; + + runtime.block_on(async { + yield_now().await; + }); + assert!(!child_ran.load(Ordering::Acquire)); + + node.resume(); + runtime + .block_on(child) + .expect("child spawned from node-bound block_on should complete"); + assert!(child_ran.load(Ordering::Acquire)); + } + + #[test] + fn nested_block_on_restores_outer_node_context() { + let mut runtime = Runtime::new(15); + let main = runtime.handle(); + let node = runtime.create_node().name("outer").build(); + let node_handle = node.handle(); + let child_ran = Arc::new(AtomicBool::new(false)); + + let outer = node.spawn({ + let main = main.clone(); + let node = node.clone(); + let child_ran = Arc::clone(&child_ran); + async move { + node_handle.block_on(async { + yield_now().await; + }); + + let child = main.spawn(async move { + child_ran.store(true, Ordering::Release); + }); + node.pause(); + Spawned(child) + } + }); + let child = runtime.block_on(outer).expect("outer task should complete").0; + + runtime.block_on(async { + yield_now().await; + }); + assert!(!child_ran.load(Ordering::Acquire)); + + node.resume(); + runtime + .block_on(child) + .expect("ambient spawn should inherit restored outer node"); + assert!(child_ran.load(Ordering::Acquire)); + } + + #[test] + #[should_panic(expected = "node 1 cannot spawn task on node 2")] + fn nested_block_on_rejects_cross_node_reentry() { + let mut runtime = Runtime::new(16); + let node_a = runtime.create_node().name("a").build(); + let node_b = runtime.create_node().name("b").build(); + let node_b_handle = node_b.handle(); + + let task = node_a.spawn(async move { + node_b_handle.block_on(async {}); + }); + + let _ = runtime.block_on(task); + } + #[test] fn block_on_returns_while_background_task_stays_ready() { let mut runtime = Runtime::new(10); diff --git a/crates/runtime-core/src/sim/executor/task.rs b/crates/runtime-core/src/sim/executor/task.rs index 043c672b833..d3cf63489b8 100644 --- a/crates/runtime-core/src/sim/executor/task.rs +++ b/crates/runtime-core/src/sim/executor/task.rs @@ -18,9 +18,9 @@ use super::TaskMeta; /// - The executor holds the `Runnable` (not visible here). pub struct JoinHandle { // async_task::FallibleTask owns a shared heap-allocated cell that holds the - // future, output, task metadata, and waker. `None` means the executor - // intentionally dropped the runnable before polling it. - pub(crate) task: async_task::FallibleTask, TaskMeta>, + // future, output, task metadata, and waker. `None` means this handle was + // detached and can no longer be awaited. + pub(crate) task: Option, TaskMeta>>, // Clone of the same AbortHandle that Abortable holds inside the task. pub(crate) abort: AbortHandle, } @@ -32,9 +32,10 @@ impl JoinHandle { } /// Drop the join handle without cancelling the task. - pub fn detach(self) { - // async_task::Task::detach makes Drop a no-op; the future keeps running. - self.task.detach(); + pub fn detach(mut self) { + if let Some(task) = self.task.take() { + task.detach(); + } } /// Poll the underlying async_task::Task for its output. @@ -42,7 +43,8 @@ impl JoinHandle { pub fn poll_join(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { // FallibleTask lets node crash discard stale runnables without panicking // when their JoinHandle is awaited. - match Pin::new(&mut self.task).poll(cx) { + let task = self.task.as_mut().expect("detached JoinHandle cannot be polled"); + match Pin::new(task).poll(cx) { Poll::Ready(Some(result)) => Poll::Ready(result), Poll::Ready(None) => Poll::Ready(Err(JoinError)), Poll::Pending => Poll::Pending, @@ -58,6 +60,14 @@ impl Future for JoinHandle { } } +impl Drop for JoinHandle { + fn drop(&mut self) { + if let Some(task) = self.task.take() { + task.detach(); + } + } +} + /// Two-phase cancellation for a simulated task. /// /// [`AbortHandle`] and [`Abortable`] work together: 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..1ce3dd36804 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::Handle::current(), Handle::Simulation(_))); + } + + assert!(matches!(crate::Handle::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"); }); } }