fix(dst): poll operator bodies on the seeded runtime under DST

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()
+}

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;

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<Installed<dyn Clock, ThreadLocal<dyn Clock>>>,
     _ids: Option<Installed<dyn IdSource, ThreadLocal<dyn IdSource>>>,
+    _blocking: Option<Installed<dyn Placement, ThreadLocal<dyn Placement>>>,
 }
 
 impl InstalledEnvironment {
@@ -79,6 +81,7 @@ impl InstalledEnvironment {
             process,
             _clock: None,
             _ids: None,
+            _blocking: None,
         }
     }
 }
@@ -150,6 +153,11 @@ pub fn run_universe<E: UniverseEnvironment, S: UniverseScenario<E::Resources>>(
                     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;

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() {

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;