memkit_async/

channel.rs

//! Zero-copy async channels.

use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

/// Zero-copy channel for transferring ownership between tasks.
pub struct MkAsyncChannel<T> {
    inner: Arc<ChannelInner<T>>,
}

struct ChannelInner<T> {
    queue: Mutex<VecDeque<T>>,
    capacity: usize,
    len: AtomicUsize,
    closed: AtomicBool,
}

impl<T> MkAsyncChannel<T> {
    /// Create a new bounded channel.
    pub fn bounded(capacity: usize) -> Self {
        Self {
            inner: Arc::new(ChannelInner {
                queue: Mutex::new(VecDeque::with_capacity(capacity)),
                capacity,
                len: AtomicUsize::new(0),
                closed: AtomicBool::new(false),
            }),
        }
    }

    /// Create sender and receiver handles.
    pub fn split(self) -> (MkAsyncSender<T>, MkAsyncReceiver<T>) {
        (
            MkAsyncSender { inner: Arc::clone(&self.inner) },
            MkAsyncReceiver { inner: self.inner },
        )
    }

    /// Get the channel capacity.
    pub fn capacity(&self) -> usize {
        self.inner.capacity
    }

    /// Get the current length.
    pub fn len(&self) -> usize {
        self.inner.len.load(Ordering::Relaxed)
    }

    /// Check if empty.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// Sender half of an async channel.
pub struct MkAsyncSender<T> {
    inner: Arc<ChannelInner<T>>,
}

impl<T> MkAsyncSender<T> {
    /// Send a value through the channel.
    pub async fn send(&self, value: T) -> Result<(), SendError<T>> {
        if self.inner.closed.load(Ordering::Acquire) {
            return Err(SendError::Closed(value));
        }

        loop {
            let len = self.inner.len.load(Ordering::Acquire);
            if len < self.inner.capacity {
                let mut queue = self.inner.queue.lock().unwrap();
                if queue.len() < self.inner.capacity {
                    queue.push_back(value);
                    self.inner.len.fetch_add(1, Ordering::Release);
                    return Ok(());
                }
            }
            
            // Yield and retry
            YieldOnce::new().await;
            
            if self.inner.closed.load(Ordering::Acquire) {
                return Err(SendError::Closed(value));
            }
        }
    }

    /// Try to send without waiting.
    pub fn try_send(&self, value: T) -> Result<(), SendError<T>> {
        if self.inner.closed.load(Ordering::Acquire) {
            return Err(SendError::Closed(value));
        }

        let mut queue = self.inner.queue.lock().unwrap();
        if queue.len() < self.inner.capacity {
            queue.push_back(value);
            self.inner.len.fetch_add(1, Ordering::Release);
            Ok(())
        } else {
            Err(SendError::Full(value))
        }
    }

    /// Close the channel.
    pub fn close(&self) {
        self.inner.closed.store(true, Ordering::Release);
    }
}

impl<T> Clone for MkAsyncSender<T> {
    fn clone(&self) -> Self {
        Self { inner: Arc::clone(&self.inner) }
    }
}

/// Receiver half of an async channel.
pub struct MkAsyncReceiver<T> {
    inner: Arc<ChannelInner<T>>,
}

impl<T> MkAsyncReceiver<T> {
    /// Receive a value from the channel.
    pub async fn recv(&self) -> Option<T> {
        loop {
            if let Some(value) = self.try_recv() {
                return Some(value);
            }
            
            if self.inner.closed.load(Ordering::Acquire) && self.inner.len.load(Ordering::Acquire) == 0 {
                return None;
            }
            
            YieldOnce::new().await;
        }
    }

    /// Try to receive without waiting.
    pub fn try_recv(&self) -> Option<T> {
        let mut queue = self.inner.queue.lock().unwrap();
        if let Some(value) = queue.pop_front() {
            self.inner.len.fetch_sub(1, Ordering::Release);
            Some(value)
        } else {
            None
        }
    }

    /// Check if the channel is closed.
    pub fn is_closed(&self) -> bool {
        self.inner.closed.load(Ordering::Acquire)
    }
}

/// Error when sending fails.
#[derive(Debug)]
pub enum SendError<T> {
    Full(T),
    Closed(T),
}

impl<T> SendError<T> {
    /// Get the value that failed to send.
    pub fn into_inner(self) -> T {
        match self {
            SendError::Full(v) | SendError::Closed(v) => v,
        }
    }
}

/// Yield once to the runtime.
struct YieldOnce(bool);

impl YieldOnce {
    fn new() -> Self {
        Self(false)
    }
}

impl Future for YieldOnce {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.0 {
            Poll::Ready(())
        } else {
            self.0 = true;
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_channel_sync() {
        let channel: MkAsyncChannel<u32> = MkAsyncChannel::bounded(3);
        let (tx, rx) = channel.split();
        
        tx.try_send(1).unwrap();
        tx.try_send(2).unwrap();
        tx.try_send(3).unwrap();
        assert!(tx.try_send(4).is_err());
        
        assert_eq!(rx.try_recv(), Some(1));
        assert_eq!(rx.try_recv(), Some(2));
        assert_eq!(rx.try_recv(), Some(3));
        assert_eq!(rx.try_recv(), None);
    }
}