asdfkkkkkkkkk

Cargo.toml

@@ -48,4 +48,4 @@ cfg-if = "1.0.0"
 async-std = "1.13.0"
 tracing-test = "0.2.5"
 tracing-tracy = "0.11.4"
-distaff = {path = "distaff"}
+distaff = {path = "distaff", features = ["tracing"]}

benches/join.rs

@@ -56,7 +56,7 @@ mod tree {
     }
 }
 
-const TREE_SIZE: usize = 16;
+const TREE_SIZE: usize = 8;
 
 #[bench]
 fn join_melange(b: &mut Bencher) {
@@ -191,11 +191,7 @@ fn join_distaff(b: &mut Bencher) {
     let pool = ThreadPool::new();
     let tree = tree::Tree::new(TREE_SIZE, 1u32);
 
-    fn sum<'scope, 'env>(
-        tree: &tree::Tree<u32>,
-        node: usize,
-        scope: &'scope Scope<'scope, 'env>,
-    ) -> u32 {
+    fn sum<'scope, 'env>(tree: &tree::Tree<u32>, node: usize, scope: Scope<'scope, 'env>) -> u32 {
         let node = tree.get(node);
         let (l, r) = scope.join(
             |s| node.left.map(|node| sum(tree, node, s)).unwrap_or_default(),
@@ -210,10 +206,12 @@ fn join_distaff(b: &mut Bencher) {
     }
 
     b.iter(move || {
-        pool.scope(|s| {
+        let sum = pool.scope(|s| {
             let sum = sum(&tree, tree.root().unwrap(), s);
-            // eprintln!("{sum}");
             assert_ne!(sum, 0);
+            sum
         });
+        std::hint::black_box(sum);
     });
+    eprintln!("Done with distaff join");
 }

distaff/Cargo.toml

@@ -3,13 +3,22 @@ name = "distaff"
 version = "0.1.0"
 edition = "2024"
 
+[profile.bench]
+debug = true
+
+[profile.release]
+debug = true
+
 [features]
 default = []
+tracing = ["dep:tracing"]
 std = []
+metrics = []
 
 [dependencies]
 parking_lot = {version = "0.12.3"}
-tracing = "0.1.40"
+atomic-wait = "1.1.0"
+tracing = {version = "0.1", optional = true}
 parking_lot_core = "0.9.10"
 crossbeam-utils = "0.8.21"
 either = "1.15.0"
@@ -17,6 +26,14 @@ either = "1.15.0"
 async-task = "4.7.1"
 
 [dev-dependencies]
-tracing-test = "0.2.5"
-tracing-tracy = "0.11.4"
-futures = "0.3"
+tracing-test = {version = "0.2"}
+tracing-tracy = {version = "0.11"}
+futures = "0.3"
+
+divan = "0.1.14"
+rayon = "1.10.0"
+chili = {path = "../../chili"}
+
+[[bench]]
+name = "overhead"
+harness = false

distaff/benches/overhead.rs

@@ -0,0 +1,119 @@
+use distaff::{Scope, ThreadPool};
+use divan::Bencher;
+
+struct Node {
+    val: u64,
+    left: Option<Box<Node>>,
+    right: Option<Box<Node>>,
+}
+
+impl Node {
+    pub fn tree(layers: usize) -> Self {
+        Self {
+            val: 1,
+            left: (layers != 1).then(|| Box::new(Self::tree(layers - 1))),
+            right: (layers != 1).then(|| Box::new(Self::tree(layers - 1))),
+        }
+    }
+}
+
+const LAYERS: &[usize] = &[10, 24];
+fn nodes() -> impl Iterator<Item = (usize, usize)> {
+    LAYERS.iter().map(|&l| (l, (1 << l) - 1))
+}
+
+#[divan::bench(args = nodes())]
+fn no_overhead(bencher: Bencher, nodes: (usize, usize)) {
+    fn join_no_overhead<A, B, RA, RB>(scope: Scope<'_, '_>, a: A, b: B) -> (RA, RB)
+    where
+        A: FnOnce(Scope<'_, '_>) -> RA + Send,
+        B: FnOnce(Scope<'_, '_>) -> RB + Send,
+        RA: Send,
+        RB: Send,
+    {
+        (a(scope), b(scope))
+    }
+
+    #[inline]
+    fn sum(node: &Node, scope: Scope<'_, '_>) -> u64 {
+        let (left, right) = join_no_overhead(
+            scope,
+            |s| node.left.as_deref().map(|n| sum(n, s)).unwrap_or_default(),
+            |s| node.right.as_deref().map(|n| sum(n, s)).unwrap_or_default(),
+        );
+
+        node.val + left + right
+    }
+
+    let tree = Node::tree(nodes.0);
+    let pool = ThreadPool::global();
+
+    bencher.bench_local(move || {
+        pool.scope(|scope| {
+            assert_eq!(sum(&tree, scope), nodes.1 as u64);
+        });
+    });
+}
+
+#[divan::bench(args = nodes())]
+fn distaff_overhead(bencher: Bencher, nodes: (usize, usize)) {
+    fn sum<'scope, 'env>(node: &Node, scope: Scope<'scope, 'env>) -> u64 {
+        let (left, right) = scope.join(
+            |s| node.left.as_deref().map(|n| sum(n, s)).unwrap_or_default(),
+            |s| node.right.as_deref().map(|n| sum(n, s)).unwrap_or_default(),
+        );
+
+        node.val + left + right
+    }
+
+    let tree = Node::tree(nodes.0);
+    let pool = ThreadPool::global();
+
+    bencher.bench_local(move || {
+        pool.scope(|scope| {
+            assert_eq!(sum(&tree, scope), nodes.1 as u64);
+        });
+    });
+}
+
+#[divan::bench(args = nodes())]
+fn rayon_overhead(bencher: Bencher, nodes: (usize, usize)) {
+    fn sum(node: &Node) -> u64 {
+        let (left, right) = rayon::join(
+            || node.left.as_deref().map(sum).unwrap_or_default(),
+            || node.right.as_deref().map(sum).unwrap_or_default(),
+        );
+
+        node.val + left + right
+    }
+
+    let tree = Node::tree(nodes.0);
+
+    bencher.bench_local(move || {
+        assert_eq!(sum(&tree), nodes.1 as u64);
+    });
+}
+
+#[divan::bench(args = nodes())]
+fn chili_overhead(bencher: Bencher, nodes: (usize, usize)) {
+    use chili::Scope;
+    fn sum(node: &Node, scope: &mut Scope<'_>) -> u64 {
+        let (left, right) = scope.join(
+            |s| node.left.as_deref().map(|n| sum(n, s)).unwrap_or_default(),
+            |s| node.right.as_deref().map(|n| sum(n, s)).unwrap_or_default(),
+        );
+
+        node.val + left + right
+    }
+
+    let tree = Node::tree(nodes.0);
+    let mut scope = Scope::global();
+
+    bencher.bench_local(move || {
+        assert_eq!(sum(&tree, &mut scope), nodes.1 as u64);
+    });
+}
+
+fn main() {
+    divan::main();
+}

distaff/src/channel.rs

@@ -0,0 +1,230 @@
+// This file is taken from [`chili`]
+
+use std::{
+    cell::UnsafeCell,
+    ptr::NonNull,
+    sync::{
+        Arc,
+        atomic::{AtomicU8, AtomicU32, Ordering},
+    },
+    thread,
+};
+
+enum State {
+    Pending,
+    Waiting,
+    Ready,
+    Taken,
+}
+
+// taken from `std`
+#[derive(Debug)]
+#[repr(transparent)]
+pub struct Parker {
+    mutex: AtomicU32,
+}
+
+impl Parker {
+    const PARKED: u32 = u32::MAX;
+    const EMPTY: u32 = 0;
+    const NOTIFIED: u32 = 1;
+    pub fn new() -> Self {
+        Self {
+            mutex: AtomicU32::new(Self::EMPTY),
+        }
+    }
+
+    pub fn is_parked(&self) -> bool {
+        self.mutex.load(Ordering::Acquire) == Self::PARKED
+    }
+
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all, fields(this = self as *const Self as usize)))]
+    pub fn park(&self) {
+        if self.mutex.fetch_sub(1, Ordering::Acquire) == Self::NOTIFIED {
+            // The thread was notified, so we can return immediately.
+            return;
+        }
+
+        loop {
+            atomic_wait::wait(&self.mutex, Self::PARKED);
+
+            // We check whether we were notified or woke up spuriously with
+            // acquire ordering in order to make-visible any writes made by the
+            // thread that notified us.
+            if self.mutex.swap(Self::EMPTY, Ordering::Acquire) == Self::NOTIFIED {
+                // The thread was notified, so we can return immediately.
+                return;
+            } else {
+                // spurious wakeup, so we need to re-park.
+                continue;
+            }
+        }
+    }
+
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all, fields(this = self as *const Self as usize)))]
+    pub fn unpark(&self) {
+        // write with Release ordering to ensure that any writes made by this
+        // thread are made-available to the unparked thread.
+        if self.mutex.swap(Self::NOTIFIED, Ordering::Release) == Self::PARKED {
+            // The thread was parked, so we need to notify it.
+            atomic_wait::wake_one(&self.mutex);
+        } else {
+            // The thread was not parked, so we don't need to do anything.
+        }
+    }
+}
+
+#[derive(Debug)]
+#[repr(C)]
+struct Channel<T = ()> {
+    state: AtomicU8,
+    /// Can only be written only by the `Receiver` and read by the `Sender` if
+    /// `state` is `State::Waiting`.
+    waiting_thread: NonNull<Parker>,
+    /// Can only be written only by the `Sender` and read by the `Receiver` if
+    /// `state` is `State::Ready`.
+    val: UnsafeCell<Option<Box<thread::Result<T>>>>,
+}
+
+impl<T> Channel<T> {
+    fn new(waiting_thread: NonNull<Parker>) -> Self {
+        Self {
+            state: AtomicU8::new(State::Pending as u8),
+            waiting_thread,
+            val: UnsafeCell::new(None),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct Receiver<T = ()>(Arc<Channel<T>>);
+
+impl<T: Send> Receiver<T> {
+    pub fn is_empty(&self) -> bool {
+        self.0.state.load(Ordering::Acquire) != State::Ready as u8
+    }
+
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub fn wait(&self) {
+        loop {
+            match self.0.state.compare_exchange(
+                State::Pending as u8,
+                State::Waiting as u8,
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            ) {
+                Ok(_) => {
+                    // SAFETY:
+                    // The `waiting_thread` is set to the current thread's parker
+                    // before we park it.
+                    unsafe {
+                        let thread = self.0.waiting_thread.as_ref();
+                        thread.park();
+                    }
+
+                    // we might have been woken up because of a shared job.
+                    // In that case, we need to check the state again.
+                    if self
+                        .0
+                        .state
+                        .compare_exchange(
+                            State::Waiting as u8,
+                            State::Pending as u8,
+                            Ordering::AcqRel,
+                            Ordering::Acquire,
+                        )
+                        .is_ok()
+                    {
+                        continue;
+                    } else {
+                        // The state was changed to `State::Ready` by the `Sender`, so we can return.
+                        break;
+                    }
+                }
+                Err(state) if state == State::Ready as u8 => {
+                    // The channel is ready, so we can return immediately.
+                    return;
+                }
+                _ => {
+                    panic!("Receiver is already waiting or consumed.");
+                }
+            }
+        }
+    }
+
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub fn poll(&self) -> Option<thread::Result<T>> {
+        if self
+            .0
+            .state
+            .compare_exchange(
+                State::Ready as u8,
+                State::Taken as u8,
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            )
+            .is_ok()
+        {
+            unsafe { Some(self.take()) }
+        } else {
+            None
+        }
+    }
+
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub fn recv(self) -> thread::Result<T> {
+        self.wait();
+
+        // SAFETY:
+        // To arrive here, either `state` is `State::Ready` or the above
+        // `compare_exchange` succeeded, the thread was parked and then
+        // unparked by the `Sender` *after* the `state` was set to
+        // `State::Ready`.
+        //
+        // In either case, this thread now has unique access to `val`.
+        unsafe { self.take() }
+    }
+
+    unsafe fn take(&self) -> thread::Result<T> {
+        assert_eq!(
+            self.0.state.swap(State::Taken as u8, Ordering::Acquire),
+            State::Ready as u8
+        );
+
+        let result = unsafe { (*self.0.val.get()).take().map(|b| *b).unwrap() };
+
+        result
+    }
+}
+
+#[derive(Debug)]
+#[repr(transparent)]
+pub struct Sender<T = ()>(Arc<Channel<T>>);
+
+impl<T: Send> Sender<T> {
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub fn send(self, val: thread::Result<T>) {
+        // SAFETY:
+        // Only this thread can write to `val` and none can read it
+        // yet.
+        unsafe {
+            *self.0.val.get() = Some(Box::new(val));
+        }
+
+        if self.0.state.swap(State::Ready as u8, Ordering::AcqRel) == State::Waiting as u8 {
+            // SAFETY:
+            // A `Receiver` already wrote its thread to `waiting_thread`
+            // *before* setting the `state` to `State::Waiting`.
+            unsafe {
+                let thread = self.0.waiting_thread.as_ref();
+                thread.unpark();
+            }
+        }
+    }
+}
+
+pub fn channel<T: Send>(thread: NonNull<Parker>) -> (Sender<T>, Receiver<T>) {
+    let channel = Arc::new(Channel::new(thread));
+
+    (Sender(channel.clone()), Receiver(channel))
+}

distaff/src/context.rs

@@ -1,5 +1,5 @@
 use std::{
-    ptr::{self, NonNull},
+    ptr::NonNull,
     sync::{
         Arc, OnceLock,
         atomic::{AtomicBool, Ordering},
@@ -9,38 +9,16 @@ use std::{
 use alloc::collections::BTreeMap;
 
 use async_task::Runnable;
-use crossbeam_utils::CachePadded;
 use parking_lot::{Condvar, Mutex};
 
 use crate::{
+    channel::{Parker, Sender},
     heartbeat::HeartbeatList,
-    job::{HeapJob, JobSender, QueuedJob as Job, StackJob},
-    latch::{AsCoreLatch, MutexLatch, NopLatch, WorkerLatch},
+    job::{HeapJob, Job2 as Job, SharedJob, StackJob},
+    util::DropGuard,
     workerthread::{HeartbeatThread, WorkerThread},
 };
 
-pub struct Heartbeat {
-    pub latch: MutexLatch,
-}
-
-impl Heartbeat {
-    pub fn new() -> NonNull<CachePadded<Self>> {
-        let ptr = Box::new(CachePadded::new(Self {
-            latch: MutexLatch::new(),
-        }));
-
-        Box::into_non_null(ptr)
-    }
-
-    pub fn is_pending(&self) -> bool {
-        self.latch.as_core_latch().poll_heartbeat()
-    }
-
-    pub fn is_sleeping(&self) -> bool {
-        self.latch.as_core_latch().is_sleeping()
-    }
-}
-
 pub struct Context {
     shared: Mutex<Shared>,
     pub shared_job: Condvar,
@@ -49,14 +27,14 @@ pub struct Context {
 }
 
 pub(crate) struct Shared {
-    pub jobs: BTreeMap<usize, NonNull<Job>>,
-    injected_jobs: Vec<NonNull<Job>>,
+    pub jobs: BTreeMap<usize, SharedJob>,
+    injected_jobs: Vec<SharedJob>,
 }
 
 unsafe impl Send for Shared {}
 
 impl Shared {
-    pub fn pop_job(&mut self) -> Option<NonNull<Job>> {
+    pub fn pop_job(&mut self) -> Option<SharedJob> {
         // this is unlikely, so make the function cold?
         // TODO: profile this
         if !self.injected_jobs.is_empty() {
@@ -68,18 +46,20 @@ impl Shared {
     }
 
     #[cold]
-    unsafe fn pop_injected_job(&mut self) -> NonNull<Job> {
+    unsafe fn pop_injected_job(&mut self) -> SharedJob {
         self.injected_jobs.pop().unwrap()
     }
 }
 
 impl Context {
-    #[inline]
     pub(crate) fn shared(&self) -> parking_lot::MutexGuard<'_, Shared> {
         self.shared.lock()
     }
 
     pub fn new_with_threads(num_threads: usize) -> Arc<Self> {
+        #[cfg(feature = "tracing")]
+        tracing::trace!("Creating context with {} threads", num_threads);
+
         let this = Arc::new(Self {
             shared: Mutex::new(Shared {
                 jobs: BTreeMap::new(),
@@ -90,8 +70,6 @@ impl Context {
             heartbeats: HeartbeatList::new(),
         });
 
-        tracing::trace!("Creating thread pool with {} threads", num_threads);
-
         // Create a barrier to synchronize the worker threads and the heartbeat thread
         let barrier = Arc::new(std::sync::Barrier::new(num_threads + 2));
 
@@ -104,8 +82,7 @@ impl Context {
                 .spawn(move || {
                     let worker = Box::new(WorkerThread::new_in(ctx));
 
-                    barrier.wait();
-                    worker.run();
+                    worker.run(barrier);
                 })
                 .expect("Failed to spawn worker thread");
         }
@@ -117,8 +94,7 @@ impl Context {
             std::thread::Builder::new()
                 .name("heartbeat-thread".to_string())
                 .spawn(move || {
-                    barrier.wait();
-                    HeartbeatThread::new(ctx).run();
+                    HeartbeatThread::new(ctx).run(barrier);
                 })
                 .expect("Failed to spawn heartbeat thread");
         }
@@ -147,7 +123,7 @@ impl Context {
         GLOBAL_CONTEXT.get_or_init(|| Self::new())
     }
 
-    pub fn inject_job(&self, job: NonNull<Job>) {
+    pub fn inject_job(&self, job: SharedJob) {
         let mut shared = self.shared.lock();
         shared.injected_jobs.push(job);
 
@@ -157,17 +133,20 @@ impl Context {
         }
     }
 
-    // caller should hold the shared lock while calling this
+    /// caller should hold the shared lock while calling this
     pub unsafe fn notify_job_shared(&self) {
-        if let Some((i, sender)) = self
-            .heartbeats
-            .inner()
+        let heartbeats = self.heartbeats.inner();
+        if let Some((i, sender)) = heartbeats
             .iter()
             .find(|(_, heartbeat)| heartbeat.is_waiting())
+            .or_else(|| heartbeats.iter().next())
         {
+            _ = i;
+            #[cfg(feature = "tracing")]
             tracing::trace!("Notifying worker thread {} about job sharing", i);
             sender.wake();
         } else {
+            #[cfg(feature = "tracing")]
             tracing::warn!("No worker found to notify about job sharing");
         }
     }
@@ -181,21 +160,13 @@ impl Context {
         // current thread is not in the same context, create a job and inject it into the other thread's context, then wait while working on our jobs.
 
         // SAFETY: we are waiting on this latch in this thread.
-        let job = StackJob::new(
-            move || {
-                let worker = WorkerThread::current_ref()
-                    .expect("WorkerThread::run_in_worker called outside of worker thread");
+        let job = StackJob::new(move |worker: &WorkerThread| f(worker));
 
-                f(worker)
-            },
-            NopLatch,
-        );
+        let job = Job::from_stackjob(&job);
 
-        let job = Job::from_stackjob(&job, worker.heartbeat.raw_latch());
+        self.inject_job(job.share(Some(worker.heartbeat.parker())));
 
-        self.inject_job(Into::into(&job));
-
-        let t = worker.wait_until_queued_job(&job).unwrap();
+        let t = worker.wait_until_shared_job(&job).unwrap();
 
         crate::util::unwrap_or_panic(t)
     }
@@ -207,47 +178,47 @@ impl Context {
         T: Send,
     {
         // current thread isn't a worker thread, create job and inject into context
-        let latch = WorkerLatch::new();
+        let parker = Parker::new();
 
-        let job = StackJob::new(
-            move || {
-                let worker = WorkerThread::current_ref()
-                    .expect("WorkerThread::run_in_worker called outside of worker thread");
+        let job = StackJob::new(move |worker: &WorkerThread| f(worker));
 
-                f(worker)
-            },
-            NopLatch,
-        );
+        let job = Job::from_stackjob(&job);
 
-        let job = Job::from_stackjob(&job, &raw const latch);
+        self.inject_job(job.share(Some(&parker)));
 
-        self.inject_job(Into::into(&job));
-        let recv = unsafe { job.as_receiver::<T>() };
+        let recv = job.take_receiver().unwrap();
 
-        crate::util::unwrap_or_panic(latch.wait_until(|| recv.poll()))
+        crate::util::unwrap_or_panic(recv.recv())
     }
 
     /// Run closure in this context.
-    #[tracing::instrument(level = "trace", skip_all)]
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
     pub fn run_in_worker<T, F>(self: &Arc<Self>, f: F) -> T
     where
         T: Send,
         F: FnOnce(&WorkerThread) -> T + Send,
     {
+        let _guard = DropGuard::new(|| {
+            #[cfg(feature = "tracing")]
+            tracing::trace!("run_in_worker: finished");
+        });
         match WorkerThread::current_ref() {
             Some(worker) => {
                 // check if worker is in the same context
                 if Arc::ptr_eq(&worker.context, self) {
+                    #[cfg(feature = "tracing")]
                     tracing::trace!("run_in_worker: current thread");
                     f(worker)
                 } else {
                     // current thread is a worker for a different context
+                    #[cfg(feature = "tracing")]
                     tracing::trace!("run_in_worker: cross-context");
                     self.run_in_worker_cross(worker, f)
                 }
             }
             None => {
                 // current thread is not a worker for any context
+                #[cfg(feature = "tracing")]
                 tracing::trace!("run_in_worker: inject into context");
                 self.run_in_worker_cold(f)
             }
@@ -260,9 +231,10 @@ impl Context {
     where
         F: FnOnce() + Send + 'static,
     {
-        let job = Job::from_heapjob(Box::new(HeapJob::new(f)), ptr::null());
+        let job = Job::from_heapjob(HeapJob::new(|_: &WorkerThread| f()));
+        #[cfg(feature = "tracing")]
         tracing::trace!("Context::spawn: spawning job: {:?}", job);
-        self.inject_job(job);
+        self.inject_job(job.share(None));
     }
 
     pub fn spawn_future<T, F>(self: &Arc<Self>, future: F) -> async_task::Task<T>
@@ -272,24 +244,16 @@ impl Context {
     {
         let schedule = move |runnable: Runnable| {
             #[align(8)]
-            unsafe fn harness<T>(this: *const (), job: *const JobSender, _: *const WorkerLatch) {
+            unsafe fn harness<T>(_: &WorkerThread, this: NonNull<()>, _: Option<Sender>) {
                 unsafe {
-                    let runnable =
-                        Runnable::<()>::from_raw(NonNull::new_unchecked(this.cast_mut()));
+                    let runnable = Runnable::<()>::from_raw(this);
                     runnable.run();
-
-                    // SAFETY: job was turned into raw
-                    drop(Box::from_raw(job.cast::<JobSender<T>>().cast_mut()));
                 }
             }
 
-            let job = Box::into_non_null(Box::new(Job::from_harness(
-                harness::<T>,
-                runnable.into_raw(),
-                ptr::null(),
-            )));
+            let job = Job::<T>::from_harness(harness::<T>, runnable.into_raw());
 
-            self.inject_job(job);
+            self.inject_job(job.share(None));
         };
 
         let (runnable, task) = unsafe { async_task::spawn_unchecked(future, schedule) };
@@ -324,12 +288,10 @@ where
 mod tests {
     use std::sync::atomic::AtomicU8;
 
-    use tracing_test::traced_test;
-
     use super::*;
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn run_in_worker() {
         let ctx = Context::global_context().clone();
         let result = ctx.run_in_worker(|_| 42);
@@ -337,7 +299,7 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn context_spawn_future() {
         let ctx = Context::global_context().clone();
         let task = ctx.spawn_future(async { 42 });
@@ -348,7 +310,7 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn context_spawn_async() {
         let ctx = Context::global_context().clone();
         let task = ctx.spawn_async(|| async { 42 });
@@ -359,7 +321,7 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn context_spawn() {
         let ctx = Context::global_context().clone();
         let counter = Arc::new(AtomicU8::new(0));
@@ -379,27 +341,25 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn inject_job_and_wake_worker() {
         let ctx = Context::new_with_threads(1);
         let counter = Arc::new(AtomicU8::new(0));
 
-        let waker = WorkerLatch::new();
+        let parker = Parker::new();
 
-        let job = StackJob::new(
-            {
-                let counter = counter.clone();
-                move || {
-                    tracing::info!("Job running");
-                    counter.fetch_add(1, Ordering::SeqCst);
+        let job = StackJob::new({
+            let counter = counter.clone();
+            move |_: &WorkerThread| {
+                #[cfg(feature = "tracing")]
+                tracing::info!("Job running");
+                counter.fetch_add(1, Ordering::SeqCst);
 
-                    42
-                }
-            },
-            NopLatch,
-        );
+                42
+            }
+        });
 
-        let job = Job::from_stackjob(&job, &raw const waker);
+        let job = Job::from_stackjob(&job);
 
         // wait for the worker to sleep
         std::thread::sleep(std::time::Duration::from_millis(100));
@@ -412,11 +372,11 @@ mod tests {
                 assert!(heartbeat.is_waiting());
             });
 
-        ctx.inject_job(Into::into(&job));
+        ctx.inject_job(job.share(Some(&parker)));
 
         // Wait for the job to be executed
-        let recv = unsafe { job.as_receiver::<i32>() };
-        let result = waker.wait_until(|| recv.poll());
+        let recv = job.take_receiver().unwrap();
+        let result = recv.recv();
         let result = crate::util::unwrap_or_panic(result);
         assert_eq!(result, 42);
         assert_eq!(counter.load(Ordering::SeqCst), 1);

distaff/src/heartbeat.rs

@@ -12,7 +12,7 @@ use std::{
 
 use parking_lot::Mutex;
 
-use crate::latch::WorkerLatch;
+use crate::channel::Parker;
 
 #[derive(Debug, Clone)]
 pub struct HeartbeatList {
@@ -27,6 +27,8 @@ impl HeartbeatList {
     }
 
     pub fn notify_nth(&self, n: usize) {
+        #[cfg(feature = "tracing")]
+        tracing::trace!("notifying worker-{}", n);
         self.inner.lock().notify_nth(n);
     }
 
@@ -125,13 +127,13 @@ impl Drop for OwnedHeartbeatReceiver {
 
 #[derive(Debug)]
 pub struct Heartbeat {
-    ptr: NonNull<(AtomicBool, WorkerLatch)>,
+    ptr: NonNull<(AtomicBool, Parker)>,
     i: u64,
 }
 
 #[derive(Debug)]
 pub struct HeartbeatReceiver {
-    ptr: NonNull<(AtomicBool, WorkerLatch)>,
+    ptr: NonNull<(AtomicBool, Parker)>,
     i: u64,
 }
 
@@ -149,7 +151,7 @@ impl Drop for Heartbeat {
 
 #[derive(Debug)]
 pub struct HeartbeatSender {
-    ptr: NonNull<(AtomicBool, WorkerLatch)>,
+    ptr: NonNull<(AtomicBool, Parker)>,
     pub last_heartbeat: Instant,
 }
 
@@ -161,7 +163,7 @@ impl Heartbeat {
         // `AtomicBool` is `Sync` and `Send`, so it can be safely shared between threads.
         let ptr = NonNull::new(Box::into_raw(Box::new((
             AtomicBool::new(true),
-            WorkerLatch::new(),
+            Parker::new(),
         ))))
         .unwrap();
         Self { ptr, i }
@@ -200,15 +202,7 @@ impl HeartbeatReceiver {
         }
     }
 
-    pub fn wait(&self) {
-        unsafe { self.ptr.as_ref().1.wait() };
-    }
-
-    pub fn raw_latch(&self) -> *const WorkerLatch {
-        unsafe { &raw const self.ptr.as_ref().1 }
-    }
-
-    pub fn latch(&self) -> &WorkerLatch {
+    pub fn parker(&self) -> &Parker {
         unsafe { &self.ptr.as_ref().1 }
     }
 
@@ -226,13 +220,13 @@ impl HeartbeatSender {
         // SAFETY:
         // `AtomicBool` is `Sync` and `Send`, so it can be safely shared between threads.
         unsafe { self.ptr.as_ref().0.store(true, Ordering::Relaxed) };
-        self.last_heartbeat = Instant::now();
+        // self.last_heartbeat = Instant::now();
     }
 
     pub fn is_waiting(&self) -> bool {
-        unsafe { self.ptr.as_ref().1.is_waiting() }
+        unsafe { self.ptr.as_ref().1.is_parked() }
     }
     pub fn wake(&self) {
-        unsafe { self.ptr.as_ref().1.wake() };
+        unsafe { self.ptr.as_ref().1.unpark() };
     }
 }

distaff/src/join.rs

@@ -1,41 +1,48 @@
+#[cfg(feature = "metrics")]
+use std::sync::atomic::Ordering;
+
 use std::{hint::cold_path, sync::Arc};
 
 use crate::{
     context::Context,
-    job::{QueuedJob as Job, StackJob},
-    latch::NopLatch,
+    job::{
+        Job2 as Job, StackJob,
+        traits::{InlineJob, IntoJob},
+    },
     workerthread::WorkerThread,
 };
 
 impl WorkerThread {
-    #[inline]
-    #[tracing::instrument(level = "trace", skip_all)]
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
     fn join_seq<A, B, RA, RB>(&self, a: A, b: B) -> (RA, RB)
     where
-        A: FnOnce() -> RA,
-        B: FnOnce() -> RB,
+        A: FnOnce(&WorkerThread) -> RA,
+        B: FnOnce(&WorkerThread) -> RB,
     {
-        let span = tracing::trace_span!("join_seq");
-        let _guard = span.enter();
-
-        let rb = b();
-        let ra = a();
+        let rb = b(self);
+        let ra = a(self);
 
         (ra, rb)
     }
 
+    pub(crate) fn join_heartbeat_every<A, B, RA, RB>(&self, a: A, b: B) -> (RA, RB)
+    where
+        A: FnOnce(&WorkerThread) -> RA + Send,
+        B: FnOnce(&WorkerThread) -> RB,
+        RA: Send,
+    {
+        // self.join_heartbeat_every_inner::<A, B, RA, RB, 2>(a, b)
+        self.join_heartbeat(a, b)
+    }
+
     /// This function must be called from a worker thread.
-    #[inline]
-    #[tracing::instrument(level = "trace", skip_all)]
-    pub(crate) fn join_heartbeat_every<A, B, RA, RB, const TIMES: usize>(
-        &self,
-        a: A,
-        b: B,
-    ) -> (RA, RB)
+    #[allow(dead_code)]
+    #[inline(always)]
+    fn join_heartbeat_every_inner<A, B, RA, RB, const TIMES: usize>(&self, a: A, b: B) -> (RA, RB)
     where
         RA: Send,
-        A: FnOnce() -> RA + Send,
-        B: FnOnce() -> RB,
+        A: FnOnce(&WorkerThread) -> RA + Send,
+        B: FnOnce(&WorkerThread) -> RB,
     {
         // SAFETY: each worker is only ever used by one thread, so this is safe.
         let count = self.join_count.get();
@@ -48,10 +55,6 @@ impl WorkerThread {
         // SAFETY: this function runs in a worker thread, so we can access the queue safely.
         if count == 0 || queue_len < 3 {
             cold_path();
-            tracing::trace!(
-                queue_len = queue_len,
-                "join_heartbeat_every: using heartbeat join",
-            );
             self.join_heartbeat(a, b)
         } else {
             self.join_seq(a, b)
@@ -59,53 +62,162 @@ impl WorkerThread {
     }
 
     /// This function must be called from a worker thread.
-    #[inline]
-    #[tracing::instrument(level = "trace", skip_all)]
-    fn join_heartbeat<A, B, RA, RB>(&self, a: A, b: B) -> (RA, RB)
+    #[allow(dead_code)]
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub(crate) fn join_heartbeat2_every<A, B, RA, RB, const TIMES: usize>(
+        &self,
+        a: A,
+        b: B,
+    ) -> (RA, RB)
     where
         RA: Send,
-        A: FnOnce() -> RA + Send,
-        B: FnOnce() -> RB,
+        A: InlineJob<RA> + Copy,
+        B: FnOnce(&WorkerThread) -> RB,
+    {
+        // SAFETY: each worker is only ever used by one thread, so this is safe.
+        let count = self.join_count.get();
+        let queue_len = unsafe { self.queue.as_ref_unchecked().len() };
+        self.join_count.set(count.wrapping_add(1) % TIMES as u8);
+
+        // TODO: add counter to job queue, check for low job count to decide whether to use heartbeat or seq.
+        // see: chili
+
+        // SAFETY: this function runs in a worker thread, so we can access the queue safely.
+        if count == 0 || queue_len < 3 {
+            cold_path();
+            self.join_heartbeat2(a, b)
+        } else {
+            (a.run_inline(self), b(self))
+        }
+    }
+
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub(crate) fn join_heartbeat2<RA, A, B, RB>(&self, a: A, b: B) -> (RA, RB)
+    where
+        B: FnOnce(&WorkerThread) -> RB,
+        A: InlineJob<RA> + Copy,
+        RA: Send,
     {
         use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
 
-        let a = StackJob::new(a, NopLatch);
-        let job = Job::from_stackjob(&a, self.heartbeat.raw_latch());
+        #[cfg(feature = "metrics")]
+        self.metrics.num_joins.fetch_add(1, Ordering::Relaxed);
+
+        let job = a.into_job();
 
         self.push_back(&job);
 
         self.tick();
 
-        let rb = match catch_unwind(AssertUnwindSafe(|| b())) {
+        let rb = match catch_unwind(AssertUnwindSafe(|| b(self))) {
             Ok(val) => val,
             Err(payload) => {
+                #[cfg(feature = "tracing")]
                 tracing::debug!("join_heartbeat: b panicked, waiting for a to finish");
                 cold_path();
+
                 // if b panicked, we need to wait for a to finish
-                self.wait_until_latch(&job);
+                let mut receiver = job.take_receiver();
+                self.wait_until_pred(|| match &receiver {
+                    Some(recv) => recv.poll().is_some(),
+                    None => {
+                        receiver = job.take_receiver();
+                        false
+                    }
+                });
+
                 resume_unwind(payload);
             }
         };
 
-        let ra = if !job.is_shared() {
-            tracing::trace!("join_heartbeat: job is not shared, running a() inline");
-            // we pushed the job to the back of the queue, any `join`s called by `b` on this worker thread will have already popped their job, or seen it be executed.
-            self.pop_back();
-
-            // a is allowed to panic here, because we already finished b.
-            unsafe { a.unwrap()() }
-        } else {
-            match self.wait_until_queued_job(&job) {
+        let ra = if let Some(recv) = job.take_receiver() {
+            match self.wait_until_recv(recv) {
                 Some(t) => crate::util::unwrap_or_panic(t),
                 None => {
+                    #[cfg(feature = "tracing")]
                     tracing::trace!(
                         "join_heartbeat: job was shared, but reclaimed, running a() inline"
                     );
                     // the job was shared, but not yet stolen, so we get to run the
                     // job inline
-                    unsafe { a.unwrap()() }
+                    a.run_inline(self)
                 }
             }
+        } else {
+            self.pop_back();
+
+            // SAFETY: we just popped the job from the queue, so it is safe to unwrap.
+            #[cfg(feature = "tracing")]
+            tracing::trace!("join_heartbeat: job was not shared, running a() inline");
+            a.run_inline(self)
+        };
+
+        (ra, rb)
+    }
+
+    /// This function must be called from a worker thread.
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    fn join_heartbeat<A, B, RA, RB>(&self, a: A, b: B) -> (RA, RB)
+    where
+        RA: Send,
+        A: FnOnce(&WorkerThread) -> RA + Send,
+        B: FnOnce(&WorkerThread) -> RB,
+    {
+        use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
+
+        #[cfg(feature = "metrics")]
+        self.metrics.num_joins.fetch_add(1, Ordering::Relaxed);
+
+        let a = StackJob::new(a);
+        let job = Job::from_stackjob(&a);
+
+        self.push_back(&job);
+
+        self.tick();
+
+        let rb = match catch_unwind(AssertUnwindSafe(|| b(self))) {
+            Ok(val) => val,
+            Err(payload) => {
+                #[cfg(feature = "tracing")]
+                tracing::debug!("join_heartbeat: b panicked, waiting for a to finish");
+                cold_path();
+
+                // if b panicked, we need to wait for a to finish
+                let mut receiver = job.take_receiver();
+                self.wait_until_pred(|| match &receiver {
+                    Some(recv) => recv.poll().is_some(),
+                    None => {
+                        receiver = job.take_receiver();
+                        false
+                    }
+                });
+
+                resume_unwind(payload);
+            }
+        };
+
+        let ra = if let Some(recv) = job.take_receiver() {
+            match self.wait_until_recv(recv) {
+                Some(t) => crate::util::unwrap_or_panic(t),
+                None => {
+                    #[cfg(feature = "tracing")]
+                    tracing::trace!(
+                        "join_heartbeat: job was shared, but reclaimed, running a() inline"
+                    );
+                    // the job was shared, but not yet stolen, so we get to run the
+                    // job inline
+                    unsafe { a.unwrap()(self) }
+                }
+            }
+        } else {
+            self.pop_back();
+
+            unsafe {
+                // SAFETY: we just popped the job from the queue, so it is safe to unwrap.
+                #[cfg(feature = "tracing")]
+                tracing::trace!("join_heartbeat: job was not shared, running a() inline");
+                a.unwrap()(self)
+            }
         };
 
         (ra, rb)
@@ -113,7 +225,6 @@ impl WorkerThread {
 }
 
 impl Context {
-    #[inline]
     pub fn join<A, B, RA, RB>(self: &Arc<Self>, a: A, b: B) -> (RA, RB)
     where
         A: FnOnce() -> RA + Send,
@@ -122,12 +233,13 @@ impl Context {
         RB: Send,
     {
         // SAFETY: join_heartbeat_every is safe to call from a worker thread.
-        self.run_in_worker(move |worker| worker.join_heartbeat_every::<_, _, _, _, 64>(a, b))
+        self.run_in_worker(move |worker| {
+            worker.join_heartbeat_every::<_, _, _, _>(|_| a(), |_| b())
+        })
     }
 }
 
 /// run two closures potentially in parallel, in the global threadpool.
-#[allow(dead_code)]
 pub fn join<A, B, RA, RB>(a: A, b: B) -> (RA, RB)
 where
     A: FnOnce() -> RA + Send,

distaff/src/latch.rs

@@ -2,19 +2,11 @@ use core::{
     marker::PhantomData,
     sync::atomic::{AtomicUsize, Ordering},
 };
-use std::{
-    cell::UnsafeCell,
-    mem,
-    ops::DerefMut,
-    sync::{
-        Arc,
-        atomic::{AtomicPtr, AtomicU8},
-    },
-};
+use std::sync::atomic::{AtomicPtr, AtomicU8};
 
 use parking_lot::{Condvar, Mutex};
 
-use crate::{WorkerThread, context::Context};
+use crate::channel::Parker;
 
 pub trait Latch {
     unsafe fn set_raw(this: *const Self);
@@ -199,21 +191,20 @@ impl Probe for NopLatch {
 
 pub struct CountLatch {
     count: AtomicUsize,
-    inner: AtomicPtr<WorkerLatch>,
+    inner: AtomicPtr<Parker>,
 }
 
 impl CountLatch {
     #[inline]
-    pub const fn new(inner: *const WorkerLatch) -> Self {
+    pub const fn new(inner: *const Parker) -> Self {
         Self {
             count: AtomicUsize::new(0),
-            inner: AtomicPtr::new(inner as *mut WorkerLatch),
+            inner: AtomicPtr::new(inner as *mut Parker),
         }
     }
 
-    pub fn set_inner(&self, inner: *const WorkerLatch) {
-        self.inner
-            .store(inner as *mut WorkerLatch, Ordering::Relaxed);
+    pub fn set_inner(&self, inner: *const Parker) {
+        self.inner.store(inner as *mut Parker, Ordering::Relaxed);
     }
 
     pub fn count(&self) -> usize {
@@ -238,11 +229,12 @@ impl Latch for CountLatch {
     unsafe fn set_raw(this: *const Self) {
         unsafe {
             if (&*this).count.fetch_sub(1, Ordering::Relaxed) == 1 {
+                #[cfg(feature = "tracing")]
                 tracing::trace!("CountLatch set_raw: count was 1, setting inner latch");
                 // If the count was 1, we need to set the inner latch.
                 let inner = (*this).inner.load(Ordering::Relaxed);
                 if !inner.is_null() {
-                    (&*inner).wake();
+                    (&*inner).unpark();
                 }
             }
         }
@@ -341,20 +333,36 @@ impl AsCoreLatch for MutexLatch {
 pub struct WorkerLatch {
     // this boolean is set when the worker is waiting.
     mutex: Mutex<bool>,
-    condvar: AtomicUsize,
+    condvar: Condvar,
 }
 
 impl WorkerLatch {
     pub fn new() -> Self {
         Self {
             mutex: Mutex::new(false),
-            condvar: AtomicUsize::new(0),
+            condvar: Condvar::new(),
         }
     }
-    pub fn lock(&self) {
-        mem::forget(self.mutex.lock());
+
+    #[cfg_attr(
+        feature = "tracing",
+        tracing::instrument(
+            level = "trace",
+            skip_all,
+            fields(this = self as *const Self as usize)
+        )
+    )]
+    pub fn lock(&self) -> parking_lot::MutexGuard<'_, bool> {
+        #[cfg(feature = "tracing")]
+        tracing::trace!("aquiring mutex..");
+        let guard = self.mutex.lock();
+        #[cfg(feature = "tracing")]
+        tracing::trace!("mutex acquired.");
+
+        guard
     }
-    pub fn unlock(&self) {
+
+    pub unsafe fn force_unlock(&self) {
         unsafe {
             self.mutex.force_unlock();
         }
@@ -362,133 +370,43 @@ impl WorkerLatch {
 
     pub fn wait(&self) {
         let condvar = &self.condvar;
-        let mut guard = self.mutex.lock();
+        let mut guard = self.lock();
 
         Self::wait_internal(condvar, &mut guard);
     }
 
-    fn wait_internal(condvar: &AtomicUsize, guard: &mut parking_lot::MutexGuard<'_, bool>) {
-        let mutex = parking_lot::MutexGuard::mutex(guard);
-        let key = condvar as *const _ as usize;
-        let lock_addr = mutex as *const _ as usize;
-        let mut requeued = false;
-
-        let state = unsafe { AtomicUsize::from_ptr(condvar as *const _ as *mut usize) };
-
+    fn wait_internal(condvar: &Condvar, guard: &mut parking_lot::MutexGuard<'_, bool>) {
         **guard = true; // set the mutex to true to indicate that the worker is waiting
-
-        unsafe {
-            parking_lot_core::park(
-                key,
-                || {
-                    let old = state.load(Ordering::Relaxed);
-                    if old == 0 {
-                        state.store(lock_addr, Ordering::Relaxed);
-                    } else if old != lock_addr {
-                        return false;
-                    }
-
-                    true
-                },
-                || {
-                    mutex.force_unlock();
-                },
-                |k, was_last_thread| {
-                    requeued = k != key;
-                    if !requeued && was_last_thread {
-                        state.store(0, Ordering::Relaxed);
-                    }
-                },
-                parking_lot_core::DEFAULT_PARK_TOKEN,
-                None,
-            );
-        }
-        // relock
-
-        let mut new = mutex.lock();
-        mem::swap(&mut new, guard);
-        mem::forget(new); // forget the new guard to avoid dropping it
-
-        **guard = false; // reset the mutex to false after waking up
+        //condvar.wait_for(guard, std::time::Duration::from_micros(100));
+        condvar.wait(guard);
+        **guard = false;
     }
 
-    fn wait_with_lock_internal<T>(&self, other: &mut parking_lot::MutexGuard<'_, T>) {
-        let key = &self.condvar as *const _ as usize;
-        let lock_addr = &self.mutex as *const _ as usize;
-        let mut requeued = false;
-
-        let mut guard = self.mutex.lock();
-
-        let state = unsafe { AtomicUsize::from_ptr(&self.condvar as *const _ as *mut usize) };
-
-        *guard = true; // set the mutex to true to indicate that the worker is waiting
-
-        unsafe {
-            let token = parking_lot_core::park(
-                key,
-                || {
-                    let old = state.load(Ordering::Relaxed);
-                    if old == 0 {
-                        state.store(lock_addr, Ordering::Relaxed);
-                    } else if old != lock_addr {
-                        return false;
-                    }
-
-                    true
-                },
-                || {
-                    drop(guard); // drop the guard to release the lock
-                    parking_lot::MutexGuard::mutex(&other).force_unlock();
-                },
-                |k, was_last_thread| {
-                    requeued = k != key;
-                    if !requeued && was_last_thread {
-                        state.store(0, Ordering::Relaxed);
-                    }
-                },
-                parking_lot_core::DEFAULT_PARK_TOKEN,
-                None,
-            );
-
-            tracing::trace!(
-                "WorkerLatch wait_with_lock_internal: unparked with token {:?}",
-                token
-            );
-        }
-        // relock
-        let mut other2 = parking_lot::MutexGuard::mutex(&other).lock();
-        tracing::trace!("WorkerLatch wait_with_lock_internal: relocked other");
-
-        // because `other` is logically unlocked, we swap it with `other2` and then forget `other2`
-        core::mem::swap(&mut *other2, &mut *other);
-        core::mem::forget(other2);
-
-        let mut guard = self.mutex.lock();
-        tracing::trace!("WorkerLatch wait_with_lock_internal: relocked self");
-
-        *guard = false; // reset the mutex to false after waking up
-    }
-
-    #[tracing::instrument(level = "trace", skip(other))]
-    pub fn wait_with_lock<T>(&self, other: &mut parking_lot::MutexGuard<'_, T>) {
-        self.wait_with_lock_internal(other);
-    }
-
-    pub fn wait_with_lock_while<T, F>(&self, other: &mut parking_lot::MutexGuard<'_, T>, mut f: F)
+    #[cfg_attr(
+        feature = "tracing",
+        tracing::instrument(level = "trace", skip_all, fields(
+        this = self as *const Self as usize,
+    )))]
+    pub fn wait_unless<F>(&self, mut f: F)
     where
-        F: FnMut(&mut T) -> bool,
+        F: FnMut() -> bool,
     {
-        while f(other.deref_mut()) {
-            self.wait_with_lock_internal(other);
+        let mut guard = self.lock();
+        if !f() {
+            Self::wait_internal(&self.condvar, &mut guard);
         }
     }
 
-    #[tracing::instrument(level = "trace", skip(f))]
+    #[cfg_attr(
+        feature = "tracing",
+        tracing::instrument(level = "trace", skip_all, fields(
+            this = self as *const Self as usize,
+    )))]
     pub fn wait_until<F, T>(&self, mut f: F) -> T
     where
         F: FnMut() -> Option<T>,
     {
-        let mut guard = self.mutex.lock();
+        let mut guard = self.lock();
         loop {
             if let Some(result) = f() {
                 return result;
@@ -501,14 +419,15 @@ impl WorkerLatch {
         *self.mutex.lock()
     }
 
-    #[tracing::instrument(level = "trace")]
+    #[cfg_attr(
+        feature = "tracing",
+tracing::instrument(level = "trace", skip_all, fields(
+        this = self as *const Self as usize,
+    )))]
     fn notify(&self) {
-        let key = &self.condvar as *const _ as usize;
-
-        unsafe {
-            let n = parking_lot_core::unpark_all(key, parking_lot_core::DEFAULT_UNPARK_TOKEN);
-            tracing::trace!("WorkerLatch notify_one: unparked {} threads", n);
-        }
+        let n = self.condvar.notify_all();
+        #[cfg(feature = "tracing")]
+        tracing::trace!("WorkerLatch notify: notified {} threads", n);
     }
 
     pub fn wake(&self) {
@@ -518,70 +437,10 @@ impl WorkerLatch {
 
 #[cfg(test)]
 mod tests {
-    use std::{ptr, sync::Barrier};
-
-    use tracing_test::traced_test;
+    use std::{ptr, sync::Arc};
 
     use super::*;
 
-    #[test]
-    #[cfg_attr(not(miri), traced_test)]
-    fn worker_latch() {
-        let latch = Arc::new(WorkerLatch::new());
-        let barrier = Arc::new(Barrier::new(2));
-        let mutex = Arc::new(parking_lot::Mutex::new(false));
-
-        let count = Arc::new(AtomicUsize::new(0));
-
-        let thread = std::thread::spawn({
-            let latch = latch.clone();
-            let mutex = mutex.clone();
-            let barrier = barrier.clone();
-            let count = count.clone();
-
-            move || {
-                tracing::info!("Thread waiting on barrier");
-                let mut guard = mutex.lock();
-                barrier.wait();
-
-                tracing::info!("Thread waiting on latch");
-                latch.wait_with_lock(&mut guard);
-                count.fetch_add(1, Ordering::Relaxed);
-                tracing::info!("Thread woke up from latch");
-                barrier.wait();
-                tracing::info!("Thread finished waiting on barrier");
-                count.fetch_add(1, Ordering::Relaxed);
-            }
-        });
-
-        assert!(!latch.is_waiting(), "Latch should not be waiting yet");
-        barrier.wait();
-        tracing::info!("Main thread finished waiting on barrier");
-        // lock mutex and notify the thread that isn't yet waiting.
-        {
-            let guard = mutex.lock();
-            tracing::info!("Main thread acquired mutex, waking up thread");
-            assert!(latch.is_waiting(), "Latch should be waiting now");
-
-            latch.wake();
-            tracing::info!("Main thread woke up thread");
-        }
-        assert_eq!(count.load(Ordering::Relaxed), 0, "Count should still be 0");
-        barrier.wait();
-        assert_eq!(
-            count.load(Ordering::Relaxed),
-            1,
-            "Count should be 1 after waking up"
-        );
-
-        thread.join().expect("Thread should join successfully");
-        assert_eq!(
-            count.load(Ordering::Relaxed),
-            2,
-            "Count should be 2 after thread has finished"
-        );
-    }
-
     #[test]
     fn test_atomic_latch() {
         let latch = AtomicLatch::new();
@@ -645,7 +504,7 @@ mod tests {
     }
 
     #[test]
-    #[traced_test]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn mutex_latch() {
         let latch = Arc::new(MutexLatch::new());
         assert!(!latch.probe());
@@ -657,8 +516,10 @@ mod tests {
         // Test wait functionality
         let latch_clone = latch.clone();
         let handle = std::thread::spawn(move || {
+            #[cfg(feature = "tracing")]
             tracing::info!("Thread waiting on latch");
             latch_clone.wait_and_reset();
+            #[cfg(feature = "tracing")]
             tracing::info!("Thread woke up from latch");
         });
 
@@ -666,8 +527,10 @@ mod tests {
         std::thread::sleep(std::time::Duration::from_millis(100));
         assert!(!latch.probe());
 
+        #[cfg(feature = "tracing")]
         tracing::info!("Setting latch from main thread");
         latch.set();
+        #[cfg(feature = "tracing")]
         tracing::info!("Latch set, joining waiting thread");
         handle.join().expect("Thread should join successfully");
     }

distaff/src/lib.rs

@@ -8,16 +8,21 @@
     box_as_ptr,
     box_vec_non_null,
     strict_provenance_atomic_ptr,
+    btree_extract_if,
+    likely_unlikely,
     let_chains
 )]
 
 extern crate alloc;
 
+mod channel;
 mod context;
 mod heartbeat;
 mod job;
 mod join;
 mod latch;
+#[cfg(feature = "metrics")]
+mod metrics;
 mod scope;
 mod threadpool;
 pub mod util;

distaff/src/metrics.rs

@@ -0,0 +1,12 @@
+use std::sync::atomic::AtomicU32;
+
+#[derive(Debug, Default)]
+pub(crate) struct WorkerMetrics {
+    pub(crate) num_jobs_shared: AtomicU32,
+    pub(crate) num_heartbeats: AtomicU32,
+    pub(crate) num_joins: AtomicU32,
+    pub(crate) num_jobs_reclaimed: AtomicU32,
+    pub(crate) num_jobs_executed: AtomicU32,
+    pub(crate) num_jobs_stolen: AtomicU32,
+    pub(crate) num_sent_to_self: AtomicU32,
+}

distaff/src/scope.rs

@@ -1,19 +1,25 @@
 use std::{
     any::Any,
     marker::PhantomData,
+    panic::{AssertUnwindSafe, catch_unwind},
+    pin::{self, Pin},
     ptr::{self, NonNull},
     sync::{
         Arc,
-        atomic::{AtomicPtr, Ordering},
+        atomic::{AtomicPtr, AtomicUsize, Ordering},
     },
 };
 
 use async_task::Runnable;
 
 use crate::{
+    channel::Sender,
     context::Context,
-    job::{HeapJob, JobSender, QueuedJob as Job},
-    latch::{CountLatch, WorkerLatch},
+    job::{
+        HeapJob, Job2 as Job,
+        traits::{InlineJob, IntoJob},
+    },
+    latch::{CountLatch, Probe},
     util::{DropGuard, SendPtr},
     workerthread::WorkerThread,
 };
@@ -29,6 +35,83 @@ use crate::{
 // - when a join() job finishes, it's latch is set
 // - when we wait for a join() job, we loop over the latch until it is set
 
+// a Scope must keep track of:
+// - The number of async jobs spawned, which is used to determine when the scope
+// is complete.
+// - A panic box, which is set when a job panics and is used to resume the panic
+// when the scope is completed.
+// - The Parker of the worker on which the scope was created, which is signaled
+// when the last outstanding async job finishes.
+// - The current worker thread in order to avoid having to query the
+// thread-local storage.
+
+struct ScopeInner {
+    outstanding_jobs: AtomicUsize,
+    parker: NonNull<crate::channel::Parker>,
+    panic: AtomicPtr<Box<dyn Any + Send + 'static>>,
+}
+
+unsafe impl Send for ScopeInner {}
+unsafe impl Sync for ScopeInner {}
+
+#[derive(Clone, Copy)]
+pub struct Scope<'scope, 'env: 'scope> {
+    inner: SendPtr<ScopeInner>,
+    worker: SendPtr<WorkerThread>,
+    _scope: PhantomData<&'scope mut &'scope ()>,
+    _env: PhantomData<&'env mut &'env ()>,
+}
+
+impl ScopeInner {
+    fn from_worker(worker: &WorkerThread) -> Self {
+        Self {
+            outstanding_jobs: AtomicUsize::new(0),
+            parker: worker.heartbeat.parker().into(),
+            panic: AtomicPtr::new(ptr::null_mut()),
+        }
+    }
+
+    fn increment(&self) {
+        self.outstanding_jobs.fetch_add(1, Ordering::Relaxed);
+    }
+
+    fn decrement(&self) {
+        if self.outstanding_jobs.fetch_sub(1, Ordering::Relaxed) == 1 {
+            unsafe {
+                self.parker.as_ref().unpark();
+            }
+        }
+    }
+
+    fn panicked(&self, err: Box<dyn Any + Send + 'static>) {
+        unsafe {
+            let err = Box::into_raw(Box::new(err));
+            if !self
+                .panic
+                .compare_exchange(ptr::null_mut(), err, Ordering::AcqRel, Ordering::Acquire)
+                .is_ok()
+            {
+                // someone else already set the panic, so we drop the error
+                _ = Box::from_raw(err);
+            }
+        }
+    }
+
+    fn maybe_propagate_panic(&self) {
+        let err = self.panic.swap(ptr::null_mut(), Ordering::AcqRel);
+
+        if err.is_null() {
+            return;
+        } else {
+            // SAFETY: we have exclusive access to the panic error, so we can safely resume it.
+            unsafe {
+                let err = *Box::from_raw(err);
+                std::panic::resume_unwind(err);
+            }
+        }
+    }
+}
+
 // find below a sketch of an unbalanced tree:
 //              []
 //             /  \
@@ -51,7 +134,7 @@ use crate::{
 // - another thread sharing a job
 // - the heartbeat waking up the worker // does this make sense? if the thread was sleeping, it didn't have any work to share.
 
-pub struct Scope<'scope, 'env: 'scope> {
+pub struct Scope2<'scope, 'env: 'scope> {
     // latch to wait on before the scope finishes
     job_counter: CountLatch,
     // local threadpool
@@ -65,7 +148,7 @@ pub struct Scope<'scope, 'env: 'scope> {
 
 pub fn scope<'env, F, R>(f: F) -> R
 where
-    F: for<'scope> FnOnce(&'scope Scope<'scope, 'env>) -> R + Send,
+    F: for<'scope> FnOnce(Scope<'scope, 'env>) -> R + Send,
     R: Send,
 {
     scope_with_context(Context::global_context(), f)
@@ -73,41 +156,25 @@ where
 
 pub fn scope_with_context<'env, F, R>(context: &Arc<Context>, f: F) -> R
 where
-    F: for<'scope> FnOnce(&'scope Scope<'scope, 'env>) -> R + Send,
+    F: for<'scope> FnOnce(Scope<'scope, 'env>) -> R + Send,
     R: Send,
 {
     context.run_in_worker(|worker| {
         // SAFETY: we call complete() after creating this scope, which
         // ensures that any jobs spawned from the scope exit before the
         // scope closes.
-        let this = unsafe { Scope::from_context(context.clone()) };
-        this.complete(worker, || f(&this))
+        let inner = pin::pin!(ScopeInner::from_worker(worker));
+        let this = Scope::<'_, 'env>::new(worker, inner.as_ref());
+        this.complete(|| f(this))
     })
 }
 
 impl<'scope, 'env> Scope<'scope, 'env> {
-    #[tracing::instrument(level = "trace", skip_all)]
-    fn wait_for_jobs(&self, worker: &WorkerThread) {
-        self.job_counter.set_inner(worker.heartbeat.raw_latch());
-        if self.job_counter.count() > 0 {
-            tracing::trace!("waiting for {} jobs to finish.", self.job_counter.count());
-            tracing::trace!(
-                "thread id: {:?}, jobs: {:?}",
-                worker.heartbeat.index(),
-                unsafe { worker.queue.as_ref_unchecked() }
-            );
-
-            // set worker index in the job counter
-            worker.wait_until_latch(&self.job_counter);
-        }
-    }
-
     /// should be called from within a worker thread.
-    #[tracing::instrument(level = "trace", skip_all)]
-    fn complete<F, R>(&self, worker: &WorkerThread, f: F) -> R
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    fn complete<F, R>(&self, f: F) -> R
     where
-        F: FnOnce() -> R + Send,
-        R: Send,
+        F: FnOnce() -> R,
     {
         use std::panic::{AssertUnwindSafe, catch_unwind};
 
@@ -119,81 +186,97 @@ impl<'scope, 'env> Scope<'scope, 'env> {
             }
         };
 
-        self.wait_for_jobs(worker);
-        self.maybe_propagate_panic();
+        self.wait_for_jobs();
+        let inner = self.inner();
+        inner.maybe_propagate_panic();
 
         // SAFETY: if result panicked, we would have propagated the panic above.
         result.unwrap()
     }
 
-    /// resumes the panic if one happened in this scope.
-    #[tracing::instrument(level = "trace", skip_all)]
-    fn maybe_propagate_panic(&self) {
-        let err_ptr = self.panic.load(Ordering::Relaxed);
-        if !err_ptr.is_null() {
-            unsafe {
-                let err = Box::from_raw(err_ptr);
-                std::panic::resume_unwind(*err);
-            }
-        }
-    }
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    fn wait_for_jobs(&self) {
+        #[cfg(feature = "tracing")]
+        tracing::trace!(
+            "waiting for {} jobs to finish.",
+            self.inner().outstanding_jobs.load(Ordering::Relaxed)
+        );
 
-    /// stores the first panic that happened in this scope.
-    #[tracing::instrument(level = "trace", skip_all)]
-    fn panicked(&self, err: Box<dyn Any + Send + 'static>) {
-        tracing::debug!("panicked in scope, storing error: {:?}", err);
-        self.panic.load(Ordering::Relaxed).is_null().then(|| {
-            use core::mem::ManuallyDrop;
-            let mut boxed = ManuallyDrop::new(Box::new(err));
-
-            let err_ptr: *mut Box<dyn Any + Send + 'static> = &mut **boxed;
-            if self
-                .panic
-                .compare_exchange(
-                    ptr::null_mut(),
-                    err_ptr,
-                    Ordering::SeqCst,
-                    Ordering::Relaxed,
-                )
-                .is_ok()
-            {
-                // we successfully set the panic, no need to drop
-            } else {
-                // drop the error, someone else already set it
-                _ = ManuallyDrop::into_inner(boxed);
-            }
+        self.worker().wait_until_pred(|| {
+            // SAFETY: we are in a worker thread, so the inner is valid.
+            let count = self.inner().outstanding_jobs.load(Ordering::Relaxed);
+            #[cfg(feature = "tracing")]
+            tracing::trace!("waiting for {} jobs to finish.", count);
+            count == 0
         });
     }
 
-    pub fn spawn<F>(&'scope self, f: F)
-    where
-        F: FnOnce(&'scope Self) + Send,
-    {
-        self.job_counter.increment();
-
-        let this = SendPtr::new_const(self).unwrap();
-
-        let job = Job::from_heapjob(
-            Box::new(HeapJob::new(move || unsafe {
-                use std::panic::{AssertUnwindSafe, catch_unwind};
-                if let Err(payload) = catch_unwind(AssertUnwindSafe(|| f(this.as_ref()))) {
-                    this.as_unchecked_ref().panicked(payload);
-                }
-                this.as_unchecked_ref().job_counter.decrement();
-            })),
-            ptr::null(),
-        );
-
-        tracing::trace!("allocated heapjob");
-
-        WorkerThread::current_ref()
-            .expect("spawn is run in workerthread.")
-            .push_front(job.as_ptr());
-
-        tracing::trace!("leaked heapjob");
+    fn inner(&self) -> &ScopeInner {
+        unsafe { self.inner.as_ref() }
     }
 
-    pub fn spawn_future<T, F>(&'scope self, future: F) -> async_task::Task<T>
+    /// stores the first panic that happened in this scope.
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    fn panicked(&self, err: Box<dyn Any + Send + 'static>) {
+        self.inner().panicked(err);
+    }
+
+    pub fn spawn<F>(&self, f: F)
+    where
+        F: FnOnce(Self) + Send,
+    {
+        struct SpawnedJob<F> {
+            f: F,
+            inner: SendPtr<ScopeInner>,
+        }
+
+        impl<F> SpawnedJob<F> {
+            fn new<'scope, 'env, T>(f: F, inner: SendPtr<ScopeInner>) -> Job
+            where
+                F: FnOnce(Scope<'scope, 'env>) -> T + Send,
+                'env: 'scope,
+                T: Send,
+            {
+                Job::from_harness(
+                    Self::harness,
+                    Box::into_non_null(Box::new(Self { f, inner })).cast(),
+                )
+            }
+
+            unsafe fn harness<'scope, 'env, T>(
+                worker: &WorkerThread,
+                this: NonNull<()>,
+                _: Option<Sender>,
+            ) where
+                F: FnOnce(Scope<'scope, 'env>) -> T + Send,
+                'env: 'scope,
+                T: Send,
+            {
+                let Self { f, inner } =
+                    unsafe { *Box::<SpawnedJob<F>>::from_non_null(this.cast()) };
+                let scope = unsafe { Scope::<'scope, 'env>::new_unchecked(worker, inner) };
+
+                // SAFETY: we are in a worker thread, so the inner is valid.
+                (f)(scope);
+            }
+        }
+
+        self.inner().increment();
+        let job = SpawnedJob::new(
+            move |scope| {
+                if let Err(payload) = catch_unwind(AssertUnwindSafe(|| f(scope))) {
+                    scope.inner().panicked(payload);
+                }
+
+                scope.inner().decrement();
+            },
+            self.inner,
+        );
+
+        self.context().inject_job(job.share(None));
+    }
+
+    pub fn spawn_future<T, F>(&self, future: F) -> async_task::Task<T>
     where
         F: Future<Output = T> + Send + 'scope,
         T: Send + 'scope,
@@ -202,9 +285,9 @@ impl<'scope, 'env> Scope<'scope, 'env> {
     }
 
     #[allow(dead_code)]
-    pub fn spawn_async<T, Fut, Fn>(&'scope self, f: Fn) -> async_task::Task<T>
+    pub fn spawn_async<T, Fut, Fn>(&self, f: Fn) -> async_task::Task<T>
     where
-        Fn: FnOnce(&'scope Self) -> Fut + Send + 'scope,
+        Fn: FnOnce(Self) -> Fut + Send + 'scope,
         Fut: Future<Output = T> + Send + 'scope,
         T: Send + 'scope,
     {
@@ -212,51 +295,43 @@ impl<'scope, 'env> Scope<'scope, 'env> {
     }
 
     #[inline]
-    fn spawn_async_internal<T, Fut, Fn>(&'scope self, f: Fn) -> async_task::Task<T>
+    fn spawn_async_internal<T, Fut, Fn>(&self, f: Fn) -> async_task::Task<T>
     where
-        Fn: FnOnce(&'scope Self) -> Fut + Send + 'scope,
+        Fn: FnOnce(Self) -> Fut + Send + 'scope,
         Fut: Future<Output = T> + Send + 'scope,
         T: Send + 'scope,
     {
-        self.job_counter.increment();
+        self.inner().increment();
 
-        let this = SendPtr::new_const(self).unwrap();
-        // let this = SendPtr::new_const(&self.job_counter).unwrap();
+        // TODO: make sure this worker lasts long enough for the
+        // reference to remain valid for the duration of the future.
+        let scope = unsafe { Self::new_unchecked(self.worker.as_ref(), self.inner) };
 
         let future = async move {
-            // SAFETY: this is valid until we decrement the job counter.
-            unsafe {
-                let _guard = DropGuard::new(move || {
-                    this.as_unchecked_ref().job_counter.decrement();
-                });
-                // TODO: handle panics here
-                f(this.as_ref()).await
-            }
+            let _guard = DropGuard::new(move || {
+                scope.inner().decrement();
+            });
+
+            // TODO: handle panics here
+            f(scope).await
         };
 
         let schedule = move |runnable: Runnable| {
             #[align(8)]
-            unsafe fn harness(this: *const (), job: *const JobSender, _: *const WorkerLatch) {
+            unsafe fn harness(_: &WorkerThread, this: NonNull<()>, _: Option<Sender>) {
                 unsafe {
-                    let runnable =
-                        Runnable::<()>::from_raw(NonNull::new_unchecked(this.cast_mut()));
+                    let runnable = Runnable::<()>::from_raw(this.cast());
                     runnable.run();
-
-                    // SAFETY: job was turned into raw
-                    drop(Box::from_raw(job.cast_mut()));
                 }
             }
 
-            let job = Box::into_raw(Box::new(Job::from_harness(
-                harness,
-                runnable.into_raw(),
-                ptr::null(),
-            )));
+            let job = Job::<()>::from_harness(harness, runnable.into_raw());
 
             // casting into Job<()> here
-            WorkerThread::current_ref()
-                .expect("spawn_async_internal is run in workerthread.")
-                .push_front(job);
+            self.context().inject_job(job.share(None));
+            // WorkerThread::current_ref()
+            //     .expect("spawn_async_internal is run in workerthread.")
+            //     .push_front(job);
         };
 
         let (runnable, task) = unsafe { async_task::spawn_unchecked(future, schedule) };
@@ -266,51 +341,176 @@ impl<'scope, 'env> Scope<'scope, 'env> {
         task
     }
 
-    #[inline]
-    pub fn join<A, B, RA, RB>(&'scope self, a: A, b: B) -> (RA, RB)
+    pub fn join<A, B, RA, RB>(&self, a: A, b: B) -> (RA, RB)
     where
         RA: Send,
-        RB: Send,
-        A: FnOnce(&'scope Self) -> RA + Send,
-        B: FnOnce(&'scope Self) -> RB + Send,
+        A: FnOnce(Self) -> RA + Send,
+        B: FnOnce(Self) -> RB,
     {
-        let worker = WorkerThread::current_ref().expect("join is run in workerthread.");
-        let this = SendPtr::new_const(self).unwrap();
+        use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
+        use std::{
+            cell::UnsafeCell,
+            mem::{self, ManuallyDrop},
+        };
 
-        worker.join_heartbeat_every::<_, _, _, _, 64>(
+        let worker = self.worker();
+
+        struct ScopeJob<F> {
+            f: UnsafeCell<ManuallyDrop<F>>,
+            inner: SendPtr<ScopeInner>,
+        }
+
+        impl<F> ScopeJob<F> {
+            fn new(f: F, inner: SendPtr<ScopeInner>) -> Self {
+                Self {
+                    f: UnsafeCell::new(ManuallyDrop::new(f)),
+                    inner,
+                }
+            }
+
+            fn into_job<'scope, 'env, T>(&self) -> Job<T>
+            where
+                F: FnOnce(Scope<'scope, 'env>) -> T + Send,
+                'env: 'scope,
+                T: Send,
             {
-                let this = this;
-                move || a(unsafe { this.as_ref() })
-            },
+                Job::from_harness(Self::harness, NonNull::from(self).cast())
+            }
+
+            unsafe fn unwrap(&self) -> F {
+                unsafe { ManuallyDrop::take(&mut *self.f.get()) }
+            }
+
+            unsafe fn harness<'scope, 'env, T>(
+                worker: &WorkerThread,
+                this: NonNull<()>,
+                sender: Option<Sender>,
+            ) where
+                F: FnOnce(Scope<'scope, 'env>) -> T + Send,
+                'env: 'scope,
+                T: Send,
             {
-                let this = this;
-                move || b(unsafe { this.as_ref() })
-            },
-        )
+                let this: &ScopeJob<F> = unsafe { this.cast().as_ref() };
+                let f = unsafe { this.unwrap() };
+                let scope = unsafe { Scope::<'scope, 'env>::new_unchecked(worker, this.inner) };
+                let sender: Sender<T> = unsafe { mem::transmute(sender) };
+
+                // SAFETY: we are in a worker thread, so the inner is valid.
+                sender.send(catch_unwind(AssertUnwindSafe(|| f(scope))));
+            }
+        }
+
+        impl<'scope, 'env, F, T> IntoJob<T> for &ScopeJob<F>
+        where
+            F: FnOnce(Scope<'scope, 'env>) -> T + Send,
+            'env: 'scope,
+            T: Send,
+        {
+            fn into_job(self) -> Job<T> {
+                self.into_job()
+            }
+        }
+
+        impl<'scope, 'env, F, T> InlineJob<T> for &ScopeJob<F>
+        where
+            F: FnOnce(Scope<'scope, 'env>) -> T + Send,
+            'env: 'scope,
+            T: Send,
+        {
+            fn run_inline(self, worker: &WorkerThread) -> T {
+                unsafe { self.unwrap()(Scope::<'scope, 'env>::new_unchecked(worker, self.inner)) }
+            }
+        }
+
+        return worker
+            .join_heartbeat2_every::<_, _, _, _, 64>(&ScopeJob::new(a, self.inner), |_| b(*self));
+
+        // let stack = ScopeJob::new(a, self.inner);
+        // let job = ScopeJob::into_job(&stack);
+
+        // worker.push_back(&job);
+
+        // worker.tick();
+
+        // let rb = match catch_unwind(AssertUnwindSafe(|| b(*self))) {
+        //     Ok(val) => val,
+        //     Err(payload) => {
+        //         #[cfg(feature = "tracing")]
+        //         tracing::debug!("join_heartbeat: b panicked, waiting for a to finish");
+        //         std::hint::cold_path();
+
+        //         // if b panicked, we need to wait for a to finish
+        //         let mut receiver = job.take_receiver();
+        //         worker.wait_until_pred(|| match &receiver {
+        //             Some(recv) => recv.poll().is_some(),
+        //             None => {
+        //                 receiver = job.take_receiver();
+        //                 false
+        //             }
+        //         });
+
+        //         resume_unwind(payload);
+        //     }
+        // };
+
+        // let ra = if let Some(recv) = job.take_receiver() {
+        //     match worker.wait_until_recv(recv) {
+        //         Some(t) => crate::util::unwrap_or_panic(t),
+        //         None => {
+        //             #[cfg(feature = "tracing")]
+        //             tracing::trace!(
+        //                 "join_heartbeat: job was shared, but reclaimed, running a() inline"
+        //             );
+        //             // the job was shared, but not yet stolen, so we get to run the
+        //             // job inline
+        //             unsafe { stack.unwrap()(*self) }
+        //         }
+        //     }
+        // } else {
+        //     worker.pop_back();
+
+        //     unsafe {
+        //         // SAFETY: we just popped the job from the queue, so it is safe to unwrap.
+        //         #[cfg(feature = "tracing")]
+        //         tracing::trace!("join_heartbeat: job was not shared, running a() inline");
+        //         stack.unwrap()(*self)
+        //     }
+        // };
+
+        // (ra, rb)
     }
 
-    unsafe fn from_context(context: Arc<Context>) -> Self {
+    fn new(worker: &WorkerThread, inner: Pin<&'scope ScopeInner>) -> Self {
+        // SAFETY: we are creating a new scope, so the inner is valid.
+        unsafe { Self::new_unchecked(worker, SendPtr::new_const(&*inner).unwrap()) }
+    }
+
+    unsafe fn new_unchecked(worker: &WorkerThread, inner: SendPtr<ScopeInner>) -> Self {
         Self {
-            context,
-            job_counter: CountLatch::new(ptr::null()),
-            panic: AtomicPtr::new(ptr::null_mut()),
+            inner,
+            worker: SendPtr::new_const(worker).unwrap(),
             _scope: PhantomData,
             _env: PhantomData,
         }
     }
+
+    pub fn context(&self) -> &Arc<Context> {
+        unsafe { &self.worker.as_ref().context }
+    }
+    pub fn worker(&self) -> &WorkerThread {
+        unsafe { self.worker.as_ref() }
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use std::sync::atomic::AtomicU8;
 
-    use tracing_test::traced_test;
-
     use super::*;
     use crate::ThreadPool;
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn scope_spawn_sync() {
         let pool = ThreadPool::new_with_threads(1);
         let count = Arc::new(AtomicU8::new(0));
@@ -325,7 +525,7 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn scope_join_one() {
         let pool = ThreadPool::new_with_threads(1);
 
@@ -338,11 +538,11 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn scope_join_many() {
         let pool = ThreadPool::new_with_threads(1);
 
-        fn sum<'scope, 'env>(scope: &'scope Scope<'scope, 'env>, n: usize) -> usize {
+        fn sum<'scope, 'env>(scope: Scope<'scope, 'env>, n: usize) -> usize {
             if n == 0 {
                 return 0;
             }
@@ -360,7 +560,7 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn scope_spawn_future() {
         let pool = ThreadPool::new_with_threads(1);
         let mut x = 0;
@@ -376,7 +576,7 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn scope_spawn_many() {
         let pool = ThreadPool::new_with_threads(1);
         let count = Arc::new(AtomicU8::new(0));

distaff/src/threadpool.rs

@@ -8,8 +8,8 @@ pub struct ThreadPool {
 
 impl Drop for ThreadPool {
     fn drop(&mut self) {
-        // Ensure that the context is properly cleaned up when the thread pool is dropped.
-        self.context.set_should_exit();
+        // TODO: Ensure that the context is properly cleaned up when the thread pool is dropped.
+        // self.context.set_should_exit();
     }
 }
 
@@ -25,9 +25,14 @@ impl ThreadPool {
         Self { context }
     }
 
+    pub fn global() -> Self {
+        let context = Context::global_context().clone();
+        Self { context }
+    }
+
     pub fn scope<'env, F, R>(&self, f: F) -> R
     where
-        F: for<'scope> FnOnce(&'scope Scope<'scope, 'env>) -> R + Send,
+        F: for<'scope> FnOnce(Scope<'scope, 'env>) -> R + Send,
         R: Send,
     {
         scope_with_context(&self.context, f)
@@ -53,17 +58,16 @@ impl ThreadPool {
 
 #[cfg(test)]
 mod tests {
-    use tracing_test::traced_test;
-
     use super::*;
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn pool_spawn_borrow() {
         let pool = ThreadPool::new_with_threads(1);
         let mut x = 0;
         pool.scope(|scope| {
             scope.spawn(|_| {
+                #[cfg(feature = "tracing")]
                 tracing::info!("Incrementing x");
                 x += 1;
             });
@@ -72,7 +76,7 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn pool_spawn_future() {
         let pool = ThreadPool::new_with_threads(1);
         let mut x = 0;
@@ -89,7 +93,7 @@ mod tests {
     }
 
     #[test]
-    #[cfg_attr(not(miri), traced_test)]
+    #[cfg_attr(all(not(miri), feature = "tracing"), tracing_test::traced_test)]
     fn pool_join() {
         let pool = ThreadPool::new_with_threads(1);
         let (a, b) = pool.join(|| 3 + 4, || 5 * 6);

distaff/src/util.rs

@@ -94,8 +94,7 @@ impl<T> SendPtr<T> {
         unsafe { Self::new_unchecked(ptr.cast_mut()) }
     }
 
-    pub unsafe fn as_unchecked_ref(&self) -> &T {
-        // SAFETY: `self.0` is a valid non-null pointer.
+    pub(crate) unsafe fn as_ref(&self) -> &T {
         unsafe { self.0.as_ref() }
     }
 }
@@ -182,7 +181,6 @@ impl<T, const BITS: u8> TaggedAtomicPtr<T, BITS> {
     }
 
     /// returns tag
-    #[inline]
     #[allow(dead_code)]
     pub fn compare_exchange_tag(
         &self,
@@ -201,7 +199,6 @@ impl<T, const BITS: u8> TaggedAtomicPtr<T, BITS> {
     }
 
     /// returns tag
-    #[inline]
     pub fn compare_exchange_weak_tag(
         &self,
         old: usize,

distaff/src/workerthread.rs

@@ -1,18 +1,20 @@
+#[cfg(feature = "metrics")]
+use std::sync::atomic::Ordering;
+
 use std::{
     cell::{Cell, UnsafeCell},
-    hint::cold_path,
     ptr::NonNull,
-    sync::Arc,
+    sync::{Arc, Barrier},
     time::Duration,
 };
 
 use crossbeam_utils::CachePadded;
 
 use crate::{
-    context::{Context, Heartbeat},
+    channel::Receiver,
+    context::Context,
     heartbeat::OwnedHeartbeatReceiver,
-    job::{JobQueue as JobList, JobResult, QueuedJob as Job, QueuedJob, StackJob},
-    latch::{AsCoreLatch, CoreLatch, Probe, WorkerLatch},
+    job::{Job2 as Job, JobQueue as JobList, SharedJob},
     util::DropGuard,
 };
 
@@ -21,6 +23,9 @@ pub struct WorkerThread {
     pub(crate) queue: UnsafeCell<JobList>,
     pub(crate) heartbeat: OwnedHeartbeatReceiver,
     pub(crate) join_count: Cell<u8>,
+
+    #[cfg(feature = "metrics")]
+    pub(crate) metrics: CachePadded<crate::metrics::WorkerMetrics>,
 }
 
 thread_local! {
@@ -36,13 +41,17 @@ impl WorkerThread {
             queue: UnsafeCell::new(JobList::new()),
             heartbeat,
             join_count: Cell::new(0),
+            #[cfg(feature = "metrics")]
+            metrics: CachePadded::new(crate::metrics::WorkerMetrics::default()),
         }
     }
 }
 
 impl WorkerThread {
-    #[tracing::instrument(level = "trace", skip_all)]
-    pub fn run(self: Box<Self>) {
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all, fields(
+        worker = self.heartbeat.index(),
+    )))]
+    pub fn run(self: Box<Self>, barrier: Arc<Barrier>) {
         let this = Box::into_raw(self);
         unsafe {
             Self::set_current(this);
@@ -54,16 +63,24 @@ impl WorkerThread {
             Self::drop_in_place(this);
         });
 
+        #[cfg(feature = "tracing")]
         tracing::trace!("WorkerThread::run: starting worker thread");
 
+        barrier.wait();
         unsafe {
             (&*this).run_inner();
         }
 
+        #[cfg(feature = "metrics")]
+        unsafe {
+            eprintln!("{:?}", (&*this).metrics);
+        }
+
+        #[cfg(feature = "tracing")]
         tracing::trace!("WorkerThread::run: worker thread finished");
     }
 
-    #[tracing::instrument(level = "trace", skip_all)]
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
     fn run_inner(&self) {
         let mut job = None;
         'outer: loop {
@@ -86,55 +103,69 @@ impl WorkerThread {
 }
 
 impl WorkerThread {
-    pub(crate) fn find_work(&self) -> Option<NonNull<Job>> {
-        self.find_work_inner().left()
-    }
-
     /// Looks for work in the local queue, then in the shared context, and if no
     /// work is found, waits for the thread to be notified of a new job, after
     /// which it returns `None`.
     /// The caller should then check for `should_exit` to determine if the
     /// thread should exit, or look for work again.
-    #[tracing::instrument(level = "trace", skip_all)]
-    pub(crate) fn find_work_or_wait(&self) -> Option<NonNull<Job>> {
-        match self.find_work_inner() {
-            either::Either::Left(job) => {
-                return Some(job);
-            }
-            either::Either::Right(mut guard) => {
-                // no jobs found, wait for a heartbeat or a new job
-                tracing::trace!("WorkerThread::find_work_or_wait: waiting for new job");
-                self.heartbeat.latch().wait_with_lock(&mut guard);
-                tracing::trace!("WorkerThread::find_work_or_wait: woken up from wait");
-                None
-            }
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub(crate) fn find_work_or_wait(&self) -> Option<SharedJob> {
+        if let Some(job) = self.find_work() {
+            return Some(job);
         }
+
+        #[cfg(feature = "tracing")]
+        tracing::trace!("waiting for new job");
+        self.heartbeat.parker().park();
+        #[cfg(feature = "tracing")]
+        tracing::trace!("woken up from wait");
+        None
     }
 
-    #[inline]
-    fn find_work_inner(
-        &self,
-    ) -> either::Either<NonNull<Job>, parking_lot::MutexGuard<'_, crate::context::Shared>> {
-        // first check the local queue for jobs
-        if let Some(job) = self.pop_front() {
-            tracing::trace!("WorkerThread::find_work_inner: found local job: {:?}", job);
-            return either::Either::Left(job);
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub(crate) fn find_work_or_wait_unless<F>(&self, mut pred: F) -> Option<SharedJob>
+    where
+        F: FnMut() -> bool,
+    {
+        if let Some(job) = self.find_work() {
+            return Some(job);
         }
+        // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+        // Check the predicate while holding the lock. This is very important,
+        // because the lock must be held when notifying us of the result of a
+        // job we scheduled.
+        // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+        // no jobs found, wait for a heartbeat or a new job
+        #[cfg(feature = "tracing")]
+        tracing::trace!(worker = self.heartbeat.index(), "waiting for new job");
+        if !pred() {
+            self.heartbeat.parker().park();
+        }
+        #[cfg(feature = "tracing")]
+        tracing::trace!(worker = self.heartbeat.index(), "woken up from wait");
 
-        // then check the shared context for jobs
+        None
+    }
+
+    fn find_work(&self) -> Option<SharedJob> {
         let mut guard = self.context.shared();
 
         if let Some(job) = guard.pop_job() {
+            #[cfg(feature = "metrics")]
+            self.metrics.num_jobs_stolen.fetch_add(1, Ordering::Relaxed);
+            #[cfg(feature = "tracing")]
             tracing::trace!("WorkerThread::find_work_inner: found shared job: {:?}", job);
-            return either::Either::Left(job);
+            return Some(job);
         }
 
-        either::Either::Right(guard)
+        None
     }
 
-    #[inline(always)]
     pub(crate) fn tick(&self) {
         if self.heartbeat.take() {
+            #[cfg(feature = "metrics")]
+            self.metrics.num_heartbeats.fetch_add(1, Ordering::Relaxed);
+            #[cfg(feature = "tracing")]
             tracing::trace!(
                 "received heartbeat, thread id: {:?}",
                 self.heartbeat.index()
@@ -143,11 +174,10 @@ impl WorkerThread {
         }
     }
 
-    #[inline]
-    #[tracing::instrument(level = "trace", skip(self))]
-    fn execute(&self, job: NonNull<Job>) {
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    fn execute(&self, job: SharedJob) {
+        unsafe { SharedJob::execute(job, self) };
         self.tick();
-        unsafe { Job::execute(job.as_ptr()) };
     }
 
     #[cold]
@@ -156,10 +186,17 @@ impl WorkerThread {
 
         if !guard.jobs.contains_key(&self.heartbeat.id()) {
             if let Some(job) = self.pop_back() {
-                Job::set_shared(unsafe { job.as_ref() });
+                #[cfg(feature = "tracing")]
                 tracing::trace!("heartbeat: sharing job: {:?}", job);
-                guard.jobs.insert(self.heartbeat.id(), job);
+
+                #[cfg(feature = "metrics")]
+                self.metrics.num_jobs_shared.fetch_add(1, Ordering::Relaxed);
+
                 unsafe {
+                    guard.jobs.insert(
+                        self.heartbeat.id(),
+                        job.as_ref().share(Some(self.heartbeat.parker())),
+                    );
                     // SAFETY: we are holding the lock on the shared context.
                     self.context.notify_job_shared();
                 }
@@ -169,29 +206,23 @@ impl WorkerThread {
 }
 
 impl WorkerThread {
-    #[inline]
     pub fn pop_back(&self) -> Option<NonNull<Job>> {
         unsafe { self.queue.as_mut_unchecked().pop_back() }
     }
 
-    #[inline]
-    pub fn push_back(&self, job: *const Job) {
-        unsafe { self.queue.as_mut_unchecked().push_back(job) }
+    pub fn push_back<T>(&self, job: *const Job<T>) {
+        unsafe { self.queue.as_mut_unchecked().push_back(job.cast()) }
+    }
+    pub fn push_front<T>(&self, job: *const Job<T>) {
+        unsafe { self.queue.as_mut_unchecked().push_front(job.cast()) }
     }
 
-    #[inline]
     pub fn pop_front(&self) -> Option<NonNull<Job>> {
         unsafe { self.queue.as_mut_unchecked().pop_front() }
     }
-
-    #[inline]
-    pub fn push_front(&self, job: *const Job) {
-        unsafe { self.queue.as_mut_unchecked().push_front(job) }
-    }
 }
 
 impl WorkerThread {
-    #[inline]
     pub fn current_ref<'a>() -> Option<&'a Self> {
         unsafe { (*WORKER.with(UnsafeCell::get)).map(|ptr| ptr.as_ref()) }
     }
@@ -242,9 +273,11 @@ impl HeartbeatThread {
         Self { ctx }
     }
 
-    #[tracing::instrument(level = "trace", skip(self))]
-    pub fn run(self) {
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self)))]
+    pub fn run(self, barrier: Arc<Barrier>) {
+        #[cfg(feature = "tracing")]
         tracing::trace!("new heartbeat thread {:?}", std::thread::current());
+        barrier.wait();
 
         let mut i = 0;
         loop {
@@ -273,113 +306,88 @@ impl HeartbeatThread {
 }
 
 impl WorkerThread {
-    #[tracing::instrument(level = "trace", skip(self))]
-    pub fn wait_until_queued_job<T>(
-        &self,
-        job: *const QueuedJob,
-    ) -> Option<std::thread::Result<T>> {
-        let recv = unsafe { (*job).as_receiver::<T>() };
-        // we've already checked that the job was popped from the queue
-        // check if shared job is our job
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self)))]
+    pub fn wait_until_shared_job<T: Send>(&self, job: &Job<T>) -> Option<std::thread::Result<T>> {
+        let recv = (*job).take_receiver().unwrap();
 
-        // if let Some(shared_job) = self.context.shared().jobs.remove(&self.heartbeat.id()) {
-        //     if core::ptr::eq(shared_job.as_ptr(), job as *const Job as _) {
-        //         // this is the job we are looking for, so we want to
-        //         // short-circuit and call it inline
-        //         tracing::trace!(
-        //             thread = self.heartbeat.index(),
-        //             "reclaiming shared job: {:?}",
-        //             shared_job
-        //         );
+        let mut out = recv.poll();
 
-        //         return None;
-        //     } else {
-        //         // this isn't the job we are looking for, but we still need to
-        //         // execute it
-        //         tracing::trace!(
-        //             thread = self.heartbeat.index(),
-        //             "executing reclaimed shared job: {:?}",
-        //             shared_job
-        //         );
-
-        //         unsafe { Job::execute(shared_job.as_ptr()) };
-        //     }
-        // }
-
-        loop {
-            match recv.poll() {
-                Some(t) => {
-                    return Some(t);
-                }
-                None => {
-                    cold_path();
-
-                    // check local jobs before locking shared context
-                    if let Some(job) = self.find_work_or_wait() {
-                        tracing::trace!(
-                            "thread {:?} executing local job: {:?}",
-                            self.heartbeat.index(),
-                            job
-                        );
-                        unsafe {
-                            Job::execute(job.as_ptr());
-                        }
-                        tracing::trace!(
-                            "thread {:?} finished local job: {:?}",
-                            self.heartbeat.index(),
-                            job
-                        );
-                        continue;
-                    }
+        while std::hint::unlikely(out.is_none()) {
+            if let Some(job) = self.find_work() {
+                unsafe {
+                    SharedJob::execute(job, self);
                 }
             }
+
+            out = recv.poll();
         }
+
+        out
     }
 
-    #[tracing::instrument(level = "trace", skip_all)]
-    pub fn wait_until_latch<L>(&self, latch: &L)
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub fn wait_until_recv<T: Send>(&self, recv: Receiver<T>) -> Option<std::thread::Result<T>> {
+        if self
+            .context
+            .shared()
+            .jobs
+            .remove(&self.heartbeat.id())
+            .is_some()
+        {
+            #[cfg(feature = "tracing")]
+            tracing::trace!("reclaiming shared job");
+            return None;
+        }
+
+        while recv.is_empty() {
+            if let Some(job) = self.find_work() {
+                unsafe {
+                    SharedJob::execute(job, self);
+                }
+                continue;
+            }
+
+            recv.wait();
+        }
+
+        Some(recv.recv())
+    }
+
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
+    pub fn wait_until_pred<F>(&self, mut pred: F)
     where
-        L: Probe,
+        F: FnMut() -> bool,
     {
-        if !latch.probe() {
-            tracing::trace!("thread {:?} waiting on latch", self.heartbeat.index());
-            self.wait_until_latch_cold(latch);
+        if !pred() {
+            #[cfg(feature = "tracing")]
+            tracing::trace!("thread {:?} waiting on predicate", self.heartbeat.index());
+            self.wait_until_latch_cold(pred);
         }
     }
 
     #[cold]
-    fn wait_until_latch_cold<L>(&self, latch: &L)
+    fn wait_until_latch_cold<F>(&self, mut pred: F)
     where
-        L: Probe,
+        F: FnMut() -> bool,
     {
         if let Some(shared_job) = self.context.shared().jobs.remove(&self.heartbeat.id()) {
+            #[cfg(feature = "tracing")]
             tracing::trace!(
                 "thread {:?} reclaiming shared job: {:?}",
                 self.heartbeat.index(),
                 shared_job
             );
-            unsafe { Job::execute(shared_job.as_ptr()) };
+            unsafe { SharedJob::execute(shared_job, self) };
         }
 
         // do the usual thing and wait for the job's latch
         // do the usual thing??? chatgipity really said this..
-        while !latch.probe() {
+        while !pred() {
             // check local jobs before locking shared context
-            if let Some(job) = self.find_work_or_wait() {
-                tracing::trace!(
-                    "thread {:?} executing local job: {:?}",
-                    self.heartbeat.index(),
-                    job
-                );
+            if let Some(job) = self.find_work() {
                 unsafe {
-                    Job::execute(job.as_ptr());
+                    SharedJob::execute(job, self);
                 }
-                tracing::trace!(
-                    "thread {:?} finished local job: {:?}",
-                    self.heartbeat.index(),
-                    job
-                );
-                continue;
             }
         }
     }

examples/join.rs

@@ -62,10 +62,11 @@ fn join_pool(tree_size: usize) {
 
 fn join_distaff(tree_size: usize) {
     use distaff::*;
-    let pool = ThreadPool::new();
+    let pool = ThreadPool::new_with_threads(6);
+
     let tree = Tree::new(tree_size, 1);
 
-    fn sum<'scope, 'env>(tree: &Tree<u32>, node: usize, scope: &'scope Scope<'scope, 'env>) -> u32 {
+    fn sum<'scope, 'env>(tree: &Tree<u32>, node: usize, scope: Scope<'scope, 'env>) -> u32 {
         let node = tree.get(node);
         let (l, r) = scope.join(
             |s| node.left.map(|node| sum(tree, node, s)).unwrap_or_default(),
@@ -81,11 +82,13 @@ fn join_distaff(tree_size: usize) {
         node.leaf + l + r
     }
 
-    let sum = pool.scope(|s| {
-        let sum = sum(&tree, tree.root().unwrap(), s);
-        sum
-    });
-    std::hint::black_box(sum);
+    for _ in 0..1000 {
+        let sum = pool.scope(|s| {
+            let sum = sum(&tree, tree.root().unwrap(), s);
+            sum
+        });
+        std::hint::black_box(sum);
+    }
 }
 
 fn join_chili(tree_size: usize) {
@@ -131,11 +134,15 @@ fn join_rayon(tree_size: usize) {
 }
 
 fn main() {
-    // use tracing_subscriber::layer::SubscriberExt;
-    // tracing::subscriber::set_global_default(
-    //     tracing_subscriber::registry().with(tracing_tracy::TracyLayer::default()),
-    // )
-    // .expect("Failed to set global default subscriber");
+    //tracing_subscriber::fmt::init();
+    use tracing_subscriber::layer::SubscriberExt;
+    tracing::subscriber::set_global_default(
+        tracing_subscriber::registry().with(tracing_tracy::TracyLayer::default()),
+    )
+    .expect("Failed to set global default subscriber");
+
+    eprintln!("Press Enter to start profiling...");
+    std::io::stdin().read_line(&mut String::new()).unwrap();
 
     let size = std::env::args()
         .nth(2)
@@ -159,6 +166,6 @@ fn main() {
     }
 
     eprintln!("Done!");
-    // wait for user input before exiting
+    // // wait for user input before exiting
     // std::io::stdin().read_line(&mut String::new()).unwrap();
 }