Fix reference counting in DispatchSemaphoreGuard::release

Author: andyleiserson

crates/dispatch2/src/semaphore.rs

@@ -26,31 +26,74 @@ impl DispatchSemaphore {
         let result = Self::wait(self, timeout);
 
         match result {
-            0 => Ok(DispatchSemaphoreGuard(self.retain())),
+            0 => Ok(DispatchSemaphoreGuard(ManuallyDrop::new(self.retain()))),
             _ => Err(WaitError::Timeout),
         }
     }
 }
 
 /// Dispatch semaphore guard.
+///
+/// The semaphore will be signaled when the guard is dropped, or when [`release`] is called.
+///
+/// [`release`]: DispatchSemaphoreGuard::release
+//
+// See `release` for discussion of the `ManuallyDrop`.
 #[derive(Debug)]
-pub struct DispatchSemaphoreGuard(DispatchRetained<DispatchSemaphore>);
+pub struct DispatchSemaphoreGuard(ManuallyDrop<DispatchRetained<DispatchSemaphore>>);
 
 impl DispatchSemaphoreGuard {
     /// Release the [`DispatchSemaphore`].
     pub fn release(self) -> bool {
-        let this = ManuallyDrop::new(self);
-
-        // SAFETY: DispatchSemaphore cannot be null.
-        let result = this.0.signal();
-
-        result != 0
+        // We suppress `Drop` for the guard because that would signal the sempahore again.
+        // The inner `DispatchRetained` is wrapped in `ManuallyDrop` so that it can be
+        // separated from the guard and dropped normally.
+        let mut this = ManuallyDrop::new(self);
+        // SAFETY: The guard is being consumed; the `ManuallyDrop` contents will not be used again.
+        let semaphore = unsafe { ManuallyDrop::take(&mut this.0) };
+        semaphore.signal() != 0
     }
 }
 
 impl Drop for DispatchSemaphoreGuard {
     fn drop(&mut self) {
-        // SAFETY: DispatchSemaphore cannot be null.
-        self.0.signal();
+        // SAFETY: The guard is being dropped; the `ManuallyDrop` contents will not be used again.
+        let semaphore = unsafe { ManuallyDrop::take(&mut self.0) };
+        semaphore.signal();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[allow(unused_imports)]
+    use super::*;
+
+    #[test]
+    #[cfg(feature = "objc2")]
+    #[cfg_attr(
+        not(target_vendor = "apple"),
+        ignore = "only Apple's libdisptch is interoperable with `objc2`"
+    )]
+    fn acquire_release() {
+        fn retain_count(semaphore: &DispatchSemaphore) -> usize {
+            // SAFETY: semaphore is valid and the method signature is correct.
+            unsafe { objc2::msg_send![semaphore, retainCount] }
+        }
+
+        let semaphore = DispatchSemaphore::new(1);
+
+        assert_eq!(retain_count(&semaphore), 1);
+        {
+            let _guard = semaphore.try_acquire(DispatchTime::NOW).unwrap();
+            assert_eq!(retain_count(&semaphore), 2);
+        }
+        assert_eq!(retain_count(&semaphore), 1);
+        {
+            let guard = semaphore.try_acquire(DispatchTime::NOW).unwrap();
+            assert_eq!(retain_count(&semaphore), 2);
+            guard.release();
+            assert_eq!(retain_count(&semaphore), 1);
+        }
+        assert_eq!(retain_count(&semaphore), 1);
     }
 }