diff --git a/src/storage.rs b/src/storage.rs index 443b53221..f63981e4f 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -25,6 +25,8 @@ pub struct StorageHandle { impl Clone for StorageHandle { fn clone(&self) -> Self { + *self.coordinate.clones.lock() += 1; + Self { zalsa_impl: self.zalsa_impl.clone(), coordinate: CoordinateDrop(Arc::clone(&self.coordinate)), @@ -51,7 +53,7 @@ impl StorageHandle { Self { zalsa_impl: Arc::new(Zalsa::new::(event_callback, jars)), coordinate: CoordinateDrop(Arc::new(Coordinate { - coordinate_lock: Mutex::default(), + clones: Mutex::new(1), cvar: Default::default(), })), phantom: PhantomData, @@ -93,6 +95,17 @@ impl Drop for Storage { } } +struct Coordinate { + /// Counter of the number of clones of actor. Begins at 1. + /// Incremented when cloned, decremented when dropped. + clones: Mutex, + cvar: Condvar, +} + +// We cannot panic while holding a lock to `clones: Mutex` and therefore we cannot enter an +// inconsistent state. +impl RefUnwindSafe for Coordinate {} + impl Default for Storage { fn default() -> Self { Self::new(None) @@ -155,15 +168,12 @@ impl Storage { .zalsa_impl .event(&|| Event::new(EventKind::DidSetCancellationFlag)); - let mut coordinate_lock = self.handle.coordinate.coordinate_lock.lock(); - let zalsa = loop { - if Arc::strong_count(&self.handle.zalsa_impl) == 1 { - // SAFETY: The strong count is 1, and we never create any weak pointers, - // so we have a unique reference. - break unsafe { &mut *(Arc::as_ptr(&self.handle.zalsa_impl).cast_mut()) }; - } - coordinate_lock = self.handle.coordinate.cvar.wait(coordinate_lock); - }; + let mut clones = self.handle.coordinate.clones.lock(); + while *clones != 1 { + clones = self.handle.coordinate.cvar.wait(clones); + } + // The ref count on the `Arc` should now be 1 + let zalsa = Arc::get_mut(&mut self.handle.zalsa_impl).unwrap(); // cancellation is done, so reset the flag zalsa.runtime_mut().reset_cancellation_flag(); zalsa @@ -250,16 +260,6 @@ impl Clone for Storage { } } -/// A simplified `WaitGroup`, this is used together with `Arc` as the actual counter -struct Coordinate { - coordinate_lock: Mutex<()>, - cvar: Condvar, -} - -// We cannot panic while holding a lock to `clones: Mutex` and therefore we cannot enter an -// inconsistent state. -impl RefUnwindSafe for Coordinate {} - struct CoordinateDrop(Arc); impl std::ops::Deref for CoordinateDrop { @@ -272,6 +272,7 @@ impl std::ops::Deref for CoordinateDrop { impl Drop for CoordinateDrop { fn drop(&mut self) { + *self.0.clones.lock() -= 1; self.0.cvar.notify_all(); } }