diff --git a/src/buffer/array.rs b/src/buffer/array.rs index 6785373..7413976 100644 --- a/src/buffer/array.rs +++ b/src/buffer/array.rs @@ -30,7 +30,7 @@ impl Buffer for ArrayBuffer { #[inline(always)] fn at(&self, idx: usize) -> *const UnsafeCell> { - &self.0[idx % CAP] as * const _ + &self.0[idx % CAP] as *const _ } } diff --git a/src/buffer/dynamic.rs b/src/buffer/dynamic.rs index c091c3b..51d042c 100644 --- a/src/buffer/dynamic.rs +++ b/src/buffer/dynamic.rs @@ -6,7 +6,7 @@ use super::Buffer; /// Holds data allocated from the heap at run time pub struct DynamicBuffer { - items: Box<[UnsafeCell>]> + items: Box<[UnsafeCell>]>, } impl DynamicBuffer { @@ -18,7 +18,7 @@ impl DynamicBuffer { let mut vec = Vec::with_capacity(size); unsafe { vec.set_len(size) }; Ok(DynamicBuffer { - items: vec.into_boxed_slice() + items: vec.into_boxed_slice(), }) } else { Err("Buffer size must be greater than 0") @@ -46,7 +46,7 @@ impl Buffer for DynamicBuffer { /// faster runtime performance due to the use of a mask instead of modulus /// when computing buffer indexes. pub struct DynamicBufferP2 { - items: Box<[UnsafeCell>]> + items: Box<[UnsafeCell>]>, } impl DynamicBufferP2 { @@ -61,7 +61,7 @@ impl DynamicBufferP2 { let mut vec = Vec::with_capacity(size); unsafe { vec.set_len(size) }; vec.into_boxed_slice() - } + }, }), _ => Err("Buffer size must be a power of two"), } diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs index a7f7265..1252bd2 100644 --- a/src/buffer/mod.rs +++ b/src/buffer/mod.rs @@ -34,7 +34,7 @@ impl> Buffer for Box { (**self).size() } - fn at(&self, idx: usize) -> * const UnsafeCell> { + fn at(&self, idx: usize) -> *const UnsafeCell> { (**self).at(idx) } } diff --git a/src/lib.rs b/src/lib.rs index aa29177..5eabee5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -142,20 +142,30 @@ impl std::error::Error for TryPopError {} /// The consumer end of the queue allows for sending data. `Producer` is /// always `Send`, but is only `Sync` for multi-producer (MPSC, MPMC) queues. pub trait Producer { + /// Check if this channel is closed. + fn is_closed(&self) -> bool; + /// Add value to front of the queue. This method will block if the queue /// is currently full. + /// If the channel is closed, this function will continue succeeding until + /// the queue becomes full. fn push(&self, value: T) -> Result<(), PushError>; /// Attempt to add a value to the front of the queue. If the value was /// added successfully, `None` will be returned. If unsuccessful, `value` /// will be returned. An unsuccessful push indicates that the queue was /// full. + /// If the channel is closed, this function will continue succeeding until + /// the queue becomes full. fn try_push(&self, value: T) -> Result<(), TryPushError>; } /// The consumer end of the queue allows for receiving data. `Consumer` is /// always `Send`, but is only `Sync` for multi-consumer (SPMC, MPMC) queues. pub trait Consumer { + /// Check if this channel is closed. + fn is_closed(&self) -> bool; + /// Remove value from the end of the queue. This method will block if the /// queue is currently empty. fn pop(&self) -> Result; diff --git a/src/mpmc.rs b/src/mpmc.rs index 7c851e7..a9bfcaa 100644 --- a/src/mpmc.rs +++ b/src/mpmc.rs @@ -34,7 +34,7 @@ unsafe impl> Sync for MPMCConsumer {} impl> Clone for MPMCConsumer { fn clone(&self) -> Self { - self.queue.consumers.fetch_add(1, Ordering::Release); + self.queue.consumers.fetch_add(1, Ordering::Relaxed); MPMCConsumer { queue: self.queue.clone(), } @@ -50,7 +50,7 @@ unsafe impl> Sync for MPMCProducer {} impl> Clone for MPMCProducer { fn clone(&self) -> Self { - self.queue.producers.fetch_add(1, Ordering::Release); + self.queue.producers.fetch_add(1, Ordering::Relaxed); MPMCProducer { queue: self.queue.clone(), } @@ -102,21 +102,24 @@ impl> Drop for MPMCQueue { } impl> Producer for MPMCProducer { + fn is_closed(&self) -> bool { + self.queue.consumers.load(Ordering::Relaxed) == 0 + } + fn push(&self, value: T) -> Result<(), PushError> { let q = &self.queue; - let head = q.head.next.fetch_add(1, Ordering::Relaxed); + loop { - if q.consumers.load(Ordering::Acquire) == 0 { - return Err(PushError::Disconnected(value)); - } else if q.tail.curr.load(Ordering::Acquire) + q.buf.size() > head { + if q.tail.curr.load(Ordering::Acquire) + q.buf.size() > head { break; + } else if q.consumers.load(Ordering::Relaxed) == 0 { + return Err(PushError::Disconnected(value)); } spin_loop(); } unsafe { buf_write(&q.buf, head, value) }; - while q.head.curr.load(Ordering::Relaxed) < head { spin_loop(); } @@ -128,43 +131,50 @@ impl> Producer for MPMCProducer { let q = &self.queue; loop { let head = q.head.curr.load(Ordering::Relaxed); - if q.consumers.load(Ordering::Acquire) == 0 { - return Err(TryPushError::Disconnected(value)); - } else if q.tail.curr.load(Ordering::Acquire) + q.buf.size() <= head { - return Err(TryPushError::Full(value)); - } else { - let next = head + 1; - if q.head - .next - .compare_exchange_weak(head, next, Ordering::Acquire, Ordering::Acquire) - .is_ok() - { - unsafe { buf_write(&q.buf, head, value) }; - q.head.curr.store(next, Ordering::Release); - return Ok(()); - } + let head_plus_one = head + 1; + + if q.tail.curr.load(Ordering::Acquire) + q.buf.size() <= head { + // buffer is full, check whether it's closed. + // relaxed is fine since Consumer.drop does an acquire/release on .tail + return if q.consumers.load(Ordering::Relaxed) == 0 { + Err(TryPushError::Disconnected(value)) + } else { + Err(TryPushError::Full(value)) + }; + } else if q + .head + .next + .compare_exchange_weak(head, head_plus_one, Ordering::Acquire, Ordering::Acquire) + .is_ok() + { + unsafe { buf_write(&q.buf, head, value) }; + q.head.curr.store(head_plus_one, Ordering::Release); + return Ok(()); } } } } impl> Consumer for MPMCConsumer { + fn is_closed(&self) -> bool { + self.queue.producers.load(Ordering::Relaxed) == 0 + } + fn pop(&self) -> Result { let q = &self.queue; - let tail = q.tail.next.fetch_add(1, Ordering::Relaxed); let tail_plus_one = tail + 1; + loop { - if tail_plus_one <= q.head.curr.load(Ordering::Acquire) { + if q.head.curr.load(Ordering::Acquire) >= tail_plus_one { break; - } else if q.producers.load(Ordering::Acquire) == 0 { + } else if q.producers.load(Ordering::Relaxed) == 0 { return Err(PopError::Disconnected); } spin_loop(); } let v = unsafe { buf_read(&q.buf, tail) }; - while q.tail.curr.load(Ordering::Relaxed) < tail { spin_loop(); } @@ -177,12 +187,15 @@ impl> Consumer for MPMCConsumer { loop { let tail = q.tail.curr.load(Ordering::Relaxed); let tail_plus_one = tail + 1; - if tail_plus_one > q.head.curr.load(Ordering::Acquire) { - if q.producers.load(Ordering::Acquire) > 0 { - return Err(TryPopError::Empty); + + if q.head.curr.load(Ordering::Acquire) < tail_plus_one { + // buffer is empty, check whether it's closed. + // relaxed is fine since Producer.drop does an acquire/release on .head + return if q.producers.load(Ordering::Relaxed) == 0 { + Err(TryPopError::Disconnected) } else { - return Err(TryPopError::Disconnected); - } + Err(TryPopError::Empty) + }; } else if q .tail .next @@ -199,13 +212,17 @@ impl> Consumer for MPMCConsumer { impl> Drop for MPMCProducer { fn drop(&mut self) { - self.queue.producers.fetch_sub(1, Ordering::Release); + self.queue.producers.fetch_sub(1, Ordering::Relaxed); + // Acquire/Release .head to ensure other threads see new .closed + self.queue.head.curr.fetch_add(0, Ordering::AcqRel); } } impl> Drop for MPMCConsumer { fn drop(&mut self) { - self.queue.consumers.fetch_sub(1, Ordering::Release); + self.queue.consumers.fetch_sub(1, Ordering::Relaxed); + // Acquire/Release .tail to ensure other threads see new .closed + self.queue.tail.curr.fetch_add(0, Ordering::AcqRel); } } @@ -308,9 +325,11 @@ mod test { assert_eq!(c.pop(), Err(PopError::Disconnected)); assert_eq!(c.try_pop(), Err(TryPopError::Disconnected)); - let (p, c) = mpmc_queue(DynamicBuffer::new(32).unwrap()); + let (p, c) = mpmc_queue(DynamicBuffer::new(2).unwrap()); p.push(1).unwrap(); std::mem::drop(c); + assert!(p.is_closed()); + p.push(1).unwrap(); assert_eq!(p.push(2), Err(PushError::Disconnected(2))); assert_eq!(p.try_push(2), Err(TryPushError::Disconnected(2))); diff --git a/src/mpsc.rs b/src/mpsc.rs index 9b7d1f5..99ae670 100644 --- a/src/mpsc.rs +++ b/src/mpsc.rs @@ -6,6 +6,7 @@ //! the `MPSCProducer` is `Send` and `Sync` while the `MPSCConsumer` is `Send` //! and `!Sync`. +use std::cell::Cell; use std::hint::spin_loop; use std::marker::PhantomData; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; @@ -28,7 +29,10 @@ struct MPSCQueue> { /// Consumer end of the queue. Implements the trait `Consumer`. pub struct MPSCConsumer> { queue: Arc>, - _not_sync: PhantomData> + /// A copy of `queue.head` for quick access. + /// + /// This value can be stale and sometimes needs to be resynchronized with `queue.head`. + cached_head: Cell, } /// Producer end of the queue. Implements the trait `Producer`. @@ -75,7 +79,10 @@ pub fn mpsc_queue>(buf: B) -> (MPSCProducer, MPSCConsumer< MPSCProducer { queue: queue.clone(), }, - MPSCConsumer { queue, _not_sync: PhantomData }, + MPSCConsumer { + queue, + cached_head: Cell::new(0), + }, ) } @@ -90,21 +97,24 @@ impl> Drop for MPSCQueue { } impl> Producer for MPSCProducer { + fn is_closed(&self) -> bool { + !self.queue.consumer.load(Ordering::Relaxed) + } + fn push(&self, value: T) -> Result<(), PushError> { let q = &self.queue; - let head = q.head.next.fetch_add(1, Ordering::Relaxed); + loop { - if !q.consumer.load(Ordering::Acquire) { - return Err(PushError::Disconnected(value)); - } else if q.tail.load(Ordering::Acquire) + q.buf.size() > head { + if q.tail.load(Ordering::Acquire) + q.buf.size() > head { break; + } else if !q.consumer.load(Ordering::Relaxed) { + return Err(PushError::Disconnected(value)); } spin_loop(); } unsafe { buf_write(&q.buf, head, value) }; - while q.head.curr.load(Ordering::Relaxed) < head { spin_loop(); } @@ -116,43 +126,54 @@ impl> Producer for MPSCProducer { let q = &self.queue; loop { let head = q.head.curr.load(Ordering::Relaxed); - if !q.consumer.load(Ordering::Acquire) { - return Err(TryPushError::Disconnected(value)); - } else if q.tail.load(Ordering::Acquire) + q.buf.size() <= head { - return Err(TryPushError::Full(value)); - } else { - let next = head + 1; - if q.head - .next - .compare_exchange_weak(head, next, Ordering::Acquire, Ordering::Acquire) - .is_ok() - { - unsafe { buf_write(&q.buf, head, value) }; - q.head.curr.store(next, Ordering::Release); - return Ok(()); - } + let head_plus_one = head + 1; + + if q.tail.load(Ordering::Acquire) + q.buf.size() <= head { + // buffer is full, check whether it's closed. + // relaxed is fine since Consumer.drop does an acquire/release on .tail + return if !q.consumer.load(Ordering::Relaxed) { + Err(TryPushError::Disconnected(value)) + } else { + Err(TryPushError::Full(value)) + }; + } else if q + .head + .next + .compare_exchange_weak(head, head_plus_one, Ordering::Acquire, Ordering::Acquire) + .is_ok() + { + unsafe { buf_write(&q.buf, head, value) }; + q.head.curr.store(head_plus_one, Ordering::Release); + return Ok(()); } } } } impl> Consumer for MPSCConsumer { + fn is_closed(&self) -> bool { + Arc::strong_count(&self.queue) < 2 + } + fn pop(&self) -> Result { let q = &self.queue; - let tail = q.tail.load(Ordering::Relaxed); let tail_plus_one = tail + 1; - loop { - if tail_plus_one <= q.head.curr.load(Ordering::Acquire) { - break; - } else if Arc::strong_count(&self.queue) < 2 { - return Err(PopError::Disconnected); + + if self.cached_head.get() < tail_plus_one { + loop { + let head = q.head.curr.load(Ordering::Acquire); + if head >= tail_plus_one { + self.cached_head.set(head); + break; + } else if Arc::strong_count(q) < 2 { + return Err(PopError::Disconnected); + } + spin_loop(); } - spin_loop(); } let v = unsafe { buf_read(&q.buf, tail) }; - q.tail.store(tail_plus_one, Ordering::Release); Ok(v) } @@ -162,23 +183,29 @@ impl> Consumer for MPSCConsumer { let tail = q.tail.load(Ordering::Relaxed); let tail_plus_one = tail + 1; - if tail_plus_one > q.head.curr.load(Ordering::Acquire) { - if Arc::strong_count(&self.queue) > 1 { - Err(TryPopError::Empty) - } else { - Err(TryPopError::Disconnected) + if self.cached_head.get() < tail_plus_one { + let head = q.head.curr.load(Ordering::Acquire); + if head < tail_plus_one { + return if Arc::strong_count(q) < 2 { + Err(TryPopError::Disconnected) + } else { + Err(TryPopError::Empty) + }; } - } else { - let v = unsafe { buf_read(&q.buf, tail) }; - q.tail.store(tail_plus_one, Ordering::Release); - Ok(v) + self.cached_head.set(head); } + + let v = unsafe { buf_read(&q.buf, tail) }; + q.tail.store(tail_plus_one, Ordering::Release); + Ok(v) } } impl> Drop for MPSCConsumer { fn drop(&mut self) { - self.queue.consumer.store(false, Ordering::Release); + self.queue.consumer.store(false, Ordering::Relaxed); + // Acquire/Release .tail to ensure other threads see new .closed + self.queue.tail.fetch_add(0, Ordering::AcqRel); } } @@ -276,9 +303,11 @@ mod test { assert_eq!(c.pop(), Err(PopError::Disconnected)); assert_eq!(c.try_pop(), Err(TryPopError::Disconnected)); - let (p, c) = mpsc_queue(DynamicBuffer::new(32).unwrap()); + let (p, c) = mpsc_queue(DynamicBuffer::new(2).unwrap()); p.push(1).unwrap(); std::mem::drop(c); + assert!(p.is_closed()); + p.push(1).unwrap(); assert_eq!(p.push(2), Err(PushError::Disconnected(2))); assert_eq!(p.try_push(2), Err(TryPushError::Disconnected(2))); diff --git a/src/spmc.rs b/src/spmc.rs index ddee6da..ea3b2fa 100644 --- a/src/spmc.rs +++ b/src/spmc.rs @@ -6,6 +6,7 @@ //! `SPMCProducer` is `Send` and `!Sync` while `SPMCConsumer` is `Send` and //! `Sync`. +use std::cell::Cell; use std::hint::spin_loop; use std::marker::PhantomData; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; @@ -43,7 +44,10 @@ impl> Clone for SPMCConsumer { /// Producer end of the queue. Implements the trait `Producer`. pub struct SPMCProducer> { queue: Arc>, - _not_sync: PhantomData> + /// A copy of `queue.tail` for quick access. + /// + /// This value can be stale and sometimes needs to be resynchronized with `queue.tail`. + cached_tail: Cell, } /// Creates a new SPMC queue @@ -74,7 +78,7 @@ pub fn spmc_queue>(buf: B) -> (SPMCProducer, SPMCConsumer< ( SPMCProducer { queue: queue.clone(), - _not_sync: PhantomData, + cached_tail: Cell::new(0), }, SPMCConsumer { queue }, ) @@ -91,17 +95,25 @@ impl> Drop for SPMCQueue { } impl> Producer for SPMCProducer { + fn is_closed(&self) -> bool { + Arc::strong_count(&self.queue) < 2 + } + fn push(&self, value: T) -> Result<(), PushError> { let q = &self.queue; let head = q.head.load(Ordering::Relaxed); - loop { - if Arc::strong_count(&self.queue) < 2 { - return Err(PushError::Disconnected(value)); - } else if q.tail.curr.load(Ordering::Acquire) + q.buf.size() > head { - break; + if self.cached_tail.get() + q.buf.size() <= head { + loop { + let tail = q.tail.curr.load(Ordering::Acquire); + if tail + q.buf.size() > head { + self.cached_tail.set(tail); + break; + } else if Arc::strong_count(q) < 2 { + return Err(PushError::Disconnected(value)); + } + spin_loop(); } - spin_loop(); } unsafe { buf_write(&q.buf, head, value) }; @@ -112,35 +124,45 @@ impl> Producer for SPMCProducer { fn try_push(&self, value: T) -> Result<(), TryPushError> { let q = &self.queue; let head = q.head.load(Ordering::Relaxed); - if Arc::strong_count(&self.queue) < 2 { - Err(TryPushError::Disconnected(value)) - } else if q.tail.curr.load(Ordering::Acquire) + q.buf.size() <= head { - Err(TryPushError::Full(value)) - } else { - unsafe { buf_write(&q.buf, head, value) }; - q.head.store(head + 1, Ordering::Release); - Ok(()) + + if self.cached_tail.get() + q.buf.size() <= head { + let tail = q.tail.curr.load(Ordering::Acquire); + if tail + q.buf.size() <= head { + return if Arc::strong_count(q) < 2 { + Err(TryPushError::Disconnected(value)) + } else { + Err(TryPushError::Full(value)) + }; + } + self.cached_tail.set(tail); } + + unsafe { buf_write(&q.buf, head, value) }; + q.head.store(head + 1, Ordering::Release); + Ok(()) } } impl> Consumer for SPMCConsumer { + fn is_closed(&self) -> bool { + !self.queue.producer.load(Ordering::Relaxed) + } + fn pop(&self) -> Result { let q = &self.queue; - let tail = q.tail.next.fetch_add(1, Ordering::Relaxed); let tail_plus_one = tail + 1; + loop { - if tail_plus_one <= q.head.load(Ordering::Acquire) { + if q.head.load(Ordering::Acquire) >= tail_plus_one { break; - } else if !q.producer.load(Ordering::Acquire) { + } else if !q.producer.load(Ordering::Relaxed) { return Err(PopError::Disconnected); } spin_loop(); } let v = unsafe { buf_read(&q.buf, tail) }; - while q.tail.curr.load(Ordering::Relaxed) < tail { spin_loop(); } @@ -153,12 +175,15 @@ impl> Consumer for SPMCConsumer { loop { let tail = q.tail.curr.load(Ordering::Relaxed); let tail_plus_one = tail + 1; - if tail_plus_one > q.head.load(Ordering::Acquire) { - if q.producer.load(Ordering::Acquire) { - return Err(TryPopError::Empty); + + if q.head.load(Ordering::Acquire) < tail_plus_one { + // buffer is empty, check whether it's closed. + // relaxed is fine since Producer.drop does an acquire/release on .head + return if !q.producer.load(Ordering::Relaxed) { + Err(TryPopError::Disconnected) } else { - return Err(TryPopError::Disconnected); - } + Err(TryPopError::Empty) + }; } else if q .tail .next @@ -175,7 +200,9 @@ impl> Consumer for SPMCConsumer { impl> Drop for SPMCProducer { fn drop(&mut self) { - self.queue.producer.store(false, Ordering::Release); + self.queue.producer.store(false, Ordering::Relaxed); + // Acquire/Release .head to ensure other threads see new .closed + self.queue.head.fetch_add(0, Ordering::AcqRel); } } @@ -272,9 +299,11 @@ mod test { assert_eq!(c.pop(), Err(PopError::Disconnected)); assert_eq!(c.try_pop(), Err(TryPopError::Disconnected)); - let (p, c) = spmc_queue(DynamicBuffer::new(32).unwrap()); + let (p, c) = spmc_queue(DynamicBuffer::new(2).unwrap()); p.push(1).unwrap(); std::mem::drop(c); + assert!(p.is_closed()); + p.push(1).unwrap(); assert_eq!(p.push(2), Err(PushError::Disconnected(2))); assert_eq!(p.try_push(2), Err(TryPushError::Disconnected(2))); diff --git a/src/spsc.rs b/src/spsc.rs index 71be1f5..cba2026 100644 --- a/src/spsc.rs +++ b/src/spsc.rs @@ -5,6 +5,7 @@ //! In other words, both the `SPSCProducer` and `SPSCConsumer` are `Send` and //! `!Sync`. +use std::cell::Cell; use std::hint::spin_loop; use std::marker::PhantomData; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -26,13 +27,19 @@ struct SPSCQueue> { /// Consumer end of the queue. Implements the trait `Consumer`. pub struct SPSCConsumer> { queue: Arc>, - _not_sync: PhantomData> + /// A copy of `queue.head` for quick access. + /// + /// This value can be stale and sometimes needs to be resynchronized with `queue.head`. + cached_head: Cell, } /// Producer end of the queue. Implements the trait `Producer`. pub struct SPSCProducer> { queue: Arc>, - _not_sync: PhantomData> + /// A copy of `queue.tail` for quick access. + /// + /// This value can be stale and sometimes needs to be resynchronized with `queue.tail`. + cached_tail: Cell, } /// Creates a new SPSC queue @@ -62,9 +69,12 @@ pub fn spsc_queue>(buf: B) -> (SPSCProducer, SPSCConsumer< ( SPSCProducer { queue: queue.clone(), - _not_sync: PhantomData, + cached_tail: Cell::new(0), + }, + SPSCConsumer { + queue, + cached_head: Cell::new(0), }, - SPSCConsumer { queue, _not_sync: PhantomData }, ) } @@ -79,17 +89,25 @@ impl> Drop for SPSCQueue { } impl> Producer for SPSCProducer { + fn is_closed(&self) -> bool { + Arc::strong_count(&self.queue) < 2 + } + fn push(&self, value: T) -> Result<(), PushError> { let q = &self.queue; let head = q.head.load(Ordering::Relaxed); - loop { - if Arc::strong_count(&self.queue) < 2 { - return Err(PushError::Disconnected(value)); - } else if q.tail.load(Ordering::Acquire) + q.buf.size() > head { - break; + if self.cached_tail.get() + q.buf.size() <= head { + loop { + let tail = q.tail.load(Ordering::Acquire); + if tail + q.buf.size() > head { + self.cached_tail.set(tail); + break; + } else if Arc::strong_count(q) < 2 { + return Err(PushError::Disconnected(value)); + } + spin_loop(); } - spin_loop(); } unsafe { buf_write(&q.buf, head, value) }; @@ -100,35 +118,49 @@ impl> Producer for SPSCProducer { fn try_push(&self, value: T) -> Result<(), TryPushError> { let q = &self.queue; let head = q.head.load(Ordering::Relaxed); - if Arc::strong_count(&self.queue) < 2 { - Err(TryPushError::Disconnected(value)) - } else if q.tail.load(Ordering::Acquire) + q.buf.size() <= head { - Err(TryPushError::Full(value)) - } else { - unsafe { buf_write(&q.buf, head, value) }; - q.head.store(head + 1, Ordering::Release); - Ok(()) + + if self.cached_tail.get() + q.buf.size() <= head { + let tail = q.tail.load(Ordering::Acquire); + if tail + q.buf.size() <= head { + return if Arc::strong_count(q) < 2 { + Err(TryPushError::Disconnected(value)) + } else { + Err(TryPushError::Full(value)) + }; + } + self.cached_tail.set(tail); } + + unsafe { buf_write(&q.buf, head, value) }; + q.head.store(head + 1, Ordering::Release); + Ok(()) } } impl> Consumer for SPSCConsumer { + fn is_closed(&self) -> bool { + Arc::strong_count(&self.queue) < 2 + } + fn pop(&self) -> Result { let q = &self.queue; - let tail = q.tail.load(Ordering::Relaxed); let tail_plus_one = tail + 1; - loop { - if tail_plus_one <= q.head.load(Ordering::Acquire) { - break; - } else if Arc::strong_count(q) < 2 { - return Err(PopError::Disconnected); + + if self.cached_head.get() < tail_plus_one { + loop { + let head = q.head.load(Ordering::Acquire); + if head >= tail_plus_one { + self.cached_head.set(head); + break; + } else if Arc::strong_count(q) < 2 { + return Err(PopError::Disconnected); + } + spin_loop(); } - spin_loop(); } let v = unsafe { buf_read(&q.buf, tail) }; - q.tail.store(tail_plus_one, Ordering::Release); Ok(v) } @@ -138,17 +170,21 @@ impl> Consumer for SPSCConsumer { let tail = q.tail.load(Ordering::Relaxed); let tail_plus_one = tail + 1; - if tail_plus_one > q.head.load(Ordering::Acquire) { - if Arc::strong_count(q) > 1 { - Err(TryPopError::Empty) - } else { - Err(TryPopError::Disconnected) + if self.cached_head.get() < tail_plus_one { + let head = q.head.load(Ordering::Acquire); + if head < tail_plus_one { + return if Arc::strong_count(q) < 2 { + Err(TryPopError::Disconnected) + } else { + Err(TryPopError::Empty) + }; } - } else { - let v = unsafe { buf_read(&q.buf, tail) }; - q.tail.store(tail_plus_one, Ordering::Release); - Ok(v) + self.cached_head.set(head); } + + let v = unsafe { buf_read(&q.buf, tail) }; + q.tail.store(tail_plus_one, Ordering::Release); + Ok(v) } } @@ -235,9 +271,11 @@ mod test { assert_eq!(c.pop(), Err(PopError::Disconnected)); assert_eq!(c.try_pop(), Err(TryPopError::Disconnected)); - let (p, c) = spsc_queue(DynamicBuffer::new(3).unwrap()); + let (p, c) = spsc_queue(DynamicBuffer::new(2).unwrap()); p.push(1).unwrap(); std::mem::drop(c); + assert!(p.is_closed()); + p.push(1).unwrap(); assert_eq!(p.push(2), Err(PushError::Disconnected(2))); assert_eq!(p.try_push(2), Err(TryPushError::Disconnected(2)));