diff --git a/crates/omnigraph-core/src/dst_blocking.rs b/crates/omnigraph-core/src/dst_blocking.rs new file mode 100644 index 000000000..835e3d83e --- /dev/null +++ b/crates/omnigraph-core/src/dst_blocking.rs @@ -0,0 +1,47 @@ +//! The blocking-work seam: where a graph operator's body is polled, sibling +//! of `dst_gate`, `dst_clock` and `dst_ids`. +//! +//! Uninstalled (default, production): the engine polls each operator body on +//! a `spawn_blocking` thread, so CPU-heavy work stays off the async workers. +//! Installed (a harness thread running a seeded current-thread runtime): the +//! body is polled on that runtime, as one more task the seed orders. Two +//! bodies that charge one memory pool then interleave the same way on every +//! replay, so a refusal's reported size is a function of the seed. + +#[cfg(feature = "dst")] +use omnigraph_seams::{Behavior, Op}; + +/// What the blocking seam holds: installing it is the whole signal. +#[cfg(feature = "dst")] +pub trait Placement: Behavior {} + +/// The one placement: bodies run on the installing thread's runtime. +#[cfg(feature = "dst")] +pub struct OnSeededRuntime; + +#[cfg(feature = "dst")] +impl Behavior for OnSeededRuntime {} + +#[cfg(feature = "dst")] +impl Placement for OnSeededRuntime {} + +#[cfg(feature = "dst")] +omnigraph_seams::thread_local_seam! { + /// The blocking seam. Install [`OnSeededRuntime`] on the harness thread. + pub static BLOCKING: dyn Placement = ("blocking", Op::Unreachable); +} + +/// Whether operator bodies on this thread run on its runtime instead of a +/// blocking thread. Without the `dst` feature this inlines to `false`. +#[cfg(not(feature = "dst"))] +#[inline(always)] +pub fn on_runtime() -> bool { + false +} + +/// Whether operator bodies on this thread run on its runtime instead of a +/// blocking thread. +#[cfg(feature = "dst")] +pub fn on_runtime() -> bool { + BLOCKING.with(|_| ()).is_some() +} diff --git a/crates/omnigraph-core/src/lib.rs b/crates/omnigraph-core/src/lib.rs index a3a57c2cd..deaa545ef 100644 --- a/crates/omnigraph-core/src/lib.rs +++ b/crates/omnigraph-core/src/lib.rs @@ -13,6 +13,7 @@ pub mod branch_control; pub mod branch_names; pub mod dataset_index; +pub mod dst_blocking; pub mod dst_clock; pub mod dst_gate; pub mod dst_ids; diff --git a/crates/omnigraph-dst/src/environment.rs b/crates/omnigraph-dst/src/environment.rs index 247fdfe59..abb921b17 100644 --- a/crates/omnigraph-dst/src/environment.rs +++ b/crates/omnigraph-dst/src/environment.rs @@ -4,6 +4,7 @@ use std::future::Future; use futures::FutureExt; +use omnigraph::dst_blocking::Placement; use omnigraph::dst_clock::Clock; use omnigraph::dst_ids::IdSource; use omnigraph::seams::{Installed, ThreadLocal}; @@ -71,6 +72,7 @@ struct InstalledEnvironment { process: UniverseProcess, _clock: Option>>, _ids: Option>>, + _blocking: Option>>, } impl InstalledEnvironment { @@ -79,6 +81,7 @@ impl InstalledEnvironment { process, _clock: None, _ids: None, + _blocking: None, } } } @@ -150,6 +153,11 @@ pub fn run_universe>( installed._clock = Some(omnigraph::dst_clock::CLOCK.install( std::sync::Arc::new(omnigraph::dst_clock::LogicalClock::default()), )); + // Operator bodies run on this seeded runtime, not on + // blocking threads the seed does not order. + installed._blocking = Some(omnigraph::dst_blocking::BLOCKING.install( + std::sync::Arc::new(omnigraph::dst_blocking::OnSeededRuntime), + )); let setup = std::panic::AssertUnwindSafe(async { environment.setup().await }) .catch_unwind() .await; diff --git a/crates/omnigraph/src/engine/operators/memory.rs b/crates/omnigraph/src/engine/operators/memory.rs index a9afa763b..02a9fef66 100644 --- a/crates/omnigraph/src/engine/operators/memory.rs +++ b/crates/omnigraph/src/engine/operators/memory.rs @@ -353,6 +353,25 @@ impl WorkMemory { }); let wake = Arc::new(BlockingWake::default()); let elapsed = MetricBuilder::new(&self.metrics).elapsed_compute(0); + if crate::dst_blocking::on_runtime() { + // A seeded runtime polls the body as one of its own tasks, so the + // seed orders its charges against every other body's. The + // checkpoint hook is keyed by thread and every body shares this + // one, so none is installed: a checkpoint only checks. + let _owner = owner; + let mut guard = crate::instrumentation::query_blocking_work_guard(); + guard.started(); + let result = std::future::poll_fn(|context| { + if let Err(error) = self.check() { + return std::task::Poll::Ready(Err(error)); + } + let _timer = elapsed.timer(); + work.as_mut().poll(context) + }) + .await; + cancellation.disarm(); + return result; + } loop { self.check()?; let memory = Arc::clone(self); @@ -1386,6 +1405,46 @@ mod tests { }); } + /// Under the DST blocking seam a body, and a body it nests, runs on the + /// installing thread's runtime, which the seed orders; without the seam + /// it runs on a blocking thread, which the seed does not. Two such + /// threads charging one pool made a csv load's refusal report a + /// different size on a DST replay of one seed. Which thread polls a body + /// is invisible to a case. + #[cfg(feature = "dst")] + #[test] + fn blocking_bodies_run_on_the_seeded_runtime_under_the_dst_seam_issue_1160() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let runtime_thread = std::thread::current().id(); + runtime.block_on(async { + let (_, memory) = test_memory(1_048_576); + let off_seam = memory + .blocking(|_| async { Ok(std::thread::current().id()) }) + .await + .unwrap(); + assert_ne!(off_seam, runtime_thread, "no seam: a blocking thread"); + let _seam = crate::dst_blocking::BLOCKING + .install(Arc::new(crate::dst_blocking::OnSeededRuntime)); + let on_seam = memory + .blocking(|_| async { Ok(std::thread::current().id()) }) + .await + .unwrap(); + assert_eq!(on_seam, runtime_thread, "seam: the seeded runtime"); + let nested = memory + .blocking(move |memory| async move { + memory + .blocking(|_| async { Ok(std::thread::current().id()) }) + .await + }) + .await + .unwrap(); + assert_eq!(nested, runtime_thread, "a nested body too"); + }); + } + /// Aborts a query task mid-flight and reads the pool; a case cannot cancel a query. #[test] fn cancelling_pending_work_drops_its_future_and_reservations() { diff --git a/crates/omnigraph/src/lib.rs b/crates/omnigraph/src/lib.rs index e19e72b7a..de713e171 100644 --- a/crates/omnigraph/src/lib.rs +++ b/crates/omnigraph/src/lib.rs @@ -16,6 +16,10 @@ pub mod changes; mod core_tests; pub mod db; #[cfg(feature = "dst")] +pub use omnigraph_core::dst_blocking; +#[cfg(not(feature = "dst"))] +pub(crate) use omnigraph_core::dst_blocking; +#[cfg(feature = "dst")] pub use omnigraph_core::dst_clock; #[cfg(not(feature = "dst"))] pub(crate) use omnigraph_core::dst_clock;