Waking the other side on async consumer / producer close

async/src/halves.rs

@@ -149,7 +149,8 @@ where
  self.base.rb().is_closed()
  }
  fn close(&self) {
- self.base.rb().close()
+ self.base.rb().close();
+ self.base.rb().wake_consumer();
  }
 }
 impl<R: RbRef> AsyncProducer for AsyncProd<R>
@@ -181,7 +182,8 @@ where
  self.base.rb().is_closed()
  }
  fn close(&self) {
- self.base.rb().close()
+ self.base.rb().close();
+ self.base.rb().wake_producer();
  }
 }
 impl<R: RbRef> AsyncConsumer for AsyncCons<R>

async/src/rb.rs

@@ -81,7 +81,14 @@ impl<S: Storage> AsyncConsumer for AsyncRb<S> {
  self.write.register(waker);
  }
 }
-impl<S: Storage> AsyncRingBuffer for AsyncRb<S> {}
+impl<S: Storage> AsyncRingBuffer for AsyncRb<S> {
+ fn wake_consumer(&self) {
+ self.write.wake()
+ }
+ fn wake_producer(&self) {
+ self.read.wake()
+ }
+}
 
 impl<S: Storage> SplitRef for AsyncRb<S> {
  type RefProd<'a> = AsyncProd<&'a Self> where Self: 'a;

async/src/tests.rs

@@ -1,6 +1,7 @@
 use crate::{async_transfer, traits::*, AsyncHeapRb};
 use core::sync::atomic::{AtomicUsize, Ordering};
 use futures::task::{noop_waker_ref, AtomicWaker};
+use std::sync::Arc;
 use std::{vec, vec::Vec};
 
 #[test]
@@ -168,3 +169,52 @@ fn wait() {
  },
  );
 }
+
+#[test]
+fn drop_close_prod() {
+ let (prod, mut cons) = AsyncHeapRb::<usize>::new(1).split();
+ let stage = Arc::new(AtomicUsize::new(0));
+ let stage_clone = stage.clone();
+ let t0 = std::thread::spawn(move || {
+ execute!(async {
+ drop(prod);
+ assert_eq!(stage.fetch_add(1, Ordering::SeqCst), 0);
+ });
+ });
+ let t1 = std::thread::spawn(move || {
+ execute!(async {
+ cons.wait_occupied(1).await;
+ assert_eq!(stage_clone.fetch_add(1, Ordering::SeqCst), 1);
+ assert!(cons.is_closed());
+ });
+ });
+ t0.join().unwrap();
+ t1.join().unwrap();
+}
+
+#[test]
+fn drop_close_cons() {
+ let (mut prod, mut cons) = AsyncHeapRb::<usize>::new(1).split();
+ let stage = Arc::new(AtomicUsize::new(0));
+ let stage_clone = stage.clone();
+ let t0 = std::thread::spawn(move || {
+ execute!(async {
+ prod.push(0).await.unwrap();
+ assert_eq!(stage.fetch_add(1, Ordering::SeqCst), 0);
+
+ prod.wait_vacant(1).await;
+ assert_eq!(stage.fetch_add(1, Ordering::SeqCst), 3);
+ assert!(prod.is_closed());
+ });
+ });
+ let t1 = std::thread::spawn(move || {
+ execute!(async {
+ cons.wait_occupied(1).await;
+ assert_eq!(stage_clone.fetch_add(1, Ordering::SeqCst), 1);
+ drop(cons);
+ assert_eq!(stage_clone.fetch_add(1, Ordering::SeqCst), 2);
+ });
+ });
+ t0.join().unwrap();
+ t1.join().unwrap();
+}

async/src/traits/mod.rs

@@ -1,11 +1,11 @@
 pub mod consumer;
 pub mod observer;
 pub mod producer;
+pub mod ring_buffer;
 
 pub use consumer::AsyncConsumer;
 pub use observer::AsyncObserver;
 pub use producer::AsyncProducer;
-
-pub trait AsyncRingBuffer: ringbuf::traits::RingBuffer + AsyncProducer + AsyncConsumer {}
+pub use ring_buffer::AsyncRingBuffer;
 
 pub use ringbuf::traits::*;

async/src/traits/ring_buffer.rs

@@ -0,0 +1,8 @@
+use crate::consumer::AsyncConsumer;
+use crate::producer::AsyncProducer;
+use ringbuf::traits::RingBuffer;
+
+pub trait AsyncRingBuffer: RingBuffer + AsyncProducer + AsyncConsumer {
+ fn wake_producer(&self);
+ fn wake_consumer(&self);
+}

src/transfer.rs

@@ -6,11 +6,7 @@ use crate::{consumer::Consumer, producer::Producer};
 /// `count` is the number of items being moved, if `None` - as much as possible items will be moved.
 ///
 /// Returns number of items been moved.
-pub fn transfer<T, C: Consumer<Item = T>, P: Producer<Item = T>>(
- src: &mut C,
- dst: &mut P,
- count: Option<usize>,
-) -> usize {
+pub fn transfer<T, C: Consumer<Item = T>, P: Producer<Item = T>>(src: &mut C, dst: &mut P, count: Option<usize>) -> usize {
  let (src_left, src_right) = src.occupied_slices();
  let (dst_left, dst_right) = dst.vacant_slices_mut();
  let src_iter = src_left.iter().chain(src_right.iter());