Restaged repo, allocator and runtime implemented, ioring-backed async fs/net/channel/timer primitives

lib/runtime/src/channel/mpsc.rs

@@ -0,0 +1,575 @@
+use std::collections::VecDeque;
+use std::future::poll_fn;
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+use std::task::{Context, Poll};
+
+use crate::op::completion::{CompletionFuture, CompletionHandle};
+use crate::sys::linux::channel::runtime_waiter;
+
+pub fn channel<T: Send + 'static>(capacity: usize) -> (Sender<T>, Receiver<T>) {
+    assert!(capacity > 0, "bounded channels require capacity > 0");
+    let shared = Arc::new(Mutex::new(State::new(Some(capacity))));
+    (
+        Sender {
+            shared: Arc::clone(&shared),
+        },
+        Receiver { shared },
+    )
+}
+
+pub fn unbounded_channel<T: Send + 'static>() -> (UnboundedSender<T>, Receiver<T>) {
+    let shared = Arc::new(Mutex::new(State::new(None)));
+    (
+        UnboundedSender {
+            shared: Arc::clone(&shared),
+        },
+        Receiver { shared },
+    )
+}
+
+pub struct Sender<T: Send + 'static> {
+    shared: Arc<Mutex<State<T>>>,
+}
+
+pub struct UnboundedSender<T: Send + 'static> {
+    shared: Arc<Mutex<State<T>>>,
+}
+
+pub struct Receiver<T: Send + 'static> {
+    shared: Arc<Mutex<State<T>>>,
+}
+
+struct State<T: Send + 'static> {
+    queue: VecDeque<T>,
+    capacity: Option<usize>,
+    sender_count: usize,
+    receiver_closed: bool,
+    recv_waiter: Option<CompletionHandle<Option<T>>>,
+    send_waiters: VecDeque<SendWaiter<T>>,
+    next_waiter_id: usize,
+}
+
+struct SendWaiter<T: Send + 'static> {
+    id: usize,
+    value: T,
+    handle: CompletionHandle<Result<(), SendError<T>>>,
+}
+
+#[derive(Debug, Eq, PartialEq)]
+pub struct SendError<T>(pub T);
+
+#[derive(Debug, Eq, PartialEq)]
+pub enum TrySendError<T> {
+    Full(T),
+    Closed(T),
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum TryRecvError {
+    Empty,
+    Disconnected,
+}
+
+impl<T: Send + 'static> State<T> {
+    fn new(capacity: Option<usize>) -> Self {
+        Self {
+            queue: VecDeque::new(),
+            capacity,
+            sender_count: 1,
+            receiver_closed: false,
+            recv_waiter: None,
+            send_waiters: VecDeque::new(),
+            next_waiter_id: 1,
+        }
+    }
+
+    fn try_send_now(&mut self, value: T) -> Result<(), TrySendError<T>> {
+        if self.receiver_closed {
+            return Err(TrySendError::Closed(value));
+        }
+
+        if let Some(waiter) = self.recv_waiter.take() {
+            waiter.complete(Some(value));
+            return Ok(());
+        }
+
+        if self
+            .capacity
+            .is_some_and(|capacity| self.queue.len() >= capacity)
+        {
+            return Err(TrySendError::Full(value));
+        }
+
+        self.queue.push_back(value);
+        Ok(())
+    }
+
+    fn enqueue_send_waiter(
+        &mut self,
+        value: T,
+        handle: CompletionHandle<Result<(), SendError<T>>>,
+    ) -> usize {
+        let id = self.next_waiter_id;
+        self.next_waiter_id = self.next_waiter_id.wrapping_add(1);
+        self.send_waiters
+            .push_back(SendWaiter { id, value, handle });
+        id
+    }
+
+    fn remove_send_waiter(&mut self, waiter_id: usize) -> bool {
+        let Some(index) = self
+            .send_waiters
+            .iter()
+            .position(|waiter| waiter.id == waiter_id)
+        else {
+            return false;
+        };
+        self.send_waiters.remove(index);
+        true
+    }
+
+    fn pump_senders(&mut self) {
+        loop {
+            if self.receiver_closed {
+                self.fail_pending_senders();
+                break;
+            }
+
+            let has_capacity = self
+                .capacity
+                .is_none_or(|capacity| self.queue.len() < capacity);
+            if !has_capacity {
+                break;
+            }
+
+            let Some(waiter) = self.send_waiters.pop_front() else {
+                break;
+            };
+
+            if let Some(receiver) = self.recv_waiter.take() {
+                receiver.complete(Some(waiter.value));
+            } else {
+                self.queue.push_back(waiter.value);
+            }
+            waiter.handle.complete(Ok(()));
+        }
+
+        if self.queue.is_empty()
+            && self.sender_count == 0
+            && let Some(waiter) = self.recv_waiter.take()
+        {
+            waiter.complete(None);
+        }
+    }
+
+    fn fail_pending_senders(&mut self) {
+        while let Some(waiter) = self.send_waiters.pop_front() {
+            waiter.handle.complete(Err(SendError(waiter.value)));
+        }
+    }
+
+    fn close_receiver(&mut self) {
+        self.receiver_closed = true;
+        self.fail_pending_senders();
+        if self.queue.is_empty()
+            && let Some(waiter) = self.recv_waiter.take()
+        {
+            waiter.complete(None);
+        }
+    }
+
+    fn drop_sender(&mut self) {
+        self.sender_count = self.sender_count.saturating_sub(1);
+        if self.sender_count == 0
+            && self.queue.is_empty()
+            && let Some(waiter) = self.recv_waiter.take()
+        {
+            waiter.complete(None);
+        }
+    }
+}
+
+impl<T: Send + 'static> Clone for Sender<T> {
+    fn clone(&self) -> Self {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .sender_count += 1;
+        Self {
+            shared: Arc::clone(&self.shared),
+        }
+    }
+}
+
+impl<T: Send + 'static> Clone for UnboundedSender<T> {
+    fn clone(&self) -> Self {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .sender_count += 1;
+        Self {
+            shared: Arc::clone(&self.shared),
+        }
+    }
+}
+
+impl<T: Send + 'static> Sender<T> {
+    pub async fn send(&self, value: T) -> Result<(), SendError<T>> {
+        let mut value = Some(value);
+        let mut wait = None;
+        poll_fn(|cx| self.poll_send(cx, &mut value, &mut wait)).await
+    }
+
+    pub fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .try_send_now(value)
+    }
+
+    pub fn is_closed(&self) -> bool {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .receiver_closed
+    }
+
+    fn poll_send(
+        &self,
+        cx: &mut Context<'_>,
+        value_slot: &mut Option<T>,
+        wait: &mut Option<CompletionFuture<Result<(), SendError<T>>>>,
+    ) -> Poll<Result<(), SendError<T>>> {
+        if let Some(future) = wait.as_mut() {
+            match Pin::new(future).poll(cx) {
+                Poll::Ready(result) => {
+                    wait.take();
+                    Poll::Ready(result)
+                }
+                Poll::Pending => Poll::Pending,
+            }
+        } else {
+            let mut state = self
+                .shared
+                .lock()
+                .expect("mpsc state should not be poisoned");
+            match state.try_send_now(value_slot.take().expect("send value should be present")) {
+                Ok(()) => Poll::Ready(Ok(())),
+                Err(TrySendError::Closed(value)) => Poll::Ready(Err(SendError(value))),
+                Err(TrySendError::Full(returned)) => {
+                    drop(state);
+                    let (future, handle) = runtime_waiter::<Result<(), SendError<T>>>();
+                    let state_shared = Arc::clone(&self.shared);
+                    let registration = {
+                        let mut state = state_shared
+                            .lock()
+                            .expect("mpsc state should not be poisoned");
+                        match state.try_send_now(returned) {
+                            Ok(()) => Ok(None),
+                            Err(TrySendError::Closed(value)) => Err(SendError(value)),
+                            Err(TrySendError::Full(value)) => {
+                                Ok(Some(state.enqueue_send_waiter(value, handle.clone())))
+                            }
+                        }
+                    };
+                    match registration {
+                        Ok(None) => {
+                            handle.complete(Ok(()));
+                            *wait = Some(future);
+                            self.poll_send(cx, value_slot, wait)
+                        }
+                        Err(error) => {
+                            handle.complete(Err(error));
+                            *wait = Some(future);
+                            self.poll_send(cx, value_slot, wait)
+                        }
+                        Ok(Some(waiter_id)) => {
+                            let cancel_shared = Arc::clone(&self.shared);
+                            let cancel_handle = handle.clone();
+                            handle.set_cancel(move || {
+                                let mut state = cancel_shared
+                                    .lock()
+                                    .expect("mpsc state should not be poisoned");
+                                let _ = state.remove_send_waiter(waiter_id);
+                                drop(state);
+                                cancel_handle.finish(None);
+                            });
+                            *wait = Some(future);
+                            self.poll_send(cx, value_slot, wait)
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+impl<T: Send + 'static> UnboundedSender<T> {
+    pub fn send(&self, value: T) -> Result<(), SendError<T>> {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .try_send_now(value)
+            .map_err(|error| match error {
+                TrySendError::Full(value) | TrySendError::Closed(value) => SendError(value),
+            })
+    }
+
+    pub fn is_closed(&self) -> bool {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .receiver_closed
+    }
+}
+
+impl<T: Send + 'static> Receiver<T> {
+    pub async fn recv(&mut self) -> Option<T> {
+        let mut wait = None;
+        poll_fn(|cx| self.poll_recv(cx, &mut wait)).await
+    }
+
+    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
+        let mut state = self
+            .shared
+            .lock()
+            .expect("mpsc state should not be poisoned");
+        if let Some(value) = state.queue.pop_front() {
+            state.pump_senders();
+            Ok(value)
+        } else if state.sender_count == 0 || state.receiver_closed {
+            Err(TryRecvError::Disconnected)
+        } else {
+            Err(TryRecvError::Empty)
+        }
+    }
+
+    pub fn close(&mut self) {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .close_receiver();
+    }
+
+    pub fn is_closed(&self) -> bool {
+        let state = self
+            .shared
+            .lock()
+            .expect("mpsc state should not be poisoned");
+        state.receiver_closed || state.sender_count == 0
+    }
+
+    fn poll_recv(
+        &mut self,
+        cx: &mut Context<'_>,
+        wait: &mut Option<CompletionFuture<Option<T>>>,
+    ) -> Poll<Option<T>> {
+        if let Some(future) = wait.as_mut() {
+            match Pin::new(future).poll(cx) {
+                Poll::Ready(result) => {
+                    wait.take();
+                    Poll::Ready(result)
+                }
+                Poll::Pending => Poll::Pending,
+            }
+        } else {
+            let (future, handle) = runtime_waiter::<Option<T>>();
+            let cancel_shared = Arc::clone(&self.shared);
+            let cancel_handle = handle.clone();
+            handle.set_cancel(move || {
+                let mut state = cancel_shared
+                    .lock()
+                    .expect("mpsc state should not be poisoned");
+                let _ = state.recv_waiter.take();
+                drop(state);
+                cancel_handle.finish(None);
+            });
+
+            {
+                let mut state = self
+                    .shared
+                    .lock()
+                    .expect("mpsc state should not be poisoned");
+                if let Some(value) = state.queue.pop_front() {
+                    state.pump_senders();
+                    handle.complete(Some(value));
+                } else if state.receiver_closed || state.sender_count == 0 {
+                    handle.complete(None);
+                } else {
+                    assert!(
+                        state.recv_waiter.is_none(),
+                        "only one mpsc receive operation may wait at a time"
+                    );
+                    state.recv_waiter = Some(handle.clone());
+                }
+            }
+
+            *wait = Some(future);
+            self.poll_recv(cx, wait)
+        }
+    }
+}
+
+impl<T: Send + 'static> Drop for Sender<T> {
+    fn drop(&mut self) {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .drop_sender();
+    }
+}
+
+impl<T: Send + 'static> Drop for UnboundedSender<T> {
+    fn drop(&mut self) {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .drop_sender();
+    }
+}
+
+impl<T: Send + 'static> Drop for Receiver<T> {
+    fn drop(&mut self) {
+        self.shared
+            .lock()
+            .expect("mpsc state should not be poisoned")
+            .close_receiver();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::{Arc, Mutex};
+    use std::time::Duration;
+
+    use crate::time::sleep;
+    use crate::{queue_future, queue_task, run, spawn_worker};
+
+    use super::{TryRecvError, TrySendError, channel, unbounded_channel};
+
+    #[test]
+    fn bounded_channel_applies_backpressure() {
+        let log = Arc::new(Mutex::new(Vec::<String>::new()));
+        let log_for_task = Arc::clone(&log);
+
+        queue_task(move || {
+            let (sender, mut receiver) = channel(1);
+            let log_for_sender = Arc::clone(&log_for_task);
+            let log_for_receiver = Arc::clone(&log_for_task);
+
+            queue_future(async move {
+                sender
+                    .send("first")
+                    .await
+                    .expect("first send should succeed");
+                log_for_sender
+                    .lock()
+                    .unwrap()
+                    .push("sent first".to_string());
+                sender
+                    .send("second")
+                    .await
+                    .expect("second send should succeed");
+                log_for_sender
+                    .lock()
+                    .unwrap()
+                    .push("sent second".to_string());
+            });
+
+            queue_future(async move {
+                sleep(Duration::from_millis(5)).await;
+                let first = receiver.recv().await.expect("first recv should succeed");
+                log_for_receiver
+                    .lock()
+                    .unwrap()
+                    .push(format!("received {first}"));
+                let second = receiver.recv().await.expect("second recv should succeed");
+                log_for_receiver
+                    .lock()
+                    .unwrap()
+                    .push(format!("received {second}"));
+            });
+        });
+        run();
+
+        let log = log.lock().unwrap();
+        let sent_first = log.iter().position(|entry| entry == "sent first").unwrap();
+        let received_first = log
+            .iter()
+            .position(|entry| entry == "received first")
+            .unwrap();
+        let sent_second = log.iter().position(|entry| entry == "sent second").unwrap();
+        let received_second = log
+            .iter()
+            .position(|entry| entry == "received second")
+            .unwrap();
+
+        assert!(
+            sent_first < received_first,
+            "first send should happen before first recv"
+        );
+        assert!(
+            received_first < sent_second,
+            "second send should not complete before capacity is freed"
+        );
+        assert!(
+            received_first < received_second,
+            "receiver should observe messages in FIFO order"
+        );
+    }
+
+    #[test]
+    fn unbounded_channel_moves_messages_across_worker_threads() {
+        let log = Arc::new(Mutex::new(Vec::new()));
+        let log_for_task = Arc::clone(&log);
+
+        queue_task(move || {
+            let (sender, mut receiver) = unbounded_channel::<String>();
+            let worker_sender = sender.clone();
+            let log_for_receiver = Arc::clone(&log_for_task);
+
+            let _worker = spawn_worker(
+                move || {
+                    queue_task(move || {
+                        worker_sender
+                            .send("worker boot".into())
+                            .expect("worker boot send should succeed");
+                        worker_sender
+                            .send("worker done".into())
+                            .expect("worker done send should succeed");
+                    });
+                },
+                || {},
+            );
+            drop(sender);
+
+            queue_future(async move {
+                while let Some(message) = receiver.recv().await {
+                    log_for_receiver.lock().unwrap().push(message);
+                }
+            });
+        });
+        run();
+
+        assert_eq!(
+            log.lock().unwrap().as_slice(),
+            ["worker boot", "worker done"]
+        );
+    }
+
+    #[test]
+    fn try_send_try_recv_and_close_semantics_work() {
+        let (sender, mut receiver) = channel(1);
+        sender
+            .try_send(1usize)
+            .expect("initial send should succeed");
+        assert_eq!(sender.try_send(2usize), Err(TrySendError::Full(2)));
+        assert_eq!(receiver.try_recv(), Ok(1));
+        assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty));
+        receiver.close();
+        assert!(sender.is_closed(), "sender should observe closed receiver");
+        assert_eq!(sender.try_send(3usize), Err(TrySendError::Closed(3)));
+        assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected));
+    }
+}