Check queue readiness and capacity before waiting

bbqueue/src/prod_cons/framed.rs

@@ -162,14 +162,28 @@ where
 {
     /// Wait for the given write grant to become available
     ///
-    /// If `sz` is larger than the storage buffer, this method will never
-    /// return.
-    ///
     /// The returned grant can be used to write up to `sz` bytes, though
     /// a smaller size may be committed. Dropping the grant without calling
     /// commit means that no data will be made visible to the consumer.
-    pub async fn wait_grant(&self, sz: H) -> FramedGrantW<Q, H> {
-        self.bbq.not.wait_for_not_full(|| self.grant(sz).ok()).await
+    ///
+    /// # Errors
+    ///
+    /// Returns [`WriteGrantError::InsufficientSize`] when a frame of `sz` cannot
+    /// fit the buffer even when empty. Releasing every outstanding read grant
+    /// would not satisfy such a request, so it is reported rather than awaited.
+    pub async fn wait_grant(&self, sz: H) -> Result<FramedGrantW<Q, H>, WriteGrantError> {
+        let needed = sz
+            .into()
+            .checked_add(core::mem::size_of::<H>())
+            .ok_or(WriteGrantError::InsufficientSize)?;
+        if needed > self.capacity() {
+            return Err(WriteGrantError::InsufficientSize);
+        }
+
+        if let Ok(grant) = self.grant(sz) {
+            return Ok(grant);
+        }
+        Ok(self.bbq.not.wait_for_not_full(|| self.grant(sz).ok()).await)
     }
 }
 
@@ -247,6 +261,9 @@ where
     ///
     /// The returned grant must be released to free the space in the buffer.
     pub async fn wait_read(&self) -> FramedGrantR<Q, H> {
+        if let Ok(grant) = self.read() {
+            return grant;
+        }
         self.bbq.not.wait_for_not_empty(|| self.read().ok()).await
     }
 }
@@ -413,3 +430,42 @@ where
     H: LenHeader + Send,
 {
 }
+
+#[cfg(all(test, feature = "maitake-sync-0_3"))]
+mod async_tests {
+    use core::{
+        future::Future,
+        pin::pin,
+        task::{Context, Poll, Waker},
+    };
+
+    use crate::{
+        BBQueue,
+        prod_cons::framed::FramedProducer,
+        traits::{
+            bbqhdl::BbqHandle,
+            coordination::{WriteGrantError, cas::AtomicCoord},
+            notifier::maitake::MaiNotSpsc,
+            storage::Inline,
+        },
+    };
+
+    type AsyncRing = BBQueue<Inline<64>, AtomicCoord, MaiNotSpsc>;
+
+    /// A frame too large for an empty buffer can never be granted, so the wait
+    /// resolves instead of parking on a consumer that cannot help.
+    #[test]
+    fn wait_grant_rejects_a_frame_larger_than_the_buffer() {
+        let bbq: AsyncRing = BBQueue::new();
+        let prod: FramedProducer<&AsyncRing, usize> = BbqHandle::framed_producer::<usize>(&&bbq);
+
+        // The payload alone fits; the header pushes it past capacity.
+        let sz = prod.capacity() - 1;
+
+        let mut wait = pin!(prod.wait_grant(sz));
+        let Poll::Ready(res) = wait.as_mut().poll(&mut Context::from_waker(Waker::noop())) else {
+            panic!("wait_grant parked on a request no consumer can satisfy");
+        };
+        assert_eq!(res.err(), Some(WriteGrantError::InsufficientSize));
+    }
+}

bbqueue/src/prod_cons/mod.rs

@@ -18,3 +18,76 @@
 
 pub mod framed;
 pub mod stream;
+
+#[cfg(test)]
+mod tests {
+    use core::sync::atomic::{AtomicUsize, Ordering};
+
+    use const_init::ConstInit;
+
+    use crate::{
+        BBQueue,
+        traits::{
+            coordination::cas::AtomicCoord,
+            notifier::{AsyncNotifier, Notifier},
+            storage::Inline,
+        },
+    };
+
+    static WAIT_CALLS: AtomicUsize = AtomicUsize::new(0);
+
+    struct CountingNotifier;
+
+    impl ConstInit for CountingNotifier {
+        const INIT: Self = Self;
+    }
+
+    impl Notifier for CountingNotifier {
+        fn wake_one_consumer(&self) {}
+
+        fn wake_one_producer(&self) {}
+    }
+
+    impl AsyncNotifier for CountingNotifier {
+        async fn wait_for_not_empty<T, F: FnMut() -> Option<T>>(&self, mut f: F) -> T {
+            WAIT_CALLS.fetch_add(1, Ordering::Relaxed);
+            f().expect("test queue has readable data")
+        }
+
+        async fn wait_for_not_full<T, F: FnMut() -> Option<T>>(&self, mut f: F) -> T {
+            WAIT_CALLS.fetch_add(1, Ordering::Relaxed);
+            f().expect("test queue has writable space")
+        }
+    }
+
+    #[tokio::test]
+    async fn ready_waits_do_not_arm_notifier() {
+        WAIT_CALLS.store(0, Ordering::Relaxed);
+
+        let stream: BBQueue<Inline<64>, AtomicCoord, CountingNotifier> = BBQueue::new();
+        let stream_producer = stream.stream_producer();
+        let stream_consumer = stream.stream_consumer();
+
+        stream_producer
+            .wait_grant_exact(1)
+            .await
+            .expect("1 byte fits a 64 byte buffer")
+            .commit(1);
+        stream_consumer.wait_read().await.release(1);
+        stream_producer.wait_grant_max_remaining(1).await.commit(1);
+        stream_consumer.wait_read().await.release(1);
+
+        let framed: BBQueue<Inline<64>, AtomicCoord, CountingNotifier> = BBQueue::new();
+        let framed_producer = framed.framed_producer();
+        let framed_consumer = framed.framed_consumer();
+
+        framed_producer
+            .wait_grant(1)
+            .await
+            .expect("a 1 byte frame fits a 64 byte buffer")
+            .commit(1);
+        framed_consumer.wait_read().await.release();
+
+        assert_eq!(WAIT_CALLS.load(Ordering::Relaxed), 0);
+    }
+}

bbqueue/src/prod_cons/stream.rs

@@ -134,6 +134,9 @@ where
 {
     /// Wait for a grant of any size, up to `max`, to become available
     pub async fn wait_grant_max_remaining(&self, max: usize) -> StreamGrantW<Q> {
+        if let Ok(grant) = self.grant_max_remaining(max) {
+            return grant;
+        }
         self.bbq
             .not
             .wait_for_not_full(|| self.grant_max_remaining(max).ok())
@@ -142,12 +145,24 @@ where
 
     /// Wait for a grant of EXACTLY `sz` to become available.
     ///
-    /// If `sz` exceeds the capacity of the buffer, this method will never return.
-    pub async fn wait_grant_exact(&self, sz: usize) -> StreamGrantW<Q> {
-        self.bbq
+    /// # Errors
+    ///
+    /// Returns [`WriteGrantError::InsufficientSize`] when `sz` cannot fit the
+    /// buffer even when empty. Releasing every outstanding read grant would not
+    /// satisfy such a request, so it is reported rather than awaited.
+    pub async fn wait_grant_exact(&self, sz: usize) -> Result<StreamGrantW<Q>, WriteGrantError> {
+        if sz > self.capacity() {
+            return Err(WriteGrantError::InsufficientSize);
+        }
+
+        if let Ok(grant) = self.grant_exact(sz) {
+            return Ok(grant);
+        }
+        Ok(self
+            .bbq
             .not
             .wait_for_not_full(|| self.grant_exact(sz).ok())
-            .await
+            .await)
     }
 }
 
@@ -193,6 +208,9 @@ where
 {
     /// Wait for any read data to become available
     pub async fn wait_read(&self) -> StreamGrantR<Q> {
+        if let Ok(grant) = self.read() {
+            return grant;
+        }
         self.bbq.not.wait_for_not_empty(|| self.read().ok()).await
     }
 }
@@ -322,3 +340,40 @@ where
 }
 
 unsafe impl<Q: BbqHandle + Send> Send for StreamGrantR<Q> {}
+
+#[cfg(all(test, feature = "maitake-sync-0_3"))]
+mod async_tests {
+    use core::{
+        future::Future,
+        pin::pin,
+        task::{Context, Poll, Waker},
+    };
+
+    use crate::{
+        BBQueue,
+        prod_cons::stream::StreamProducer,
+        traits::{
+            coordination::{WriteGrantError, cas::AtomicCoord},
+            notifier::maitake::MaiNotSpsc,
+            storage::Inline,
+        },
+    };
+
+    type AsyncRing = BBQueue<Inline<64>, AtomicCoord, MaiNotSpsc>;
+
+    /// A request larger than the whole buffer can never be granted, so the wait
+    /// resolves instead of parking on a consumer that cannot help.
+    #[test]
+    fn wait_grant_exact_rejects_a_size_larger_than_the_buffer() {
+        let bbq: AsyncRing = BBQueue::new();
+        let prod: StreamProducer<&AsyncRing> = bbq.stream_producer();
+
+        let sz = prod.capacity() + 1;
+
+        let mut wait = pin!(prod.wait_grant_exact(sz));
+        let Poll::Ready(res) = wait.as_mut().poll(&mut Context::from_waker(Waker::noop())) else {
+            panic!("wait_grant_exact parked on a request no consumer can satisfy");
+        };
+        assert_eq!(res.err(), Some(WriteGrantError::InsufficientSize));
+    }
+}