From 940c6812224df85e3d704e0db8f4607c2cb1fe69 Mon Sep 17 00:00:00 2001
From: Janis <janis@nirgendwo.xyz>
Date: Fri, 20 Jun 2025 12:16:51 +0200
Subject: [PATCH] threadlocal count for join/heatbeat

---
 src/praetor/mod.rs | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/praetor/mod.rs b/src/praetor/mod.rs
index ae267a1..55c9e08 100644
--- a/src/praetor/mod.rs
+++ b/src/praetor/mod.rs
@@ -1013,7 +1013,7 @@ mod job {
 
 use std::{
     any::Any,
-    cell::UnsafeCell,
+    cell::{Cell, UnsafeCell},
     collections::BTreeMap,
     future::Future,
     hint::cold_path,
@@ -1084,8 +1084,6 @@ struct WorkerThread {
 }
 
 pub struct Scope<'scope> {
-    // counter for join_every_* function
-    join_count: AtomicUsize,
     // latch to wait on before the scope finishes
     job_counter: JobCounter,
     // local threadpool
@@ -1572,9 +1570,20 @@ impl<'scope> Scope<'scope> {
         A: FnOnce(&Self) -> RA + Send,
         B: FnOnce(&Self) -> RB + Send,
     {
-        let count = self.join_count.load(Ordering::Relaxed);
-        self.join_count
-            .store(count.wrapping_add(1) % TIMES, Ordering::Relaxed);
+        thread_local! {
+            static JOIN_COUNT: Cell<usize> = Cell::new(0);
+        }
+
+        // a threadlocal counter is much faster than a sync atomic counter
+        let count = JOIN_COUNT.with(|count| {
+            count.set(count.get().wrapping_add(1) % TIMES);
+            count.get()
+        });
+
+        // let count = self.join_count.load(Ordering::Relaxed);
+        // self.join_count
+        //     .store(count.wrapping_add(1) % TIMES, Ordering::Relaxed);
+
         // let count = self
         //     .join_count
         //     .update(Ordering::Relaxed, Ordering::Relaxed, |n| {
@@ -1645,7 +1654,6 @@ impl<'scope> Scope<'scope> {
     fn from_context(ctx: Arc<Context>) -> Self {
         Self {
             context: ctx,
-            join_count: AtomicUsize::new(0),
             job_counter: JobCounter::default(),
             panic: AtomicPtr::new(ptr::null_mut()),
             _pd: PhantomData,