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>, } +impl PartialEq for Owner { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.inner, &other.inner) + } +} + thread_local! { static OWNER: RefCell> = 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> StoredValue { /// assert_eq!(last, None); /// assert_eq!(length_fn(), None); /// ``` + #[track_caller] pub fn try_with_value(&self, fun: impl FnOnce(&T) -> U) -> Option { S::try_with(self.node, fun) } @@ -347,6 +348,7 @@ impl> StoredValue { /// assert_eq!(sum, 6); /// assert_eq!(length_fn(), 3); /// ``` + #[track_caller] pub fn with_value(&self, fun: impl FnOnce(&T) -> U) -> U { self.try_with_value(fun) .unwrap_or_else(unwrap_signal!(self))