#[cfg(feature = "hydration")]
use hydration_context::SharedContext;
use or_poisoned::OrPoisoned;
use rustc_hash::FxHashMap;
use std::{
any::{Any, TypeId},
cell::RefCell,
fmt::Debug,
mem,
sync::{Arc, RwLock, Weak},
};
mod arc_stored_value;
mod arena;
mod arena_item;
mod context;
mod storage;
mod stored_value;
use self::arena::Arena;
pub use arc_stored_value::ArcStoredValue;
#[cfg(feature = "sandboxed-arenas")]
pub use arena::sandboxed::Sandboxed;
#[cfg(feature = "sandboxed-arenas")]
use arena::ArenaMap;
use arena::NodeId;
pub use arena_item::*;
pub use context::*;
pub use storage::*;
#[allow(deprecated)] pub use stored_value::{store_value, FromLocal, StoredValue};
#[derive(Debug, Clone, Default)]
#[must_use]
pub struct Owner {
pub(crate) inner: Arc<RwLock<OwnerInner>>,
#[cfg(feature = "hydration")]
pub(crate) shared_context: Option<Arc<dyn SharedContext + Send + Sync>>,
}
impl Owner {
fn downgrade(&self) -> WeakOwner {
WeakOwner {
inner: Arc::downgrade(&self.inner),
#[cfg(feature = "hydration")]
shared_context: self.shared_context.as_ref().map(Arc::downgrade),
}
}
}
#[derive(Clone)]
struct WeakOwner {
inner: Weak<RwLock<OwnerInner>>,
#[cfg(feature = "hydration")]
shared_context: Option<Weak<dyn SharedContext + Send + Sync>>,
}
impl WeakOwner {
fn upgrade(&self) -> Option<Owner> {
self.inner.upgrade().map(|inner| {
#[cfg(feature = "hydration")]
let shared_context =
self.shared_context.as_ref().and_then(|sc| sc.upgrade());
Owner {
inner,
#[cfg(feature = "hydration")]
shared_context,
}
})
}
}
impl PartialEq for Owner {
fn eq(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.inner, &other.inner)
}
}
thread_local! {
static OWNER: RefCell<Option<WeakOwner>> = Default::default();
}
impl Owner {
pub fn debug_id(&self) -> usize {
Arc::as_ptr(&self.inner) as usize
}
pub fn ancestry(&self) -> Vec<usize> {
let mut ancestors = Vec::new();
let mut curr_parent = self
.inner
.read()
.or_poisoned()
.parent
.as_ref()
.and_then(|n| n.upgrade());
while let Some(parent) = curr_parent {
ancestors.push(Arc::as_ptr(&parent) as usize);
curr_parent = parent
.read()
.or_poisoned()
.parent
.as_ref()
.and_then(|n| n.upgrade());
}
ancestors
}
pub fn new() -> Self {
#[cfg(not(feature = "hydration"))]
let parent = OWNER.with(|o| {
o.borrow()
.as_ref()
.and_then(|o| o.upgrade())
.map(|o| Arc::downgrade(&o.inner))
});
#[cfg(feature = "hydration")]
let (parent, shared_context) = OWNER
.with(|o| {
o.borrow().as_ref().and_then(|o| o.upgrade()).map(|o| {
(Some(Arc::downgrade(&o.inner)), o.shared_context.clone())
})
})
.unwrap_or((None, None));
let this = Self {
inner: Arc::new(RwLock::new(OwnerInner {
parent: parent.clone(),
nodes: Default::default(),
contexts: Default::default(),
cleanups: Default::default(),
children: Default::default(),
#[cfg(feature = "sandboxed-arenas")]
arena: parent
.as_ref()
.and_then(|parent| parent.upgrade())
.map(|parent| parent.read().or_poisoned().arena.clone())
.unwrap_or_default(),
paused: false,
})),
#[cfg(feature = "hydration")]
shared_context,
};
if let Some(parent) = parent.and_then(|n| n.upgrade()) {
parent
.write()
.or_poisoned()
.children
.push(Arc::downgrade(&this.inner));
}
this
}
#[cfg(feature = "hydration")]
#[track_caller]
pub fn new_root(
shared_context: Option<Arc<dyn SharedContext + Send + Sync>>,
) -> Self {
let this = Self {
inner: Arc::new(RwLock::new(OwnerInner {
parent: None,
nodes: Default::default(),
contexts: Default::default(),
cleanups: Default::default(),
children: Default::default(),
#[cfg(feature = "sandboxed-arenas")]
arena: Default::default(),
paused: false,
})),
#[cfg(feature = "hydration")]
shared_context,
};
this.set();
this
}
pub fn child(&self) -> Self {
let parent = Some(Arc::downgrade(&self.inner));
let mut inner = self.inner.write().or_poisoned();
#[cfg(feature = "sandboxed-arenas")]
let arena = inner.arena.clone();
let paused = inner.paused;
let child = Self {
inner: Arc::new(RwLock::new(OwnerInner {
parent,
nodes: Default::default(),
contexts: Default::default(),
cleanups: Default::default(),
children: Default::default(),
#[cfg(feature = "sandboxed-arenas")]
arena,
paused,
})),
#[cfg(feature = "hydration")]
shared_context: self.shared_context.clone(),
};
inner.children.push(Arc::downgrade(&child.inner));
child
}
pub fn set(&self) {
OWNER.with_borrow_mut(|owner| *owner = Some(self.downgrade()));
#[cfg(feature = "sandboxed-arenas")]
Arena::set(&self.inner.read().or_poisoned().arena);
}
pub fn with<T>(&self, fun: impl FnOnce() -> T) -> T {
fn inner_1(self_: &Owner) -> Option<WeakOwner> {
let prev = {
OWNER.with(|o| (*o.borrow_mut()).replace(self_.downgrade()))
};
#[cfg(feature = "sandboxed-arenas")]
Arena::set(&self_.inner.read().or_poisoned().arena);
prev
}
let prev = inner_1(self);
let val = fun();
fn inner_2(prev: Option<WeakOwner>) {
OWNER.with(|o| {
*o.borrow_mut() = prev;
});
}
inner_2(prev);
val
}
pub fn with_cleanup<T>(&self, fun: impl FnOnce() -> T) -> T {
self.cleanup();
self.with(fun)
}
pub fn cleanup(&self) {
self.inner.cleanup();
}
pub fn on_cleanup(fun: impl FnOnce() + Send + Sync + 'static) {
if let Some(owner) = Owner::current() {
let mut inner = owner.inner.write().or_poisoned();
#[cfg(feature = "sandboxed-arenas")]
let fun = {
let arena = Arc::clone(&inner.arena);
move || {
Arena::set(&arena);
fun()
}
};
inner.cleanups.push(Box::new(fun));
}
}
fn register(&self, node: NodeId) {
self.inner.write().or_poisoned().nodes.push(node);
}
pub fn current() -> Option<Owner> {
OWNER.with(|o| o.borrow().as_ref().and_then(|n| n.upgrade()))
}
#[cfg(feature = "hydration")]
pub fn shared_context(
&self,
) -> Option<Arc<dyn SharedContext + Send + Sync>> {
self.shared_context.clone()
}
pub fn unset(self) {
OWNER.with_borrow_mut(|owner| {
if owner.as_ref().and_then(|n| n.upgrade()) == Some(self) {
mem::take(owner);
}
})
}
#[cfg(feature = "hydration")]
pub fn current_shared_context(
) -> Option<Arc<dyn SharedContext + Send + Sync>> {
OWNER.with(|o| {
o.borrow()
.as_ref()
.and_then(|o| o.upgrade())
.and_then(|current| current.shared_context.clone())
})
}
#[cfg(feature = "hydration")]
pub fn with_hydration<T>(fun: impl FnOnce() -> T + 'static) -> T {
fn inner<T>(fun: Box<dyn FnOnce() -> T>) -> T {
provide_context(IsHydrating(true));
let sc = OWNER.with_borrow(|o| {
o.as_ref()
.and_then(|o| o.upgrade())
.and_then(|current| current.shared_context.clone())
});
match sc {
None => fun(),
Some(sc) => {
let prev = sc.get_is_hydrating();
sc.set_is_hydrating(true);
let value = fun();
sc.set_is_hydrating(prev);
value
}
}
}
inner(Box::new(fun))
}
#[cfg(feature = "hydration")]
pub fn with_no_hydration<T>(fun: impl FnOnce() -> T + 'static) -> T {
fn inner<T>(fun: Box<dyn FnOnce() -> T>) -> T {
provide_context(IsHydrating(false));
let sc = OWNER.with_borrow(|o| {
o.as_ref()
.and_then(|o| o.upgrade())
.and_then(|current| current.shared_context.clone())
});
match sc {
None => fun(),
Some(sc) => {
let prev = sc.get_is_hydrating();
sc.set_is_hydrating(false);
let value = fun();
sc.set_is_hydrating(prev);
value
}
}
}
inner(Box::new(fun))
}
pub fn pause(&self) {
let mut stack = Vec::with_capacity(16);
stack.push(Arc::downgrade(&self.inner));
while let Some(curr) = stack.pop() {
if let Some(curr) = curr.upgrade() {
let mut curr = curr.write().or_poisoned();
curr.paused = true;
stack.extend(curr.children.iter().map(Weak::clone));
}
}
}
pub fn paused(&self) -> bool {
self.inner.read().or_poisoned().paused
}
pub fn resume(&self) {
let mut stack = Vec::with_capacity(16);
stack.push(Arc::downgrade(&self.inner));
while let Some(curr) = stack.pop() {
if let Some(curr) = curr.upgrade() {
let mut curr = curr.write().or_poisoned();
curr.paused = false;
stack.extend(curr.children.iter().map(Weak::clone));
}
}
}
}
#[doc(hidden)]
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct IsHydrating(pub bool);
pub fn on_cleanup(fun: impl FnOnce() + Send + Sync + 'static) {
Owner::on_cleanup(fun)
}
#[derive(Default)]
pub(crate) struct OwnerInner {
pub parent: Option<Weak<RwLock<OwnerInner>>>,
nodes: Vec<NodeId>,
pub contexts: FxHashMap<TypeId, Box<dyn Any + Send + Sync>>,
pub cleanups: Vec<Box<dyn FnOnce() + Send + Sync>>,
pub children: Vec<Weak<RwLock<OwnerInner>>>,
#[cfg(feature = "sandboxed-arenas")]
arena: Arc<RwLock<ArenaMap>>,
paused: bool,
}
impl Debug for OwnerInner {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OwnerInner")
.field("parent", &self.parent)
.field("nodes", &self.nodes)
.field("contexts", &self.contexts)
.field("cleanups", &self.cleanups.len())
.finish()
}
}
impl Drop for OwnerInner {
fn drop(&mut self) {
for child in std::mem::take(&mut self.children) {
if let Some(child) = child.upgrade() {
child.cleanup();
}
}
for cleanup in mem::take(&mut self.cleanups) {
cleanup();
}
let nodes = mem::take(&mut self.nodes);
if !nodes.is_empty() {
#[cfg(not(feature = "sandboxed-arenas"))]
Arena::with_mut(|arena| {
for node in nodes {
_ = arena.remove(node);
}
});
#[cfg(feature = "sandboxed-arenas")]
{
let mut arena = self.arena.write().or_poisoned();
for node in nodes {
_ = arena.remove(node);
}
}
}
}
}
trait Cleanup {
fn cleanup(&self);
}
impl Cleanup for RwLock<OwnerInner> {
fn cleanup(&self) {
let (cleanups, nodes, children) = {
let mut lock = self.write().or_poisoned();
(
mem::take(&mut lock.cleanups),
mem::take(&mut lock.nodes),
mem::take(&mut lock.children),
)
};
for child in children {
if let Some(child) = child.upgrade() {
child.cleanup();
}
}
for cleanup in cleanups {
cleanup();
}
if !nodes.is_empty() {
#[cfg(not(feature = "sandboxed-arenas"))]
Arena::with_mut(|arena| {
for node in nodes {
_ = arena.remove(node);
}
});
#[cfg(feature = "sandboxed-arenas")]
{
let arena = self.read().or_poisoned().arena.clone();
let mut arena = arena.write().or_poisoned();
for node in nodes {
_ = arena.remove(node);
}
}
}
}
}