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) {
344 OWNER.with_borrow_mut(|owner| {
345 if owner.as_ref().and_then(|n| n.upgrade()) == Some(self) {
346 mem::take(owner);
347 }
348 })
349 }
350
351 #[cfg(feature = "hydration")]
353 pub fn current_shared_context(
354 ) -> Option<Arc<dyn SharedContext + Send + Sync>> {
355 OWNER.with(|o| {
356 o.borrow()
357 .as_ref()
358 .and_then(|o| o.upgrade())
359 .and_then(|current| current.shared_context.clone())
360 })
361 }
362
363 #[cfg(feature = "hydration")]
366 pub fn with_hydration<T>(fun: impl FnOnce() -> T + 'static) -> T {
367 fn inner<T>(fun: Box<dyn FnOnce() -> T>) -> T {
368 provide_context(IsHydrating(true));
369
370 let sc = OWNER.with_borrow(|o| {
371 o.as_ref()
372 .and_then(|o| o.upgrade())
373 .and_then(|current| current.shared_context.clone())
374 });
375 match sc {
376 None => fun(),
377 Some(sc) => {
378 let prev = sc.get_is_hydrating();
379 sc.set_is_hydrating(true);
380 let value = fun();
381 sc.set_is_hydrating(prev);
382 value
383 }
384 }
385 }
386
387 inner(Box::new(fun))
388 }
389
390 #[cfg(feature = "hydration")]
392 pub fn with_no_hydration<T>(fun: impl FnOnce() -> T + 'static) -> T {
393 fn inner<T>(fun: Box<dyn FnOnce() -> T>) -> T {
394 provide_context(IsHydrating(false));
395
396 let sc = OWNER.with_borrow(|o| {
397 o.as_ref()
398 .and_then(|o| o.upgrade())
399 .and_then(|current| current.shared_context.clone())
400 });
401 match sc {
402 None => fun(),
403 Some(sc) => {
404 let prev = sc.get_is_hydrating();
405 sc.set_is_hydrating(false);
406 let value = fun();
407 sc.set_is_hydrating(prev);
408 value
409 }
410 }
411 }
412
413 inner(Box::new(fun))
414 }
415
416 pub fn pause(&self) {
423 let mut stack = Vec::with_capacity(16);
424 stack.push(Arc::downgrade(&self.inner));
425 while let Some(curr) = stack.pop() {
426 if let Some(curr) = curr.upgrade() {
427 let mut curr = curr.write().or_poisoned();
428 curr.paused = true;
429 stack.extend(curr.children.iter().map(Weak::clone));
430 }
431 }
432 }
433
434 pub fn paused(&self) -> bool {
436 self.inner.read().or_poisoned().paused
437 }
438
439 pub fn resume(&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 = false;
452 stack.extend(curr.children.iter().map(Weak::clone));
453 }
454 }
455 }
456}
457
458#[doc(hidden)]
459#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
460pub struct IsHydrating(pub bool);
461
462pub fn on_cleanup(fun: impl FnOnce() + Send + Sync + 'static) {
470 Owner::on_cleanup(fun)
471}
472
473#[derive(Default)]
474pub(crate) struct OwnerInner {
475 pub parent: Option<Weak<RwLock<OwnerInner>>>,
476 nodes: Vec<NodeId>,
477 pub contexts: FxHashMap<TypeId, Box<dyn Any + Send + Sync>>,
478 pub cleanups: Vec<Box<dyn FnOnce() + Send + Sync>>,
479 pub children: Vec<Weak<RwLock<OwnerInner>>>,
480 #[cfg(feature = "sandboxed-arenas")]
481 arena: Arc<RwLock<ArenaMap>>,
482 paused: bool,
483}
484
485impl Debug for OwnerInner {
486 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
487 f.debug_struct("OwnerInner")
488 .field("parent", &self.parent)
489 .field("nodes", &self.nodes)
490 .field("contexts", &self.contexts)
491 .field("cleanups", &self.cleanups.len())
492 .finish()
493 }
494}
495
496impl Drop for OwnerInner {
497 fn drop(&mut self) {
498 for child in std::mem::take(&mut self.children) {
499 if let Some(child) = child.upgrade() {
500 child.cleanup();
501 }
502 }
503
504 for cleanup in mem::take(&mut self.cleanups) {
505 cleanup();
506 }
507
508 let nodes = mem::take(&mut self.nodes);
509 if !nodes.is_empty() {
510 #[cfg(not(feature = "sandboxed-arenas"))]
511 Arena::with_mut(|arena| {
512 for node in nodes {
513 _ = arena.remove(node);
514 }
515 });
516 #[cfg(feature = "sandboxed-arenas")]
517 {
518 let mut arena = self.arena.write().or_poisoned();
519 for node in nodes {
520 _ = arena.remove(node);
521 }
522 }
523 }
524 }
525}
526
527trait Cleanup {
528 fn cleanup(&self);
529}
530
531impl Cleanup for RwLock<OwnerInner> {
532 fn cleanup(&self) {
533 let (cleanups, nodes, children) = {
534 let mut lock = self.write().or_poisoned();
535 (
536 mem::take(&mut lock.cleanups),
537 mem::take(&mut lock.nodes),
538 mem::take(&mut lock.children),
539 )
540 };
541 for child in children {
542 if let Some(child) = child.upgrade() {
543 child.cleanup();
544 }
545 }
546 for cleanup in cleanups {
547 cleanup();
548 }
549
550 if !nodes.is_empty() {
551 #[cfg(not(feature = "sandboxed-arenas"))]
552 Arena::with_mut(|arena| {
553 for node in nodes {
554 _ = arena.remove(node);
555 }
556 });
557 #[cfg(feature = "sandboxed-arenas")]
558 {
559 let arena = self.read().or_poisoned().arena.clone();
560 let mut arena = arena.write().or_poisoned();
561 for node in nodes {
562 _ = arena.remove(node);
563 }
564 }
565 }
566 }
567}