A fully async mailbox

src/mailbox.rs

@@ -1,5 +1,6 @@
 use std::{
 	future::Future,
+	pin::Pin,
 	sync::Mutex,
 	task::{Context, Poll, Waker},
 };
@@ -8,45 +9,70 @@ pub struct Mailbox<T>(Mutex<MailboxInner<T>>);
 
 impl<T> Mailbox<T> {
 	pub fn new() -> Self {
-		Mailbox(Mutex::new(MailboxInner {
-			value: None,
-			waker: None,
-		}))
+		Mailbox(Mutex::new(MailboxInner::Empty(None)))
 	}
 }
 
-struct MailboxInner<T> {
-	value: Option<T>,
-	waker: Option<Waker>,
+enum MailboxInner<T> {
+	Empty(Option<Waker>),
+	Full(Option<Waker>, T),
 }
 
-struct MailboxFuture<'a, T>(&'a Mailbox<T>);
-impl<'a, T> Future for MailboxFuture<'a, T> {
+struct MailboxPut<'a, T>(&'a Mailbox<T>, Option<T>);
+impl<'a, T> Unpin for MailboxPut<'a, T> {}
+impl<'a, T> Future for MailboxPut<'a, T> {
+	type Output = ();
+
+	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+		let mut mb = self.0 .0.lock().unwrap();
+		match &mut *mb {
+			MailboxInner::Empty(ref mut waker) => {
+				let waker = waker.take();
+				*mb = MailboxInner::Full(None, self.1.take().expect("Mailbox future overpooled!"));
+				drop(mb);
+				if let Some(waker) = waker {
+					waker.wake();
+				}
+				Poll::Ready(())
+			}
+			MailboxInner::Full(ref mut waker, _) => {
+				*waker = Some(cx.waker().clone());
+				Poll::Pending
+			}
+		}
+	}
+}
+
+struct MailboxGet<'a, T>(&'a Mailbox<T>);
+impl<'a, T> Future for MailboxGet<'a, T> {
 	type Output = T;
 
-	fn poll(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
 		let mut mb = self.0 .0.lock().unwrap();
-		if let Some(value) = mb.value.take() {
-			return Poll::Ready(value);
-		}
-		mb.waker = Some(cx.waker().clone());
+		let old = std::mem::replace(&mut *mb, MailboxInner::Empty(None));
+		match old {
+			MailboxInner::Empty(_) => {
+				*mb = MailboxInner::Empty(Some(cx.waker().clone()));
 				Poll::Pending
 			}
+			MailboxInner::Full(waker, value) => {
+				drop(mb);
+				if let Some(waker) = waker {
+					waker.wake();
+				}
+				Poll::Ready(value)
+			}
+		}
+	}
 }
 
 impl<T> Mailbox<T> {
-	pub fn put(&self, value: T) {
-		let mut mb = self.0.lock().unwrap();
-		mb.value = Some(value);
-		let Some(waker) = mb.waker.take() else {
-			return;
-		};
-		drop(mb);
-		waker.wake();
+	pub fn put(&self, value: T) -> impl Future<Output = ()> + '_ {
+		MailboxPut(&self, Some(value))
 	}
 
 	pub fn get(&self) -> impl Future<Output = T> + '_ {
-		MailboxFuture(&self)
+		MailboxGet(&self)
 	}
 }
 
@@ -78,7 +104,7 @@ mod tests {
 	#[test]
 	fn test_mailbox_once() {
 		let mb = Mailbox::<i32>::new();
-		mb.put(42);
+		pollster::block_on(mb.put(42));
 		assert_eq!(pollster::block_on(mb.get()), 42);
 		assert!(!ready(mb.get()));
 	}
@@ -87,16 +113,16 @@ mod tests {
 	fn test_mailbox_once_oor() {
 		let mb = Mailbox::<i32>::new();
 		let f = mb.get();
-		mb.put(42);
+		pollster::block_on(mb.put(42));
 		assert_eq!(pollster::block_on(f), 42);
 	}
 
 	#[test]
 	fn test_mailbox_twice_no_wait() {
 		let mb = Mailbox::<i32>::new();
-		mb.put(42);
-		mb.put(13);
-		assert_eq!(pollster::block_on(mb.get()), 13);
+		pollster::block_on(mb.put(42));
+		assert!(!ready(mb.put(13)));
+		assert_eq!(pollster::block_on(mb.get()), 42);
 		assert!(!ready(mb.get()));
 	}
 
@@ -104,10 +130,10 @@ mod tests {
 	fn test_mailbox_twice_in_order() {
 		let mb = Mailbox::<i32>::new();
 		let f = mb.get();
-		mb.put(42);
+		pollster::block_on(mb.put(42));
 		assert_eq!(pollster::block_on(f), 42);
 		let f = mb.get();
-		mb.put(13);
+		pollster::block_on(mb.put(13));
 		assert_eq!(pollster::block_on(f), 13);
 		assert!(!ready(mb.get()));
 	}