perf(thread_aware)!: partition relocation storage per affinity

[Sander Saares (sandersaares)]

[Copilot speaking]

Motivation

thread_aware::Arc<T, S>::relocate is the hot path of the crate: a thread-per-core runtime calls it on every cross-core handoff, for every Arc<_, PerCore> reachable in the relocated object graph. It took an unconditional exclusive RwLock write lock on every call, even though the overwhelmingly common case is a read-only lookup of an already-populated slot. Every relocation of every clone of a value therefore serialized on one lock, turning a read-only lookup into a process-wide serialization point. (Azure DevOps 7747020.)

Changes

Per-affinity partitioning. The slot table is no longer guarded by a single lock. Each affinity owns its own cache-line-padded RwLock, and the table is sized once to the strategy's fixed slot count, so there is no growth path and no table-wide lock. Relocations into different affinities now touch different locks on different cache lines and never contend: the hit path is a shared read of one slot, and a miss takes only that slot's exclusive lock. The two-stage probe (shared read, escalate to exclusive on a miss) is retained per slot.

Before, a fanout that handed work to every core funnelled every relocation through one lock and throughput was flat regardless of core count. After, it scales with cores:

shape	before	after
16 workers, distinct affinities	2.72 us / 5.8 Melem/s	200 ns / 80 Melem/s
32 workers, distinct affinities	5.50 us / 5.8 Melem/s	290 ns / 108 Melem/s

The subject is a five-layer object tree, the realistic case where a message reaches several thread-aware nodes. Single-threaded relocation and the deref/use path are unchanged; the latter never touched the storage lock in any design.

Deadlock-safe source handling. A miss also records the value the Arc moved away from into the source affinity's slot, so a later relocation back into that affinity finds it instead of re-materializing. That write happens under its own lock and never while the destination lock is held, so two threads relocating in opposite directions cannot each wait for the lock the other holds.

Slot locks never poison. The only caller code run while a slot lock is held is the relocation factory on the miss path. It runs under catch_unwind, which drops the guard before resuming the unwind, so a panicking factory leaves the slot empty and still usable rather than poisoned.

Constructible storage handle. The per-affinity table is crate-private (SlotTable); the public storage::Storage handle exposes an affinity-keyed contract — build with Storage::new, seed slots with insert, read them back with get — without revealing the locking or slot layout, so the representation can change without breaking callers. This lets a caller prepare per-affinity values up front and hand the storage to Arc::from_storage. A Strategy must report the same slot count for every affinity that shares a storage; the built-in strategies do, since processor and memory-region counts are machine properties.

Benchmarks. A concurrent benchmark measures the cost of one relocation while every core relocates at once, batching many relocations behind a single barrier release so the fixed synchronization cost falls into the regression intercept. hit_path and miss_path cover the single-threaded branches under both Criterion and Callgrind.

Structure

Arc<T, S> clones ── share ──> Storage<T, S>            (public, opaque handle)
                                  |
                                  +-- SlotTable: [ CachePadded<RwLock<Option<Arc<T>>>> ; N ]
                                                   slot 0    slot 1   ...  slot N-1
                                                   (one independently-locked slot per affinity)

[Copilot]

Pull request overview

This PR optimizes thread_aware::Arc<T, S>::relocate by replacing the single global lock with per-affinity, cache-line-padded slot locks, reducing cross-core contention on the hot relocation path. It also adds documentation and benchmark coverage to validate and track the new behavior/perf characteristics.

Changes:

• Replace the shared storage lock with an affinity-partitioned SlotTable (OnceLock<Box<[CachePadded<RwLock<...>>]>>) and update relocation logic to use a read-probe + write-on-miss per-slot scheme.
• Add concurrency/panic-safety tests to validate “materialize once” behavior and that slot locks are not left poisoned after a factory panic.
• Add implementation documentation and new Criterion/Callgrind benchmarks for hit/miss and concurrent relocation scenarios.

Reviewed changes

Copilot reviewed 8 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
crates/thread_aware/src/cell/mod.rs	Updates Arc to use the new opaque storage handle and revises relocation to two-stage per-slot locking with unwind safety.
crates/thread_aware/src/cell/storage.rs	Introduces SlotTable and new Storage wrapper (per-affinity slot locks, OnceLock sizing, counting helpers).
crates/thread_aware/src/cell/tests.rs	Updates tests to match new storage shape; adds new concurrency and panic/unpoisoning tests.
crates/thread_aware/docs/implementation.md	New design/locking and benchmarking rationale documentation.
crates/thread_aware/Cargo.toml	Adds crossbeam-utils dependency and registers new benchmark targets.
crates/thread_aware/benches/thread_aware_relocate.rs	New Criterion wall-clock benchmark suite for relocation (hit/miss/concurrent).
crates/thread_aware/benches/thread_aware_relocate_cg.rs	New Callgrind benchmark suite for relocation hit/miss paths.
Cargo.toml	Adds workspace dependency entry for crossbeam-utils.
Cargo.lock	Locks crossbeam-utils for the workspace.

Suppressed comments (1)

crates/thread_aware/src/cell/storage.rs:82

• Because the slot array is initialized once, it’s worth asserting (at least in debug) that Strategy::count is consistent across affinities. Without this, a custom Strategy that varies count() could cause a hard-to-diagnose panic on indexing later.

    /// Returns the slot array, sizing it on first use to hold every affinity.
    fn slots(&self, affinity: Affinity) -> &[Slot<T>] {
        self.slots.get_or_init(|| {
            (0..S::count(affinity))
                .map(|_| CachePadded::new(RwLock::new(None)))
                .collect::<Vec<_>>()
                .into_boxed_slice()
        })
    }

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.0%. Comparing base (9877935) to head (5d17d48).
⚠️ Report is 10 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##             main     #681     +/-   ##
=========================================
  Coverage   100.0%   100.0%             
=========================================
  Files         503      543     +40     
  Lines       57406    60474   +3068     
=========================================
+ Hits        57406    60474   +3068     

Flag	Coverage Δ
linux	?
linux-arm	?
scheduled	?
windows	?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

✅ Version increments look sufficient

cargo semver-checks compared the 1 crate(s) this PR publishes against their previous version-bump commit in git history. Every version increment is sufficient for the detected API changes.

Crate	Baseline	Baseline commit	This PR	Minimum required	Status
thread_aware	0.9.0	7c185b4	0.10.0	0.9.1	✅ ok

This check is informational and does not block the merge.

View the check run

[Copilot]

Pull request overview

Copilot reviewed 8 out of 9 changed files in this pull request and generated no new comments.

Suppressed comments (2)

crates/thread_aware/src/cell/mod.rs:536

• Arc::from_storage is public, but storage::Storage is effectively opaque outside the crate (its constructor and methods are pub(crate)), and Arc does not expose a way to obtain the backing Storage handle. As a result, downstream callers can’t practically use from_storage. Consider exposing a Storage handle accessor on Arc (or alternatively making Storage constructible/manipulable publicly, or making from_storage crate-private if it’s intended to be internal-only).

    /// Creates a new `Arc` from the given storage and the current affinity.
    ///
    /// If the resulting `Arc` is transferred to an affinity which does not have data in the storage,
    /// it will behave like a [`sync::Arc`].
    ///
    /// # Panics
    /// Panics if the storage does not contain data for the current affinity.
    pub fn from_storage(storage: sync::Arc<Storage<T, S>>, current_affinity: Affinity) -> Self {
        let value = storage.get_clone(current_affinity).expect("No data found for the current affinity");

crates/thread_aware/src/cell/tests.rs:764

• In test code, prefer .unwrap() over .expect(...) (the backtrace is typically sufficient, and it keeps tests consistent with the repo’s guidance). Here RACERS is a const so the message isn’t adding useful diagnostics.

        let values = racers.into_iter().map(|racer| racer.join().unwrap()).collect::<Vec<_>>();
        let (first, rest) = values.split_first().expect("RACERS is nonzero");

[Copilot]

Pull request overview

Copilot reviewed 9 out of 10 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (2)

crates/thread_aware/src/cell/storage.rs:40

• The NEVER_POISONED contract is documented as unconditional, but SlotTable<T, _> is generic and performs T: Clone operations while holding a slot lock (e.g., in get_clone / count_where). A panicking Clone would poison the lock, contradicting the docs.

If the intent is that this guarantee is relied on only for Storage (which stores sync::Arc<T>, whose clone doesn’t panic), consider scoping the documentation accordingly so it stays accurate for SlotTable’s generic use.

/// A slot lock is never left poisoned, so acquiring it never fails.
///
/// Poisoning happens when a thread panics while holding the lock and the guard is
/// dropped during unwinding. The operations run under a slot lock — cloning,
/// storing and comparing the reference-counted handle it holds — cannot unwind,

crates/thread_aware/src/cell/mod.rs:595

• strong_count can underflow under concurrent relocation because it reads raw first and then counts internal references without any global synchronization; new internal references may be published between those reads, making internal > raw and wrapping the subtraction to a huge usize.

Since strong_count is inherently racy, prefer making it robust by taking the internal count first and using saturating_sub (or checked_sub with a fallback) to avoid wraparound.

    pub fn strong_count(this: &Self) -> usize {
        let raw = sync::Arc::strong_count(&this.value);
        let internal = this.storage.count_where(|stored| sync::Arc::ptr_eq(stored, &this.value));
        raw - internal

[Copilot]

Pull request overview

Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (3)

crates/thread_aware/src/cell/mod.rs:596

• Arc::strong_count subtracts an internal-slot count from the raw Arc strong count. With the new per-slot locking, count_where is explicitly racy under concurrent relocation, so internal can transiently exceed the previously-captured raw, causing an integer underflow and returning a huge number. Use a saturating subtraction (or otherwise guard against underflow) so this remains safe under races.

    pub fn strong_count(this: &Self) -> usize {
        let raw = sync::Arc::strong_count(&this.value);
        let internal = this.storage.count_where(|stored| sync::Arc::ptr_eq(stored, &this.value));
        raw - internal
    }

crates/thread_aware/benches/thread_aware_relocate.rs:46

• This benchmark imports many_cpus::SystemHardware, but the crate's many_cpus dependency is only enabled by the threads feature (not by std), and the bench target is currently declared with required-features = ["std"]. As written, the bench won't compile when built with only std. Prefer using std::thread::available_parallelism() (std-only) and drop the many_cpus dependency here.

use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main};
use many_cpus::SystemHardware;
use thread_aware::affinity::{Affinity, pinned_affinities};
use thread_aware::{Arc, PerCore, ThreadAware, Unaware};

crates/thread_aware/benches/thread_aware_relocate.rs:428

• After removing many_cpus, compute the processor count using std::thread::available_parallelism() so the concurrent group still sizes its worker pool correctly without requiring the threads feature.

    let saturated = SystemHardware::current().processors().len();

[Copilot]

Pull request overview

Copilot reviewed 9 out of 10 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

crates/thread_aware/Cargo.toml:140

• This bench also uses #[derive(ThreadAware)] in its Linux-only module, so it won’t compile on Linux when std is enabled without derive. Adding derive to required-features will make the target auto-skip in that feature set.

[[bench]]
name = "thread_aware_relocate_cg"
harness = false
required-features = ["std"]

crates/thread_aware/src/cell/storage.rs:230

• The unit-test module here is missing #[cfg_attr(coverage_nightly, coverage(off))]. This repo uses a 100% coverage gate and consistently marks #[cfg(test)] mod tests blocks with coverage(off) so the test harness itself doesn’t count against coverage (e.g., crates/thread_aware/src/registry.rs:218-220).

#[cfg(test)]
mod tests {

[Copilot]

Pull request overview

Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (1)

crates/thread_aware/src/cell/storage.rs:19

• Grammar in the public trait docs: use “an affinity-aware manner” (vowel sound) instead of “a affinity-aware manner”.

/// A strategy for storing data in a affinity-aware manner.

[Sander Saares (sandersaares)]

[Copilot speaking]

Published 3 findings. No finding follows up on an existing discussion thread.

See diagnostics

Diagnostic	Value
Cache	Miss