use crate::{
QueueError,
owned::{MpmcQueue, QueueBuilder},
traits::{QueueConsumer, QueueFactory, QueueProducer},
};
use std::{fmt, marker::PhantomData, sync::Arc};
pub struct PointerQueue<T, I = u32, const N: usize = 0>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
queue: Arc<MpmcQueue<usize, I, N>>,
_phantom: PhantomData<T>,
}
impl<T, I, const N: usize> fmt::Debug for PointerQueue<T, I, N>
where
T: Send + Sync + fmt::Debug,
I: Copy + Into<u128> + fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PointerQueue")
.field("capacity", &self.capacity())
.field("len", &self.len())
.field("is_empty", &self.is_empty())
.finish()
}
}
#[derive(Debug, Clone)]
pub struct PointerQueueBuilder<T, I = u32>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
inner: QueueBuilder<usize, I>,
_phantom: PhantomData<T>,
}
impl<T, I> Default for PointerQueueBuilder<T, I>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
fn default() -> Self {
Self::new()
}
}
impl<T, I> PointerQueueBuilder<T, I>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
pub const fn new() -> Self {
Self {
inner: QueueBuilder::new(),
_phantom: PhantomData,
}
}
#[must_use]
pub const fn capacity(mut self, cap: usize) -> Self {
self.inner = self.inner.capacity(cap);
self
}
pub fn build(self) -> Result<Arc<PointerQueue<T, I>>, QueueError> {
let queue = self.inner.build()?;
Ok(Arc::new(PointerQueue {
queue,
_phantom: PhantomData,
}))
}
pub fn build_static<const N: usize>(self) -> Result<Arc<PointerQueue<T, I, N>>, QueueError> {
let queue = self.inner.build_static::<N>()?;
Ok(Arc::new(PointerQueue {
queue,
_phantom: PhantomData,
}))
}
pub fn channels(self) -> Result<(PointerProducer<T, I>, PointerConsumer<T, I>), QueueError> {
let queue = self.build()?;
Ok((queue.producer(), queue.consumer()))
}
pub fn channels_static<const N: usize>(
self,
) -> Result<(PointerProducer<T, I, N>, PointerConsumer<T, I, N>), QueueError> {
let queue = self.build_static::<N>()?;
Ok((queue.producer(), queue.consumer()))
}
}
pub const fn pointer_queue<T>() -> PointerQueueBuilder<T, u32>
where
T: Send + Sync,
{
PointerQueueBuilder::new()
}
pub const fn pointer_queue_with_index<T, I>() -> PointerQueueBuilder<T, I>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
PointerQueueBuilder::new()
}
impl<T, I, const N: usize> PointerQueue<T, I, N>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
pub fn new(capacity: usize) -> Result<Self, QueueError> {
let queue = Arc::new(MpmcQueue::new(capacity)?);
Ok(Self {
queue,
_phantom: PhantomData,
})
}
pub fn capacity(&self) -> usize {
self.queue.capacity()
}
pub fn len(&self) -> usize {
self.queue.len()
}
pub fn is_empty(&self) -> bool {
self.queue.is_empty()
}
pub fn is_full(&self) -> bool {
self.queue.is_full()
}
pub fn try_push(&self, arc: Arc<T>) -> Result<(), (Arc<T>, QueueError)> {
let raw_ptr = Arc::into_raw(arc) as usize;
match self.queue.try_push(raw_ptr) {
Ok(()) => Ok(()),
Err((ptr, err)) => {
let recovered_arc = unsafe { Arc::from_raw(ptr as *const T) };
Err((recovered_arc, err))
},
}
}
pub fn push(&self, arc: Arc<T>) -> Result<(), QueueError> {
let raw_ptr = Arc::into_raw(arc) as usize;
self.queue.push(raw_ptr)
}
pub fn try_pop(&self) -> Result<Arc<T>, QueueError> {
match self.queue.try_pop() {
Ok(raw_ptr) => {
let arc = unsafe { Arc::from_raw(raw_ptr as *const T) };
Ok(arc)
},
Err(e) => Err(e),
}
}
pub fn pop(&self) -> Result<Arc<T>, QueueError> {
match self.queue.pop() {
Ok(raw_ptr) => {
let arc = unsafe { Arc::from_raw(raw_ptr as *const T) };
Ok(arc)
},
Err(e) => Err(e),
}
}
pub fn peek(&self) -> Result<Arc<T>, QueueError> {
match self.queue.peek() {
Ok(raw_ptr) => {
let temp_arc = unsafe { Arc::from_raw(raw_ptr as *const T) };
let cloned_arc = Arc::clone(&temp_arc);
let _ = Arc::into_raw(temp_arc);
Ok(cloned_arc)
},
Err(e) => Err(e),
}
}
}
pub type PointerProducer<T, I = u32, const N: usize = 0> = PointerProducerHandle<T, I, N>;
pub type PointerConsumer<T, I = u32, const N: usize = 0> = PointerConsumerHandle<T, I, N>;
#[derive(Debug)]
pub struct PointerProducerHandle<T, I = u32, const N: usize = 0>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
queue: Arc<PointerQueue<T, I, N>>,
}
impl<T, I, const N: usize> Clone for PointerProducerHandle<T, I, N>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
fn clone(&self) -> Self {
Self {
queue: self.queue.clone(),
}
}
}
impl<T, I, const N: usize> QueueProducer<Arc<T>> for PointerProducerHandle<T, I, N>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
fn try_push(&self, value: Arc<T>) -> Result<(), (Arc<T>, QueueError)> {
self.queue.try_push(value)
}
fn push(&self, arc: Arc<T>) -> Result<(), QueueError> {
self.queue.push(arc)
}
fn push_with_seq(&self, arc: Arc<T>) -> Result<usize, QueueError> {
let raw_ptr = Arc::into_raw(arc) as usize;
self.queue.queue.push_impl(raw_ptr, true)
}
}
#[derive(Debug)]
pub struct PointerConsumerHandle<T, I = u32, const N: usize = 0>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
queue: Arc<PointerQueue<T, I, N>>,
}
impl<T, I, const N: usize> Clone for PointerConsumerHandle<T, I, N>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
fn clone(&self) -> Self {
Self {
queue: self.queue.clone(),
}
}
}
impl<T, I, const N: usize> QueueConsumer<Arc<T>> for PointerConsumerHandle<T, I, N>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
fn try_pop(&self) -> Result<Arc<T>, QueueError> {
self.queue.try_pop()
}
fn pop(&self) -> Result<Arc<T>, QueueError> {
self.queue.pop()
}
fn pop_with_seq(&self) -> Result<(Arc<T>, usize), QueueError> {
match self.queue.queue.pop_impl(true) {
Ok((raw_ptr, idx)) => {
let arc = unsafe { Arc::from_raw(raw_ptr as *const T) };
Ok((arc, idx))
},
Err(e) => Err(e),
}
}
fn peek(&self) -> Result<Arc<T>, QueueError> {
self.queue.peek()
}
fn peek_with_seq(&self) -> Result<(Arc<T>, usize), QueueError> {
let seq = self.queue.queue.len();
let arc = self.queue.peek()?;
Ok((arc, seq))
}
fn pop_if<F>(&self, mut predicate: F) -> Result<Arc<T>, QueueError>
where
F: FnMut(&Arc<T>, usize) -> bool,
{
let (peeked_arc, seq) = self.peek_with_seq()?;
if predicate(&peeked_arc, seq) {
self.pop()
} else {
Err(QueueError::Empty)
}
}
fn consume<F>(&self, mut consumer: F) -> usize
where
F: FnMut(Arc<T>, usize) -> bool,
{
let mut count = 0;
while let Ok((arc, seq)) = self.pop_with_seq() {
count += 1;
if consumer(arc, seq) {
break;
}
}
count
}
fn is_empty(&self) -> bool {
self.queue.is_empty()
}
fn size(&self) -> usize {
self.queue.len()
}
}
impl<T, I, const N: usize> QueueFactory<Arc<T>> for Arc<PointerQueue<T, I, N>>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
type Producer = PointerProducerHandle<T, I, N>;
type Consumer = PointerConsumerHandle<T, I, N>;
fn producer(&self) -> Self::Producer {
PointerProducerHandle {
queue: self.clone(),
}
}
fn consumer(&self) -> Self::Consumer {
PointerConsumerHandle {
queue: self.clone(),
}
}
}
unsafe impl<T, I, const N: usize> Send for PointerQueue<T, I, N>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
}
unsafe impl<T, I, const N: usize> Sync for PointerQueue<T, I, N>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
}
impl<T, I, const N: usize> Drop for PointerQueue<T, I, N>
where
T: Send + Sync,
I: Copy + Into<u128>,
{
fn drop(&mut self) {
while self.try_pop().is_ok() {}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::{
sync::atomic::{AtomicUsize, Ordering},
time::Instant,
};
#[derive(Debug, Clone, PartialEq)]
struct LargeData {
id: u64,
data: Vec<u8>,
name: String,
}
impl LargeData {
fn new(id: u64, size: usize) -> Self {
Self {
id,
data: vec![0u8; size],
name: format!("item_{id}"),
}
}
}
#[test]
fn test_builder_pattern() {
let queue = pointer_queue::<LargeData>().capacity(16).build().unwrap();
assert_eq!(queue.capacity(), 16);
assert!(queue.is_empty());
}
#[test]
fn test_channels() {
let (producer, consumer) = pointer_queue::<LargeData>().capacity(8).channels().unwrap();
let data1 = Arc::new(LargeData::new(1, 1024));
let data2 = Arc::new(LargeData::new(2, 2048));
producer.push(data1.clone()).unwrap();
producer.push(data2.clone()).unwrap();
let popped1 = consumer.pop().unwrap();
let popped2 = consumer.pop().unwrap();
assert_eq!(*popped1, *data1);
assert_eq!(*popped2, *data2);
assert!(consumer.is_empty());
}
#[test]
fn test_try_operations() {
let queue = pointer_queue::<LargeData>().capacity(2).build().unwrap();
let data1 = Arc::new(LargeData::new(1, 128));
let data2 = Arc::new(LargeData::new(2, 128));
let data3 = Arc::new(LargeData::new(3, 128));
assert!(queue.try_push(data1.clone()).is_ok());
assert!(queue.try_push(data2).is_ok());
assert!(queue.is_full());
match queue.try_push(data3.clone()) {
Err((returned_arc, QueueError::Full)) => {
assert_eq!(*returned_arc, *data3);
},
_ => panic!("Expected full queue error"),
}
let popped = queue.try_pop().unwrap();
assert_eq!(*popped, *data1);
}
#[test]
fn test_peek() {
let queue = pointer_queue::<LargeData>().capacity(4).build().unwrap();
let data = Arc::new(LargeData::new(99, 256));
queue.push(data.clone()).unwrap();
let peeked = queue.peek().unwrap();
assert_eq!(*peeked, *data);
let popped = queue.pop().unwrap();
assert_eq!(*popped, *data);
}
#[test]
fn test_reference_counting() {
let queue = pointer_queue::<LargeData>().capacity(4).build().unwrap();
let data = Arc::new(LargeData::new(123, 64));
assert_eq!(Arc::strong_count(&data), 1);
queue.push(data.clone()).unwrap();
assert_eq!(Arc::strong_count(&data), 2);
let peeked = queue.peek().unwrap();
assert_eq!(Arc::strong_count(&data), 3);
drop(peeked);
assert_eq!(Arc::strong_count(&data), 2);
let popped = queue.pop().unwrap();
assert_eq!(Arc::strong_count(&data), 2);
drop(popped);
assert_eq!(Arc::strong_count(&data), 1);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn stress_test() {
const CAPACITY: usize = 128;
const PRODUCERS: usize = 2;
const CONSUMERS: usize = 2;
const ITEMS_PER_PRODUCER: usize = 10_000;
let (producer, consumer) = pointer_queue::<LargeData>()
.capacity(CAPACITY)
.channels()
.unwrap();
let total_items = PRODUCERS * ITEMS_PER_PRODUCER;
let consumed_count = Arc::new(AtomicUsize::new(0));
let mut consumer_handles = Vec::new();
for _ in 0..CONSUMERS {
let consumer = consumer.clone();
let consumed_clone = consumed_count.clone();
let handle = tokio::task::spawn(async move {
loop {
if consumed_clone.load(Ordering::SeqCst) >= total_items {
break;
}
match consumer.try_pop() {
Ok(_data) => {
consumed_clone.fetch_add(1, Ordering::SeqCst);
},
Err(QueueError::Empty) => {
tokio::task::yield_now().await;
},
Err(e) => panic!("Unexpected error: {e:?}"),
}
}
});
consumer_handles.push(handle);
}
let mut producer_handles = Vec::new();
let start = Instant::now();
for producer_id in 0..PRODUCERS {
let producer = producer.clone();
let handle = tokio::task::spawn(async move {
for item_id in 0..ITEMS_PER_PRODUCER {
let data = Arc::new(LargeData::new(
(producer_id * ITEMS_PER_PRODUCER + item_id) as u64,
64,
));
loop {
match producer.try_push(data.clone()) {
Ok(()) => break,
Err((_, QueueError::Full)) => {
tokio::task::yield_now().await;
},
Err((_, e)) => panic!("Unexpected error: {e:?}"),
}
}
}
});
producer_handles.push(handle);
}
for handle in producer_handles {
handle.await.unwrap();
}
while consumed_count.load(Ordering::SeqCst) < total_items {
tokio::time::sleep(tokio::time::Duration::from_millis(1)).await;
}
for handle in consumer_handles {
handle.await.unwrap();
}
let elapsed = start.elapsed();
let throughput = (total_items as f64) / elapsed.as_secs_f64();
println!(
"Pointer queue stress test: {PRODUCERS} producers, {CONSUMERS} consumers, {ITEMS_PER_PRODUCER} items each = {total_items} total in {elapsed:?} ({throughput:.0} ops/sec)"
);
assert_eq!(consumed_count.load(Ordering::SeqCst), total_items);
}
}