From 57bd343f4a1103ca323694234b79fca379629a46 Mon Sep 17 00:00:00 2001
From: Greg Johnston <greg.johnston@gmail.com>
Date: Sat, 7 Sep 2024 07:01:54 -0400
Subject: [PATCH] fix: remove Owner from thread-local when it is ordinarily
 dropped, to ensure cleanup (closes #2942) (#2944)

---
 integrations/utils/src/lib.rs            |  2 +-
 reactive_graph/src/owner.rs              | 15 +++++++++++++++
 reactive_graph/src/owner/stored_value.rs |  2 ++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/integrations/utils/src/lib.rs b/integrations/utils/src/lib.rs
index 5bf1e61f5..86ed2a239 100644
--- a/integrations/utils/src/lib.rs
+++ b/integrations/utils/src/lib.rs
@@ -59,7 +59,7 @@ pub trait ExtendResponse: Sized {
                     // drop the owner, cleaning up the reactive runtime,
                     // once the stream is over
                     .chain(once(async move {
-                        drop(owner);
+                        owner.unset();
                         Default::default()
                     })),
             ));
diff --git a/reactive_graph/src/owner.rs b/reactive_graph/src/owner.rs
index 338cd956d..e17e649cf 100644
--- a/reactive_graph/src/owner.rs
+++ b/reactive_graph/src/owner.rs
@@ -55,6 +55,12 @@ pub struct Owner {
     pub(crate) shared_context: Option<Arc<dyn SharedContext + Send + Sync>>,
 }
 
+impl PartialEq for Owner {
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.inner, &other.inner)
+    }
+}
+
 thread_local! {
     static OWNER: RefCell<Option<Owner>> = Default::default();
 }
@@ -241,6 +247,15 @@ impl Owner {
         self.shared_context.clone()
     }
 
+    /// Removes this from its state as the thread-local owner and drops it.
+    pub fn unset(self) {
+        OWNER.with_borrow_mut(|owner| {
+            if owner.as_ref() == Some(&self) {
+                mem::take(owner);
+            }
+        })
+    }
+
     /// Returns the current [`SharedContext`], if any.
     #[cfg(feature = "hydration")]
     pub fn current_shared_context(
diff --git a/reactive_graph/src/owner/stored_value.rs b/reactive_graph/src/owner/stored_value.rs
index 6d5fcbb37..91aca2851 100644
--- a/reactive_graph/src/owner/stored_value.rs
+++ b/reactive_graph/src/owner/stored_value.rs
@@ -313,6 +313,7 @@ impl<T, S: Storage<T>> StoredValue<T, S> {
     /// assert_eq!(last, None);
     /// assert_eq!(length_fn(), None);
     /// ```
+    #[track_caller]
     pub fn try_with_value<U>(&self, fun: impl FnOnce(&T) -> U) -> Option<U> {
         S::try_with(self.node, fun)
     }
@@ -347,6 +348,7 @@ impl<T, S: Storage<T>> StoredValue<T, S> {
     /// assert_eq!(sum, 6);
     /// assert_eq!(length_fn(), 3);
     /// ```
+    #[track_caller]
     pub fn with_value<U>(&self, fun: impl FnOnce(&T) -> U) -> U {
         self.try_with_value(fun)
             .unwrap_or_else(unwrap_signal!(self))