ArcCell: wait-free get()

src/sync/arc_cell.rs

@@ -5,48 +5,56 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 
 /// A type providing atomic storage and retrieval of an `Arc<T>`.
 #[derive(Debug)]
-pub struct ArcCell<T>(AtomicUsize, PhantomData<Arc<T>>);
+pub struct ArcCell<T> {
+    ptr: AtomicUsize,
+    sem: AtomicUsize,
+    _marker: PhantomData<Arc<T>>,
+}
 
 impl<T> Drop for ArcCell<T> {
     fn drop(&mut self) {
-        self.take();
+        unsafe { mem::transmute::<_, Arc<T>>(self.ptr.load(Ordering::Relaxed)); }
     }
 }
 
 impl<T> ArcCell<T> {
     /// Creates a new `ArcCell`.
     pub fn new(t: Arc<T>) -> ArcCell<T> {
-        ArcCell(AtomicUsize::new(unsafe { mem::transmute(t) }), PhantomData)
-    }
-
-    fn take(&self) -> Arc<T> {
-        loop {
-            match self.0.swap(0, Ordering::Acquire) {
-                0 => {}
-                n => return unsafe { mem::transmute(n) }
-            }
+        ArcCell {
+            ptr: AtomicUsize::new(unsafe { mem::transmute(t) }),
+            sem: AtomicUsize::new(0),
+            _marker: PhantomData,
         }
     }
 
-    fn put(&self, t: Arc<T>) {
-        debug_assert_eq!(self.0.load(Ordering::SeqCst), 0);
-        self.0.store(unsafe { mem::transmute(t) }, Ordering::Release);
+    /// Create a new `ArcCell` from the given `Arc` interior.
+    pub fn with_val(v: T) -> ArcCell<T> {
+        ArcCell {
+            ptr: AtomicUsize::new(unsafe { mem::transmute(Arc::new(v)) }),
+            sem: AtomicUsize::new(0),
+            _marker: PhantomData,
+        }
     }
 
     /// Stores a new value in the `ArcCell`, returning the previous
     /// value.
     pub fn set(&self, t: Arc<T>) -> Arc<T> {
-        let old = self.take();
-        self.put(t);
-        old
+        unsafe {
+            let t: usize = mem::transmute(t);
+            let old: Arc<T> = mem::transmute(self.ptr.swap(t, Ordering::Acquire));
+            while self.sem.load(Ordering::Relaxed) > 0 {}
+            old
+        }
     }
 
     /// Returns a copy of the value stored by the `ArcCell`.
     pub fn get(&self) -> Arc<T> {
-        let t = self.take();
+        self.sem.fetch_add(1, Ordering::Relaxed);
+        let t: Arc<T> = unsafe { mem::transmute(self.ptr.load(Ordering::SeqCst)) };
         // NB: correctness here depends on Arc's clone impl not panicking
         let out = t.clone();
-        self.put(t);
+        self.sem.fetch_sub(1, Ordering::Relaxed);
+        mem::forget(t);
         out
     }
 }