subetha_cxc/
blocking_mpmc_ring.rs

1//! `BlockingMpmcRing`: composed-SPSC MPMC grid with cross-process
2//! futex-shaped `send_blocking` / `recv_blocking`.
3//!
4//! Wraps [`crate::mpmc_ring::SharedRingMpmc`] (N independent Lamport
5//! SPSC rings, statically partitioned across M consumers
6//! round-robin: consumer `m` owns rings `m, m + M, m + 2M, ...`).
7//!
8//! Wake routing:
9//! - Each producer parks on its own ring's `producer_waker[i]` when
10//!   the ring is full. The owning consumer wakes that specific
11//!   waker after popping from ring `i`.
12//! - Each consumer parks on its own `consumer_waker[m]` when every
13//!   ring in its subset is empty. A producer publishing to ring `i`
14//!   wakes `consumer_waker[i % M]` only.
15//!
16//! Per-subset shared `total_published[m]` counter drives the wake
17//! seq so consumer parks at `total_published[m] + 1`.
18//!
19//! See [`crate::cross_process_waker`] for the wake protocol and
20//! [`crate::blocking_spsc_ring::BlockingSpscRing`] /
21//! [`crate::blocking_mpsc_ring::BlockingMpscRing`] for the simpler
22//! shapes.
23
24use std::cell::Cell;
25use std::marker::PhantomData;
26use std::path::{Path, PathBuf};
27use std::sync::Arc;
28use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
29use std::time::{Duration, Instant};
30
31use crate::blocking_spsc_ring::BlockingError;
32use crate::cross_process_waker::{
33    CrossProcessWaker, MAX_WAITERS_DEFAULT, WakerError,
34};
35use crate::shared_ring::RingError;
36use crate::spsc_ring::SpscRingCore;
37
38const PRE_PARK_SPIN: u32 = 32;
39
40/// Factory for an MPMC grid of N SPSC rings partitioned across M
41/// consumers, every blocking call backed by a cross-process waker.
42pub struct BlockingMpmcRing;
43
44/// One producer handle. Sole writer to one underlying SPSC ring.
45/// `!Sync + !Clone + Send`: one producer per thread.
46pub struct BlockingMpmcProducer {
47    ring: Arc<SpscRingCore>,
48    own_waker: Arc<CrossProcessWaker>,
49    /// Consumer waker for the subset that owns this ring
50    /// (`subset = ring_index % n_consumers`).
51    consumer_waker: Arc<CrossProcessWaker>,
52    /// Shared total-published counter for the subset that owns this
53    /// ring (drives the consumer-side wake seq).
54    subset_total_published: Arc<AtomicU64>,
55    _not_sync: PhantomData<Cell<()>>,
56}
57
58/// One consumer handle. Sole drainer of an assigned subset of
59/// producer rings (round-robin from the factory).
60/// `!Sync + !Clone + Send`.
61pub struct BlockingMpmcConsumer {
62    /// Producer rings in this consumer's subset.
63    rings: Vec<Arc<SpscRingCore>>,
64    /// Per-ring producer-side wakers (parallel to `rings`).
65    producer_wakers: Vec<Arc<CrossProcessWaker>>,
66    /// This consumer's own waker.
67    own_waker: Arc<CrossProcessWaker>,
68    /// Counter shared with all producers in this consumer's subset.
69    subset_total_published: Arc<AtomicU64>,
70    next_drain: AtomicUsize,
71    _not_sync: PhantomData<Cell<()>>,
72}
73
74impl BlockingMpmcRing {
75    /// In-process grid: N rings, M consumer subsets, anon-mapped.
76    pub fn create_anon_grid(
77        n_producers: usize,
78        n_consumers: usize,
79        capacity: usize,
80    ) -> Result<(Vec<BlockingMpmcProducer>, Vec<BlockingMpmcConsumer>), BlockingError>
81    {
82        assert!(n_consumers >= 1, "n_consumers must be >= 1");
83        assert!(
84            n_producers >= n_consumers,
85            "n_producers ({n_producers}) must be >= n_consumers ({n_consumers})",
86        );
87        let mut rings = Vec::with_capacity(n_producers);
88        let mut producer_wakers = Vec::with_capacity(n_producers);
89        for _ in 0..n_producers {
90            rings.push(Arc::new(
91                SpscRingCore::create_anon(capacity).map_err(BlockingError::from)?,
92            ));
93            producer_wakers.push(Arc::new(
94                CrossProcessWaker::create_anon(MAX_WAITERS_DEFAULT)
95                    .map_err(BlockingError::from)?,
96            ));
97        }
98        let mut consumer_wakers = Vec::with_capacity(n_consumers);
99        let mut subset_total_published = Vec::with_capacity(n_consumers);
100        for _ in 0..n_consumers {
101            consumer_wakers.push(Arc::new(
102                CrossProcessWaker::create_anon(MAX_WAITERS_DEFAULT)
103                    .map_err(BlockingError::from)?,
104            ));
105            subset_total_published.push(Arc::new(AtomicU64::new(0)));
106        }
107        Ok(build_grid(
108            rings,
109            producer_wakers,
110            consumer_wakers,
111            subset_total_published,
112            n_consumers,
113        ))
114    }
115
116    /// File-backed grid. Path layout:
117    ///   `<prefix>.ring.{i}.bin`  - SPSC ring for producer `i`
118    ///   `<prefix>.pw.{i}.bin`    - producer-side waker for ring `i`
119    ///   `<prefix>.cw.{m}.bin`    - consumer-side waker for subset `m`
120    pub fn create_grid(
121        path_prefix: impl AsRef<Path>,
122        n_producers: usize,
123        n_consumers: usize,
124        capacity: usize,
125    ) -> Result<(Vec<BlockingMpmcProducer>, Vec<BlockingMpmcConsumer>), BlockingError>
126    {
127        assert!(n_consumers >= 1, "n_consumers must be >= 1");
128        assert!(n_producers >= n_consumers, "n_producers must be >= n_consumers");
129        let base = path_prefix.as_ref().to_path_buf();
130        let mut rings = Vec::with_capacity(n_producers);
131        let mut producer_wakers = Vec::with_capacity(n_producers);
132        for i in 0..n_producers {
133            rings.push(Arc::new(
134                SpscRingCore::create(ring_path(&base, i), capacity)
135                    .map_err(BlockingError::from)?,
136            ));
137            producer_wakers.push(Arc::new(
138                CrossProcessWaker::create(pw_path(&base, i), MAX_WAITERS_DEFAULT)
139                    .map_err(BlockingError::from)?,
140            ));
141        }
142        let mut consumer_wakers = Vec::with_capacity(n_consumers);
143        let mut subset_total_published = Vec::with_capacity(n_consumers);
144        for m in 0..n_consumers {
145            consumer_wakers.push(Arc::new(
146                CrossProcessWaker::create(cw_path(&base, m), MAX_WAITERS_DEFAULT)
147                    .map_err(BlockingError::from)?,
148            ));
149            subset_total_published.push(Arc::new(AtomicU64::new(0)));
150        }
151        Ok(build_grid(
152            rings,
153            producer_wakers,
154            consumer_wakers,
155            subset_total_published,
156            n_consumers,
157        ))
158    }
159
160    /// Open an existing file-backed grid.
161    pub fn open_grid(
162        path_prefix: impl AsRef<Path>,
163        n_producers: usize,
164        n_consumers: usize,
165        expected_capacity: usize,
166    ) -> Result<(Vec<BlockingMpmcProducer>, Vec<BlockingMpmcConsumer>), BlockingError>
167    {
168        assert!(n_consumers >= 1, "n_consumers must be >= 1");
169        assert!(n_producers >= n_consumers, "n_producers must be >= n_consumers");
170        let base = path_prefix.as_ref().to_path_buf();
171        let mut rings = Vec::with_capacity(n_producers);
172        let mut producer_wakers = Vec::with_capacity(n_producers);
173        for i in 0..n_producers {
174            rings.push(Arc::new(
175                SpscRingCore::open(ring_path(&base, i), expected_capacity)
176                    .map_err(BlockingError::from)?,
177            ));
178            producer_wakers.push(Arc::new(
179                CrossProcessWaker::open(pw_path(&base, i), MAX_WAITERS_DEFAULT)
180                    .map_err(BlockingError::from)?,
181            ));
182        }
183        let mut consumer_wakers = Vec::with_capacity(n_consumers);
184        let mut subset_total_published = Vec::with_capacity(n_consumers);
185        for m in 0..n_consumers {
186            consumer_wakers.push(Arc::new(
187                CrossProcessWaker::open(cw_path(&base, m), MAX_WAITERS_DEFAULT)
188                    .map_err(BlockingError::from)?,
189            ));
190            subset_total_published.push(Arc::new(AtomicU64::new(0)));
191        }
192        Ok(build_grid(
193            rings,
194            producer_wakers,
195            consumer_wakers,
196            subset_total_published,
197            n_consumers,
198        ))
199    }
200}
201
202fn ring_path(base: &Path, i: usize) -> PathBuf {
203    let mut s = base.as_os_str().to_owned();
204    s.push(format!(".ring.{i}.bin"));
205    PathBuf::from(s)
206}
207
208fn pw_path(base: &Path, i: usize) -> PathBuf {
209    let mut s = base.as_os_str().to_owned();
210    s.push(format!(".pw.{i}.bin"));
211    PathBuf::from(s)
212}
213
214fn cw_path(base: &Path, m: usize) -> PathBuf {
215    let mut s = base.as_os_str().to_owned();
216    s.push(format!(".cw.{m}.bin"));
217    PathBuf::from(s)
218}
219
220fn build_grid(
221    rings: Vec<Arc<SpscRingCore>>,
222    producer_wakers: Vec<Arc<CrossProcessWaker>>,
223    consumer_wakers: Vec<Arc<CrossProcessWaker>>,
224    subset_total_published: Vec<Arc<AtomicU64>>,
225    n_consumers: usize,
226) -> (Vec<BlockingMpmcProducer>, Vec<BlockingMpmcConsumer>) {
227    let producers: Vec<BlockingMpmcProducer> = rings
228        .iter()
229        .zip(producer_wakers.iter())
230        .enumerate()
231        .map(|(i, (r, pw))| {
232            let subset = i % n_consumers;
233            BlockingMpmcProducer {
234                ring: Arc::clone(r),
235                own_waker: Arc::clone(pw),
236                consumer_waker: Arc::clone(&consumer_wakers[subset]),
237                subset_total_published: Arc::clone(&subset_total_published[subset]),
238                _not_sync: PhantomData,
239            }
240        })
241        .collect();
242
243    let mut consumer_rings: Vec<Vec<Arc<SpscRingCore>>> =
244        (0..n_consumers).map(|_| Vec::new()).collect();
245    let mut consumer_pwakers: Vec<Vec<Arc<CrossProcessWaker>>> =
246        (0..n_consumers).map(|_| Vec::new()).collect();
247    for (i, (r, pw)) in rings.iter().zip(producer_wakers.iter()).enumerate() {
248        let subset = i % n_consumers;
249        consumer_rings[subset].push(Arc::clone(r));
250        consumer_pwakers[subset].push(Arc::clone(pw));
251    }
252    let consumers: Vec<BlockingMpmcConsumer> = consumer_rings
253        .into_iter()
254        .zip(consumer_pwakers)
255        .zip(consumer_wakers)
256        .zip(subset_total_published)
257        .map(|(((subset_rings, subset_pwakers), own_waker), total)| {
258            BlockingMpmcConsumer {
259                rings: subset_rings,
260                producer_wakers: subset_pwakers,
261                own_waker,
262                subset_total_published: total,
263                next_drain: AtomicUsize::new(0),
264                _not_sync: PhantomData,
265            }
266        })
267        .collect();
268
269    (producers, consumers)
270}
271
272impl BlockingMpmcProducer {
273    /// Non-blocking push. On success, advances subset-shared
274    /// counter + fires a wake at the consumer waker.
275    #[inline]
276    pub fn try_push(&self, payload: &[u8]) -> Result<(), RingError> {
277        let r = self.ring.try_push(payload);
278        if r.is_ok() {
279            let new_seq =
280                self.subset_total_published.fetch_add(1, Ordering::Release) + 1;
281            self.consumer_waker.wake_up_to(new_seq);
282        }
283        r
284    }
285
286    /// Block until either push succeeds or `timeout` elapses.
287    pub fn send_blocking(
288        &self,
289        payload: &[u8],
290        timeout: Option<Duration>,
291    ) -> Result<(), BlockingError> {
292        let deadline = timeout.map(|d| Instant::now() + d);
293        loop {
294            match self.try_push(payload) {
295                Ok(()) => return Ok(()),
296                Err(RingError::Full) => {}
297                Err(e) => return Err(BlockingError::Ring(e)),
298            }
299            for _ in 0..PRE_PARK_SPIN {
300                if self.ring.try_push(payload).is_ok() {
301                    let new_seq = self
302                        .subset_total_published
303                        .fetch_add(1, Ordering::Release)
304                        + 1;
305                    self.consumer_waker.wake_up_to(new_seq);
306                    return Ok(());
307                }
308                std::hint::spin_loop();
309            }
310            let target = self.ring.tail() + 1;
311            let token = self.own_waker.try_park(target)?;
312            if self.ring.try_push(payload).is_ok() {
313                self.own_waker.release(token);
314                let new_seq = self
315                    .subset_total_published
316                    .fetch_add(1, Ordering::Release)
317                    + 1;
318                self.consumer_waker.wake_up_to(new_seq);
319                return Ok(());
320            }
321            let remaining = match deadline {
322                None => None,
323                Some(d) => {
324                    let now = Instant::now();
325                    if now >= d {
326                        self.own_waker.release(token);
327                        return Err(BlockingError::Timeout);
328                    }
329                    Some(d - now)
330                }
331            };
332            match self.own_waker.wait(token, remaining) {
333                Ok(()) => continue,
334                Err(WakerError::Timeout) => return Err(BlockingError::Timeout),
335                Err(e) => return Err(BlockingError::from(e)),
336            }
337        }
338    }
339
340    pub fn capacity(&self) -> usize { self.ring.capacity() }
341    pub fn head(&self) -> u64 { self.ring.head() }
342}
343
344impl BlockingMpmcConsumer {
345    /// Non-blocking pop, round-robin across this consumer's subset.
346    pub fn try_pop(&self, out: &mut [u8]) -> Result<usize, RingError> {
347        self.try_pop_inner(out)
348    }
349
350    /// Block until either a pop succeeds or `timeout` elapses.
351    pub fn recv_blocking(
352        &self,
353        out: &mut [u8],
354        timeout: Option<Duration>,
355    ) -> Result<usize, BlockingError> {
356        let deadline = timeout.map(|d| Instant::now() + d);
357        loop {
358            match self.try_pop_inner(out) {
359                Ok(n) => return Ok(n),
360                Err(RingError::Empty) => {}
361                Err(e) => return Err(BlockingError::Ring(e)),
362            }
363            for _ in 0..PRE_PARK_SPIN {
364                if let Ok(n) = self.try_pop_inner(out) {
365                    return Ok(n);
366                }
367                std::hint::spin_loop();
368            }
369            let target = self.subset_total_published.load(Ordering::Acquire) + 1;
370            let token = self.own_waker.try_park(target)?;
371            if let Ok(n) = self.try_pop_inner(out) {
372                self.own_waker.release(token);
373                return Ok(n);
374            }
375            let remaining = match deadline {
376                None => None,
377                Some(d) => {
378                    let now = Instant::now();
379                    if now >= d {
380                        self.own_waker.release(token);
381                        return Err(BlockingError::Timeout);
382                    }
383                    Some(d - now)
384                }
385            };
386            match self.own_waker.wait(token, remaining) {
387                Ok(()) => continue,
388                Err(WakerError::Timeout) => return Err(BlockingError::Timeout),
389                Err(e) => return Err(BlockingError::from(e)),
390            }
391        }
392    }
393
394    #[inline]
395    fn try_pop_inner(&self, out: &mut [u8]) -> Result<usize, RingError> {
396        let n = self.rings.len();
397        let start = self.next_drain.load(Ordering::Relaxed);
398        for i in 0..n {
399            let idx = (start + i) % n;
400            if let Ok(bytes) = self.rings[idx].try_pop(out) {
401                self.next_drain.store((idx + 1) % n, Ordering::Relaxed);
402                let tail = self.rings[idx].tail();
403                self.producer_wakers[idx].wake_up_to(tail);
404                return Ok(bytes);
405            }
406        }
407        Err(RingError::Empty)
408    }
409
410    /// Number of rings in this consumer's subset.
411    pub fn n_rings(&self) -> usize { self.rings.len() }
412
413    /// Approximate pending items across this consumer's subset.
414    pub fn approx_subset_len(&self) -> usize {
415        self.rings.iter().map(|r| r.approx_len()).sum()
416    }
417}
418
419#[cfg(test)]
420mod tests {
421    use super::*;
422    use std::thread;
423
424    #[test]
425    fn round_trip_4p_2c_anon() {
426        let (producers, consumers) =
427            BlockingMpmcRing::create_anon_grid(4, 2, 8).expect("create");
428        const PER_PROD: u64 = 25;
429        let total: u64 = PER_PROD * 4;
430
431        let prod_handles: Vec<_> = producers
432            .into_iter()
433            .enumerate()
434            .map(|(pid, p)| {
435                thread::spawn(move || {
436                    for i in 0..PER_PROD {
437                        let val = (pid as u64) * 1_000_000 + i;
438                        let mut payload = [0u8; 56];
439                        payload[..8].copy_from_slice(&val.to_le_bytes());
440                        p.send_blocking(&payload, Some(Duration::from_secs(5)))
441                            .expect("send");
442                    }
443                })
444            })
445            .collect();
446
447        let cons_handles: Vec<_> = consumers
448            .into_iter()
449            .map(|c| {
450                let per_consumer = (total / 2) as usize;
451                thread::spawn(move || {
452                    let mut buf = [0u8; 64];
453                    let mut got: Vec<u64> = Vec::with_capacity(per_consumer);
454                    for _ in 0..per_consumer {
455                        c.recv_blocking(&mut buf, Some(Duration::from_secs(5)))
456                            .expect("recv");
457                        got.push(u64::from_le_bytes(buf[..8].try_into().unwrap()));
458                    }
459                    got
460                })
461            })
462            .collect();
463
464        for h in prod_handles {
465            h.join().unwrap();
466        }
467        let mut all: Vec<u64> = cons_handles
468            .into_iter()
469            .flat_map(|h| h.join().unwrap())
470            .collect();
471        all.sort_unstable();
472
473        let mut expected: Vec<u64> = Vec::with_capacity(total as usize);
474        for pid in 0..4u64 {
475            for i in 0..PER_PROD {
476                expected.push(pid * 1_000_000 + i);
477            }
478        }
479        expected.sort_unstable();
480        assert_eq!(all, expected, "every item delivered exactly once");
481    }
482
483    #[test]
484    fn recv_blocking_returns_timeout() {
485        let (_producers, consumers) =
486            BlockingMpmcRing::create_anon_grid(2, 2, 4).expect("create");
487        let mut buf = [0u8; 64];
488        let t0 = Instant::now();
489        let err = consumers[0].recv_blocking(&mut buf, Some(Duration::from_millis(60)));
490        assert_eq!(err, Err(BlockingError::Timeout));
491        assert!(t0.elapsed() >= Duration::from_millis(50));
492    }
493}
subetha_cxc/blocking_mpmc_ring.rs

subetha_cxc/
blocking_mpmc_ring.rs
