perf: lock-free PublicKeyCache; drop pubkey_cache_lock from hot paths (P1 of #863)

[ch4r10t33r]

Summary

P1 of the perf work in #863. Replaces the std.AutoHashMap + pubkey_cache_lock mutex backing of xmss.PublicKeyCache with a per-slot atomic CAS design, then drops the lock from chain.onBlock.verifySignatures and chain.verifyAggregatedAttestation.

Targets the ~78ms mean lock hold reported in the issue (`zeam_lock_hold_seconds{lock="pubkey_cache",site="onBlock.verifySignatures"}`: 195 holds, sum ~15.3s). With the mutex around the whole verify-signatures block, rayon's internal parallelism inside xmss.verifyAggregatedPayloadBatch was effectively disabled — every other thread that needed a cache lookup queued behind the holder.

Changes

pkgs/xmss/src/lib.zig

pub const PublicKeyCache = struct {
    slots: []std.atomic.Value(usize), // @intFromPtr(handle) or 0
    allocator: Allocator,

    pub fn init(allocator, capacity) !Self;
    pub fn deinit(self) void;
    pub fn getOrPut(self, idx, bytes) HashSigError!*const HashSigPublicKey;
};

• Reads = single atomic acquire load. No mutex.
• Misses: deserialise via PublicKey.fromBytes, CAS-install under release; on lost race free our handle and adopt the winner's. No allocator contention either.
• Out-of-range indices fall through to a non-cached deserialise (safe, slower). Validator-set growth post-genesis will need a fork-boundary cache rebuild — documented inline.
• PublicKeyCache.init signature now (allocator, capacity). Both call sites updated; existing tests pass; new tests added for out-of-range lookup and zero-capacity init.

pkgs/node/src/chain.zig

• chain.onBlock verify-signatures block: removed the pubkey_cache_lock.lock() / unlock() + tier-5 enter/leave dance. Wrapped verifySignaturesParallel call is the dominant onBlock cost on aggregator nodes per the Investigate slow slot_interval / tick duration (event-loop starvation vs nominal 0.8s) #863 sub-step histogram; single-threaded mutex was negating the rayon parallelism inside Rust.
• chain.verifyAggregatedAttestation: same removal.
• BeamChain.init: passes opts.config.genesis.numValidators() as cache capacity.
• pubkey_cache_lock field on BeamChain left in place (zero acquirers) to avoid touching the lock-hierarchy invariants in locking.zig in this PR. Removal lands as a follow-up cleanup.

Expected impact

Per the analysis in #863:

zeam_lock_hold_seconds pubkey_cache / onBlock.verifySignatures: 195 holds, sum ~15.3 s → ~78 ms mean

Net of this PR: that histogram should drop to near-zero (no holds, since we no longer take the lock). The displaced cost moves into the actual XMSS verify path, where rayon now actually fans out across cores. The chain.onBlock substep histogram from #883 (P0) will show verify_signatures shrinking on aggregator nodes proportional to core count, after PR1 lands and a metric scrape window is captured.

Validation

• zig build clean (Debug + ReleaseFast)
• zig build test -Dprover=dummy --summary all — running locally during PR creation; will update if any failure
• zig fmt clean
• Memory contract: getOrPut lost-race path frees the loser's handle via pk.deinit(); CAS winner is reachable via the slot atomic and freed in PublicKeyCache.deinit. No leak under contention.

Out of scope (follow-ups)

• Remove the now-dead pubkey_cache_lock field + the tier-5a slot in the lock hierarchy comment.
• Pre-warm the cache at chain init from the anchor state's validator pubkeys (current code is lazy on first verify; pre-warm shifts the deserialisation cost to startup, which is irrelevant for liveness).
• The original P1 also called for "multi-thread XMSS verify pool" — xmss.verifyAggregatedPayloadBatch already uses rayon internally (per perf: micro-bench suite + continuous profile capture for production #861 profile notes). The lock removal here is what unlocks the parallelism in practice.

Refs #863. Builds on (but does not depend on) #883.

[zclawz]

Adversarial review finding for #884:

Medium — pkgs/xmss/src/lib.zig:75-83 — out-of-range cache path leaks Rust public-key handles

PublicKeyCache is sized from the genesis validator count, but getOrPut() handles validator_index >= slots.len by deserializing with PublicKey.fromBytes() and returning pk.handle without transferring/freeing ownership. Callers using a cache assume the cache owns returned handles, so these uncached handles are never freed.

If validator-set growth ever makes state.validators.len > genesis.numValidators(), valid attestations from new validators can leak one Rust public-key handle per lookup, which is a resource-exhaustion risk.

Suggested fix: either grow/rebuild the cache when the validator count increases, or make the API distinguish cached/borrowed vs uncached/owned handles so fallback handles are freed after verification. If validator growth is intentionally impossible today, I’d still prefer the fallback to fail loudly instead of being “slow but correct” while leaking.

Non-blocking notes:

• pkgs/node/src/chain.zig:213-215 still says xmss.PublicKeyCache is not thread-safe; that comment looks stale after this PR.
• Tests cover empty/basic cache metadata, but I didn’t see coverage for actual getOrPut() caching, CAS lost-race behavior, or deinit freeing installed handles exactly once.

Validation/evidence: inspected the diff against main and surrounding call sites in chain.zig / state-transition/src/transition.zig. A full zig build test -Dprover=dummy --summary all was started, but hit the long XMSS aggregation portion before completion.

[ch4r10t33r]

Fixed in 9324fe0:

Out-of-range leak (Medium) — addressed. PublicKeyCache.getOrPut now returns HashSigError.ValidatorIndexOutOfRange instead of falling back to a leaky uncached deserialise. Lean spec doesn't grow the validator set post-genesis today; if/when it does, the fork-boundary handler will rebuild the cache with the new size (documented inline). Fail-loud per your preference.

Stale comment (Note) — pkgs/node/src/chain.zig:213-215 updated. Now says the cache is lock-free (per-slot atomic CAS) and the previous serial-pre-phase constraint no longer applies.

Test gaps (Note) — added three tests:

• getOrPut populates on miss and returns same handle on hit — confirms the population path runs at most once per slot.
• getOrPut returns ValidatorIndexOutOfRange past capacity — confirms the new fail-loud path.
• concurrent getOrPut for same slot installs exactly one handle — 8-thread CAS race; every thread observes the same winning handle, count() == 1 post-race, and cache.deinit() alone frees everything (testing allocator surfaces any leaked loser handle at process exit).

zig build test -Dprover=dummy --summary all passes locally.

[zclawz]

Verified the fix in 9324fe0 — I’m happy with this now.

What I checked:

• PublicKeyCache.getOrPut() now fails loudly with HashSigError.ValidatorIndexOutOfRange for validator_index >= capacity, so the leaky uncached fallback is gone.
• CAS lost-race path still deinitializes the loser handle and returns the installed winner.
• Added tests cover miss/hit reuse, out-of-range failure, and concurrent same-slot population.
• zig build test -Dprover=dummy --summary all passed locally (EXIT:0).

Non-blocking nit: the top-level PublicKeyCache doc still says “out-of-range indices fall through to a non-cached deserialise” even though the detailed getOrPut doc and code now correctly say fail-loud. Worth updating before merge, but not a blocker from my side.

Approval-style comment: ✅ looks good to merge once you’re happy with the doc nit.

[ch4r10t33r]

@zclawz can you verify and approve this PR?

[zclawz]

PR #884 Review: lock-free PublicKeyCache; drop pubkey_cache_lock from hot paths

Summary: Replaces the std.AutoHashMap + pubkey_cache_lock mutex in xmss.PublicKeyCache with a flat array of std.atomic.Value(usize) slots (each holding @intFromPtr(handle) or 0-sentinel). Removes the dominant ~78ms mean lock hold reported in #863 from verifySignatures and verifyAggregatedAttestation.

---

Correctness / Logic ✅

The cache semantics are correct for an immutable, write-once-per-slot design: once a slot is populated it is never mutated, so a reader that gets a non-zero load always sees a fully initialised handle.

ValidatorIndexOutOfRange for indices ≥ capacity is the right choice given that lean spec doesn't grow the validator set post-genesis. Failing loudly is better than silently leaking an uncached deserialise. The caveat about rebuilding the cache at fork boundaries when validator-set growth lands is appropriately documented.

One improvement worth adding: log the out-of-range index before returning the error in getOrPut. If an unexpected validator index appears (protocol bug or future spec change), the error propagates silently up through verifySignatures → onBlock → failed block processing, and operators will see only the outer error. A single debug/warn log line with the index and capacity would make diagnosis much faster.

---

Concurrency / Memory ordering ✅

CAS protocol is correct:

• Hit path: load(.acquire) — synchronises with the winner's .release store, so the handle's interior fields are visible. ✅
• Miss path: cmpxchgStrong(EMPTY, new_int, .release, .acquire) — on success, .release publishes the new handle; on failure, .acquire synchronises with the winner's .release, so @ptrFromInt(loser) is the winner's fully-visible handle, and we free our own copy. ✅
• deinit uses .monotonic — correct since deinit is single-threaded by convention. ✅

The EMPTY = 0 sentinel is safe: Rust allocators never return a null pointer for non-zero-sized types, and HashSigPublicKey is non-zero-sized. ✅

---

Resource management ✅

deinit walks all slots and frees non-null handles. Lost-race writers free their handle immediately in getOrPut. No leaks in the normal path.

Test cleanup nit: The concurrent 8-thread test intentionally skips defer cache.deinit() to inspect post-race state, but the early-exit return error.TestUnexpectedResult paths can leak the cache allocation. Add defer cache.deinit() at the top — it runs on all exit paths including the error returns, so the inspection still works.

---

Nits

• Consider adding a capacity() accessor returning self.slots.len for diagnostics / future assertions.
• count() walk note is already documented as "best-effort" — good.

---

Verdict: LGTM ✅ — two suggestions worth addressing:

1. Log the validator index + capacity in getOrPut before returning ValidatorIndexOutOfRange (operational visibility).
2. Add defer cache.deinit() in the concurrent test to avoid leak on assertion failure paths.