atomptr: implement Ref::consume and fix a mild memory leak

development/libs/atomptr/src/lib.rs

@@ -1,13 +1,13 @@
-//! A safe, strongly typed (generic) atomic pointer abstraction to build
-//! datastructures, and lock-free algorithms on top of.  Only uses
-//! `libstd`.
+//! A safe, strongly typed (generic) atomic pointer abstraction to
+//! build datastructures, and lock-free algorithms on top of.  Only
+//! uses `libstd`.
 //!
 //! The standard library contains an `AtomicPtr` type, which by itself
-//! isn't very ergonomic to use, because it deals with raw pointers.  This
-//! library assumes that types can always be heap allocated, wrapping them
-//! in a `Box<T>`, and provides a nicer (and safe!) abstraction for
-//! `std::sync::atomic::AtomicPtr`.  Using this crate is fairely
-//! self-explanatory:
+//! isn't very ergonomic to use, because it deals with raw pointers.
+//! This library assumes that types can always be heap allocated,
+//! wrapping them in a `Box<T>`, and provides a nicer (and safe!)
+//! abstraction for `std::sync::atomic::AtomicPtr`.  Using this crate
+//! is fairely self-explanatory:
 //!
 //! ```rust
 //! use atomptr::AtomPtr;
@@ -23,11 +23,9 @@
 //! ```
 //!
 //! Note that the type that is returned by `get_ref` and `swap` is
-//! `Ref<T>`, which means that the old data is not de-allocated after a
-//! swap, before this last reference goes out of scope.  You can of course
-//! always manually call `drop()` on it.
-
-
+//! `Ref<T>`, which means that the old data is not de-allocated after
+//! a swap, before this last reference goes out of scope.  You can of
+//! course always manually call `drop()` on it.
 
 use std::sync::{
     atomic::{AtomicPtr, Ordering},
@@ -40,6 +38,18 @@ pub struct Ref<T> {
     inner: Box<Arc<T>>,
 }
 
+impl<T> Ref<T> {
+    /// Consume this Ref wrapper to yield the underlying `Arc<T>`
+    ///
+    /// If you want to take ownership of the underlying type data, and
+    /// you can prove that only one strong-reference Arc exists to
+    /// this type, you can use `std::arc::Arc::try_unwrap()` to peel
+    /// the reference counter and take exclusive ownership.
+    pub fn consume(self) -> Arc<T> {
+        *self.inner
+    }
+}
+
 impl<T> Deref for Ref<T> {
     type Target = Arc<T>;
 
@@ -97,13 +107,8 @@ impl<T> AtomPtr<T> {
         let new = Self::make_raw_ptr(new);
         let prev = self.inner.swap(new, Ordering::Relaxed);
 
-        let b = unsafe { Box::from_raw(prev) };
-        let arc = Arc::clone(&*b);
-        std::mem::forget(b);
-
-        Ref {
-            inner: Box::new(arc),
-        }
+        let inner = unsafe { Box::from_raw(prev) };
+        Ref { inner }
     }
 }
 
@@ -146,3 +151,27 @@ fn swap() {
     // But the old ref is still valid
     assert_eq!(ts1, *still_ts1.as_ref());
 }
+
+#[test]
+fn take_from_swap() {
+    let ts1 = TestStruct {
+        name: "Hello 1".into(),
+    };
+
+    let ts2 = TestStruct {
+        name: "Hello 2".into(),
+    };
+
+    // Make an AtomPtr with some data
+    let ptr = AtomPtr::new(ts1.clone());
+    assert_eq!(ptr.get_ref().name, "Hello 1".to_string());
+
+    // Swap the data
+    let still_ts1 = ptr.swap(ts2);
+    assert_eq!(ptr.get_ref().name, "Hello 2".to_string());
+    assert_eq!(Arc::strong_count(&still_ts1), 1);
+
+    // We can now also take ownership of the Arc
+    let ts1_again = Arc::try_unwrap(still_ts1.consume()).unwrap();
+    assert_eq!(ts1_again, ts1);
+}