Release owned queue handles after grant completion

bbqueue/miri.sh

@@ -1,10 +1,7 @@
 #!/bin/bash
 
 # We disable isolation because we use tokio sleep
-# We disable leaks because ???
-#
-# TODO: Can we eliminate some of these limitations for testing?
-MIRIFLAGS="-Zmiri-disable-isolation -Zmiri-ignore-leaks" \
+MIRIFLAGS="-Zmiri-disable-isolation" \
     cargo +nightly miri test \
     --target x86_64-unknown-linux-gnu \
     --features=std

bbqueue/src/lib.rs

@@ -143,22 +143,21 @@ pub mod export {
 
 #[cfg(all(test, feature = "alloc"))]
 mod test {
+    use alloc::sync::Arc;
     use core::{ops::Deref, time::Duration};
 
     use crate::{
         queue::{ArcBBQueue, BBQueue},
         traits::{
             coordination::cas::AtomicCoord,
-            notifier::maitake::MaiNotSpsc,
+            notifier::{maitake::MaiNotSpsc, polling::Polling},
             storage::{BoxedSlice, Inline},
         },
     };
 
     #[cfg(all(target_has_atomic = "ptr", feature = "alloc"))]
     #[test]
     fn ux() {
-        use crate::traits::{notifier::polling::Polling, storage::BoxedSlice};
-
         static BBQ: BBQueue<Inline<64>, AtomicCoord, Polling> = BBQueue::new();
         let _ = BBQ.stream_producer();
         let _ = BBQ.stream_consumer();
@@ -177,9 +176,6 @@ mod test {
     #[cfg(target_has_atomic = "ptr")]
     #[test]
     fn smoke() {
-        use crate::traits::notifier::polling::Polling;
-        use core::ops::Deref;
-
         static BBQ: BBQueue<Inline<64>, AtomicCoord, Polling> = BBQueue::new();
         let prod = BBQ.stream_producer();
         let cons = BBQ.stream_consumer();
@@ -203,9 +199,6 @@ mod test {
     #[cfg(target_has_atomic = "ptr")]
     #[test]
     fn smoke_framed() {
-        use crate::traits::notifier::polling::Polling;
-        use core::ops::Deref;
-
         static BBQ: BBQueue<Inline<64>, AtomicCoord, Polling> = BBQueue::new();
         let prod = BBQ.framed_producer();
         let cons = BBQ.framed_consumer();
@@ -225,8 +218,6 @@ mod test {
     #[cfg(target_has_atomic = "ptr")]
     #[test]
     fn framed_misuse() {
-        use crate::traits::notifier::polling::Polling;
-
         static BBQ: BBQueue<Inline<64>, AtomicCoord, Polling> = BBQueue::new();
         let prod = BBQ.stream_producer();
         let cons = BBQ.framed_consumer();
@@ -255,6 +246,88 @@ mod test {
         assert!(cons.read().is_err());
     }
 
+    #[test]
+    fn framed_commit_releases_owned_queue_handle() {
+        let bbq: ArcBBQueue<Inline<64>, AtomicCoord, Polling> =
+            ArcBBQueue::new_with_storage(Inline::new());
+        let weak = Arc::downgrade(&bbq.0);
+        let prod = bbq.framed_producer();
+        let baseline = Arc::strong_count(&bbq.0);
+
+        let wgr = prod.grant(1).expect("space should be available");
+        assert_eq!(Arc::strong_count(&bbq.0), baseline + 1);
+        wgr.commit(1);
+        assert_eq!(Arc::strong_count(&bbq.0), baseline);
+
+        drop(prod);
+        drop(bbq);
+        assert!(weak.upgrade().is_none());
+    }
+
+    #[test]
+    fn framed_release_releases_owned_queue_handle() {
+        let bbq: ArcBBQueue<Inline<64>, AtomicCoord, Polling> =
+            ArcBBQueue::new_with_storage(Inline::new());
+        let weak = Arc::downgrade(&bbq.0);
+        let seed_prod = bbq.0.as_ref().framed_producer();
+        seed_prod
+            .grant(1)
+            .expect("space should be available")
+            .commit(1);
+        let cons = bbq.framed_consumer();
+        let baseline = Arc::strong_count(&bbq.0);
+
+        let rgr = cons.read().expect("seeded frame should be available");
+        assert_eq!(Arc::strong_count(&bbq.0), baseline + 1);
+        rgr.release();
+        assert_eq!(Arc::strong_count(&bbq.0), baseline);
+
+        drop(cons);
+        drop(bbq);
+        assert!(weak.upgrade().is_none());
+    }
+
+    #[test]
+    fn stream_commit_releases_owned_queue_handle() {
+        let bbq: ArcBBQueue<Inline<64>, AtomicCoord, Polling> =
+            ArcBBQueue::new_with_storage(Inline::new());
+        let weak = Arc::downgrade(&bbq.0);
+        let prod = bbq.stream_producer();
+        let baseline = Arc::strong_count(&bbq.0);
+
+        let wgr = prod.grant_exact(1).expect("space should be available");
+        assert_eq!(Arc::strong_count(&bbq.0), baseline + 1);
+        wgr.commit(1);
+        assert_eq!(Arc::strong_count(&bbq.0), baseline);
+
+        drop(prod);
+        drop(bbq);
+        assert!(weak.upgrade().is_none());
+    }
+
+    #[test]
+    fn stream_release_releases_owned_queue_handle() {
+        let bbq: ArcBBQueue<Inline<64>, AtomicCoord, Polling> =
+            ArcBBQueue::new_with_storage(Inline::new());
+        let weak = Arc::downgrade(&bbq.0);
+        let seed_prod = bbq.0.as_ref().stream_producer();
+        seed_prod
+            .grant_exact(1)
+            .expect("space should be available")
+            .commit(1);
+        let cons = bbq.stream_consumer();
+        let baseline = Arc::strong_count(&bbq.0);
+
+        let rgr = cons.read().expect("seeded byte should be available");
+        assert_eq!(Arc::strong_count(&bbq.0), baseline + 1);
+        rgr.release(1);
+        assert_eq!(Arc::strong_count(&bbq.0), baseline);
+
+        drop(cons);
+        drop(bbq);
+        assert!(weak.upgrade().is_none());
+    }
+
     #[tokio::test]
     async fn asink() {
         static BBQ: BBQueue<Inline<64>, AtomicCoord, MaiNotSpsc> = BBQueue::new();

bbqueue/src/prod_cons/framed.rs

@@ -4,8 +4,9 @@
 
 use core::{
     marker::PhantomData,
+    mem::ManuallyDrop,
     ops::{Deref, DerefMut},
-    ptr::NonNull,
+    ptr::{self, NonNull},
 };
 
 use crate::traits::{
@@ -277,8 +278,11 @@ where
         }
 
         self.bbq.cor.commit_inner(cap, grant_len, used_len);
-        self.bbq.not.wake_one_consumer();
-        core::mem::forget(self);
+        let this = ManuallyDrop::new(self);
+        // SAFETY: Read once; the ManuallyDrop grant is not accessed or dropped afterward.
+        let bbq = unsafe { ptr::read(&this.bbq) };
+        bbq.not.wake_one_consumer();
+        drop(bbq);
     }
 
     /// Aborts the grant, making no frame available to the consumer
@@ -359,8 +363,11 @@ where
         let hdrlen: usize = const { core::mem::size_of::<H>() };
         let used = len + hdrlen;
         self.bbq.cor.release_inner(used);
-        self.bbq.not.wake_one_producer();
-        core::mem::forget(self);
+        let this = ManuallyDrop::new(self);
+        // SAFETY: Read once; the ManuallyDrop grant is not accessed or dropped afterward.
+        let bbq = unsafe { ptr::read(&this.bbq) };
+        bbq.not.wake_one_producer();
+        drop(bbq);
     }
 
     /// Drop the grant WITHOUT releasing the message from the queue.

bbqueue/src/prod_cons/stream.rs

@@ -208,14 +208,8 @@ where
     /// Make `used` bytes available to be read.
     ///
     /// `used` is capped to the length of the grant.
-    pub fn commit(self, used: usize) {
-        let (_, cap) = unsafe { self.bbq.sto.ptr_len() };
-        let used = used.min(self.len);
-        self.bbq.cor.commit_inner(cap, self.len, used);
-        if used != 0 {
-            self.bbq.not.wake_one_consumer();
-        }
-        core::mem::forget(self);
+    pub fn commit(mut self, used: usize) {
+        self.to_commit = used;
     }
 }
 
@@ -271,13 +265,8 @@ where
     /// Make `used` bytes available for writing.
     ///
     /// `used` is capped to the length of the grant
-    pub fn release(self, used: usize) {
-        let used = used.min(self.len);
-        self.bbq.cor.release_inner(used);
-        if used != 0 {
-            self.bbq.not.wake_one_producer();
-        }
-        core::mem::forget(self);
+    pub fn release(mut self, used: usize) {
+        self.to_release = used;
     }
 }