timely/scheduling/
current.rs

1//! The operator the current thread is scheduling.
2//!
3//! Timely records the identifier of each operator for the duration of its `schedule` call, the
4//! same identifier that `ScheduleEvent` reports. Code that runs without a handle to the operator,
5//! such as a global allocator attributing allocations, can read it with [`operator_id`].
6//!
7//! The identifier names an operator within one worker. Peer workers of a computation build the
8//! same dataflows in the same order, so an identifier names the same operator on each of them,
9//! but it means nothing beyond that computation. The value assumes one computation is stepped on
10//! a thread at a time. A host that steps several workers of unrelated computations on one thread
11//! has to tell them apart itself.
12//!
13//! Work outside any operator, such as communication in `Worker::step` and dataflow construction,
14//! has no current operator. Work that happens on behalf of an operator outside its `schedule`
15//! call, for example trace maintenance driven by the host application, can claim that operator
16//! with [`enter`]. Hold the guard in a scoped binding around that work, and never across a call
17//! to `Worker::step`.
18//!
19//! Reading and writing the value does not allocate where the platform provides native
20//! thread-locals. Elsewhere, the first access on a thread may allocate.
21
22use std::cell::Cell;
23use std::marker::PhantomData;
24
25thread_local! {
26    // The identifier plus one, or zero outside any operator. Const-initialized and without
27    // `Drop`, so reading it never allocates or registers a destructor, which a global allocator
28    // relies on to avoid reentrancy.
29    static CURRENT: Cell<usize> = const { Cell::new(0) };
30}
31
32/// The identifier of the innermost operator this thread is scheduling, if any.
33#[inline]
34pub fn operator_id() -> Option<usize> {
35    CURRENT.with(|c| c.get().checked_sub(1))
36}
37
38/// Marks `id` as the operator this thread is scheduling until the returned guard drops.
39///
40/// Guards nest: dropping one restores the operator that was current when it was created. Guards
41/// must drop in reverse creation order, which holding them in scoped bindings ensures.
42#[inline]
43pub fn enter(id: usize) -> Guard {
44    Guard { previous: CURRENT.with(|c| c.replace(id + 1)), _not_send: PhantomData }
45}
46
47/// Restores the previously current operator on drop. See [`enter`].
48///
49/// A guard restores the thread it was created on, so it cannot move to another thread:
50///
51/// ```compile_fail
52/// fn assert_send<T: Send>(_: T) {}
53/// assert_send(timely::scheduling::current::enter(0));
54/// ```
55#[must_use = "the operator is current only until the guard drops"]
56pub struct Guard {
57    previous: usize,
58    // Dropped on another thread, the guard would overwrite that thread's current operator.
59    _not_send: PhantomData<*const ()>,
60}
61
62impl Drop for Guard {
63    #[inline]
64    fn drop(&mut self) {
65        CURRENT.with(|c| c.set(self.previous));
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn guards_nest_and_restore() {
75        assert_eq!(operator_id(), None);
76        {
77            let _outer = enter(3);
78            assert_eq!(operator_id(), Some(3));
79            {
80                let _inner = enter(0);
81                assert_eq!(operator_id(), Some(0));
82            }
83            assert_eq!(operator_id(), Some(3));
84        }
85        assert_eq!(operator_id(), None);
86    }
87}
timely/scheduling/current.rs

timely/scheduling/
current.rs
