From 15960da098b119a50c20f405c03ab1e8ae10aaab Mon Sep 17 00:00:00 2001
From: Jonathan Johnson <jon@khonsulabs.com>
Date: Fri, 29 Dec 2023 14:23:49 -0800
Subject: [PATCH] Fixing edge case in DynamicReader disconnect

There was a small window between when notify_all and the strong count
for the Arc is decreased that a DynamicReader could observe the strong
count still being greater than the reader count, but the drop thread
just hasn't proceeded far enough.

Now the on_disconnect is stored in an option, and its presence denotes
that the disconnect logic has not fired yet. DynamicReader now checks
on_disconnect in addition to the strong count.
---
 CHANGELOG.md |  2 ++
 src/value.rs | 29 +++++++++++++++++------------
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ee91470..ff81df2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -28,6 +28,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   contains more than one element. Previously, if rows contained widgets that
   filled the given space, this would cause the grid to calculate layouts
   incorrectly.
+- A potential edge case where a `DynamicReader` would not return after being
+  disconnected has been removed.
 
 ### Changed
 
diff --git a/src/value.rs b/src/value.rs
index 602187a..8acfa1b 100644
--- a/src/value.rs
+++ b/src/value.rs
@@ -43,7 +43,7 @@ impl<T> Dynamic<T> {
                 readers: 0,
                 wakers: Vec::new(),
                 widgets: AHashSet::new(),
-                on_disconnect: Vec::new(),
+                on_disconnect: Some(Vec::new()),
                 source_callback: CallbackHandle::default(),
             }),
             during_callback_state: Mutex::default(),
@@ -811,10 +811,10 @@ impl<T> Drop for Dynamic<T> {
         // `Dynamic`.
         if let Ok(mut state) = self.state() {
             if Arc::strong_count(&self.0) == state.readers + 1 {
-                let on_disconnect = std::mem::take(&mut state.on_disconnect);
+                let on_disconnect = state.on_disconnect.take();
                 drop(state);
 
-                for on_disconnect in on_disconnect {
+                for on_disconnect in on_disconnect.into_iter().flatten() {
                     on_disconnect.invoke(());
                 }
 
@@ -1182,7 +1182,7 @@ struct State<T> {
     windows: AHashSet<WindowHandle>,
     widgets: AHashSet<(WindowHandle, WidgetId)>,
     wakers: Vec<Waker>,
-    on_disconnect: Vec<OnceCallback>,
+    on_disconnect: Option<Vec<OnceCallback>>,
     readers: usize,
 }
 
@@ -1606,7 +1606,9 @@ impl<T> DynamicReader<T> {
             let state = self.source.state.lock().ignore_poison();
             if state.wrapped.generation != self.read_generation {
                 return true;
-            } else if state.readers == Arc::strong_count(&self.source) {
+            } else if state.readers == Arc::strong_count(&self.source)
+                || state.on_disconnect.is_none()
+            {
                 return false;
             }
             drop(state);
@@ -1619,7 +1621,8 @@ impl<T> DynamicReader<T> {
     /// Returns true if this reader still has any writers connected to it.
     #[must_use]
     pub fn connected(&self) -> bool {
-        self.source.state.lock().ignore_poison().readers < Arc::strong_count(&self.source)
+        let state = self.source.state.lock().ignore_poison();
+        state.readers < Arc::strong_count(&self.source) && state.on_disconnect.is_some()
     }
 
     /// Suspends the current async task until the contained value has been
@@ -1643,11 +1646,11 @@ impl<T> DynamicReader<T> {
     where
         OnDisconnect: FnOnce() + Send + 'static,
     {
-        self.source
-            .state()
-            .expect("deadlocked")
-            .on_disconnect
-            .push(OnceCallback::new(|()| on_disconnect()));
+        let mut state = self.source.state().expect("deadlocked");
+
+        if let Some(callbacks) = &mut state.on_disconnect {
+            callbacks.push(OnceCallback::new(|()| on_disconnect()));
+        }
     }
 }
 
@@ -1705,7 +1708,9 @@ impl<'a, T> Future for BlockUntilUpdatedFuture<'a, T> {
         let mut state = self.0.source.state().expect("deadlocked");
         if state.wrapped.generation != self.0.read_generation {
             return Poll::Ready(true);
-        } else if state.readers == Arc::strong_count(&self.0.source) {
+        } else if state.readers == Arc::strong_count(&self.0.source)
+            || state.on_disconnect.is_none()
+        {
             return Poll::Ready(false);
         }