1use 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#[derive(Debug)]
53pub enum PubSubCapacityMorphError {
54 InvalidCapacity,
56 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
75pub struct CapacityPubSubRing {
77 chain: Mutex<Vec<Arc<PubSubRing>>>,
81 capacity_atom: AtomicU64,
83 pin_generation: AtomicU64,
85 backing_source: PubSubBackingSource,
87 morph_seq: AtomicU64,
89 morph_lock: Mutex<()>,
92 warm: Mutex<Option<(usize, Arc<PubSubRing>)>>,
97 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 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 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 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 pub fn current_capacity(&self) -> usize {
180 self.capacity_atom.load(Ordering::Acquire) as usize
181 }
182
183 pub fn pin_generation(&self) -> u64 {
185 self.pin_generation.load(Ordering::Acquire)
186 }
187
188 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 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 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 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 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 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 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 pub fn warm_capacity(&self) -> Option<usize> {
327 self.warm.lock().as_ref().map(|(c, _)| *c)
328 }
329
330 pub fn warm_hits(&self) -> u64 {
332 self.warm_hits.load(Ordering::Relaxed)
333 }
334
335 pub fn clear_warm(&self) {
338 *self.warm.lock() = None;
339 }
340
341 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 pub fn ring_handle(&self) -> Arc<PubSubRing> {
361 let chain = self.chain.lock();
362 chain.last().expect("chain non-empty").clone()
363 }
364
365 pub fn chain_len(&self) -> usize {
368 self.chain.lock().len()
369 }
370
371 pub fn chain_total_capacity(&self) -> usize {
379 let chain = self.chain.lock();
380 chain.iter().map(|r| r.capacity()).sum()
381 }
382}
383
384pub struct CapacityPubSubSubscriber {
388 cap_ring: Arc<CapacityPubSubRing>,
389 backing_idx: u64,
390 position: u64,
391}
392
393impl CapacityPubSubSubscriber {
394 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 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 pub fn backing_idx(&self) -> u64 { self.backing_idx }
435
436 pub fn position(&self) -> u64 { self.position }
438}
439
440fn 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 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}