use crate::common::{
engine::{Engine, EngineSizing, validate_completed_dataset, write_file_durable},
error::{Result, VsdbError},
vsdb_freeze_base_dir, vsdb_get_base_dir,
};
use parking_lot::Mutex;
use ruc::pnk;
use serde::{Deserialize, Serialize};
use std::{
cell::RefCell,
collections::HashMap,
fmt, fs, io,
path::{Component, Path, PathBuf},
str::FromStr,
sync::{Arc, LazyLock},
time::{SystemTime, UNIX_EPOCH},
};
pub type NsId = u64;
pub const DEFAULT_NS_ID: NsId = 0;
const NS_REGISTRY_REL_PATH: &str = "__SYSTEM__/__namespaces__";
const NS_STATE_DIR_REL_PATH: &str = "__SYSTEM__/__namespace_state__";
const NS_DERIVED_DIR: &str = "__NAMESPACES__";
const DEFAULT_NS_SHARDS: usize = 4;
const DEFAULT_NS_BUDGET_MB: usize = 512;
#[derive(
Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize,
)]
#[serde(from = "InstanceIdWire")]
pub struct InstanceId {
pub map_id: u64,
pub ns: Option<NsId>,
}
#[derive(Deserialize)]
struct InstanceIdWire {
map_id: u64,
ns: Option<NsId>,
}
impl InstanceId {
pub fn new(map_id: u64, ns: NsId) -> Self {
Self {
map_id,
ns: (ns != DEFAULT_NS_ID).then_some(ns),
}
}
}
impl From<InstanceIdWire> for InstanceId {
fn from(w: InstanceIdWire) -> Self {
Self {
map_id: w.map_id,
ns: w.ns.filter(|&n| n != DEFAULT_NS_ID),
}
}
}
impl From<u64> for InstanceId {
fn from(map_id: u64) -> Self {
Self { map_id, ns: None }
}
}
impl fmt::Display for InstanceId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.ns {
None => write!(f, "{}", self.map_id),
Some(ns) => write!(f, "{}@{}", self.map_id, ns),
}
}
}
impl FromStr for InstanceId {
type Err = VsdbError;
fn from_str(s: &str) -> Result<Self> {
let parse = |v: &str, what: &str| {
v.parse::<u64>().map_err(|_| VsdbError::Decode {
detail: format!("invalid InstanceId {what}: {v:?}"),
})
};
match s.split_once('@') {
None => Ok(Self {
map_id: parse(s, "map_id")?,
ns: None,
}),
Some((m, n)) => Ok(Self {
map_id: parse(m, "map_id")?,
ns: Some(parse(n, "ns_id")?).filter(|&n| n != DEFAULT_NS_ID),
}),
}
}
}
#[derive(Clone, Debug)]
pub struct NamespaceOpts {
pub path: Option<PathBuf>,
pub shards: usize,
pub mem_budget_mb: Option<usize>,
}
impl Default for NamespaceOpts {
fn default() -> Self {
Self {
path: None,
shards: DEFAULT_NS_SHARDS,
mem_budget_mb: None,
}
}
}
#[derive(Clone, Debug)]
pub struct NsInfo {
pub id: NsId,
pub path: PathBuf,
pub pinned: bool,
pub shards: usize,
pub created_at: u64,
}
#[derive(Serialize, Deserialize, Clone)]
struct NsRecord {
id: NsId,
path: Option<String>,
shards: u32,
mem_budget_mb: Option<u64>,
created_at: u64,
}
#[derive(Serialize, Deserialize)]
struct RegistryFile {
next_id: NsId,
entries: Vec<NsRecord>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum NsLifecycleState {
Pending,
Established,
}
impl Default for RegistryFile {
fn default() -> Self {
Self {
next_id: 1,
entries: Vec::new(),
}
}
}
static REGISTRY_LOCK: Mutex<()> = Mutex::new(());
static OPEN_NAMESPACES: LazyLock<Mutex<HashMap<NsId, Namespace>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
static DEFAULT_NS: LazyLock<Namespace> = LazyLock::new(|| {
Namespace(Arc::new(NsInner {
id: DEFAULT_NS_ID,
path: vsdb_get_base_dir(),
engine: pnk!(Engine::new()),
}))
});
thread_local! {
static NS_STACK: RefCell<Vec<Namespace>> = const { RefCell::new(Vec::new()) };
}
fn registry_path() -> PathBuf {
vsdb_get_base_dir().join(NS_REGISTRY_REL_PATH)
}
fn lifecycle_path(base: &Path, id: NsId) -> PathBuf {
base.join(NS_STATE_DIR_REL_PATH).join(format!("{id:016x}"))
}
fn load_lifecycle(base: &Path, id: NsId) -> Result<Option<NsLifecycleState>> {
match fs::read(lifecycle_path(base, id)) {
Ok(bytes) if bytes.as_slice() == b"P" => Ok(Some(NsLifecycleState::Pending)),
Ok(bytes) if bytes.as_slice() == b"E" => Ok(Some(NsLifecycleState::Established)),
Ok(bytes) => Err(ns_err(format!(
"namespace {id} has corrupt lifecycle state bytes: {bytes:?}"
))),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(e.into()),
}
}
fn save_lifecycle(base: &Path, id: NsId, state: NsLifecycleState) -> Result<()> {
let byte = match state {
NsLifecycleState::Pending => b"P",
NsLifecycleState::Established => b"E",
};
let path = lifecycle_path(base, id);
fs::create_dir_all(path.parent().expect("has parent"))?;
write_file_durable(&path, byte).map_err(VsdbError::from)
}
fn remove_lifecycle(base: &Path, id: NsId) -> Result<()> {
let path = lifecycle_path(base, id);
match fs::remove_file(&path) {
Ok(()) => {
if let Some(parent) = path.parent() {
fs::File::open(parent)?.sync_all()?;
}
Ok(())
}
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(e.into()),
}
}
fn load_registry() -> Result<RegistryFile> {
match fs::read(registry_path()) {
Ok(bytes) => Ok(postcard::from_bytes(&bytes)?),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(RegistryFile::default()),
Err(e) => Err(e.into()),
}
}
fn save_registry(reg: &RegistryFile) -> Result<()> {
let path = registry_path();
fs::create_dir_all(path.parent().expect("has parent"))?;
let bytes = postcard::to_allocvec(reg)?;
write_file_durable(&path, &bytes).map_err(VsdbError::from)
}
fn resolve_root(base: &Path, rec: &NsRecord) -> PathBuf {
match &rec.path {
Some(p) => PathBuf::from(p),
None => base.join(NS_DERIVED_DIR).join(format!("{:016x}", rec.id)),
}
}
fn path_contains(outer: &Path, inner: &Path) -> bool {
let o: Vec<Component<'_>> = outer.components().collect();
let i: Vec<Component<'_>> = inner.components().collect();
i.len() >= o.len() && i[..o.len()] == o[..]
}
fn normalize_physical(p: &Path) -> PathBuf {
let mut existing = p;
let mut tail: Vec<std::ffi::OsString> = Vec::new();
loop {
match existing.canonicalize() {
Ok(mut c) => {
for seg in tail.iter().rev() {
c.push(seg);
}
return c;
}
Err(_) => match (existing.parent(), existing.file_name()) {
(Some(parent), Some(name)) => {
tail.push(name.to_owned());
existing = parent;
}
_ => return p.to_path_buf(),
},
}
}
}
fn ns_err(detail: impl Into<String>) -> VsdbError {
VsdbError::Namespace {
detail: detail.into(),
}
}
fn validate_explicit_root(
base: &Path,
reg: &RegistryFile,
candidate: &Path,
) -> Result<()> {
if !candidate.is_absolute() {
return Err(ns_err(format!(
"namespace path must be absolute: {}",
candidate.display()
)));
}
if candidate.as_os_str().to_str().is_none() {
return Err(ns_err(format!(
"namespace path must be valid UTF-8: {}",
candidate.display()
)));
}
if candidate
.components()
.any(|c| matches!(c, Component::ParentDir | Component::CurDir))
{
return Err(ns_err(format!(
"namespace path must not contain `.`/`..` components: {}",
candidate.display()
)));
}
let cand_norm = normalize_physical(candidate);
let base_norm = normalize_physical(base);
if path_contains(&base_norm, &cand_norm) || path_contains(&cand_norm, &base_norm) {
return Err(ns_err(format!(
"namespace path {} overlaps the default base dir {}",
candidate.display(),
base.display()
)));
}
for rec in ®.entries {
let other = normalize_physical(&resolve_root(base, rec));
if path_contains(&other, &cand_norm) || path_contains(&cand_norm, &other) {
return Err(ns_err(format!(
"namespace path {} overlaps namespace {}'s root",
candidate.display(),
rec.id,
)));
}
}
Ok(())
}
fn ensure_root_adoptable(root: &Path) -> Result<bool> {
match fs::read_dir(root) {
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(false),
Ok(mut entries) => {
if entries.next().is_some() {
Err(ns_err(format!(
"explicit namespace root {} already exists and is not \
empty; importing foreign data dirs is unsupported",
root.display()
)))
} else {
Ok(true)
}
}
Err(e) => Err(e.into()),
}
}
fn cleanup_failed_root(root: &Path, preexisted: bool) {
if preexisted {
if let Ok(entries) = fs::read_dir(root) {
for e in entries.flatten() {
let p = e.path();
let _ = if p.is_dir() {
fs::remove_dir_all(&p)
} else {
fs::remove_file(&p)
};
}
}
} else {
let _ = fs::remove_dir_all(root);
}
}
fn sizing_for(mem_budget_mb: Option<usize>) -> EngineSizing {
EngineSizing::from_budget_mb(mem_budget_mb.unwrap_or(DEFAULT_NS_BUDGET_MB))
}
struct NsInner {
id: NsId,
path: PathBuf,
engine: Engine,
}
#[derive(Clone)]
pub struct Namespace(Arc<NsInner>);
impl fmt::Debug for Namespace {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Namespace")
.field("id", &self.0.id)
.field("path", &self.0.path)
.finish()
}
}
impl Namespace {
pub fn default_ns() -> Namespace {
DEFAULT_NS.clone()
}
pub fn create() -> Result<Namespace> {
Self::create_with(NamespaceOpts::default())
}
pub fn create_with(opts: NamespaceOpts) -> Result<Namespace> {
vsdb_freeze_base_dir();
let base = vsdb_get_base_dir();
let shards = opts.shards.clamp(1, 64);
let _g = REGISTRY_LOCK.lock();
let mut reg = load_registry()?;
let mut root_preexisted = false;
if let Some(p) = &opts.path {
validate_explicit_root(&base, ®, p)?;
root_preexisted = ensure_root_adoptable(p)?;
}
let id = reg.next_id;
let rec = NsRecord {
id,
path: opts
.path
.as_ref()
.map(|p| p.to_str().expect("validated UTF-8").to_owned()),
shards: shards as u32,
mem_budget_mb: opts.mem_budget_mb.map(|v| v as u64),
created_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0),
};
let root = resolve_root(&base, &rec);
save_lifecycle(&base, id, NsLifecycleState::Pending)?;
reg.next_id += 1;
reg.entries.push(rec.clone());
if let Err(e) = save_registry(®) {
let _ = remove_lifecycle(&base, id);
return Err(e);
}
match open_record_locked(&base, &rec, &root) {
Ok(ns) => Ok(ns),
Err(e) => {
reg.entries.retain(|r| r.id != rec.id);
let rolled_back = save_registry(®).is_ok();
cleanup_failed_root(&root, root_preexisted);
if rolled_back {
let _ = remove_lifecycle(&base, rec.id);
}
Err(e)
}
}
}
pub fn open(id: NsId) -> Result<Namespace> {
if id == DEFAULT_NS_ID {
return Ok(Self::default_ns());
}
if let Some(ns) = OPEN_NAMESPACES.lock().get(&id) {
return Ok(ns.clone());
}
vsdb_freeze_base_dir();
let base = vsdb_get_base_dir();
let _g = REGISTRY_LOCK.lock();
if let Some(ns) = OPEN_NAMESPACES.lock().get(&id) {
return Ok(ns.clone());
}
let reg = load_registry()?;
let rec = reg.entries.iter().find(|r| r.id == id).ok_or_else(|| {
ns_err(format!(
"namespace {id} is not registered (destroyed, or from \
another universe)"
))
})?;
let root = resolve_root(&base, rec);
open_record_locked(&base, rec, &root)
}
pub fn scope<R>(&self, f: impl FnOnce() -> R) -> R {
struct PopGuard;
impl Drop for PopGuard {
fn drop(&mut self) {
NS_STACK.with(|s| {
s.borrow_mut().pop();
});
}
}
NS_STACK.with(|s| s.borrow_mut().push(self.clone()));
let _guard = PopGuard;
f()
}
pub fn current() -> Namespace {
NS_STACK
.with(|s| s.borrow().last().cloned())
.unwrap_or_else(Self::default_ns)
}
pub fn id(&self) -> NsId {
self.0.id
}
pub fn path(&self) -> &Path {
&self.0.path
}
pub fn system_dir(&self) -> PathBuf {
self.0.path.join("__SYSTEM__")
}
pub fn meta_dir(&self) -> PathBuf {
self.system_dir().join("__instance_meta__")
}
pub fn meta_path(&self, map_id: u64) -> PathBuf {
let mut p = self.meta_dir();
p.push(format!("{:016x}", map_id));
p
}
pub fn flush(&self) {
self.0.engine.flush()
}
pub fn shard_properties(&self, name: &str) -> Vec<Option<String>> {
self.0.engine.shard_properties(name)
}
#[inline(always)]
pub(crate) fn engine(&self) -> &Engine {
&self.0.engine
}
pub fn close(self) -> std::result::Result<(), (Option<Namespace>, VsdbError)> {
ns_close_impl(self.0.id, Some(self))
}
}
fn validated_shards(rec: &NsRecord) -> Result<usize> {
let shards = rec.shards as usize;
if !(1..=64).contains(&shards) {
return Err(ns_err(format!(
"registry entry for namespace {} carries an invalid shard \
count ({}); the registry file is damaged",
rec.id, rec.shards
)));
}
Ok(shards)
}
fn open_record_locked(base: &Path, rec: &NsRecord, root: &Path) -> Result<Namespace> {
let shards = validated_shards(rec)?;
let lifecycle = load_lifecycle(base, rec.id)?;
if lifecycle != Some(NsLifecycleState::Pending) {
validate_completed_dataset(root, shards, true).map_err(VsdbError::from)?;
}
let sizing = sizing_for(rec.mem_budget_mb.map(|v| v as usize));
let engine = Engine::open_at(root, shards, sizing).map_err(VsdbError::from)?;
if lifecycle != Some(NsLifecycleState::Established) {
save_lifecycle(base, rec.id, NsLifecycleState::Established)?;
}
let ns = Namespace(Arc::new(NsInner {
id: rec.id,
path: root.to_path_buf(),
engine,
}));
OPEN_NAMESPACES.lock().insert(rec.id, ns.clone());
Ok(ns)
}
pub fn vsdb_ns_list() -> Result<Vec<NsInfo>> {
vsdb_freeze_base_dir();
let base = vsdb_get_base_dir();
let _g = REGISTRY_LOCK.lock();
let reg = load_registry()?;
Ok(reg
.entries
.iter()
.map(|rec| NsInfo {
id: rec.id,
path: resolve_root(&base, rec),
pinned: rec.path.is_some(),
shards: rec.shards as usize,
created_at: rec.created_at,
})
.collect())
}
pub fn vsdb_ns_destroy(id: NsId) -> Result<()> {
if id == DEFAULT_NS_ID {
return Err(ns_err("the default namespace cannot be destroyed"));
}
vsdb_freeze_base_dir();
let base = vsdb_get_base_dir();
let _g = REGISTRY_LOCK.lock();
if OPEN_NAMESPACES.lock().contains_key(&id) {
return Err(ns_err(format!(
"namespace {id} is open in this process; destroy requires a \
not-open target"
)));
}
let mut reg = load_registry()?;
let Some(pos) = reg.entries.iter().position(|r| r.id == id) else {
return Err(ns_err(format!("namespace {id} is not registered")));
};
let root = resolve_root(&base, ®.entries[pos]);
reg.entries.remove(pos);
save_registry(®)?;
let _ = remove_lifecycle(&base, id);
match fs::remove_dir_all(&root) {
Ok(()) => Ok(()),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(e.into()),
}
}
pub fn vsdb_ns_relocate(id: NsId, new_path: impl AsRef<Path>) -> Result<()> {
if id == DEFAULT_NS_ID {
return Err(ns_err(
"the default namespace's root is the base dir; relocate it \
via VSDB_BASE_DIR / vsdb_set_base_dir before first use",
));
}
vsdb_freeze_base_dir();
let base = vsdb_get_base_dir();
let new_path = new_path.as_ref();
let _g = REGISTRY_LOCK.lock();
if OPEN_NAMESPACES.lock().contains_key(&id) {
return Err(ns_err(format!(
"namespace {id} is open in this process; relocate requires a \
not-open target"
)));
}
let mut reg = load_registry()?;
let Some(rec) = reg.entries.iter().find(|r| r.id == id) else {
return Err(ns_err(format!("namespace {id} is not registered")));
};
let rec_shards = validated_shards(rec)?;
let mut probe = reg.clone_without(id);
validate_explicit_root(&base, &probe, new_path)?;
drop(probe.entries.drain(..));
validate_completed_dataset(new_path, rec_shards, true).map_err(VsdbError::from)?;
let rec = reg
.entries
.iter_mut()
.find(|r| r.id == id)
.expect("checked above");
rec.path = Some(
new_path
.to_str()
.expect("validated UTF-8 in validate_explicit_root")
.to_owned(),
);
save_registry(®)
}
pub fn vsdb_ns_close(id: NsId) -> Result<()> {
ns_close_impl(id, None).map_err(|(_, e)| e)
}
fn ns_close_impl(
id: NsId,
caller_handle: Option<Namespace>,
) -> std::result::Result<(), (Option<Namespace>, VsdbError)> {
if id == DEFAULT_NS_ID {
return Err((
caller_handle,
ns_err("the default namespace cannot be closed"),
));
}
let _g = REGISTRY_LOCK.lock();
let ns_owned = {
let mut open = OPEN_NAMESPACES.lock();
let Some(entry) = open.get(&id) else {
return Err((
caller_handle,
ns_err(format!("namespace {id} is not open in this process")),
));
};
if let Some(h) = &caller_handle {
debug_assert!(Arc::ptr_eq(&entry.0, &h.0));
}
let accounted = 1 + usize::from(caller_handle.is_some());
let others = Arc::strong_count(&entry.0) - accounted;
if others > 0 {
let qualifier = if caller_handle.is_some() {
"other "
} else {
""
};
return Err((
caller_handle,
ns_err(format!(
"namespace {id} still has {others} {qualifier}live handle(s); \
drop every collection handle and `Namespace` clone first"
)),
));
}
drop(caller_handle);
open.remove(&id).expect("present: checked above")
};
let inner = Arc::try_unwrap(ns_owned.0)
.unwrap_or_else(|_| unreachable!("count was 1 under both locks"));
inner.engine.close().map_err(|e| (None, VsdbError::from(e)))
}
impl RegistryFile {
fn clone_without(&self, id: NsId) -> RegistryFile {
RegistryFile {
next_id: self.next_id,
entries: self
.entries
.iter()
.filter(|r| r.id != id)
.cloned()
.collect(),
}
}
}
pub(crate) fn flush_all_open() {
let handles: Vec<Namespace> = OPEN_NAMESPACES.lock().values().cloned().collect();
for ns in handles {
ns.flush();
}
}