use std::cell::RefCell;
use std::marker::PhantomData;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use crate::Vm;
thread_local! {
static SCOPED_PERSISTENT_STATE_ROOT: RefCell<Option<PathBuf>> = const { RefCell::new(None) };
}
#[derive(Clone, Copy, Debug)]
pub struct PersistentStateRoot<'a>(&'a Path);
impl<'a> PersistentStateRoot<'a> {
#[must_use]
pub fn new(path: &'a Path) -> Self {
Self(path)
}
fn as_path(self) -> &'a Path {
self.0
}
}
#[derive(Debug)]
#[must_use = "retain this guard for the isolated VM execution"]
pub struct ScopedPersistentStateRoot {
previous: Option<PathBuf>,
_not_send: PhantomData<Rc<()>>,
}
pub fn scope_persistent_state_root(root: PersistentStateRoot<'_>) -> ScopedPersistentStateRoot {
let previous =
SCOPED_PERSISTENT_STATE_ROOT.with(|slot| slot.replace(Some(root.as_path().to_path_buf())));
ScopedPersistentStateRoot {
previous,
_not_send: PhantomData,
}
}
impl Drop for ScopedPersistentStateRoot {
fn drop(&mut self) {
SCOPED_PERSISTENT_STATE_ROOT.with(|slot| {
slot.replace(self.previous.take());
});
}
}
pub(crate) fn current_persistent_state_root() -> Option<PathBuf> {
SCOPED_PERSISTENT_STATE_ROOT.with(|slot| slot.borrow().clone())
}
pub fn register_persistent_state_builtins_at_root(
vm: &mut Vm,
base_dir: &Path,
state_root: PersistentStateRoot<'_>,
pipeline_name: &str,
) {
let state_root = state_root.as_path();
crate::store::register_store_builtins_at_state_root(vm, state_root);
crate::metadata::register_metadata_builtins_at_state_root(vm, base_dir, state_root);
crate::checkpoint::register_checkpoint_builtins_at_state_root(vm, state_root, pipeline_name);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn scoped_persistent_state_root_restores_nested_owner() {
let outer = Path::new("/isolated/outer/.harn");
let inner = Path::new("/isolated/inner/.harn");
assert_eq!(current_persistent_state_root(), None);
let outer_guard = scope_persistent_state_root(PersistentStateRoot::new(outer));
assert_eq!(current_persistent_state_root().as_deref(), Some(outer));
{
let _inner_guard = scope_persistent_state_root(PersistentStateRoot::new(inner));
assert_eq!(current_persistent_state_root().as_deref(), Some(inner));
}
assert_eq!(current_persistent_state_root().as_deref(), Some(outer));
drop(outer_guard);
assert_eq!(current_persistent_state_root(), None);
}
}