1#[cfg(feature = "hydration")]
4use hydration_context::SharedContext;
5use or_poisoned::OrPoisoned;
6use rustc_hash::FxHashMap;
7use std::{
8 any::{Any, TypeId},
9 cell::RefCell,
10 fmt::Debug,
11 mem,
12 sync::{Arc, RwLock, Weak},
13};
14
15mod arc_stored_value;
16mod arena;
17mod arena_item;
18mod context;
19mod storage;
20mod stored_value;
21use self::arena::Arena;
22pub use arc_stored_value::ArcStoredValue;
23#[cfg(feature = "sandboxed-arenas")]
24pub use arena::sandboxed::Sandboxed;
25#[cfg(feature = "sandboxed-arenas")]
26use arena::ArenaMap;
27use arena::NodeId;
28pub use arena_item::*;
29pub use context::*;
30pub use storage::*;
31#[allow(deprecated)] pub use stored_value::{store_value, FromLocal, StoredValue};
33
34#[derive(Debug, Clone, Default)]
56#[must_use]
57pub struct Owner {
58 pub(crate) inner: Arc<RwLock<OwnerInner>>,
59 #[cfg(feature = "hydration")]
60 pub(crate) shared_context: Option<Arc<dyn SharedContext + Send + Sync>>,
61}
62
63impl Owner {
64 pub fn downgrade(&self) -> WeakOwner {
66 WeakOwner {
67 inner: Arc::downgrade(&self.inner),
68 #[cfg(feature = "hydration")]
69 shared_context: self.shared_context.as_ref().map(Arc::downgrade),
70 }
71 }
72}
73
74#[derive(Clone)]
79pub struct WeakOwner {
80 inner: Weak<RwLock<OwnerInner>>,
81 #[cfg(feature = "hydration")]
82 shared_context: Option<Weak<dyn SharedContext + Send + Sync>>,
83}
84
85impl WeakOwner {
86 pub fn upgrade(&self) -> Option<Owner> {
90 self.inner.upgrade().map(|inner| {
91 #[cfg(feature = "hydration")]
92 let shared_context =
93 self.shared_context.as_ref().and_then(|sc| sc.upgrade());
94 Owner {
95 inner,
96 #[cfg(feature = "hydration")]
97 shared_context,
98 }
99 })
100 }
101}
102
103impl PartialEq for Owner {
104 fn eq(&self, other: &Self) -> bool {
105 Arc::ptr_eq(&self.inner, &other.inner)
106 }
107}
108
109thread_local! {
110 static OWNER: RefCell<Option<WeakOwner>> = Default::default();
111}
112
113impl Owner {
114 pub fn debug_id(&self) -> usize {
119 Arc::as_ptr(&self.inner) as usize
120 }
121
122 pub fn ancestry(&self) -> Vec<usize> {
127 let mut ancestors = Vec::new();
128 let mut curr_parent = self
129 .inner
130 .read()
131 .or_poisoned()
132 .parent
133 .as_ref()
134 .and_then(|n| n.upgrade());
135 while let Some(parent) = curr_parent {
136 ancestors.push(Arc::as_ptr(&parent) as usize);
137 curr_parent = parent
138 .read()
139 .or_poisoned()
140 .parent
141 .as_ref()
142 .and_then(|n| n.upgrade());
143 }
144 ancestors
145 }
146
147 pub fn new() -> Self {
149 #[cfg(not(feature = "hydration"))]
150 let parent = OWNER.with(|o| {
151 o.borrow()
152 .as_ref()
153 .and_then(|o| o.upgrade())
154 .map(|o| Arc::downgrade(&o.inner))
155 });
156 #[cfg(feature = "hydration")]
157 let (parent, shared_context) = OWNER
158 .with(|o| {
159 o.borrow().as_ref().and_then(|o| o.upgrade()).map(|o| {
160 (Some(Arc::downgrade(&o.inner)), o.shared_context.clone())
161 })
162 })
163 .unwrap_or((None, None));
164 let this = Self {
165 inner: Arc::new(RwLock::new(OwnerInner {
166 parent: parent.clone(),
167 nodes: Default::default(),
168 contexts: Default::default(),
169 cleanups: Default::default(),
170 children: Default::default(),
171 #[cfg(feature = "sandboxed-arenas")]
172 arena: parent
173 .as_ref()
174 .and_then(|parent| parent.upgrade())
175 .map(|parent| parent.read().or_poisoned().arena.clone())
176 .unwrap_or_default(),
177 paused: false,
178 })),
179 #[cfg(feature = "hydration")]
180 shared_context,
181 };
182 if let Some(parent) = parent.and_then(|n| n.upgrade()) {
183 parent
184 .write()
185 .or_poisoned()
186 .children
187 .push(Arc::downgrade(&this.inner));
188 }
189 this
190 }
191
192 #[cfg(feature = "hydration")]
198 #[track_caller]
199 pub fn new_root(
200 shared_context: Option<Arc<dyn SharedContext + Send + Sync>>,
201 ) -> Self {
202 let this = Self {
203 inner: Arc::new(RwLock::new(OwnerInner {
204 parent: None,
205 nodes: Default::default(),
206 contexts: Default::default(),
207 cleanups: Default::default(),
208 children: Default::default(),
209 #[cfg(feature = "sandboxed-arenas")]
210 arena: Default::default(),
211 paused: false,
212 })),
213 #[cfg(feature = "hydration")]
214 shared_context,
215 };
216 this.set();
217 this
218 }
219
220 pub fn parent(&self) -> Option<Owner> {
226 self.inner
227 .read()
228 .or_poisoned()
229 .parent
230 .as_ref()
231 .and_then(|p| p.upgrade())
232 .map(|inner| Owner {
233 inner,
234 #[cfg(feature = "hydration")]
235 shared_context: self.shared_context.clone(),
236 })
237 }
238
239 pub fn child(&self) -> Self {
241 let parent = Some(Arc::downgrade(&self.inner));
242 let mut inner = self.inner.write().or_poisoned();
243 #[cfg(feature = "sandboxed-arenas")]
244 let arena = inner.arena.clone();
245 let paused = inner.paused;
246 let child = Self {
247 inner: Arc::new(RwLock::new(OwnerInner {
248 parent,
249 nodes: Default::default(),
250 contexts: Default::default(),
251 cleanups: Default::default(),
252 children: Default::default(),
253 #[cfg(feature = "sandboxed-arenas")]
254 arena,
255 paused,
256 })),
257 #[cfg(feature = "hydration")]
258 shared_context: self.shared_context.clone(),
259 };
260 inner.children.push(Arc::downgrade(&child.inner));
261 child
262 }
263
264 pub fn set(&self) {
266 OWNER.with_borrow_mut(|owner| *owner = Some(self.downgrade()));
267 #[cfg(feature = "sandboxed-arenas")]
268 Arena::set(&self.inner.read().or_poisoned().arena);
269 }
270
271 pub fn with<T>(&self, fun: impl FnOnce() -> T) -> T {
273 fn inner_1(self_: &Owner) -> Option<WeakOwner> {
275 let prev = OWNER.with_borrow_mut(|o| o.replace(self_.downgrade()));
276 #[cfg(feature = "sandboxed-arenas")]
277 Arena::set(&self_.inner.read().or_poisoned().arena);
278 prev
279 }
280 let prev = inner_1(self);
281
282 let val = fun();
283
284 fn inner_2(prev: Option<WeakOwner>) {
286 OWNER.with_borrow_mut(|o| *o = prev);
287 }
288 inner_2(prev);
289
290 val
291 }
292
293 pub fn with_cleanup<T>(&self, fun: impl FnOnce() -> T) -> T {
295 self.cleanup();
296 self.with(fun)
297 }
298
299 pub fn cleanup(&self) {
304 self.inner.cleanup();
305 }
306
307 pub fn on_cleanup(fun: impl FnOnce() + Send + Sync + 'static) {
313 if let Some(owner) = Owner::current() {
314 let mut inner = owner.inner.write().or_poisoned();
315
316 #[cfg(feature = "sandboxed-arenas")]
317 let fun = {
318 let arena = Arc::clone(&inner.arena);
319 move || {
320 Arena::set(&arena);
321 fun()
322 }
323 };
324
325 inner.cleanups.push(Box::new(fun));
326 }
327 }
328
329 fn register(&self, node: NodeId) {
330 self.inner.write().or_poisoned().nodes.push(node);
331 }
332
333 pub fn current() -> Option<Owner> {
335 OWNER.with(|o| o.borrow().as_ref().and_then(|n| n.upgrade()))
336 }
337
338 #[cfg(feature = "hydration")]
340 pub fn shared_context(
341 &self,
342 ) -> Option<Arc<dyn SharedContext + Send + Sync>> {
343 self.shared_context.clone()
344 }
345
346 pub fn unset(self) {
350 OWNER.with_borrow_mut(|owner| {
351 if owner.as_ref().and_then(|n| n.upgrade()) == Some(self) {
352 mem::take(owner);
353 }
354 })
355 }
356
357 pub fn unset_with_forced_cleanup(self) {
361 OWNER.with_borrow_mut(|owner| {
362 if owner
363 .as_ref()
364 .and_then(|n| n.upgrade())
365 .map(|o| o == self)
366 .unwrap_or(false)
367 {
368 mem::take(owner);
369 }
370 });
371 self.cleanup();
372 }
373
374 #[cfg(feature = "hydration")]
376 pub fn current_shared_context(
377 ) -> Option<Arc<dyn SharedContext + Send + Sync>> {
378 OWNER.with(|o| {
379 o.borrow()
380 .as_ref()
381 .and_then(|o| o.upgrade())
382 .and_then(|current| current.shared_context.clone())
383 })
384 }
385
386 #[cfg(feature = "hydration")]
389 pub fn with_hydration<T>(fun: impl FnOnce() -> T + 'static) -> T {
390 fn inner<T>(fun: Box<dyn FnOnce() -> T>) -> T {
391 provide_context(IsHydrating(true));
392
393 let sc = OWNER.with_borrow(|o| {
394 o.as_ref()
395 .and_then(|o| o.upgrade())
396 .and_then(|current| current.shared_context.clone())
397 });
398 match sc {
399 None => fun(),
400 Some(sc) => {
401 let prev = sc.get_is_hydrating();
402 sc.set_is_hydrating(true);
403 let value = fun();
404 sc.set_is_hydrating(prev);
405 value
406 }
407 }
408 }
409
410 inner(Box::new(fun))
411 }
412
413 #[cfg(feature = "hydration")]
415 pub fn with_no_hydration<T>(fun: impl FnOnce() -> T + 'static) -> T {
416 fn inner<T>(fun: Box<dyn FnOnce() -> T>) -> T {
417 provide_context(IsHydrating(false));
418
419 let sc = OWNER.with_borrow(|o| {
420 o.as_ref()
421 .and_then(|o| o.upgrade())
422 .and_then(|current| current.shared_context.clone())
423 });
424 match sc {
425 None => fun(),
426 Some(sc) => {
427 let prev = sc.get_is_hydrating();
428 sc.set_is_hydrating(false);
429 let value = fun();
430 sc.set_is_hydrating(prev);
431 value
432 }
433 }
434 }
435
436 inner(Box::new(fun))
437 }
438
439 pub fn pause(&self) {
446 let mut stack = Vec::with_capacity(16);
447 stack.push(Arc::downgrade(&self.inner));
448 while let Some(curr) = stack.pop() {
449 if let Some(curr) = curr.upgrade() {
450 let mut curr = curr.write().or_poisoned();
451 curr.paused = true;
452 stack.extend(curr.children.iter().map(Weak::clone));
453 }
454 }
455 }
456
457 pub fn paused(&self) -> bool {
459 self.inner.read().or_poisoned().paused
460 }
461
462 pub fn resume(&self) {
469 let mut stack = Vec::with_capacity(16);
470 stack.push(Arc::downgrade(&self.inner));
471 while let Some(curr) = stack.pop() {
472 if let Some(curr) = curr.upgrade() {
473 let mut curr = curr.write().or_poisoned();
474 curr.paused = false;
475 stack.extend(curr.children.iter().map(Weak::clone));
476 }
477 }
478 }
479}
480
481#[doc(hidden)]
482#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
483pub struct IsHydrating(pub bool);
484
485pub fn on_cleanup(fun: impl FnOnce() + Send + Sync + 'static) {
493 Owner::on_cleanup(fun)
494}
495
496#[derive(Default)]
497pub(crate) struct OwnerInner {
498 pub parent: Option<Weak<RwLock<OwnerInner>>>,
499 nodes: Vec<NodeId>,
500 pub contexts: FxHashMap<TypeId, Box<dyn Any + Send + Sync>>,
501 pub cleanups: Vec<Box<dyn FnOnce() + Send + Sync>>,
502 pub children: Vec<Weak<RwLock<OwnerInner>>>,
503 #[cfg(feature = "sandboxed-arenas")]
504 arena: Arc<RwLock<ArenaMap>>,
505 paused: bool,
506}
507
508impl Debug for OwnerInner {
509 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
510 f.debug_struct("OwnerInner")
511 .field("parent", &self.parent)
512 .field("nodes", &self.nodes)
513 .field("contexts", &self.contexts)
514 .field("cleanups", &self.cleanups.len())
515 .finish()
516 }
517}
518
519impl Drop for OwnerInner {
520 fn drop(&mut self) {
521 for child in std::mem::take(&mut self.children) {
522 if let Some(child) = child.upgrade() {
523 child.cleanup();
524 }
525 }
526
527 for cleanup in mem::take(&mut self.cleanups) {
528 cleanup();
529 }
530
531 let nodes = mem::take(&mut self.nodes);
532 if !nodes.is_empty() {
533 #[cfg(not(feature = "sandboxed-arenas"))]
534 Arena::with_mut(|arena| {
535 for node in nodes {
536 _ = arena.remove(node);
537 }
538 });
539 #[cfg(feature = "sandboxed-arenas")]
540 {
541 let mut arena = self.arena.write().or_poisoned();
542 for node in nodes {
543 _ = arena.remove(node);
544 }
545 }
546 }
547 }
548}
549
550trait Cleanup {
551 fn cleanup(&self);
552}
553
554impl Cleanup for RwLock<OwnerInner> {
555 fn cleanup(&self) {
556 let (cleanups, nodes, children) = {
557 let mut lock = self.write().or_poisoned();
558 (
559 mem::take(&mut lock.cleanups),
560 mem::take(&mut lock.nodes),
561 mem::take(&mut lock.children),
562 )
563 };
564 for child in children {
565 if let Some(child) = child.upgrade() {
566 child.cleanup();
567 }
568 }
569 for cleanup in cleanups {
570 cleanup();
571 }
572
573 if !nodes.is_empty() {
574 #[cfg(not(feature = "sandboxed-arenas"))]
575 Arena::with_mut(|arena| {
576 for node in nodes {
577 _ = arena.remove(node);
578 }
579 });
580 #[cfg(feature = "sandboxed-arenas")]
581 {
582 let arena = self.read().or_poisoned().arena.clone();
583 let mut arena = arena.write().or_poisoned();
584 for node in nodes {
585 _ = arena.remove(node);
586 }
587 }
588 }
589 }
590}