fix(crypto): run every proof on one dedicated prover thread

[MegaRedHand]

🗒️ Description / Motivation

devnet-5 aggregators running --prover-arena were OOM-killed about every 40 minutes (77-98 restarts in ~65 h, each at ~29 GB anon RSS against a 28 GiB memory limit). Non-aggregators, and aggregators without the arena, stayed flat.

The cause is how leanVM's arena (zk_alloc) interacts with the threads ethlambda proves from:

• The arena gives each thread that allocates during a proof its own slab, claimed once and never released, not even when the thread exits. A new proof resets the cursor but does not return the pages, so each slab stays resident at its owner's peak.
• leanVM's pool workers are spawned once and live forever, so their slabs are reused. The thread that drives a proof is the caller, and its slab is the one that fills (1-4.3 GiB each, measured from /proc/PID/pagemap; pool-worker slabs stayed near 0).
• ethlambda proved on whichever thread asked: the interval-2 aggregation worker (spawn_blocking, a thread tokio replaces once it idles) and the actor (block building, re-aggregation). Every new calling thread pinned another proof peak, up to the 24-slab cap.

What Changed

• crates/common/crypto/src/lib.rs: prove() replaces the acquire_prover() mutex. It sends the proving closure to one long-lived leanvm-prover thread and blocks the caller until the proof returns. Each of the five proving sites wraps only its aggregate call; decoding and argument conversion stay on the caller.
• crates/common/crypto/tests/arena_slabs.rs (new): with the arena on, proves once from each of 5 fresh threads and checks that zk_alloc::stats().threads (slabs handed out) does not grow after the first.
• zk_alloc is a new dev-dependency, pinned to the same leanVM rev as leanvm, because the facade does not re-export it. The comment in Cargo.toml says the two revs must stay equal: only then is it the same crate instance, with the same statics.
• crates/common/crypto/tests/common/mod.rs: the key helper arena.rs had, now shared by both arena tests.
• CLAUDE.md: a note on why all proving goes through one thread.

Correctness / Behavior Guarantees

• One proof at a time is still enforced: the single thread runs jobs in order, which is what the mutex guaranteed.
• Panics: a panicking proof is caught on the prover thread (logged with error!) and re-raised on the caller with resume_unwind, so callers see the same panic as before and later proofs still run. The mutex version recovered from poisoning for the same reason.
• Stack: the prover thread uses std's default stack, the same size as tokio's worker and blocking threads that proved until now. The real proving tests pass on it in release-fast.
• Latency: callers already blocked for the whole proof while holding the mutex; now they block on a channel instead. Proofs are not slower.
• A job must not call prove() again (the prover thread would wait on itself). No caller does this; the doc comment says so.
• No change without --prover-arena, apart from the thread the proof runs on.

Tests Added / Run

Test	main	This PR
arena_slabs (ignored, slow): slab count after 4 more proofs from fresh threads	10 → 14 ❌	10 → 10 ✅
prove_runs_every_job_on_one_thread (unit)	n/a	✅
prove_reraises_a_panic_and_keeps_proving (unit)	n/a	✅
test_aggregate_single_signature, test_type_2_merge_verify_split_round_trip, aggregates_on_the_arena_when_enabled (ignored, real proofs)		✅

make fmt && make lint
cargo test --workspace --profile release-fast --lib --bins
cargo test -p ethlambda-crypto --profile release-fast --test arena --test arena_slabs -- --ignored
cargo test -p ethlambda-crypto --profile release-fast --lib -- --ignored test_aggregate_single_signature test_type_2_merge_verify_split_round_trip

test_setup_is_idempotent no longer acquires the permit twice, since the permit is gone.

Related Issues / PRs

• Hand an exited thread's arena slab to the next thread leanEthereum/leanVM#346 fixes part of this inside leanVM: it hands an exited thread's slab to the next thread. It is merged into riscv-exploration, not main, and does not cover long-lived threads that take turns proving.
• Left over in leanVM (out of scope here): verifier threads also claim slabs. verify reaches mle_eval_par → eq_table_arena on the calling thread, which uses the arena whenever another thread has a proof's phase open. Each such slab touches only KiB, so it is not an OOM source.

✅ Verification Checklist

• Ran make fmt — clean
• Ran make lint (clippy with -D warnings) — clean
• Ran make test (cargo test --workspace --profile release-fast) — ran --lib --bins plus the crypto tests above; leanSpec spec tests not run

[github-actions]

🤖 Kimi Code Review

I'll review this PR which replaces a mutex-based prover permit with a dedicated single prover thread to prevent OOM issues from arena slab accumulation.

Summary

The change is architecturally sound: moving from Mutex<()> to a dedicated mpsc::channel-driven prover thread correctly bounds leanVM's per-thread arena slabs. However, there are several issues to address.

---

Critical Issues

1. Unbounded channel allows unbounded memory growth under load

File: crates/common/crypto/src/lib.rs, lines 131-132

let (tx, rx) = mpsc::channel::<ProverJob>();

mpsc::channel is unbounded. Under heavy load, if proofs queue faster than the prover thread processes them, this grows without bound. Given this is consensus-critical infrastructure that must remain available, a bounded channel with backpressure is safer.

Fix: Use mpsc::sync_channel(BOUND) or document why unbounded is acceptable (the prove() function blocks on reply_rx.recv(), so callers naturally backpressure, but the queue itself can still grow).

---

2. reply_tx.send(outcome) uses let _ = which silently drops send errors

File: crates/common/crypto/src/lib.rs, lines 108-109

let _ = reply_tx.send(outcome);

The comment claims "the caller waits on the reply until it arrives, so the send cannot fail." This is incorrect reasoning. The send can fail if the caller has panicked or been cancelled (e.g., tokio::task::abort), dropping reply_rx. In that case, the prover thread silently drops the result and continues, but the job's panic payload is lost.

More critically: if the caller is cancelled after the prover thread starts but before this send, the prover thread leaks the result and the caller may not observe the panic.

Fix: At minimum log the error. Better: this is actually fine for correctness (the caller is gone), but the comment should be honest. However, for panic payloads, consider whether silently dropping a panic is acceptable.

---

3. Thread spawning uses .expect() in get_or_init — panics poison the OnceLock

File: crates/common/crypto/src/lib.rs, lines 137-138

.spawn(move || rx.into_iter().for_each(|job| job()))
.expect("failed to spawn the leanVM prover thread");

If thread spawning fails (e.g., RLIMIT_NPROC, cgroups limits), the OnceLock is poisoned with an uninited state. The expect panics, and subsequent calls to prover_queue() will retry get_or_init (which is fine for OnceLock), but this crashes the process on the first failure.

Fix: Consider returning a Result from prove() and propagating thread spawn failures. Given this is consensus code, graceful degradation is preferable to crashing.

---

Security & Correctness Issues

4. AssertUnwindSafe is technically unsound if T's drop panics

File: crates/common/crypto/src/lib.rs, lines 103-104

let outcome = panic::catch_unwind(AssertUnwindSafe(job));

The AssertUnwindSafe is required because F: FnOnce() doesn't implement UnwindSafe. This is generally acceptable for a closure that's fully owned and moved. However, if T's Drop implementation panics, double-panic occurs. This is a standard Rust issue, but worth noting.

More importantly: if job captures !UnwindSafe types (like MutexGuard), unwinding could leave them in a bad state. The closure captures raw_xmss, children_native, etc. — these are plain vectors, so this is fine in practice.

Suggestion: Add a comment explaining why AssertUnwindSafe is sound (all captured types are owned, simple data).

---

5. No timeout on reply_rx.recv() — prover thread deadlock blocks caller forever

File: crates/common/crypto/src/lib.rs, lines 114-115

reply_rx.recv()
    .expect("the prover thread replies to every job")

If the prover thread deadlocks inside leanVM (bug, infinite loop), the caller blocks forever. In an async context (tokio blocking pool), this consumes a thread permanently.

Fix: Consider a generous timeout with recv_timeout and panic/restart logic, or document that leanVM's aggregate is expected to terminate.

---

Performance Issues

6. mpsc::sync_channel(1) in prove() creates a channel per proof

File: crates/common/crypto/src/lib.rs, lines 99-100

let (reply_tx, reply_rx) = mpsc::sync_channel(1);

This allocates a new channel for every proof. Given proofs are expensive, this is negligible overhead, but oneshot channels (e.g., tokio::sync::oneshot or crossbeam::channel) are more idiomatic.

Actually, re-reading: this is fine. The allocation is dwarfed by proof cost.

---

Code Quality Issues

7. Test prove_runs_every_job_on_one_thread has race potential

File: crates/common/crypto/src/lib.rs, lines 688-704

let askers: Vec<_> = (0..4)
    .map(|_| {
        thread::spawn(|| {
            let asker = thread::current().id();
            let prover = prove(|| thread::current().id());
            (asker, prover)
        })
    })
    .collect();

This test assumes all 4 threads are spawned and call prove() concurrently enough that they'd be queued. However, thread::spawn + prove() may execute sequentially enough that the first finishes before the second starts. The test still passes (same prover thread), but doesn't fully exercise the serialization path.

Fix: Add a barrier or thread::yield_now loop to ensure concurrent prove() calls. Or accept this as sufficient.

---

8. arena_slabs.rs test has misleading assertion message

File: crates/common/crypto/tests/arena_slabs.rs, lines 44-48

assert_eq!(
    zk_alloc::stats().threads,
    slabs,
    "proving from new threads claimed new slabs"
);

If this fails, it means new slabs were claimed, but the message says "claimed new slabs" which reads like a statement of fact rather than failure description. The message is actually fine grammatically ("expected X, found Y: proving from new threads claimed new slabs"), but slightly awkward.

Suggestion: "new threads unexpectedly claimed additional arena slabs" for clarity.

---

9. Missing zk_alloc dev-dependency documentation in workspace

File: Cargo.toml, lines 90-93

# leanVM's prover arena, for tests that read its slab accounting (the facade does
# not re-export it). Keep the rev equal to leanvm's: only then is this the same
# crate instance, with the same statics, as the one leanvm proves on.

Good comment, but this is a critical invariant for correctness. If these revs diverge, the test reads from a different zk_alloc static than leanVM writes to, making the test pass falsely (it sees 0 slabs, thinks everything's fine).

Suggestion: Add a build.rs or CI check that enforces leanvm.rev == zk_alloc.rev. Or use a workspace variable:

[workspace]
leanvm-rev = "48a904208d682848dac0e18ef8b01ebfc40df9ad"

[dependencies]
leanvm = { git = "...", rev = { workspace = "leanvm-rev" } }  # syntax may vary

Actually, Cargo doesn't support workspace-level variables for rev. Consider a [workspace.metadata] section and a CI lint.

---

Positive Observations

1. Good: The prove() function correctly moves expensive work (decoding) outside the prover thread, only wrapping the actual aggregate call.

2. Good: Panic isolation prevents one bad proof from killing the prover thread permanently.

3. Good: The arena_slabs.rs integration test correctly validates the core invariant in a separate process.

4. Good: Thread name "leanvm-prover" aids debugging.

5. Good: The common/mod.rs extraction removes duplication between integration tests.

---

Minor Suggestions

10. Consider std::thread::JoinHandle storage for graceful shutdown

The prover thread is leaked (no JoinHandle stored). This is fine for process lifetime, but prevents clean shutdown in tests or embedding scenarios.

11. init_leanvm comment update

The doc comment in CLAUDE.md is good, but init_leanvm's rustdoc could mention the arena/thread interaction.

---

Final Assessment

Category	Verdict
Correctness	Acceptable with fixes for items 1, 2, 3
Security	No direct vulnerabilities; resource exhaustion risk (item 1)
Performance	Neutral to slightly positive (thread creation amortized)
Maintainability	Improved with better comments on AssertUnwindSafe

Required before merge: Address item 1 (unbounded channel) or justify why it's safe. Item 2 (silent send error) should at minimum have corrected comment or logging. Item 3 (thread spawn panic) should be considered for production hardening.

---

Automated review by Kimi (Moonshot AI) · kimi-k2.6 · custom prompt

[github-actions]

🤖 Claude Code Review

Review of PR 641: run every proof on one dedicated prover thread

Verdict: The change is sound. I read the diff only and did not build or run the tests. I have no blocking concerns.

What it does: prove() in crates/common/crypto/src/lib.rs replaces the Mutex<()> permit. A single long-lived leanvm-prover thread now runs every proof, which does two things:

• It serializes proofs, as the mutex did.
• It bounds the arena's per-thread slabs, which was the OOM cause.

Correct:

• Panic handling: catch_unwind runs on the prover thread and resume_unwind re-raises on the caller. A bad proof can't kill the worker, and the mutex-poison recovery is no longer needed.
• expect calls: The queue send and the reply recv can only fail if the prover thread exits. The loop never returns and every job replies, so neither can fail.
• Lifetimes: The 'static and move closures are handled correctly at each call site. split_type_2_by_message builds ClaimSelection inside the closure, so the borrow of kept stays valid.
• Tests: The unit tests cover single-thread execution and panic recovery. tests/arena_slabs.rs checks the slab count directly. It runs in its own process, is #[ignore]d like the other slow crypto tests, and asserts slabs > 0 to catch a duplicate zk_alloc instance.
• Dependency pin: Pinning zk_alloc to the same rev as leanvm is documented, and it is a dev-dependency only.

Minor points:

1. Rev drift: The two git dependencies share a rev only by convention (Cargo.toml:90-93). If someone bumps leanvm and forgets zk_alloc, Cargo would build two crate instances. The slabs > 0 assertion would then catch it, but only in an ignored test. A [workspace.dependencies] comment on the leanvm line pointing at zk_alloc would help, or a CI check.
2. Reentrancy: A nested prove call deadlocks silently. The doc comment warns about this. A debug_assert! comparing the current thread's name or ID with the prover's would turn a hang into a panic. This is optional, since no current caller nests.
3. Memory on the prover thread: Inputs such as children_native and raw_xmss are dropped on the prover thread, and their allocations were made on the caller. With the arena enabled, check that frees across threads are safe in zk_alloc. I expect they are, since the previous code already crossed threads through tokio's blocking pool.
4. Stack size: The prover thread uses std's default 2 MiB stack. CLAUDE.md mentions stack overflows in signature work. Proving previously ran on tokio blocking threads, which also default to 2 MiB, so behavior is unchanged. A named constant or explicit stack_size would make this deliberate if leanVM's recursion ever deepens.
5. Latency: Callers block while earlier jobs finish, as they did with the mutex. There is no fairness change. The queue is unbounded, but callers block on the reply, so depth is bounded by the number of callers.

The CLAUDE.md note accurately describes the new behavior.

---

Automated review by Claude (Anthropic) · sonnet · custom prompt

[github-actions]

🤖 Codex Code Review

Looks good overall. The main change — serializing all leanVM proving onto a dedicated thread to cap arena slab growth — is sensible, and the added tests cover the key behavioral contract.

• No correctness bug jumped out in the consensus-critical paths listed; this PR is scoped to ethlambda-crypto proving orchestration rather than fork choice / attestation / STF logic.
• Security-wise, catching panics inside the prover worker and re-raising them on the caller is a good containment choice for a shared proving service.

One thing I’d still flag:

• Potential self-deadlock on re-entrant proving: prove() now always enqueues onto the single prover thread and blocks on reply_rx.recv(). If any code ever calls prove() from a closure already running on that prover thread, it deadlocks permanently rather than panicking fast. The doc comment mentions this constraint, but the implementation does not enforce it. Consider detecting the prover thread and asserting early in crates/common/crypto/src/lib.rs:102 / crates/common/crypto/src/lib.rs:132 so misuse fails loudly instead of hanging the process.

Minor maintainability note:

• The test dependency on zk_alloc being “the same crate instance” as leanvm is subtle and brittle across dependency graph changes. The comment in Cargo.toml:90 helps, but this coupling is worth preserving carefully in future refactors.

If you want, I can also turn this into a ready-to-post GitHub review comment.

---

Automated review by OpenAI Codex · gpt-5.4 · custom prompt

[pablodeymo]

Suggested change

	match reply_rx
	.recv()
	.expect("the prover thread replies to every job")
	{
	Ok(value) => value,
	Err(payload) => panic::resume_unwind(payload),
	}
	match reply_rx
	.recv()
	.expect("the prover thread replies to every job")
	.unwrap_or_else(|_| panic::resume_unwind(payload))

[pablodeymo]

Small note on the stack: 2 MiB matches the tokio threads the node proved on, but not every caller. ethlambda benchmark runs synchronously on the main thread (bin/ethlambda/src/main.rs:72, 8 MiB stack), so its proofs (benchmark/corpus.rs:155 plus the block building it drives) now drop to 2 MiB. Same for the actor's proofs under shadow-integration, where the runtime is current_thread on main.

Probably fine since production already proves on 2 MiB, but since aggregation has hit stack overflows before (the reason tests run under release-fast), maybe set an explicit .stack_size(...) here, or at least soften the "the size tokio gives its own threads" wording in the doc comment and PR body?