Skip to main content

subetha_cxc/
capacity_pubsub_ring.rs

1//! `CapacityPubSubRing`: runtime-resizable wrapper around
2//! [`PubSubRing`] that adds the capacity-axis morph to the
3//! pub/sub (1P/NC absolute-position) primitive.
4//!
5//! Sibling of [`CapacityAdaptiveRing`](crate::CapacityAdaptiveRing)
6//! and [`CapacityBroadcastRing`](crate::CapacityBroadcastRing).
7//! `CapacityPubSubRing` morphs PubSubRing's slot count at runtime
8//! under a chain-of-backings invariant: the producer publishes to
9//! the most-recent backing; subscribers carry their own
10//! (backing_idx, position) state and drain each backing in turn
11//! before advancing to the next.
12//!
13//! # Per-subscriber position tracking
14//!
15//! Pub/sub's per-subscriber position state already lives outside
16//! the ring (in `PubSubSubscriber::position` /
17//! [`SubscriberPosition`](crate::replay_positions::SubscriberPosition)),
18//! so the capacity-morph wrapper threads each subscriber through
19//! the chain of historical backings as the active one rolls
20//! forward. A [`CapacityPubSubSubscriber`] holds:
21//!
22//! - `cap_ring: Arc<CapacityPubSubRing>` to see the chain
23//! - `backing_idx: u64` - which backing in the chain we are
24//!   currently reading from
25//! - `position: u64` - position within `backings[backing_idx]`
26//!
27//! On `try_next()`, the subscriber reads at its current
28//! `(backing_idx, position)`. On `Pending` AND when not on the
29//! most-recent backing, it advances `backing_idx` and resets
30//! `position` to 0 (every stale backing's prior content drains
31//! before the subscriber crosses into the next).
32//!
33//! # Chain pruning
34//!
35//! Chain entries grow append-only across morphs. A separate
36//! `gc()` method walks the chain and drops the oldest entries
37//! whose strong count is 1 (only the chain itself holds the
38//! Arc) - i.e. no subscriber is currently reading from them. The
39//! producer continues publishing to the active end of the chain
40//! during gc; the morph guard ensures gc and morph are mutually
41//! exclusive.
42
43use std::path::{Path, PathBuf};
44use std::sync::Arc;
45use std::sync::atomic::{AtomicU64, Ordering};
46
47use parking_lot::Mutex;
48
49use crate::protocol_pubsub::{PubSubReadError, PubSubRing};
50
51/// Errors returned by capacity-morph operations on a pubsub ring.
52#[derive(Debug)]
53pub enum PubSubCapacityMorphError {
54    /// Target capacity is not a power of two, or less than 2.
55    InvalidCapacity,
56    /// I/O error during backing allocation.
57    Io(std::io::Error),
58}
59
60impl From<std::io::Error> for PubSubCapacityMorphError {
61    fn from(e: std::io::Error) -> Self { Self::Io(e) }
62}
63
64impl std::fmt::Display for PubSubCapacityMorphError {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        match self {
67            Self::InvalidCapacity => write!(f, "capacity must be pow2 >= 2"),
68            Self::Io(e) => write!(f, "io error during pubsub morph: {e}"),
69        }
70    }
71}
72
73impl std::error::Error for PubSubCapacityMorphError {}
74
75/// Runtime-resizable pubsub ring.
76pub struct CapacityPubSubRing {
77    /// Chain of all backings ever allocated, oldest-first. The
78    /// active backing is always at the end (index chain.len() -
79    /// 1). Append-only; pruning happens via [`gc`](Self::gc).
80    chain: Mutex<Vec<Arc<PubSubRing>>>,
81    /// Cached observable capacity of the active backing.
82    capacity_atom: AtomicU64,
83    /// Bumped on every morph for caller-polled pin invalidation.
84    pin_generation: AtomicU64,
85    /// Locale source for morph-allocated backings.
86    backing_source: PubSubBackingSource,
87    /// Monotonic morph counter for path / shm-name uniqueness.
88    morph_seq: AtomicU64,
89    /// Serialises morph callers (and gc) so the chain mutations
90    /// are atomic with respect to each other.
91    morph_lock: Mutex<()>,
92    /// One-slot warm cache: a fully constructed backing at a
93    /// predicted capacity, built off the morph lock by
94    /// [`prewarm`](Self::prewarm). Same design as
95    /// `CapacityAdaptiveRing`'s warm cache.
96    warm: Mutex<Option<(usize, Arc<PubSubRing>)>>,
97    /// Successful warm-cache hits consumed by `morph_capacity_to`.
98    warm_hits: AtomicU64,
99}
100
101unsafe impl Send for CapacityPubSubRing {}
102unsafe impl Sync for CapacityPubSubRing {}
103
104enum PubSubBackingSource {
105    Anon,
106    File(PathBuf),
107    Shm(String),
108}
109
110impl CapacityPubSubRing {
111    /// Anon (in-process) capacity-adaptive pubsub ring.
112    pub fn create_anon(
113        initial_capacity: usize,
114    ) -> Result<Arc<Self>, PubSubCapacityMorphError> {
115        if !initial_capacity.is_power_of_two() || initial_capacity < 2 {
116            return Err(PubSubCapacityMorphError::InvalidCapacity);
117        }
118        let ring = PubSubRing::create_anon(initial_capacity)?;
119        Ok(Arc::new(Self {
120            chain: Mutex::new(vec![Arc::new(ring)]),
121            capacity_atom: AtomicU64::new(initial_capacity as u64),
122            pin_generation: AtomicU64::new(0),
123            backing_source: PubSubBackingSource::Anon,
124            morph_seq: AtomicU64::new(0),
125            morph_lock: Mutex::new(()),
126            warm: Mutex::new(None),
127            warm_hits: AtomicU64::new(0),
128        }))
129    }
130
131    /// File-backed capacity-adaptive pubsub ring.
132    pub fn create(
133        base_path: impl AsRef<Path>,
134        initial_capacity: usize,
135    ) -> Result<Arc<Self>, PubSubCapacityMorphError> {
136        if !initial_capacity.is_power_of_two() || initial_capacity < 2 {
137            return Err(PubSubCapacityMorphError::InvalidCapacity);
138        }
139        let base = base_path.as_ref().to_path_buf();
140        let path = path_for_capacity_seq(&base, initial_capacity, 0);
141        let ring = PubSubRing::create(&path, initial_capacity)?;
142        Ok(Arc::new(Self {
143            chain: Mutex::new(vec![Arc::new(ring)]),
144            capacity_atom: AtomicU64::new(initial_capacity as u64),
145            pin_generation: AtomicU64::new(0),
146            backing_source: PubSubBackingSource::File(base),
147            morph_seq: AtomicU64::new(1),
148            morph_lock: Mutex::new(()),
149            warm: Mutex::new(None),
150            warm_hits: AtomicU64::new(0),
151        }))
152    }
153
154    /// ShmFs (named shared memory) capacity-adaptive pubsub ring.
155    pub fn create_shmfs(
156        name_prefix: &str,
157        initial_capacity: usize,
158    ) -> Result<Arc<Self>, PubSubCapacityMorphError> {
159        if !initial_capacity.is_power_of_two() || initial_capacity < 2 {
160            return Err(PubSubCapacityMorphError::InvalidCapacity);
161        }
162        let name = format!("{name_prefix}_cap_{initial_capacity}_g0");
163        let total = crate::protocol_pubsub::pubsub_ring_file_size(initial_capacity);
164        let shm = crate::shm_file::ShmFile::create_or_open_named(&name, total)?;
165        let ring = PubSubRing::create_from_shm(shm, initial_capacity)?;
166        Ok(Arc::new(Self {
167            chain: Mutex::new(vec![Arc::new(ring)]),
168            capacity_atom: AtomicU64::new(initial_capacity as u64),
169            pin_generation: AtomicU64::new(0),
170            backing_source: PubSubBackingSource::Shm(name_prefix.to_owned()),
171            morph_seq: AtomicU64::new(1),
172            morph_lock: Mutex::new(()),
173            warm: Mutex::new(None),
174            warm_hits: AtomicU64::new(0),
175        }))
176    }
177
178    /// Current capacity of the active backing.
179    pub fn current_capacity(&self) -> usize {
180        self.capacity_atom.load(Ordering::Acquire) as usize
181    }
182
183    /// Current pin generation.
184    pub fn pin_generation(&self) -> u64 {
185        self.pin_generation.load(Ordering::Acquire)
186    }
187
188    /// Publish a payload to the currently-active backing. Returns
189    /// the absolute position assigned within that backing (not
190    /// globally unique across backings - subscribers identify
191    /// items via payload contents, not position).
192    ///
193    /// The chain lock is held through the inner publish call so a
194    /// concurrent morph cannot slip in and make this publish land
195    /// in a backing that just became stale. Subscribers walk the
196    /// chain oldest-to-newest and only advance forward; if a
197    /// publish landed in a now-stale backing past where any
198    /// subscriber had already advanced, those items would be
199    /// silently lost. Holding the lock through publish prevents
200    /// that.
201    pub fn publish(&self, payload: &[u8]) -> u64 {
202        let chain = self.chain.lock();
203        chain.last().expect("chain always has at least one backing").publish(payload)
204    }
205
206    /// Subscribe to the stream from the CURRENT active backing's
207    /// current head. The subscriber drains forward from there,
208    /// crossing into newly-morphed backings as it catches up.
209    /// "From now" semantics: late joiners do NOT see history
210    /// from before they subscribed.
211    pub fn subscribe_from_now(self: &Arc<Self>) -> CapacityPubSubSubscriber {
212        let chain = self.chain.lock();
213        let backing_idx = (chain.len() - 1) as u64;
214        let active = &chain[backing_idx as usize];
215        let position = active.head();
216        drop(chain);
217        CapacityPubSubSubscriber {
218            cap_ring: Arc::clone(self),
219            backing_idx,
220            position,
221        }
222    }
223
224    /// Subscribe starting from the beginning of the OLDEST
225    /// backing currently in the chain. The subscriber drains
226    /// every item from every backing oldest-to-newest, crossing
227    /// chain entries as it catches up. Used when a subscriber
228    /// needs to replay the full available history.
229    pub fn subscribe_from_oldest(self: &Arc<Self>) -> CapacityPubSubSubscriber {
230        CapacityPubSubSubscriber {
231            cap_ring: Arc::clone(self),
232            backing_idx: 0,
233            position: 0,
234        }
235    }
236
237    /// Morph the active backing's capacity. Allocates a fresh
238    /// backing at `new_capacity`, appends it to the chain,
239    /// bumps pin_generation, and publishes the new active end.
240    /// Subscribers reading from older chain entries continue
241    /// undisturbed; they advance into the new backing
242    /// individually as their try_next catches up.
243    pub fn morph_capacity_to(
244        &self,
245        new_capacity: usize,
246    ) -> Result<(), PubSubCapacityMorphError> {
247        let _morph_guard = self.morph_lock.lock();
248
249        if !new_capacity.is_power_of_two() || new_capacity < 2 {
250            return Err(PubSubCapacityMorphError::InvalidCapacity);
251        }
252
253        let current = self.capacity_atom.load(Ordering::Acquire) as usize;
254        if current == new_capacity {
255            return Ok(());
256        }
257
258        // Warm-cache probe: a prediction matching the morph target
259        // skips allocation entirely; a mismatch stays cached and
260        // the cold path runs unchanged.
261        let warm_hit = {
262            let mut warm = self.warm.lock();
263            warm.take_if(|(cap, _)| *cap == new_capacity)
264        };
265        let new_ring = match warm_hit {
266            Some((_, ring)) => {
267                self.warm_hits.fetch_add(1, Ordering::Relaxed);
268                ring
269            }
270            None => self.build_backing(new_capacity)?,
271        };
272
273        {
274            let mut chain = self.chain.lock();
275            chain.push(new_ring);
276        }
277        self.pin_generation.fetch_add(1, Ordering::AcqRel);
278        self.capacity_atom
279            .store(new_capacity as u64, Ordering::Release);
280
281        Ok(())
282    }
283
284    /// Construct a fresh backing at `capacity`, at the wrapper's
285    /// locale, with a unique per-build name. Shared by the cold
286    /// morph path and [`prewarm`](Self::prewarm).
287    fn build_backing(
288        &self,
289        capacity: usize,
290    ) -> Result<Arc<PubSubRing>, PubSubCapacityMorphError> {
291        let seq = self.morph_seq.fetch_add(1, Ordering::AcqRel);
292        let ring = match &self.backing_source {
293            PubSubBackingSource::Anon => PubSubRing::create_anon(capacity)?,
294            PubSubBackingSource::File(base) => {
295                let path = path_for_capacity_seq(base, capacity, seq);
296                PubSubRing::create(&path, capacity)?
297            }
298            PubSubBackingSource::Shm(prefix) => {
299                let name = format!("{prefix}_cap_{capacity}_g{seq}");
300                let total = crate::protocol_pubsub::pubsub_ring_file_size(capacity);
301                let shm = crate::shm_file::ShmFile::create_or_open_named(&name, total)?;
302                PubSubRing::create_from_shm(shm, capacity)?
303            }
304        };
305        Ok(Arc::new(ring))
306    }
307
308    /// Speculatively build a backing at `capacity` into the
309    /// one-slot warm cache, off the morph lock's critical path.
310    /// The next `morph_capacity_to(capacity)` consumes it and
311    /// skips allocation. Re-prewarming the cached capacity is a
312    /// no-op; a different capacity replaces the slot.
313    pub fn prewarm(&self, capacity: usize) -> Result<(), PubSubCapacityMorphError> {
314        if !capacity.is_power_of_two() || capacity < 2 {
315            return Err(PubSubCapacityMorphError::InvalidCapacity);
316        }
317        if self.warm.lock().as_ref().map(|(c, _)| *c) == Some(capacity) {
318            return Ok(());
319        }
320        let ring = self.build_backing(capacity)?;
321        *self.warm.lock() = Some((capacity, ring));
322        Ok(())
323    }
324
325    /// Capacity currently held in the warm cache, if any.
326    pub fn warm_capacity(&self) -> Option<usize> {
327        self.warm.lock().as_ref().map(|(c, _)| *c)
328    }
329
330    /// Number of morphs that consumed a warm-cache prediction.
331    pub fn warm_hits(&self) -> u64 {
332        self.warm_hits.load(Ordering::Relaxed)
333    }
334
335    /// Drop any cached prediction, releasing its memory (and its
336    /// file / shm region for non-anon locales).
337    pub fn clear_warm(&self) {
338        *self.warm.lock() = None;
339    }
340
341    /// Garbage-collect stale backings from the front of the
342    /// chain. Drops the oldest contiguous run of backings whose
343    /// strong-count is 1 (only the chain itself holds them; no
344    /// subscriber is currently reading them). The active backing
345    /// is never dropped even when it has strong-count 1 - that
346    /// would lose the producer's target. Returns the number of
347    /// stale backings reclaimed.
348    pub fn gc(&self) -> usize {
349        let _morph_guard = self.morph_lock.lock();
350        let mut chain = self.chain.lock();
351        let mut reclaimed = 0;
352        while chain.len() > 1 && Arc::strong_count(&chain[0]) == 1 {
353            chain.remove(0);
354            reclaimed += 1;
355        }
356        reclaimed
357    }
358
359    /// Direct access to the currently-active [`PubSubRing`].
360    pub fn ring_handle(&self) -> Arc<PubSubRing> {
361        let chain = self.chain.lock();
362        chain.last().expect("chain non-empty").clone()
363    }
364
365    /// Number of backings currently in the chain (active + any
366    /// not-yet-gc'd stale entries).
367    pub fn chain_len(&self) -> usize {
368        self.chain.lock().len()
369    }
370
371    /// Sum of capacities across every backing currently in the
372    /// chain. Used by KeepAll-style producers to bound in-flight
373    /// items to actual buffering room: any item the producer
374    /// publishes is held in SOME backing until the slowest
375    /// subscriber catches up; with at most `chain_total_capacity()`
376    /// in-flight items, no backing wraps past a subscriber's
377    /// position before that subscriber drains it.
378    pub fn chain_total_capacity(&self) -> usize {
379        let chain = self.chain.lock();
380        chain.iter().map(|r| r.capacity()).sum()
381    }
382}
383
384/// Subscriber-side handle to a [`CapacityPubSubRing`]. Holds its
385/// own backing_idx + position within that backing and advances
386/// through the chain as it catches up.
387pub struct CapacityPubSubSubscriber {
388    cap_ring: Arc<CapacityPubSubRing>,
389    backing_idx: u64,
390    position: u64,
391}
392
393impl CapacityPubSubSubscriber {
394    /// Try to read the next payload. On `Ok`, the subscriber
395    /// advances its position by 1. On `Pending` at a stale
396    /// backing's head, transparently advances to the next backing
397    /// in the chain and retries; on `Pending` at the active
398    /// backing's head, returns `Pending` (no more data right
399    /// now).
400    pub fn try_next(&mut self, out: &mut [u8]) -> Result<(), PubSubReadError> {
401        loop {
402            let (backing, is_latest) = {
403                let chain = self.cap_ring.chain.lock();
404                let len = chain.len();
405                let idx = self.backing_idx as usize;
406                if idx >= len {
407                    return Err(PubSubReadError::Pending);
408                }
409                let b = Arc::clone(&chain[idx]);
410                (b, idx == len - 1)
411            };
412
413            match backing.read_at(self.position, out) {
414                Ok(()) => {
415                    self.position += 1;
416                    return Ok(());
417                }
418                Err(PubSubReadError::Pending) => {
419                    if !is_latest {
420                        // Caught up to a stale backing's head;
421                        // cross into the next entry of the chain.
422                        self.backing_idx += 1;
423                        self.position = 0;
424                        continue;
425                    }
426                    return Err(PubSubReadError::Pending);
427                }
428                err => return err,
429            }
430        }
431    }
432
433    /// Current backing chain index this subscriber is reading.
434    pub fn backing_idx(&self) -> u64 { self.backing_idx }
435
436    /// Current position within the current backing.
437    pub fn position(&self) -> u64 { self.position }
438}
439
440/// Compose the per-morph pubsub file path.
441fn path_for_capacity_seq(base: &Path, capacity: usize, seq: u64) -> PathBuf {
442    let mut s = base.as_os_str().to_owned();
443    s.push(format!(".cap_{capacity}_g{seq}.bin"));
444    PathBuf::from(s)
445}
446
447#[cfg(test)]
448mod tests {
449    use super::*;
450
451    #[test]
452    fn prewarm_hit_consumes_cache_and_subscribers_cross_chain() {
453        let ring = CapacityPubSubRing::create_anon(64).unwrap();
454        let mut sub = ring.subscribe_from_oldest();
455        ring.publish(&7u64.to_le_bytes());
456
457        ring.prewarm(256).unwrap();
458        assert_eq!(ring.warm_capacity(), Some(256));
459        ring.morph_capacity_to(256).unwrap();
460        assert_eq!(ring.warm_hits(), 1, "morph must consume the prediction");
461        assert_eq!(ring.warm_capacity(), None, "the slot is one-shot");
462        assert_eq!(ring.current_capacity(), 256);
463        assert_eq!(ring.chain_len(), 2);
464
465        // Published-pre-morph item reads from the stale chain
466        // entry; a post-morph publish lands on the warm backing
467        // and the subscriber crosses into it.
468        ring.publish(&9u64.to_le_bytes());
469        let mut out = [0u8; 64];
470        sub.try_next(&mut out).unwrap();
471        assert_eq!(u64::from_le_bytes(out[..8].try_into().unwrap()), 7);
472        sub.try_next(&mut out).unwrap();
473        assert_eq!(u64::from_le_bytes(out[..8].try_into().unwrap()), 9);
474    }
475
476    #[test]
477    fn prewarm_mismatch_stays_cached() {
478        let ring = CapacityPubSubRing::create_anon(64).unwrap();
479        ring.prewarm(512).unwrap();
480        ring.morph_capacity_to(256).unwrap();
481        assert_eq!(ring.warm_hits(), 0);
482        assert_eq!(ring.warm_capacity(), Some(512));
483        ring.morph_capacity_to(512).unwrap();
484        assert_eq!(ring.warm_hits(), 1);
485        assert_eq!(ring.warm_capacity(), None);
486    }
487
488    #[test]
489    fn prewarm_rejects_non_pow2_and_clear_drops() {
490        let ring = CapacityPubSubRing::create_anon(64).unwrap();
491        assert!(matches!(
492            ring.prewarm(100),
493            Err(PubSubCapacityMorphError::InvalidCapacity)
494        ));
495        ring.prewarm(128).unwrap();
496        ring.clear_warm();
497        assert_eq!(ring.warm_capacity(), None);
498    }
499}