Skip to main content

proxima_core/
per_core.rs

1//! Per-core sharded storage: `count` slots, each emitting thread/core routed to
2//! one to cut contention on a shared structure (e.g. one MPMC [`ring`] per core).
3//!
4//! The **primitive** is `slot(core_id)` + `count()` — pure, no_std, no ambient
5//! state: the caller says which core it is. `local()` is a **std convenience**
6//! (behind `feature = "std"`) that assigns a sticky per-thread slot via TLS; a
7//! bare-metal caller instead passes an id from its runtime (prime's per-core id,
8//! a hardware CPU id) to `slot`. Two storage tiers: heap-backed [`PerCore`]
9//! (`Vec`, `feature = "alloc"`) and inline no-alloc [`StaticPerCore`] (`[T; N]`).
10//!
11//! [`ring`]: crate::ring
12
13#[cfg(feature = "alloc")]
14use alloc::vec::Vec;
15#[cfg(feature = "std")]
16use core::sync::atomic::{AtomicUsize, Ordering};
17#[cfg(feature = "std")]
18use std::cell::Cell;
19
20#[cfg(feature = "std")]
21std::thread_local! {
22    // a per-thread monotonic ticket, NOT a pre-resolved slot index — modded by
23    // `count` at access so one ticket maps into any `PerCore` regardless of its
24    // slot count (caching `ticket % count` would index a smaller one out of
25    // bounds). Sticky for cache locality: a thread keeps landing on the same slot.
26    static CORE_TICKET: Cell<Option<usize>> = const { Cell::new(None) };
27}
28
29#[cfg(feature = "std")]
30static NEXT_CORE: AtomicUsize = AtomicUsize::new(0);
31
32// std-only routing: assign this thread a sticky ticket, map it into `count` slots.
33#[cfg(feature = "std")]
34fn local_index(count: usize) -> usize {
35    let ticket = CORE_TICKET.with(|cell| {
36        cell.get().unwrap_or_else(|| {
37            let assigned = NEXT_CORE.fetch_add(1, Ordering::Relaxed);
38            cell.set(Some(assigned));
39            assigned
40        })
41    });
42    ticket % count
43}
44
45/// Heap-backed per-core storage (`Vec`, runtime slot count). See the module docs
46/// for the `slot`/`local` split; [`StaticPerCore`] is the inline no-alloc tier.
47#[cfg(feature = "alloc")]
48pub struct PerCore<T> {
49    slots: Vec<T>,
50    count: usize,
51}
52
53#[cfg(feature = "alloc")]
54impl<T> PerCore<T> {
55    pub fn new_with(count: usize, factory: impl FnMut(usize) -> T) -> Self {
56        let slots = (0..count).map(factory).collect();
57        Self { slots, count }
58    }
59
60    pub fn from_vec(slots: Vec<T>) -> Self {
61        let count = slots.len();
62        Self { slots, count }
63    }
64
65    /// The calling thread's slot (std only). Threads spread round-robin; when
66    /// emitters outnumber slots, several share one — safe for a multi-producer
67    /// structure like [`crate::ring::Ring`], just more contended.
68    #[cfg(feature = "std")]
69    pub fn local(&self) -> &T {
70        &self.slots[local_index(self.count)]
71    }
72
73    pub fn slot(&self, index: usize) -> &T {
74        &self.slots[index]
75    }
76
77    #[must_use]
78    pub fn count(&self) -> usize {
79        self.count
80    }
81}
82
83/// Inline, no-alloc per-core storage (`[T; N]`) — the bare-metal tier of
84/// [`PerCore`]. Same `slot`/`local` routing; the buffer lives inline (no heap).
85pub struct StaticPerCore<T, const N: usize> {
86    slots: [T; N],
87}
88
89impl<T, const N: usize> StaticPerCore<T, N> {
90    pub fn new_with(factory: impl FnMut(usize) -> T) -> Self {
91        Self {
92            slots: core::array::from_fn(factory),
93        }
94    }
95
96    /// The calling thread's slot (std only); bare-metal callers use [`slot`].
97    ///
98    /// [`slot`]: StaticPerCore::slot
99    #[cfg(feature = "std")]
100    pub fn local(&self) -> &T {
101        &self.slots[local_index(N)]
102    }
103
104    pub fn slot(&self, index: usize) -> &T {
105        &self.slots[index]
106    }
107
108    #[must_use]
109    pub const fn count(&self) -> usize {
110        N
111    }
112}
113
114#[cfg(test)]
115#[allow(clippy::unwrap_used, clippy::expect_used)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn static_per_core_routes_by_index() {
121        let cores = StaticPerCore::<usize, 4>::new_with(|index| index * 10);
122        assert_eq!(cores.count(), 4);
123        assert_eq!(*cores.slot(0), 0);
124        assert_eq!(*cores.slot(3), 30);
125    }
126
127    #[cfg(feature = "alloc")]
128    #[test]
129    fn per_core_new_with_and_from_vec() {
130        let built = PerCore::new_with(3, |index| index + 1);
131        assert_eq!(built.count(), 3);
132        assert_eq!(*built.slot(2), 3);
133
134        let from = PerCore::from_vec(alloc::vec![10usize, 20]);
135        assert_eq!(from.count(), 2);
136        assert_eq!(*from.slot(1), 20);
137    }
138
139    #[cfg(feature = "std")]
140    #[test]
141    fn local_is_sticky_within_a_thread() {
142        let cores = StaticPerCore::<usize, 8>::new_with(|index| index);
143        // a thread keeps landing on the same slot (cache locality) — the ticket
144        // is assigned once and reused for every `local()` on this thread.
145        let first = *cores.local();
146        assert_eq!(*cores.local(), first);
147        assert!(first < cores.count());
148    }
149}