Use wrapping addition for JEC

rayon-core/src/join/test.rs

@@ -125,3 +125,21 @@ fn join_context_second() {
     assert!(!a_migrated);
     assert!(b_migrated);
 }
+
+#[test]
+fn join_counter_overflow() {
+    const MAX: u32 = 500_000;
+
+    let mut i = 0;
+    let mut j = 0;
+    let pool = ThreadPoolBuilder::new().num_threads(2).build().unwrap();
+
+    // Hammer on join a bunch of times -- used to hit overflow debug-assertions
+    // in JEC on 32-bit targets: https://github.com/rayon-rs/rayon/issues/797
+    for _ in 0..MAX {
+        pool.join(|| i += 1, || j += 1);
+    }
+
+    assert_eq!(i, MAX);
+    assert_eq!(j, MAX);
+}

rayon-core/src/sleep/counters.rs

@@ -126,7 +126,7 @@ impl AtomicCounters {
         loop {
             let old_value = self.load(Ordering::SeqCst);
             if increment_when(old_value.jobs_counter()) {
-                let new_value = old_value.plus(ONE_JEC);
+                let new_value = old_value.increment_jobs_counter();
                 if self.try_exchange(old_value, new_value, Ordering::SeqCst) {
                     return new_value;
                 }
@@ -221,9 +221,11 @@ impl Counters {
     }
 
     #[inline]
-    fn plus(self, word: usize) -> Counters {
+    fn increment_jobs_counter(self) -> Counters {
+        // We can freely add to JEC because it occupies the most significant bits.
+        // Thus it doesn't overflow into the other counters, just wraps itself.
         Counters {
-            word: self.word + word,
+            word: self.word.wrapping_add(ONE_JEC),
         }
     }