use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use aetherdb::{
encode_control_plane_payload, ClusterReconfigCmd, ClusterReconfigResponse, ControlPlaneCmd,
DatabaseRef, DmaRequest, DmaResponse, Key, Node, QueryResult, SafetyConfig, ShardBundle,
ShortCircuitMode, Value, WritePermit,
};
use bytes::Bytes;
use thiserror::Error;
use tracing::info;
use crate::aether_host::{self, HostConfig, HostHandle};
use crate::aether_wire::WireRoutePool;
use crate::kwt_replay::connect;
use crate::kwt_replay::{DeployMode, JtiStore, OpenError, Topology};
const MAX_MIGRATE_SESSIONS: usize = 8;
#[derive(Debug, Error)]
pub enum RuntimeError {
#[error(transparent)]
Open(#[from] OpenError),
#[error(transparent)]
Connect(#[from] connect::ConnectError),
#[error("host: {0}")]
Host(String),
#[error("AetherDB: {0}")]
Db(#[from] aetherdb::Error),
#[error("data plane not running (use `aether-runtime {{ host start … }}`)")]
HostStopped,
#[error("cluster node required for this operation")]
ClusterRequired,
#[error("topology: {0}")]
Topology(String),
}
#[derive(Clone)]
pub struct AetherRuntime {
topology: Arc<RwLock<Topology>>,
host: Arc<Mutex<Option<HostSlot>>>,
jti: Arc<JtiStore>,
primary_db_id: Arc<Mutex<Option<String>>>,
migrate_cache: Arc<Mutex<HashMap<String, MigrateCacheEntry>>>,
wire_pool: Arc<Mutex<Option<Arc<WireRoutePool>>>>,
}
enum MigrateCacheEntry {
Rows {
database_id: String,
rows: Vec<(Key, Value)>,
},
Bundle {
database_id: String,
bundle: ShardBundle,
},
}
struct HostSlot {
handle: HostHandle,
#[allow(dead_code)]
kdl: String,
}
impl AetherRuntime {
pub fn bootstrap_from_env() -> Result<Self, RuntimeError> {
let topology = Arc::new(RwLock::new(Topology::from_env()));
let jti = Arc::new(JtiStore::open_from_env()?);
let runtime = Self {
topology: Arc::clone(&topology),
host: Arc::new(Mutex::new(None)),
jti,
primary_db_id: Arc::new(Mutex::new(None)),
migrate_cache: Arc::new(Mutex::new(HashMap::new())),
wire_pool: Arc::new(Mutex::new(None)),
};
if data_plane_autostart_from_env() {
runtime.start_host_from_env()?;
}
Ok(runtime)
}
#[must_use]
pub fn jti_store(&self) -> Arc<JtiStore> {
Arc::clone(&self.jti)
}
#[must_use]
pub fn topology(&self) -> Topology {
self.topology.read().expect("topology lock poisoned").clone()
}
pub fn apply_topology(&self, next: Topology) -> Result<(), RuntimeError> {
if next.shards == 0 {
return Err(RuntimeError::Topology("shards must be >= 1".into()));
}
*self.topology.write().expect("topology lock poisoned") = next;
Ok(())
}
pub fn apply_topology_from_kdl(&self, node: &kdl::KdlNode) -> Result<(), RuntimeError> {
let mut topo = self.topology();
if let Some(mode) = node.get("mode").and_then(|v| v.as_string()) {
topo.mode = match mode.to_ascii_lowercase().as_str() {
"cluster" => DeployMode::Cluster,
"replicate" | "replay" => DeployMode::Replicate,
"local" | "embedded" => DeployMode::Local,
other => {
return Err(RuntimeError::Topology(format!("unknown mode `{other}`")));
}
};
}
if let Some(shards) = node.get("shards").and_then(|v| v.as_integer()) {
topo.shards = u32::try_from(shards).unwrap_or(1).max(1);
}
if let Some(replication) = node.get("replication").and_then(|v| v.as_integer()) {
topo.replication = u32::try_from(replication).unwrap_or(1).max(1);
}
if let Some(node_id) = node.get("node_id").and_then(|v| v.as_integer()) {
topo.node_id = u64::try_from(node_id).unwrap_or(1).max(1);
}
if let Some(query) = node.get("query").and_then(|v| v.as_string()) {
topo.query = query.to_string();
}
self.apply_topology(topo)
}
pub fn host_status(&self) -> HostStatus {
let topo = self.topology();
let guard = self.host.lock().expect("host lock poisoned");
match guard.as_ref() {
None => HostStatus {
running: false,
mode: topo.mode_label().into(),
endpoint: None,
primary_id: self.primary_db_id.lock().expect("primary lock").clone(),
},
Some(slot) => HostStatus {
running: true,
mode: slot.handle.mode_label().into(),
endpoint: slot.handle.bind_endpoint(),
primary_id: self.primary_db_id.lock().expect("primary lock").clone(),
},
}
}
#[must_use]
pub fn data_plane_expected_from_env() -> bool {
data_plane_autostart_from_env()
}
pub fn start_host_from_env(&self) -> Result<(), RuntimeError> {
let cfg = host_config_from_env().map_err(RuntimeError::Host)?;
self.start_host(cfg, None)
}
pub fn start_host(&self, cfg: HostConfig, kdl: Option<String>) -> Result<(), RuntimeError> {
if self.host.lock().expect("host lock").is_some() {
return Err(RuntimeError::Host("data plane already running".into()));
}
let topo = self.topology();
let tls_block = cfg.tls.as_ref().map_or(String::new(), |tls| {
crate::kwt_replay::kdl_build::tls_block_from_paths(&tls.cert, &tls.key, &tls.ca)
});
let host_kdl = kdl.unwrap_or_else(|| {
crate::kwt_replay::kdl_build::host_data_kdl(
&topo,
&cfg.bind,
cfg.replicate_hex.trim(),
&tls_block,
cfg.control_key_hex.as_deref(),
)
});
let handle = aether_host::open_host(&cfg, &topo).map_err(RuntimeError::Host)?;
let primary_id = format!("nautalid-primary-{}", topo.node_id);
let db = handle
.primary_database(&primary_id, &host_kdl)
.map_err(RuntimeError::Host)?;
self.jti.reload_colocated(db)?;
*self.primary_db_id.lock().expect("primary lock") = Some(primary_id);
*self.host.lock().expect("host lock") = Some(HostSlot {
handle,
kdl: host_kdl,
});
info!(mode = %topo.mode_label(), "LID data plane started in-process");
Ok(())
}
pub fn stop_host(&self) -> Result<(), RuntimeError> {
let slot = self
.host
.lock()
.expect("host lock")
.take()
.ok_or(RuntimeError::HostStopped)?;
slot.handle.shutdown().map_err(RuntimeError::Host)?;
*self.primary_db_id.lock().expect("primary lock") = None;
*self.wire_pool.lock().expect("wire pool lock") = None;
self.jti.reload_from_env()?;
info!("LID data plane stopped; JTI reverted to env connect spec");
Ok(())
}
pub fn jti_reconnect(&self, url: &str) -> Result<(), RuntimeError> {
let spec = connect::parse_url(url);
self.jti.reload_from_connect(spec)?;
Ok(())
}
pub fn jti_attach_primary(&self) -> Result<(), RuntimeError> {
let guard = self.host.lock().expect("host lock");
let slot = guard.as_ref().ok_or(RuntimeError::HostStopped)?;
let primary_id = self
.primary_db_id
.lock()
.expect("primary lock")
.clone()
.ok_or_else(|| RuntimeError::Host("no primary database id".into()))?;
let db = match &slot.handle {
HostHandle::Replicate(db) => db.clone(),
HostHandle::Cluster(node) => node
.database(&primary_id)
.ok_or_else(|| RuntimeError::Host(format!("database `{primary_id}` not found")))?,
};
self.jti.reload_colocated(db)?;
Ok(())
}
pub fn spawn_database(
&self,
id: &str,
source: &str,
preset: &str,
) -> Result<DatabaseRef, RuntimeError> {
let node = self.cluster_node()?;
node.spawn_db_str(id, source, preset).map_err(Into::into)
}
pub fn close_database(&self, id: &str) -> Result<(), RuntimeError> {
let node = self.cluster_node()?;
node.close_db(id).map_err(Into::into)
}
pub fn cluster_reconfig(&self, cmd: ClusterReconfigCmd) -> ClusterReconfigResponse {
match self.host.lock().expect("host lock").as_ref() {
Some(slot) => match &slot.handle {
HostHandle::Cluster(node) => node.apply_cluster_reconfig(cmd),
HostHandle::Replicate(_) => {
ClusterReconfigResponse::Error("cluster ops require cluster data plane".into())
}
},
None => ClusterReconfigResponse::Error(RuntimeError::HostStopped.to_string()),
}
}
pub fn control_plane(&self, cmd: ControlPlaneCmd) -> ClusterReconfigResponse {
match self.host.lock().expect("host lock").as_ref() {
Some(slot) => match &slot.handle {
HostHandle::Cluster(node) => node.apply_control_plane(cmd),
HostHandle::Replicate(_) => {
ClusterReconfigResponse::Error("control-plane ops require cluster node".into())
}
},
None => ClusterReconfigResponse::Error(RuntimeError::HostStopped.to_string()),
}
}
pub fn migrate_import_cached(
&self,
session: &str,
database_id: Option<String>,
) -> ClusterReconfigResponse {
let entry = self
.migrate_cache
.lock()
.expect("migrate cache lock")
.remove(session);
let Some(entry) = entry else {
return ClusterReconfigResponse::Error(format!("migrate session `{session}` not found"));
};
match entry {
MigrateCacheEntry::Rows {
database_id: cached_id,
rows,
} => {
let db_id = database_id.unwrap_or(cached_id);
self.control_plane(ControlPlaneCmd::MigrateImport {
database_id: db_id,
rows,
})
}
MigrateCacheEntry::Bundle {
database_id: cached_id,
bundle,
} => {
let db_id = database_id.unwrap_or(cached_id);
self.control_plane(ControlPlaneCmd::MigrateImportBundle {
database_id: db_id,
bundle,
})
}
}
}
pub fn forward_control_cmd(
&self,
target_node_id: u64,
cmd: ControlPlaneCmd,
) -> ClusterReconfigResponse {
let node = match self.cluster_node() {
Ok(node) => node,
Err(err) => return ClusterReconfigResponse::Error(err.to_string()),
};
let token = node
.control_kwt_gate()
.and_then(|gate| gate.issue_control_token("nautalid-bus", 120).ok());
let payload = encode_control_plane_payload(&cmd, token.as_deref());
node.forward_control(target_node_id, payload)
}
pub fn query_at(&self, address: &str, text: &str) -> Result<QueryResult, RuntimeError> {
let node = self.cluster_node()?;
node.query_at(address, text).map_err(Into::into)
}
pub fn database_addresses(&self) -> Result<Vec<String>, RuntimeError> {
let node = self.cluster_node()?;
Ok(node
.database_addresses()
.into_iter()
.map(|a| a.to_string())
.collect())
}
pub fn issue_write_permit(&self) -> Result<WritePermit, RuntimeError> {
let node = self.cluster_node()?;
node.issue_write_permit().map_err(Into::into)
}
pub fn short_circuit_dump(&self) -> Result<aetherdb::ShortCircuitReport, RuntimeError> {
let node = self.cluster_node()?;
node.short_circuit_dump().map_err(Into::into)
}
pub fn checkpoint_at(&self, address: &str) -> Result<u64, RuntimeError> {
let node = self.cluster_node()?;
let db = node.database_by_address(address)?;
db.checkpoint().map_err(Into::into)
}
pub fn emergency_persist_at(
&self,
address: &str,
) -> Result<(std::path::PathBuf, u64), RuntimeError> {
let node = self.cluster_node()?;
let db = node.database_by_address(address)?;
db.emergency_persist().map_err(Into::into)
}
pub fn dma_at(&self, address: &str, request: DmaRequest) -> DmaResponse {
let Ok(node) = self.cluster_node() else {
return DmaResponse::Error(RuntimeError::ClusterRequired.to_string());
};
let Ok(db) = node.database_by_address(address) else {
return DmaResponse::Error(format!("database `{address}` not found"));
};
let payload = Bytes::from(request.encode().to_vec());
db.dispatch_dma_wire(payload)
}
pub fn attach_write_token_at(&self, address: &str, token: &str) -> Result<(), RuntimeError> {
let node = self.cluster_node()?;
let db = node.database_by_address(address)?;
db.attach_write_token(token).map_err(Into::into)
}
pub fn store_migrate_export(
&self,
session: &str,
database_id: &str,
resp: &ClusterReconfigResponse,
) {
let mut cache = self.migrate_cache.lock().expect("migrate cache lock");
let entry = match resp {
ClusterReconfigResponse::MigrateData { rows } => MigrateCacheEntry::Rows {
database_id: database_id.to_string(),
rows: rows.clone(),
},
ClusterReconfigResponse::MigrateBundle { bundle } => MigrateCacheEntry::Bundle {
database_id: database_id.to_string(),
bundle: bundle.clone(),
},
_ => return,
};
if cache.len() >= MAX_MIGRATE_SESSIONS {
cache.clear();
}
cache.insert(session.to_string(), entry);
}
pub fn database_catalog_kdl(&self) -> Result<String, RuntimeError> {
let node = self.cluster_node()?;
Ok(node.database_catalog().to_kdl())
}
pub fn lifecycle_kdl(&self) -> Result<String, RuntimeError> {
let node = self.cluster_node()?;
Ok(node.instance_lifecycle_log_kdl())
}
pub fn cluster_snapshot_kdl(&self) -> ClusterReconfigResponse {
self.cluster_reconfig(ClusterReconfigCmd::GetStatus)
}
pub fn dispatch_wt_wire(
&self,
route_tag: u8,
address: Option<&str>,
payload: Bytes,
) -> Result<Bytes, String> {
match route_tag {
0 => {
let db = self.resolve_wt_database(address)?;
Ok(db.dispatch_command_wire(payload).encode())
}
1 => {
let db = self.resolve_wt_database(address)?;
Ok(db.dispatch_dma_wire(payload).encode())
}
2 => self.dispatch_control_wire(payload),
3 => self.dispatch_replicate_wire(payload),
other => Err(format!("unknown WT wire route tag {other}")),
}
}
fn dispatch_control_wire(&self, payload: Bytes) -> Result<Bytes, String> {
if let Ok(node) = self.cluster_node() {
return Ok(node.dispatch_control_wire(payload).encode());
}
self.wire_route_at_host("control", payload)
}
fn dispatch_replicate_wire(&self, payload: Bytes) -> Result<Bytes, String> {
self.wire_route_at_host("replicate", payload)
}
fn wire_route_at_host(&self, route: &str, payload: Bytes) -> Result<Bytes, String> {
let endpoint = self
.host_bind_endpoint()
.ok_or_else(|| format!("data plane not running (cannot wire `{route}`)"))?;
let pool = self.wire_pool()?;
pool.request(&endpoint, route, payload)
}
fn host_bind_endpoint(&self) -> Option<String> {
self.host
.lock()
.expect("host lock poisoned")
.as_ref()
.and_then(|slot| slot.handle.bind_endpoint())
}
pub fn ensure_wire_pool(&self) {
let _ = self.wire_pool();
}
pub(crate) fn wire_pool(&self) -> Result<Arc<WireRoutePool>, String> {
let mut guard = self.wire_pool.lock().expect("wire pool lock poisoned");
if guard.is_none() {
let runtime = self.host_engine_runtime().unwrap_or_else(|_| {
aetherdb::EngineRuntime::with_defaults().expect("aetherdb engine runtime")
});
let tls = crate::aether_tunnel::wire_client_zmq_tls().map_err(|err| err.to_string())?;
*guard = Some(Arc::new(WireRoutePool::with_tls(runtime, tls)));
}
Ok(Arc::clone(guard.as_ref().expect("wire pool")))
}
pub(crate) fn host_engine_runtime(&self) -> Result<Arc<aetherdb::EngineRuntime>, String> {
let guard = self.host.lock().expect("host lock poisoned");
let slot = guard.as_ref().ok_or_else(|| RuntimeError::HostStopped.to_string())?;
match &slot.handle {
HostHandle::Replicate(db) => Ok(db.runtime()),
HostHandle::Cluster(node) => Ok(node.runtime()),
}
}
fn resolve_wt_database(&self, address: Option<&str>) -> Result<DatabaseRef, String> {
let guard = self.host.lock().expect("host lock poisoned");
let slot = guard.as_ref().ok_or_else(|| RuntimeError::HostStopped.to_string())?;
match &slot.handle {
HostHandle::Replicate(db) => Ok(db.clone()),
HostHandle::Cluster(node) => {
if let Some(addr) = address {
return node
.database_by_address(addr)
.map_err(|err| err.to_string());
}
if let Some(id) = self
.primary_db_id
.lock()
.expect("primary lock")
.clone()
{
if let Some(db) = node.database(&id) {
return Ok(db);
}
}
let addresses = node.database_addresses();
let first = addresses
.first()
.ok_or_else(|| "no databases on cluster node".to_string())?;
node.database_by_address(&first.to_string())
.map_err(|err| err.to_string())
}
}
}
fn cluster_node(&self) -> Result<Arc<Node>, RuntimeError> {
match self.host.lock().expect("host lock").as_ref() {
Some(slot) => match &slot.handle {
HostHandle::Cluster(node) => Ok(Arc::clone(node)),
HostHandle::Replicate(_) => Err(RuntimeError::ClusterRequired),
},
None => Err(RuntimeError::HostStopped),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HostStatus {
pub running: bool,
pub mode: String,
pub endpoint: Option<String>,
pub primary_id: Option<String>,
}
fn host_config_from_env() -> Result<HostConfig, String> {
aether_host::host_config_from_env()
}
fn data_plane_autostart_from_env() -> bool {
if crate::surfaces::env_enabled("NAUTALID_AETHERDB_DATA") {
return true;
}
std::env::var("AETHERDB_HOST_BIND")
.ok()
.is_some_and(|s| !s.trim().is_empty())
|| std::env::var("AETHERDB_HOST_REPLICATE_KEY_HEX")
.ok()
.is_some_and(|s| !s.trim().is_empty())
}
pub fn format_cluster_response(resp: &ClusterReconfigResponse) -> String {
match resp {
ClusterReconfigResponse::Status(snap) => format!(
"ok status revision={} read_only={} shards={} databases={}",
snap.revision,
snap.read_only,
snap.shards,
snap.database_count
),
ClusterReconfigResponse::Applied { revision } => {
format!("ok applied revision={revision}")
}
ClusterReconfigResponse::Placement(entries) => {
let mut out = format!("ok placement count={}", entries.len());
for entry in entries.iter().take(32) {
out.push_str(&format!(" shard{}=node{}", entry.shard_id, entry.node_id));
}
out
}
ClusterReconfigResponse::Balance { revision, loads } => {
format!("ok balance revision={revision} nodes={}", loads.len())
}
ClusterReconfigResponse::MigrateData { rows } => {
format!("ok migrate_export rows={}", rows.len())
}
ClusterReconfigResponse::MigrateBundle { bundle } => {
format!(
"ok migrate_bundle kv={} vertices={} edges={}",
bundle.kv.len(),
bundle.vertices.len(),
bundle.edges.len()
)
}
ClusterReconfigResponse::Migrated { records, revision } => {
format!("ok migrated records={records} revision={revision}")
}
ClusterReconfigResponse::Resharded { records, revision } => {
format!("ok resharded records={records} revision={revision}")
}
ClusterReconfigResponse::Balanced {
migrations,
records,
revision,
} => format!(
"ok balanced migrations={migrations} records={records} revision={revision}"
),
ClusterReconfigResponse::Dispatched { target_node_id } => {
format!("ok dispatched target_node_id={target_node_id}")
}
ClusterReconfigResponse::Error(msg) => format!("err {msg}"),
}
}
pub fn format_dma_response(resp: &DmaResponse) -> String {
match resp {
DmaResponse::Value(Some(value)) => {
format!("ok value={}", crate::kwt_replay::bus_api::format_bytes_attr("v", value.as_bytes()))
}
DmaResponse::Value(None) => "ok miss".into(),
DmaResponse::Rows(rows) => format!("ok rows count={}", rows.len()),
DmaResponse::VectorHits(hits) => format!("ok vector_hits count={}", hits.len()),
DmaResponse::Query(result) => crate::kwt_replay::bus_api::format_query_result_public(result),
DmaResponse::Error(msg) => format!("err {msg}"),
}
}
pub fn parse_short_circuit_mode(raw: &str) -> Result<ShortCircuitMode, RuntimeError> {
match raw.trim().to_ascii_lowercase().as_str() {
"sync" => Ok(ShortCircuitMode::Sync),
"best_effort" | "best-effort" => Ok(ShortCircuitMode::BestEffort),
"disabled" => Ok(ShortCircuitMode::Disabled),
other => Err(RuntimeError::Topology(format!(
"unknown short_circuit mode `{other}` (sync, best_effort, disabled)"
))),
}
}
pub fn safety_config_from_kdl(node: &kdl::KdlNode) -> Result<SafetyConfig, RuntimeError> {
let persist_drive = node
.get("persist_drive")
.and_then(|v| v.as_string())
.map(str::to_string)
.filter(|s| !s.is_empty())
.ok_or_else(|| RuntimeError::Topology("persist_drive required".into()))?;
let short_circuit = node
.get("short_circuit")
.and_then(|v| v.as_string())
.map(str::to_string)
.unwrap_or_else(|| "sync".into());
Ok(SafetyConfig {
persist_drive,
short_circuit: parse_short_circuit_mode(&short_circuit)?,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn bootstrap_embedded_jti() {
let runtime = AetherRuntime::bootstrap_from_env().expect("bootstrap");
assert!(runtime.jti_store().is_enabled());
assert!(!runtime.host_status().running);
}
#[test]
fn apply_topology_updates_mode() {
let runtime = AetherRuntime::bootstrap_from_env().expect("bootstrap");
let mut topo = runtime.topology();
topo.mode = DeployMode::Cluster;
topo.shards = 8;
runtime.apply_topology(topo).expect("apply");
assert_eq!(runtime.topology().mode, DeployMode::Cluster);
assert_eq!(runtime.topology().shards, 8);
}
}