Async unbuffered channel that blocks on read and on write - rust

Question

For threaded applications, the Rust standard library provides std::sync::mpsc::sync_channel, a buffered channel which blocks on the reading end when the buffer is empty and blocks on the writing end when the buffer is full. In particular, if you set the buffer size to 0, then any read or write will block until there is a matching write or read.
For async code, there is futures::channel::mpsc::channel, but this does not have the same behavior. Here the minimum capacity is the number of senders on the channel, which is greater than 0. The sending end can still block (because it implements Sink, so we can use SinkExt::send and await it), but only after there's at least one thing already in the buffer.
I took a look to see if there were any packages that provide the functionality I'm looking for, but I could not find anything. Tokio provides lots of nice async synchronization primitives, but none of them did quite what I'm looking for. Plus, my program is going to run in the browser, so I don't think I'm able to use a runtime like Tokio. Does anyone know of a package that fits my use case? I would try to implement this myself, since this almost feels like the most minimal use case for the Sink and Stream traits, but even a minimal implementation of these traits seems like it would be really complicated. Thoughts?
Edit: here's a minimal example of what I mean:
fn main() {
let (tx, rx) = blocking_channel();
let ft = async move {
tx.send(3).await;
println!("message sent and received");
}
let fr = async move {
let msg = rx.recv().await;
println!("received {}", msg);
}
block_on(async { join!(ft, fr) });
}
In this example, whichever future runs first will yield to the other, and only print after both rx.recv and tx.send have been called. Obviously, the receiving end can only progress after tx.send has been called, but I want the less obvious behavior of the transmitting end also having to wait.

Answer 1

Interesting question. I don't think something like that already exists.
Here is a quick proof-of-concept prototype I wrote for that. It's not the prettiest, but it seems to work. There might be a better struct layout than just wrapping everything in a RefCell<Option<...>>, though. And I don't particularly like the sender_dropped and receiver_dropped variables.
Be sure to unittest it properly if used in production!!!
extern crate alloc;
use alloc::rc::Rc;
use core::cell::RefCell;
use core::pin::Pin;
use core::task::{Poll, Waker};
use futures::SinkExt;
use futures::StreamExt;
struct Pipe<T> {
send_waker: RefCell<Option<Waker>>,
receive_waker: RefCell<Option<Waker>>,
value: RefCell<Option<T>>,
sender_dropped: RefCell<bool>,
receiver_dropped: RefCell<bool>,
}
impl<T> Pipe<T> {
fn new() -> Rc<Pipe<T>> {
Rc::new(Self {
value: RefCell::new(None),
send_waker: RefCell::new(None),
receive_waker: RefCell::new(None),
sender_dropped: RefCell::new(false),
receiver_dropped: RefCell::new(false),
})
}
}
impl<T> Pipe<T> {
fn wake_sender(&self) {
if let Some(waker) = self.send_waker.replace(None) {
waker.wake();
}
}
fn wake_receiver(&self) {
if let Some(waker) = self.receive_waker.replace(None) {
waker.wake();
}
}
}
pub struct PipeSender<T> {
pipe: Rc<Pipe<T>>,
}
pub struct PipeReceiver<T> {
pipe: Rc<Pipe<T>>,
}
pub fn create_pipe<T>() -> (PipeSender<T>, PipeReceiver<T>) {
let pipe = Pipe::new();
(PipeSender { pipe: pipe.clone() }, PipeReceiver { pipe })
}
impl<T> futures::Sink<T> for PipeSender<T> {
type Error = ();
fn poll_ready(
self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> Poll<Result<(), Self::Error>> {
let result = if *self.pipe.receiver_dropped.borrow() {
Poll::Ready(Err(()))
} else if self.pipe.receive_waker.borrow().is_some() && self.pipe.value.borrow().is_none() {
Poll::Ready(Ok(()))
} else {
self.pipe.send_waker.replace(Some(cx.waker().clone()));
Poll::Pending
};
// Wake potential receiver
self.pipe.wake_receiver();
result
}
fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
let prev = self.pipe.value.replace(Some(item));
assert!(prev.is_none(), "A value got lost in the pipe.");
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_: &mut futures::task::Context<'_>,
) -> Poll<Result<(), Self::Error>> {
// Noop, start_send already completes the send
Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_: &mut std::task::Context<'_>,
) -> Poll<Result<(), Self::Error>> {
// Noop, start_send already completes the send
Poll::Ready(Ok(()))
}
}
impl<T> futures::Stream for PipeReceiver<T> {
type Item = T;
fn poll_next(
self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
let result = {
let value = self.pipe.value.replace(None);
if let Some(value) = value {
Poll::Ready(Some(value))
} else if *self.pipe.sender_dropped.borrow() {
Poll::Ready(None)
} else {
self.pipe.receive_waker.replace(Some(cx.waker().clone()));
Poll::Pending
}
};
// Wake potential sender
self.pipe.wake_sender();
result
}
}
impl<T> Drop for PipeSender<T> {
fn drop(&mut self) {
self.pipe.sender_dropped.replace(true);
self.pipe.wake_receiver();
}
}
impl<T> Drop for PipeReceiver<T> {
fn drop(&mut self) {
self.pipe.receiver_dropped.replace(true);
self.pipe.wake_sender();
}
}
#[tokio::main]
async fn main() {
use std::time::Duration;
let (mut sender, mut receiver) = create_pipe();
tokio::join!(
async move {
for i in 0..5u32 {
println!("Sending {i} ...");
if let Err(_) = sender.send(i).await {
println!("Stream closed.");
break;
}
println!("Sent {i}.");
}
println!("Sender closed.");
},
async move {
println!("Attempting to receive ...");
while let Some(val) = receiver.next().await {
println!("Received: {val}");
println!("\n=== Waiting ... ===\n");
tokio::time::sleep(Duration::from_secs(1)).await;
println!("Attempting to receive ...");
}
println!("Receiver closed.");
}
);
}
Sending 0 ...
Attempting to receive ...
Sent 0.
Sending 1 ...
Received: 0
=== Waiting ... ===
Attempting to receive ...
Sent 1.
Sending 2 ...
Received: 1
=== Waiting ... ===
Attempting to receive ...
Sent 2.
Sending 3 ...
Received: 2
=== Waiting ... ===
Attempting to receive ...
Sent 3.
Sending 4 ...
Received: 3
=== Waiting ... ===
Attempting to receive ...
Sent 4.
Sender closed.
Received: 4
=== Waiting ... ===
Attempting to receive ...
Receiver closed.