From 5546158fff4628308ebd36529d67d23ab2499b7d Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Thu, 29 Oct 2026 17:38:37 +0100 Subject: [PATCH 1/3] Record the operator a thread is scheduling Adds `scheduling::current`, a thread-local holding the global identifier of the operator whose `schedule` call is running. Code without a handle to the operator, such as a global allocator attributing allocations, reads it with `current::operator()`. Hosts claim an operator for work done on its behalf outside `schedule` with `current::enter`. The thread-local is const-initialized and has no destructor, so reading it from an allocator neither allocates nor reenters. Co-Authored-By: Claude Opus 5.5 --- timely/src/progress/subgraph.rs | 5 ++- timely/src/scheduling/current.rs | 65 ++++++++++++++++++++++++++++++++ timely/src/scheduling/mod.rs | 1 + timely/src/worker.rs | 5 ++- timely/tests/current_operator.rs | 47 +++++++++++++++++++++++ 5 files changed, 121 insertions(+), 2 deletions(-) create mode 100644 timely/src/scheduling/current.rs create mode 100644 timely/tests/current_operator.rs diff --git a/timely/src/progress/subgraph.rs b/timely/src/progress/subgraph.rs index d863132fb..d8bdf243d 100644 --- a/timely/src/progress/subgraph.rs +++ b/timely/src/progress/subgraph.rs @@ -715,7 +715,10 @@ impl PerOperatorState { l.log(crate::logging::ScheduleEvent::start(self.id)); } - let incomplete = operator.schedule(); + let incomplete = { + let _current = crate::scheduling::current::enter(self.id); + operator.schedule() + }; // Perhaps log information about the stop of the schedule call. if let Some(l) = self.logging.as_mut() { diff --git a/timely/src/scheduling/current.rs b/timely/src/scheduling/current.rs new file mode 100644 index 000000000..3fb47a450 --- /dev/null +++ b/timely/src/scheduling/current.rs @@ -0,0 +1,65 @@ +//! The operator the current thread is scheduling. +//! +//! Timely records the identifier of each operator for the duration of its `schedule` call, the +//! same identifier that `ScheduleEvent` reports. Code that runs without a handle to the operator, +//! such as a global allocator attributing allocations, can read it with [`operator`]. +//! +//! Work that happens on behalf of an operator outside its `schedule` call, for example trace +//! maintenance driven by the host application, can claim that operator with [`enter`]. + +use std::cell::Cell; + +thread_local! { + // The identifier plus one, or zero outside any operator. Const-initialized and without + // `Drop`, so reading it never allocates or registers a destructor, which a global allocator + // relies on to avoid reentrancy. + static CURRENT: Cell = const { Cell::new(0) }; +} + +/// The identifier of the innermost operator this thread is scheduling, if any. +#[inline] +pub fn operator() -> Option { + CURRENT.with(|c| c.get().checked_sub(1)) +} + +/// Marks `id` as the operator this thread is scheduling until the returned guard drops. +/// +/// Guards nest: dropping one restores the operator that was current when it was created. Guards +/// must drop in reverse creation order, which holding them in scoped bindings ensures. +#[inline] +pub fn enter(id: usize) -> Guard { + Guard { previous: CURRENT.with(|c| c.replace(id + 1)) } +} + +/// Restores the previously current operator on drop. See [`enter`]. +#[must_use = "the operator is current only until the guard drops"] +pub struct Guard { + previous: usize, +} + +impl Drop for Guard { + #[inline] + fn drop(&mut self) { + CURRENT.with(|c| c.set(self.previous)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn guards_nest_and_restore() { + assert_eq!(operator(), None); + { + let _outer = enter(3); + assert_eq!(operator(), Some(3)); + { + let _inner = enter(0); + assert_eq!(operator(), Some(0)); + } + assert_eq!(operator(), Some(3)); + } + assert_eq!(operator(), None); + } +} diff --git a/timely/src/scheduling/mod.rs b/timely/src/scheduling/mod.rs index ec46ddbf0..cf3c351df 100644 --- a/timely/src/scheduling/mod.rs +++ b/timely/src/scheduling/mod.rs @@ -1,6 +1,7 @@ //! Types and traits to activate and schedule fibers. pub mod activate; +pub mod current; pub use self::activate::{Activations, Activator, ActivateOnDrop, SyncActivator}; diff --git a/timely/src/worker.rs b/timely/src/worker.rs index f24caefca..f3715cbf4 100644 --- a/timely/src/worker.rs +++ b/timely/src/worker.rs @@ -810,7 +810,10 @@ impl Wrapper { l.log(crate::logging::ScheduleEvent::start(self.identifier)); } - let incomplete = self.operate.as_mut().map(|op| op.schedule()).unwrap_or(false); + let incomplete = { + let _current = crate::scheduling::current::enter(self.identifier); + self.operate.as_mut().map(|op| op.schedule()).unwrap_or(false) + }; if !incomplete { self.operate = None; self.resources = None; diff --git a/timely/tests/current_operator.rs b/timely/tests/current_operator.rs new file mode 100644 index 000000000..2a88d4c63 --- /dev/null +++ b/timely/tests/current_operator.rs @@ -0,0 +1,47 @@ +use std::cell::Cell; +use std::rc::Rc; + +use timely::dataflow::channels::pact::Pipeline; +use timely::dataflow::operators::generic::Operator; +use timely::dataflow::operators::{Input, Probe}; +use timely::dataflow::InputHandle; +use timely::scheduling::current; + +/// An operator observes its own global identifier while scheduled, and nothing is current +/// between steps. +#[test] +fn operator_is_current_while_scheduled() { + timely::execute_directly(|worker| { + let mut input = InputHandle::new(); + let observed = Rc::new(Cell::new(0usize)); + let expected = Rc::new(Cell::new(None)); + let probe = worker.dataflow::(|scope| { + let observed = Rc::clone(&observed); + let expected = Rc::clone(&expected); + scope + .input_from(&mut input) + .container::>() + .unary(Pipeline, "Observe", move |_cap, info| { + expected.set(Some(info.global_id)); + move |records, output| { + assert_eq!(current::operator(), Some(info.global_id)); + observed.set(observed.get() + 1); + records.for_each(|time, data| output.session(&time).give_container(data)); + } + }) + .probe() + .0 + }); + + for round in 0..3 { + input.send(round); + input.advance_to(round + 1); + while probe.less_than(&(round + 1)) { + worker.step(); + assert_eq!(current::operator(), None); + } + } + assert!(expected.get().is_some()); + assert!(observed.get() > 0); + }); +} From 9d5433c2a749533727a356e76e776f2ab492c0f9 Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Thu, 29 Oct 2026 20:48:47 +0100 Subject: [PATCH 2/3] Make the operator guard !Send, rename to operator_id, document the value's meaning Co-Authored-By: Claude Opus 5.5 --- timely/src/scheduling/current.rs | 43 +++++++++++++++++++++++++------- timely/tests/current_operator.rs | 43 +++++++++++++++++++++----------- 2 files changed, 63 insertions(+), 23 deletions(-) diff --git a/timely/src/scheduling/current.rs b/timely/src/scheduling/current.rs index 3fb47a450..bfc103787 100644 --- a/timely/src/scheduling/current.rs +++ b/timely/src/scheduling/current.rs @@ -2,12 +2,28 @@ //! //! Timely records the identifier of each operator for the duration of its `schedule` call, the //! same identifier that `ScheduleEvent` reports. Code that runs without a handle to the operator, -//! such as a global allocator attributing allocations, can read it with [`operator`]. +//! such as a global allocator attributing allocations, can read it with [`operator_id`]. +//! +//! The identifier names an operator within one worker. Peer workers of a computation build the +//! same dataflows in the same order, so an identifier names the same operator on each of them, +//! but it means nothing beyond that computation. The value assumes one computation is stepped on +//! a thread at a time. A host that steps several workers of unrelated computations on one thread +//! has to tell them apart itself. +//! +//! The current operator is whoever is running, not whoever owns the data. An allocation made by +//! one operator and handed to another, such as a buffer pushed through `Pipeline`, is attributed +//! to the operator that made it. Work outside any operator, such as communication in +//! `Worker::step` and dataflow construction, has no current operator. //! //! Work that happens on behalf of an operator outside its `schedule` call, for example trace -//! maintenance driven by the host application, can claim that operator with [`enter`]. +//! maintenance driven by the host application, can claim that operator with [`enter`]. Hold the +//! guard in a scoped binding around that work, and never across a call to `Worker::step`. +//! +//! Reading and writing the value does not allocate where the platform provides native +//! thread-locals. Elsewhere, the first access on a thread may allocate. use std::cell::Cell; +use std::marker::PhantomData; thread_local! { // The identifier plus one, or zero outside any operator. Const-initialized and without @@ -18,7 +34,7 @@ thread_local! { /// The identifier of the innermost operator this thread is scheduling, if any. #[inline] -pub fn operator() -> Option { +pub fn operator_id() -> Option { CURRENT.with(|c| c.get().checked_sub(1)) } @@ -28,13 +44,22 @@ pub fn operator() -> Option { /// must drop in reverse creation order, which holding them in scoped bindings ensures. #[inline] pub fn enter(id: usize) -> Guard { - Guard { previous: CURRENT.with(|c| c.replace(id + 1)) } + Guard { previous: CURRENT.with(|c| c.replace(id + 1)), _not_send: PhantomData } } /// Restores the previously current operator on drop. See [`enter`]. +/// +/// A guard restores the thread it was created on, so it cannot move to another thread: +/// +/// ```compile_fail +/// fn assert_send(_: T) {} +/// assert_send(timely::scheduling::current::enter(0)); +/// ``` #[must_use = "the operator is current only until the guard drops"] pub struct Guard { previous: usize, + // Dropped on another thread, the guard would overwrite that thread's current operator. + _not_send: PhantomData<*const ()>, } impl Drop for Guard { @@ -50,16 +75,16 @@ mod tests { #[test] fn guards_nest_and_restore() { - assert_eq!(operator(), None); + assert_eq!(operator_id(), None); { let _outer = enter(3); - assert_eq!(operator(), Some(3)); + assert_eq!(operator_id(), Some(3)); { let _inner = enter(0); - assert_eq!(operator(), Some(0)); + assert_eq!(operator_id(), Some(0)); } - assert_eq!(operator(), Some(3)); + assert_eq!(operator_id(), Some(3)); } - assert_eq!(operator(), None); + assert_eq!(operator_id(), None); } } diff --git a/timely/tests/current_operator.rs b/timely/tests/current_operator.rs index 2a88d4c63..9d42f7f96 100644 --- a/timely/tests/current_operator.rs +++ b/timely/tests/current_operator.rs @@ -7,27 +7,43 @@ use timely::dataflow::operators::{Input, Probe}; use timely::dataflow::InputHandle; use timely::scheduling::current; -/// An operator observes its own global identifier while scheduled, and nothing is current -/// between steps. +/// An operator observes its own global identifier while scheduled, also inside a nested region, +/// and nothing is current between steps. #[test] fn operator_is_current_while_scheduled() { timely::execute_directly(|worker| { let mut input = InputHandle::new(); - let observed = Rc::new(Cell::new(0usize)); - let expected = Rc::new(Cell::new(None)); + let outer_runs = Rc::new(Cell::new(0usize)); + let inner_runs = Rc::new(Cell::new(0usize)); let probe = worker.dataflow::(|scope| { - let observed = Rc::clone(&observed); - let expected = Rc::clone(&expected); - scope + let outer_observed = Rc::clone(&outer_runs); + let inner_observed = Rc::clone(&inner_runs); + let stream = scope .input_from(&mut input) .container::>() - .unary(Pipeline, "Observe", move |_cap, info| { - expected.set(Some(info.global_id)); + .unary(Pipeline, "Outer", move |_cap, info| { move |records, output| { - assert_eq!(current::operator(), Some(info.global_id)); - observed.set(observed.get() + 1); + assert_eq!(current::operator_id(), Some(info.global_id)); + outer_observed.set(outer_observed.get() + 1); records.for_each(|time, data| output.session(&time).give_container(data)); } + }); + scope + .region(|inner| { + use timely::dataflow::operators::{Enter, Leave}; + stream + .enter(inner) + .unary(Pipeline, "Inner", move |_cap, info| { + move |records, output| { + // The innermost operator, not the region scheduling it. + assert_eq!(current::operator_id(), Some(info.global_id)); + inner_observed.set(inner_observed.get() + 1); + records.for_each(|time, data| { + output.session(&time).give_container(data) + }); + } + }) + .leave(scope) }) .probe() .0 @@ -38,10 +54,9 @@ fn operator_is_current_while_scheduled() { input.advance_to(round + 1); while probe.less_than(&(round + 1)) { worker.step(); - assert_eq!(current::operator(), None); + assert_eq!(current::operator_id(), None); } } - assert!(expected.get().is_some()); - assert!(observed.get() > 0); + assert!(outer_runs.get() > 0 && inner_runs.get() > 0, "both operators ran"); }); } From 50a7f84ca41ec6347d5e52454bcf0fc698e7c000 Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Fri, 30 Oct 2026 11:13:36 +0100 Subject: [PATCH 3/3] Drop the ownership caveat from the current operator docs Co-Authored-By: Claude Opus 5.5 --- timely/src/scheduling/current.rs | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/timely/src/scheduling/current.rs b/timely/src/scheduling/current.rs index bfc103787..9741acdf9 100644 --- a/timely/src/scheduling/current.rs +++ b/timely/src/scheduling/current.rs @@ -10,14 +10,11 @@ //! a thread at a time. A host that steps several workers of unrelated computations on one thread //! has to tell them apart itself. //! -//! The current operator is whoever is running, not whoever owns the data. An allocation made by -//! one operator and handed to another, such as a buffer pushed through `Pipeline`, is attributed -//! to the operator that made it. Work outside any operator, such as communication in -//! `Worker::step` and dataflow construction, has no current operator. -//! -//! Work that happens on behalf of an operator outside its `schedule` call, for example trace -//! maintenance driven by the host application, can claim that operator with [`enter`]. Hold the -//! guard in a scoped binding around that work, and never across a call to `Worker::step`. +//! Work outside any operator, such as communication in `Worker::step` and dataflow construction, +//! has no current operator. Work that happens on behalf of an operator outside its `schedule` +//! call, for example trace maintenance driven by the host application, can claim that operator +//! with [`enter`]. Hold the guard in a scoped binding around that work, and never across a call +//! to `Worker::step`. //! //! Reading and writing the value does not allocate where the platform provides native //! thread-locals. Elsewhere, the first access on a thread may allocate.