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 fn downgrade(&self) -> WeakOwner {
65 WeakOwner {
66 inner: Arc::downgrade(&self.inner),
67 #[cfg(feature = "hydration")]
68 shared_context: self.shared_context.as_ref().map(Arc::downgrade),
69 }
70 }
71}
72
73#[derive(Clone)]
74struct WeakOwner {
75 inner: Weak<RwLock<OwnerInner>>,
76 #[cfg(feature = "hydration")]
77 shared_context: Option<Weak<dyn SharedContext + Send + Sync>>,
78}
79
80impl WeakOwner {
81 fn upgrade(&self) -> Option<Owner> {
82 self.inner.upgrade().map(|inner| {
83 #[cfg(feature = "hydration")]
84 let shared_context =
85 self.shared_context.as_ref().and_then(|sc| sc.upgrade());
86 Owner {
87 inner,
88 #[cfg(feature = "hydration")]
89 shared_context,
90 }
91 })
92 }
93}
94
95impl PartialEq for Owner {
96 fn eq(&self, other: &Self) -> bool {
97 Arc::ptr_eq(&self.inner, &other.inner)
98 }
99}
100
101thread_local! {
102 static OWNER: RefCell<Option<WeakOwner>> = Default::default();
103}
104
105impl Owner {
106 pub fn debug_id(&self) -> usize {
111 Arc::as_ptr(&self.inner) as usize
112 }
113
114 pub fn ancestry(&self) -> Vec<usize> {
119 let mut ancestors = Vec::new();
120 let mut curr_parent = self
121 .inner
122 .read()
123 .or_poisoned()
124 .parent
125 .as_ref()
126 .and_then(|n| n.upgrade());
127 while let Some(parent) = curr_parent {
128 ancestors.push(Arc::as_ptr(&parent) as usize);
129 curr_parent = parent
130 .read()
131 .or_poisoned()
132 .parent
133 .as_ref()
134 .and_then(|n| n.upgrade());
135 }
136 ancestors
137 }
138
139 pub fn new() -> Self {
141 #[cfg(not(feature = "hydration"))]
142 let parent = OWNER.with(|o| {
143 o.borrow()
144 .as_ref()
145 .and_then(|o| o.upgrade())
146 .map(|o| Arc::downgrade(&o.inner))
147 });
148 #[cfg(feature = "hydration")]
149 let (parent, shared_context) = OWNER
150 .with(|o| {
151 o.borrow().as_ref().and_then(|o| o.upgrade()).map(|o| {
152 (Some(Arc::downgrade(&o.inner)), o.shared_context.clone())
153 })
154 })
155 .unwrap_or((None, None));
156 let this = Self {
157 inner: Arc::new(RwLock::new(OwnerInner {
158 parent: parent.clone(),
159 nodes: Default::default(),
160 contexts: Default::default(),
161 cleanups: Default::default(),
162 children: Default::default(),
163 #[cfg(feature = "sandboxed-arenas")]
164 arena: parent
165 .as_ref()
166 .and_then(|parent| parent.upgrade())
167 .map(|parent| parent.read().or_poisoned().arena.clone())
168 .unwrap_or_default(),
169 paused: false,
170 })),
171 #[cfg(feature = "hydration")]
172 shared_context,
173 };
174 if let Some(parent) = parent.and_then(|n| n.upgrade()) {
175 parent
176 .write()
177 .or_poisoned()
178 .children
179 .push(Arc::downgrade(&this.inner));
180 }
181 this
182 }
183
184 #[cfg(feature = "hydration")]
190 #[track_caller]
191 pub fn new_root(
192 shared_context: Option<Arc<dyn SharedContext + Send + Sync>>,
193 ) -> Self {
194 let this = Self {
195 inner: Arc::new(RwLock::new(OwnerInner {
196 parent: None,
197 nodes: Default::default(),
198 contexts: Default::default(),
199 cleanups: Default::default(),
200 children: Default::default(),
201 #[cfg(feature = "sandboxed-arenas")]
202 arena: Default::default(),
203 paused: false,
204 })),
205 #[cfg(feature = "hydration")]
206 shared_context,
207 };
208 this.set();
209 this
210 }
211
212 pub fn parent(&self) -> Option<Owner> {
218 self.inner
219 .read()
220 .or_poisoned()
221 .parent
222 .as_ref()
223 .and_then(|p| p.upgrade())
224 .map(|inner| Owner {
225 inner,
226 #[cfg(feature = "hydration")]
227 shared_context: self.shared_context.clone(),
228 })
229 }
230
231 pub fn child(&self) -> Self {
233 let parent = Some(Arc::downgrade(&self.inner));
234 let mut inner = self.inner.write().or_poisoned();
235 #[cfg(feature = "sandboxed-arenas")]
236 let arena = inner.arena.clone();
237 let paused = inner.paused;
238 let child = Self {
239 inner: Arc::new(RwLock::new(OwnerInner {
240 parent,
241 nodes: Default::default(),
242 contexts: Default::default(),
243 cleanups: Default::default(),
244 children: Default::default(),
245 #[cfg(feature = "sandboxed-arenas")]
246 arena,
247 paused,
248 })),
249 #[cfg(feature = "hydration")]
250 shared_context: self.shared_context.clone(),
251 };
252 inner.children.push(Arc::downgrade(&child.inner));
253 child
254 }
255
256 pub fn set(&self) {
258 OWNER.with_borrow_mut(|owner| *owner = Some(self.downgrade()));
259 #[cfg(feature = "sandboxed-arenas")]
260 Arena::set(&self.inner.read().or_poisoned().arena);
261 }
262
263 pub fn with<T>(&self, fun: impl FnOnce() -> T) -> T {
265 fn inner_1(self_: &Owner) -> Option<WeakOwner> {
267 let prev = {
268 OWNER.with(|o| (*o.borrow_mut()).replace(self_.downgrade()))
269 };
270 #[cfg(feature = "sandboxed-arenas")]
271 Arena::set(&self_.inner.read().or_poisoned().arena);
272 prev
273 }
274 let prev = inner_1(self);
275
276 let val = fun();
277
278 fn inner_2(prev: Option<WeakOwner>) {
280 OWNER.with(|o| {
281 *o.borrow_mut() = prev;
282 });
283 }
284 inner_2(prev);
285
286 val
287 }
288
289 pub fn with_cleanup<T>(&self, fun: impl FnOnce() -> T) -> T {
291 self.cleanup();
292 self.with(fun)
293 }
294
295 pub fn cleanup(&self) {
300 self.inner.cleanup();
301 }
302
303 pub fn on_cleanup(fun: impl FnOnce() + Send + Sync + 'static) {
309 if let Some(owner) = Owner::current() {
310 let mut inner = owner.inner.write().or_poisoned();
311
312 #[cfg(feature = "sandboxed-arenas")]
313 let fun = {
314 let arena = Arc::clone(&inner.arena);
315 move || {
316 Arena::set(&arena);
317 fun()
318 }
319 };
320
321 inner.cleanups.push(Box::new(fun));
322 }
323 }
324
325 fn register(&self, node: NodeId) {
326 self.inner.write().or_poisoned().nodes.push(node);
327 }
328
329 pub fn current() -> Option<Owner> {
331 OWNER.with(|o| o.borrow().as_ref().and_then(|n| n.upgrade()))
332 }
333
334 #[cfg(feature = "hydration")]
336 pub fn shared_context(
337 &self,
338 ) -> Option<Arc<dyn SharedContext + Send + Sync>> {
339 self.shared_context.clone()
340 }
341
342 pub fn unset(self) {
346 OWNER.with_borrow_mut(|owner| {
347 if owner.as_ref().and_then(|n| n.upgrade()) == Some(self) {
348 mem::take(owner);
349 }
350 })
351 }
352
353 pub fn unset_with_forced_cleanup(self) {
357 OWNER.with_borrow_mut(|owner| {
358 if owner
359 .as_ref()
360 .and_then(|n| n.upgrade())
361 .map(|o| o == self)
362 .unwrap_or(false)
363 {
364 mem::take(owner);
365 }
366 });
367 self.cleanup();
368 }
369
370 #[cfg(feature = "hydration")]
372 pub fn current_shared_context(
373 ) -> Option<Arc<dyn SharedContext + Send + Sync>> {
374 OWNER.with(|o| {
375 o.borrow()
376 .as_ref()
377 .and_then(|o| o.upgrade())
378 .and_then(|current| current.shared_context.clone())
379 })
380 }
381
382 #[cfg(feature = "hydration")]
385 pub fn with_hydration<T>(fun: impl FnOnce() -> T + 'static) -> T {
386 fn inner<T>(fun: Box<dyn FnOnce() -> T>) -> T {
387 provide_context(IsHydrating(true));
388
389 let sc = OWNER.with_borrow(|o| {
390 o.as_ref()
391 .and_then(|o| o.upgrade())
392 .and_then(|current| current.shared_context.clone())
393 });
394 match sc {
395 None => fun(),
396 Some(sc) => {
397 let prev = sc.get_is_hydrating();
398 sc.set_is_hydrating(true);
399 let value = fun();
400 sc.set_is_hydrating(prev);
401 value
402 }
403 }
404 }
405
406 inner(Box::new(fun))
407 }
408
409 #[cfg(feature = "hydration")]
411 pub fn with_no_hydration<T>(fun: impl FnOnce() -> T + 'static) -> T {
412 fn inner<T>(fun: Box<dyn FnOnce() -> T>) -> T {
413 provide_context(IsHydrating(false));
414
415 let sc = OWNER.with_borrow(|o| {
416 o.as_ref()
417 .and_then(|o| o.upgrade())
418 .and_then(|current| current.shared_context.clone())
419 });
420 match sc {
421 None => fun(),
422 Some(sc) => {
423 let prev = sc.get_is_hydrating();
424 sc.set_is_hydrating(false);
425 let value = fun();
426 sc.set_is_hydrating(prev);
427 value
428 }
429 }
430 }
431
432 inner(Box::new(fun))
433 }
434
435 pub fn pause(&self) {
442 let mut stack = Vec::with_capacity(16);
443 stack.push(Arc::downgrade(&self.inner));
444 while let Some(curr) = stack.pop() {
445 if let Some(curr) = curr.upgrade() {
446 let mut curr = curr.write().or_poisoned();
447 curr.paused = true;
448 stack.extend(curr.children.iter().map(Weak::clone));
449 }
450 }
451 }
452
453 pub fn paused(&self) -> bool {
455 self.inner.read().or_poisoned().paused
456 }
457
458 pub fn resume(&self) {
465 let mut stack = Vec::with_capacity(16);
466 stack.push(Arc::downgrade(&self.inner));
467 while let Some(curr) = stack.pop() {
468 if let Some(curr) = curr.upgrade() {
469 let mut curr = curr.write().or_poisoned();
470 curr.paused = false;
471 stack.extend(curr.children.iter().map(Weak::clone));
472 }
473 }
474 }
475}
476
477#[doc(hidden)]
478#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
479pub struct IsHydrating(pub bool);
480
481pub fn on_cleanup(fun: impl FnOnce() + Send + Sync + 'static) {
489 Owner::on_cleanup(fun)
490}
491
492#[derive(Default)]
493pub(crate) struct OwnerInner {
494 pub parent: Option<Weak<RwLock<OwnerInner>>>,
495 nodes: Vec<NodeId>,
496 pub contexts: FxHashMap<TypeId, Box<dyn Any + Send + Sync>>,
497 pub cleanups: Vec<Box<dyn FnOnce() + Send + Sync>>,
498 pub children: Vec<Weak<RwLock<OwnerInner>>>,
499 #[cfg(feature = "sandboxed-arenas")]
500 arena: Arc<RwLock<ArenaMap>>,
501 paused: bool,
502}
503
504impl Debug for OwnerInner {
505 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
506 f.debug_struct("OwnerInner")
507 .field("parent", &self.parent)
508 .field("nodes", &self.nodes)
509 .field("contexts", &self.contexts)
510 .field("cleanups", &self.cleanups.len())
511 .finish()
512 }
513}
514
515impl Drop for OwnerInner {
516 fn drop(&mut self) {
517 for child in std::mem::take(&mut self.children) {
518 if let Some(child) = child.upgrade() {
519 child.cleanup();
520 }
521 }
522
523 for cleanup in mem::take(&mut self.cleanups) {
524 cleanup();
525 }
526
527 let nodes = mem::take(&mut self.nodes);
528 if !nodes.is_empty() {
529 #[cfg(not(feature = "sandboxed-arenas"))]
530 Arena::with_mut(|arena| {
531 for node in nodes {
532 _ = arena.remove(node);
533 }
534 });
535 #[cfg(feature = "sandboxed-arenas")]
536 {
537 let mut arena = self.arena.write().or_poisoned();
538 for node in nodes {
539 _ = arena.remove(node);
540 }
541 }
542 }
543 }
544}
545
546trait Cleanup {
547 fn cleanup(&self);
548}
549
550impl Cleanup for RwLock<OwnerInner> {
551 fn cleanup(&self) {
552 let (cleanups, nodes, children) = {
553 let mut lock = self.write().or_poisoned();
554 (
555 mem::take(&mut lock.cleanups),
556 mem::take(&mut lock.nodes),
557 mem::take(&mut lock.children),
558 )
559 };
560 for child in children {
561 if let Some(child) = child.upgrade() {
562 child.cleanup();
563 }
564 }
565 for cleanup in cleanups {
566 cleanup();
567 }
568
569 if !nodes.is_empty() {
570 #[cfg(not(feature = "sandboxed-arenas"))]
571 Arena::with_mut(|arena| {
572 for node in nodes {
573 _ = arena.remove(node);
574 }
575 });
576 #[cfg(feature = "sandboxed-arenas")]
577 {
578 let arena = self.read().or_poisoned().arena.clone();
579 let mut arena = arena.write().or_poisoned();
580 for node in nodes {
581 _ = arena.remove(node);
582 }
583 }
584 }
585 }
586}