fix: remove Owner from thread-local when it is ordinarily dropped, to ensure cleanup (closes #2942) (#2944)

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()
                     })),
             ));

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(

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))