use crate::{Consumer, Producer, SpscRingBuffer};
pub struct MpscFanIn;
impl MpscFanIn {
pub fn with_capacity<T: Send + 'static>(
producer_count: usize,
per_ring_capacity: usize,
) -> (Vec<MpscFanInProducer<T>>, MpscFanInConsumer<T>) {
assert!(producer_count >= 1, "need at least one producer");
let mut producers = Vec::with_capacity(producer_count);
let mut consumers = Vec::with_capacity(producer_count);
for _ in 0..producer_count {
let (p, c) = SpscRingBuffer::with_capacity::<T>(per_ring_capacity);
producers.push(MpscFanInProducer { inner: p });
consumers.push(c);
}
(
producers,
MpscFanInConsumer {
rings: consumers,
cursor: 0,
},
)
}
}
pub struct MpscFanInProducer<T> {
inner: Producer<T>,
}
impl<T> MpscFanInProducer<T> {
pub fn try_push(&mut self, value: T) -> Result<(), T> {
self.inner.try_push(value)
}
pub fn capacity(&self) -> usize {
self.inner.capacity()
}
}
pub struct MpscFanInConsumer<T> {
rings: Vec<Consumer<T>>,
cursor: usize,
}
impl<T> MpscFanInConsumer<T> {
pub fn try_pop(&mut self) -> Option<T> {
let n = self.rings.len();
for offset in 0..n {
let idx = (self.cursor + offset) % n;
if let Some(v) = self.rings[idx].try_pop() {
self.cursor = (idx + 1) % n;
return Some(v);
}
}
None
}
pub fn producer_count(&self) -> usize {
self.rings.len()
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;
use super::*;
#[test]
fn single_producer_acts_like_plain_spsc() {
let (mut producers, mut consumer) = MpscFanIn::with_capacity::<u32>(1, 4);
let p = &mut producers[0];
p.try_push(7).unwrap();
assert_eq!(consumer.try_pop(), Some(7));
assert_eq!(consumer.try_pop(), None);
}
#[test]
fn round_robin_visits_every_producer() {
let (mut producers, mut consumer) = MpscFanIn::with_capacity::<u32>(3, 4);
producers[0].try_push(10).unwrap();
producers[1].try_push(20).unwrap();
producers[2].try_push(30).unwrap();
let mut got = Vec::new();
for _ in 0..3 {
got.push(consumer.try_pop().unwrap());
}
got.sort();
assert_eq!(got, vec![10u32, 20, 30]);
}
#[test]
fn quiet_producer_doesnt_block_busy_one() {
let (mut producers, mut consumer) = MpscFanIn::with_capacity::<u32>(2, 16);
for i in 0..10u32 {
producers[1].try_push(i).unwrap();
}
let mut got = Vec::new();
while let Some(v) = consumer.try_pop() {
got.push(v);
}
assert_eq!(got, (0..10u32).collect::<Vec<_>>());
}
#[test]
fn try_pop_on_all_empty_returns_none() {
let (_producers, mut consumer) = MpscFanIn::with_capacity::<u32>(4, 4);
assert_eq!(consumer.try_pop(), None);
}
#[test]
fn three_producers_one_consumer_under_threads() {
let (mut producers, mut consumer) = MpscFanIn::with_capacity::<u64>(3, 256);
let per_producer = 50_000u64;
let consumed = Arc::new(AtomicUsize::new(0));
let consumed_c = consumed.clone();
let total = per_producer * 3;
let consumer_t = thread::spawn(move || {
let mut local = 0u64;
while local < total {
if let Some(_v) = consumer.try_pop() {
local += 1;
consumed_c.fetch_add(1, Ordering::Relaxed);
}
}
});
let mut handles = Vec::new();
for (i, _) in (0..3).enumerate() {
let mut producer = producers.remove(0);
handles.push(thread::spawn(move || {
for j in 0..per_producer {
let v = (i as u64) * 1_000_000 + j;
while producer.try_push(v).is_err() {
std::hint::spin_loop();
}
}
}));
}
for h in handles {
h.join().unwrap();
}
consumer_t.join().unwrap();
assert_eq!(consumed.load(Ordering::Relaxed), total as usize);
}
#[test]
fn cursor_advances_past_drained_ring() {
let (mut producers, mut consumer) = MpscFanIn::with_capacity::<u32>(2, 4);
producers[0].try_push(1).unwrap();
assert_eq!(consumer.try_pop(), Some(1));
producers[1].try_push(2).unwrap();
assert_eq!(consumer.try_pop(), Some(2));
}
}