fix rare zkalloc/rayon interraction bug

Author: TomWambsgans

Cargo.lock

@@ -5,6 +5,7 @@ edition.workspace = true
 
 [dependencies]
 libc = "0.2"
+rayon.workspace = true
 
 [lints]
 workspace = true

crates/backend/system-info/Cargo.toml

@@ -9,3 +9,36 @@ pub fn peak_rss_bytes() -> u64 {
     // ru_maxrss unit: bytes on macOS, KiB on Linux.
     if cfg!(target_os = "macos") { max } else { max * 1024 }
 }
+
+/// Number of jobs [`flush_rayon`] pushes. Must exceed
+/// `crossbeam_deque::deque::BLOCK_CAP` (currently 63 —
+/// `crossbeam-deque-0.8.6/src/deque.rs:1191`).
+const RAYON_FLUSH_JOBS: usize = 256;
+
+/// Drain rayon's internal queues so they release any storage allocated during the
+/// previous phase.
+///
+/// Rayon's global pool owns a `crossbeam_deque::Injector`, internally a linked list
+/// of fixed-size blocks (`Block` and `Injector::push` —
+/// `crossbeam-deque-0.8.6/src/deque.rs:1219` and `:1371`). A block is freed only
+/// once its last slot has been consumed.
+///
+/// `rayon::join` from a non-worker thread reaches that injector via
+/// `join` (`rayon-core-1.13.0/src/join/mod.rs:132`) ->
+/// `registry::in_worker` (`registry.rs:946`) ->
+/// `Registry::in_worker_cold` (`:517`) ->
+/// `Registry::inject` (`:428`) -> `Injector::push`.
+///
+/// Under an arena allocator that recycles memory between phases (e.g. `zk-alloc`),
+/// a block allocated *during* a phase points into a slab the next `begin_phase()`
+/// will reuse. The next push then writes a `JobRef` straight through whatever the
+/// application has placed on top, silently corrupting it.
+///
+/// Pushing more than `BLOCK_CAP` jobs while the arena is off forces the Injector                                        
+/// to allocate a fresh tail block (which lands in System), and forces workers to                                      
+/// steal the last slot of every preceding block (which destroys them).
+pub fn flush_rayon() {
+    for _ in 0..RAYON_FLUSH_JOBS {
+        rayon::join(|| {}, || {});
+    }
+}

crates/backend/system-info/src/lib.rs

@@ -7,6 +7,9 @@ description = "Bump+reset arena allocator for ZK proving workloads"
 [dependencies]
 system-info.workspace = true
 
+[dev-dependencies]
+rayon.workspace = true
+
 [target.'cfg(not(all(target_os = "linux", target_arch = "x86_64")))'.dependencies]
 libc = "0.2"
 

crates/backend/zk-alloc/Cargo.toml

@@ -106,8 +106,12 @@ pub fn begin_phase() {
 
 /// Deactivates the arena. New allocations go to the system allocator; existing arena
 /// pointers stay valid until the next `begin_phase()` resets the slabs.
+///
+/// Also calls [`system_info::flush_rayon`] to release any rayon/crossbeam storage
+/// still referencing this phase's arena memory.
 pub fn end_phase() {
     ARENA_ACTIVE.store(false, Ordering::Release);
+    system_info::flush_rayon();
 }
 
 #[cold]

crates/backend/zk-alloc/src/lib.rs

@@ -0,0 +1,26 @@
+//! Regression test for the bug prevented by `system_info::flush_rayon`.
+
+use rayon::prelude::*;
+
+#[global_allocator]
+static A: zk_alloc::ZkAllocator = zk_alloc::ZkAllocator;
+
+#[test]
+fn rayon_does_not_corrupt_zkalloc() {
+    zk_alloc::init();
+    let _: u64 = (0..1_000_000_u64).into_par_iter().sum();
+
+    zk_alloc::begin_phase();
+    for _ in 0..200 {
+        rayon::join(|| {}, || {});
+    }
+    zk_alloc::end_phase();
+
+    zk_alloc::begin_phase();
+    let canary = vec![0xAB_u8; 8192];
+    rayon::join(|| {}, || {});
+    zk_alloc::end_phase();
+
+    let pos = canary.iter().position(|&b| b != 0xAB);
+    assert!(pos.is_none(), "canary corrupted at offset {}", pos.unwrap());
+}