use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use mongreldb_cluster::bootstrap::{self, ClusterStatus, TrustConfig};
use mongreldb_cluster::network::{TlsConfig, TransportConfig, TransportSecurity};
use mongreldb_cluster::node::{ClusterError, NodeIdentity};
use mongreldb_cluster::runtime::{
GroupTiming, MetaMembership, NodeInternalRpcClient, NodeRuntime, NodeRuntimeConfig,
RuntimeError, RuntimeStatus,
};
use mongreldb_cluster::tablet::Key;
use mongreldb_log::commit_log::ExecutionControl;
use mongreldb_types::ids::{NodeId, TabletId};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::sync::Mutex;
#[derive(Clone, Debug)]
pub struct ClusterRuntimeOptions {
pub node_data: PathBuf,
pub rpc_listen: String,
pub plaintext_test: bool,
pub fast_timing: bool,
}
impl ClusterRuntimeOptions {
pub fn resolve(node_data: PathBuf, rpc_listen: Option<String>) -> Self {
let rpc_listen = rpc_listen
.or_else(|| std::env::var("MONGRELDB_CLUSTER_RPC_LISTEN").ok())
.filter(|s| !s.trim().is_empty())
.unwrap_or_else(|| "127.0.0.1:17443".to_owned());
let plaintext_test = env_flag_true("MONGRELDB_CLUSTER_PLAINTEXT_TEST");
Self {
node_data,
rpc_listen,
plaintext_test,
fast_timing: plaintext_test,
}
}
}
#[derive(Clone)]
pub struct ClusterRuntimeHandle {
inner: Arc<Mutex<Option<NodeRuntime>>>,
node_data: PathBuf,
}
#[derive(Debug)]
pub enum ClusterRuntimeError {
Cluster(ClusterError),
Runtime(RuntimeError),
Config(String),
}
impl std::fmt::Display for ClusterRuntimeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Cluster(error) => write!(f, "{error}"),
Self::Runtime(error) => write!(f, "{error}"),
Self::Config(message) => write!(f, "{message}"),
}
}
}
impl std::error::Error for ClusterRuntimeError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Cluster(error) => Some(error),
Self::Runtime(error) => Some(error),
Self::Config(_) => None,
}
}
}
impl From<ClusterError> for ClusterRuntimeError {
fn from(error: ClusterError) -> Self {
Self::Cluster(error)
}
}
impl From<RuntimeError> for ClusterRuntimeError {
fn from(error: RuntimeError) -> Self {
Self::Runtime(error)
}
}
impl From<std::io::Error> for ClusterRuntimeError {
fn from(error: std::io::Error) -> Self {
Self::Config(format!("cluster runtime I/O: {error}"))
}
}
impl ClusterRuntimeHandle {
pub async fn start(options: ClusterRuntimeOptions) -> Result<Self, ClusterRuntimeError> {
let _identity =
NodeIdentity::load(&options.node_data)?.ok_or(ClusterError::NotInitialized)?;
let status = bootstrap::cluster_status(&options.node_data)?;
let listen = resolve_listen_address(&options.rpc_listen)?;
let security = resolve_security(&options.node_data, options.plaintext_test)?;
let peers = peers_from_status(&status, &listen);
let meta = meta_membership_from_status(&status, &listen);
let config = NodeRuntimeConfig {
node_data: options.node_data.clone(),
security,
transport: transport_config(options.plaintext_test),
listen_address: listen.clone(),
rpc_address: Some(listen),
peers,
meta,
timing: options.fast_timing.then(fast_timing),
};
let runtime = NodeRuntime::start(config).await?;
Ok(Self {
inner: Arc::new(Mutex::new(Some(runtime))),
node_data: options.node_data,
})
}
pub fn node_data(&self) -> &Path {
&self.node_data
}
pub async fn is_live(&self) -> bool {
self.inner.lock().await.is_some()
}
pub async fn runtime_status_json(&self) -> Result<Value, ClusterRuntimeError> {
let guard = self.inner.lock().await;
let runtime = guard
.as_ref()
.ok_or_else(|| ClusterRuntimeError::Config("cluster runtime not running".into()))?;
Ok(runtime_status_to_json(&runtime.status()))
}
pub async fn transfer_leader(
&self,
tablet_id: TabletId,
to: NodeId,
) -> Result<Value, ClusterRuntimeError> {
let guard = self.inner.lock().await;
let runtime = guard
.as_ref()
.ok_or_else(|| ClusterRuntimeError::Config("cluster runtime not running".into()))?;
let descriptor = runtime.tablet_descriptor(tablet_id).ok_or_else(|| {
ClusterRuntimeError::Config(format!(
"this node hosts no live tablet group for tablet {tablet_id}"
))
})?;
let target = descriptor.replica_on(to).ok_or_else(|| {
ClusterRuntimeError::Config(format!("node {to} is not a replica of tablet {tablet_id}"))
})?;
let group = runtime.tablet_group(tablet_id).ok_or_else(|| {
ClusterRuntimeError::Config(format!(
"this node hosts no live tablet group for tablet {tablet_id}"
))
})?;
group
.transfer_leader(target.raft_node_id, LEADER_TIMEOUT)
.await
.map_err(|error| ClusterRuntimeError::Config(error.to_string()))?;
Ok(json!({
"command": "TRANSFER LEADER",
"tablet_id": tablet_id.to_string(),
"to": to.to_string(),
"status": "ok",
"target_raft_node_id": target.raft_node_id,
}))
}
pub async fn split_tablet(
&self,
tablet_id: TabletId,
at_key_hex: Option<String>,
) -> Result<Value, ClusterRuntimeError> {
let split_key = match at_key_hex.as_deref() {
Some(hex) => Some(parse_key_hex(hex)?),
None => None,
};
let mut guard = self.inner.lock().await;
let runtime = guard
.as_mut()
.ok_or_else(|| ClusterRuntimeError::Config("cluster runtime not running".into()))?;
if runtime.tablet_group(tablet_id).is_none() {
return Err(ClusterRuntimeError::Config(format!(
"this node hosts no live tablet group for tablet {tablet_id}"
)));
}
let control = ExecutionControl::default();
let published = runtime.split_tablet(tablet_id, split_key, &control).await?;
Ok(json!({
"command": "SPLIT TABLET",
"tablet_id": tablet_id.to_string(),
"at_key_hex": at_key_hex,
"status": "ok",
"published": true,
"left_tablet_id": published.children[0].tablet_id.to_string(),
"right_tablet_id": published.children[1].tablet_id.to_string(),
}))
}
pub async fn merge_tablets(
&self,
left: TabletId,
right: TabletId,
) -> Result<Value, ClusterRuntimeError> {
let mut guard = self.inner.lock().await;
let runtime = guard
.as_mut()
.ok_or_else(|| ClusterRuntimeError::Config("cluster runtime not running".into()))?;
for tablet_id in [left, right] {
if runtime.tablet_group(tablet_id).is_none() {
return Err(ClusterRuntimeError::Config(format!(
"this node hosts no live tablet group for tablet {tablet_id}"
)));
}
}
let control = ExecutionControl::default();
let published = runtime.merge_tablets(left, right, &control).await?;
Ok(json!({
"command": "MERGE TABLETS",
"left": left.to_string(),
"right": right.to_string(),
"status": "ok",
"published": true,
"replacement_tablet_id": published.replacement.tablet_id.to_string(),
}))
}
pub fn runtime_mutex(&self) -> Arc<Mutex<Option<NodeRuntime>>> {
Arc::clone(&self.inner)
}
pub async fn attach_internal_rpc_handler(
&self,
service_id: u32,
handler: Arc<dyn mongreldb_cluster::network::InternalRpcHandler>,
) -> Result<(), ClusterRuntimeError> {
let guard = self.inner.lock().await;
let runtime = guard
.as_ref()
.ok_or_else(|| ClusterRuntimeError::Config("cluster runtime not running".into()))?;
runtime.attach_internal_rpc_handler(service_id, handler);
Ok(())
}
pub async fn internal_rpc_client(&self) -> Result<NodeInternalRpcClient, ClusterRuntimeError> {
let guard = self.inner.lock().await;
let runtime = guard
.as_ref()
.ok_or_else(|| ClusterRuntimeError::Config("cluster runtime not running".into()))?;
Ok(runtime.internal_rpc_client())
}
pub async fn shutdown(&self) -> Result<(), ClusterRuntimeError> {
let runtime = {
let mut guard = self.inner.lock().await;
guard.take()
};
if let Some(runtime) = runtime {
runtime.shutdown().await?;
}
Ok(())
}
}
const LEADER_TIMEOUT: Duration = Duration::from_secs(15);
fn env_flag_true(name: &str) -> bool {
match std::env::var(name) {
Ok(value) => {
let value = value.trim();
value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes")
}
Err(_) => false,
}
}
fn resolve_listen_address(listen: &str) -> Result<String, ClusterRuntimeError> {
let trimmed = listen.trim();
if trimmed.is_empty() {
return Err(ClusterRuntimeError::Config(
"cluster RPC listen address is empty".into(),
));
}
if trimmed.ends_with(":0") {
let listener = std::net::TcpListener::bind(trimmed)?;
return Ok(listener.local_addr()?.to_string());
}
Ok(trimmed.to_owned())
}
fn resolve_security(
node_data: &Path,
plaintext_test: bool,
) -> Result<TransportSecurity, ClusterRuntimeError> {
if plaintext_test {
return Ok(TransportSecurity::PlaintextForTesting);
}
let trust = load_persisted_trust(node_data)?.ok_or_else(|| {
ClusterRuntimeError::Config(
"cluster trust material missing under cluster-meta/trust.json; \
run `mongreldb-server cluster init` first, or set \
MONGRELDB_CLUSTER_PLAINTEXT_TEST=1 for non-production tests"
.into(),
)
})?;
let tls = TlsConfig::from_trust(&trust).map_err(|error| {
ClusterRuntimeError::Config(format!(
"cluster trust PEMs are not usable for mTLS ({error}); \
supply real certificates or set MONGRELDB_CLUSTER_PLAINTEXT_TEST=1 \
for non-production tests"
))
})?;
Ok(TransportSecurity::Mtls(tls))
}
#[derive(Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
struct TrustEnvelope {
format_version: u32,
trust: TrustConfig,
}
fn load_persisted_trust(node_data: &Path) -> Result<Option<TrustConfig>, ClusterRuntimeError> {
let path = node_data.join("cluster-meta").join("trust.json");
if !path.exists() {
return Ok(None);
}
let bytes = std::fs::read(&path)?;
let envelope: TrustEnvelope = serde_json::from_slice(&bytes).map_err(|error| {
ClusterRuntimeError::Config(format!("cluster-meta/trust.json is corrupt: {error}"))
})?;
if envelope.format_version != 1 {
return Err(ClusterRuntimeError::Config(format!(
"cluster-meta/trust.json format version {} is unsupported",
envelope.format_version
)));
}
envelope
.trust
.validate()
.map_err(ClusterRuntimeError::from)?;
Ok(Some(envelope.trust))
}
fn peers_from_status(status: &ClusterStatus, self_listen: &str) -> Vec<(NodeId, String)> {
if !status.membership.is_empty() {
return status
.membership
.iter()
.map(|member| {
let address = if member.node_id == status.identity.node_id {
self_listen.to_owned()
} else {
member.rpc_address.clone()
};
(member.node_id, address)
})
.collect();
}
vec![(status.identity.node_id, self_listen.to_owned())]
}
fn meta_membership_from_status(
status: &ClusterStatus,
self_listen: &str,
) -> Option<MetaMembership> {
let group = status.database_group.as_ref()?;
if !group.voter_ids.contains(&status.identity.node_id) {
return None;
}
let bootstrap_voters = if group.voter_ids.len() == 1 {
Some(vec![(status.identity.node_id, self_listen.to_owned())])
} else {
None
};
Some(MetaMembership {
meta_group_id: group.raft_group_id,
bootstrap_voters,
})
}
fn transport_config(plaintext_test: bool) -> TransportConfig {
if plaintext_test {
TransportConfig {
connect_timeout: Duration::from_millis(500),
rpc_timeout: Duration::from_secs(2),
snapshot_timeout: Duration::from_secs(5),
connect_attempts: 5,
reconnect_backoff: Duration::from_millis(10),
max_frame_bytes: 16 * 1024 * 1024,
max_connections: 64,
handshake_timeout: Duration::from_secs(2),
shutdown_grace: Duration::from_secs(2),
}
} else {
TransportConfig::default()
}
}
fn fast_timing() -> GroupTiming {
GroupTiming {
heartbeat_interval: Duration::from_millis(100),
election_timeout_min: Duration::from_millis(300),
election_timeout_max: Duration::from_millis(600),
install_snapshot_timeout: Duration::from_millis(2_000),
}
}
fn parse_key_hex(text: &str) -> Result<Key, ClusterRuntimeError> {
let trimmed = text.trim();
if trimmed.is_empty() {
return Err(ClusterRuntimeError::Config("split key hex is empty".into()));
}
if !trimmed.len().is_multiple_of(2) {
return Err(ClusterRuntimeError::Config(
"split key hex must have an even number of digits".into(),
));
}
let mut bytes = Vec::with_capacity(trimmed.len() / 2);
let chars: Vec<char> = trimmed.chars().collect();
for chunk in chars.chunks(2) {
let hi = chunk[0].to_digit(16).ok_or_else(|| {
ClusterRuntimeError::Config(format!("invalid split key hex `{text}`"))
})?;
let lo = chunk[1].to_digit(16).ok_or_else(|| {
ClusterRuntimeError::Config(format!("invalid split key hex `{text}`"))
})?;
bytes.push(((hi << 4) | lo) as u8);
}
Ok(Key::from_bytes(bytes))
}
fn runtime_status_to_json(status: &RuntimeStatus) -> Value {
json!({
"live": true,
"node_id": status.identity.node_id.to_string(),
"cluster_id": status.identity.cluster_id.to_string(),
"rpc_address": status.rpc_address,
"meta_present": status.meta.is_some(),
"meta": status.meta.as_ref().map(|meta| json!({
"meta_group_id": meta.meta_group_id.to_string(),
"metadata_version": meta.metadata_version.get(),
"current_leader": meta.metrics.current_leader,
"local_raft_node_id": meta.metrics.node_id,
})),
"tablet_count": status.tablets.len(),
"tablets": status.tablets.iter().map(|tablet| json!({
"tablet_id": tablet.tablet_id.to_string(),
"raft_group_id": tablet.raft_group_id.to_string(),
"state": tablet.state.to_string(),
"replica_count": tablet.replicas.len(),
"applied_index": tablet.applied.index,
"current_leader": tablet.metrics.current_leader,
"local_raft_node_id": tablet.metrics.node_id,
})).collect::<Vec<_>>(),
})
}
pub fn cluster_node_data_from_env(cli_value: Option<String>) -> Option<PathBuf> {
cli_value
.or_else(|| std::env::var("MONGRELDB_CLUSTER_NODE_DATA").ok())
.map(|s| s.trim().to_owned())
.filter(|s| !s.is_empty())
.map(PathBuf::from)
}
#[cfg(test)]
mod tests {
use super::*;
use mongreldb_cluster::bootstrap::{cluster_init, InitRequest};
use mongreldb_cluster::node::{Locality, NodeCapacity};
use tempfile::tempdir;
const CA_PEM: &str = "-----BEGIN CERTIFICATE-----\nY2E=\n-----END CERTIFICATE-----\n";
const CERT_PEM: &str = "-----BEGIN CERTIFICATE-----\nbm9kZQ==\n-----END CERTIFICATE-----\n";
const KEY_PEM: &str = "-----BEGIN PRIVATE KEY-----\nc2VjcmV0\n-----END PRIVATE KEY-----\n";
fn bootstrap(data: &Path) -> NodeIdentity {
let mut counter = 0u64;
let mut csprng = |buf: &mut [u8]| {
for chunk in buf.chunks_mut(8) {
counter += 1;
let bytes = counter.to_le_bytes();
chunk.copy_from_slice(&bytes[..chunk.len()]);
}
Ok(())
};
let identity = NodeIdentity::load_or_create(data, &mut csprng).unwrap();
let request = InitRequest {
rpc_address: "127.0.0.1:0".to_owned(),
locality: Locality::default(),
capacity: NodeCapacity::default(),
trust: TrustConfig::from_pems(
CA_PEM.to_owned(),
CERT_PEM.to_owned(),
KEY_PEM.to_owned(),
vec![identity.node_id],
)
.unwrap(),
};
cluster_init(data, &request, &mut csprng).unwrap().identity
}
#[tokio::test]
async fn plaintext_start_status_and_missing_tablet_errors() {
let dir = tempdir().unwrap();
let identity = bootstrap(dir.path());
let listen = {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
listener.local_addr().unwrap().to_string()
};
let handle = ClusterRuntimeHandle::start(ClusterRuntimeOptions {
node_data: dir.path().to_path_buf(),
rpc_listen: listen.clone(),
plaintext_test: true,
fast_timing: true,
})
.await
.expect("runtime starts after cluster init");
assert!(handle.is_live().await);
let status = handle.runtime_status_json().await.unwrap();
assert_eq!(status["live"], true);
assert_eq!(status["node_id"], identity.node_id.to_string());
assert_eq!(status["rpc_address"], listen);
assert_eq!(status["meta_present"], true);
assert_eq!(status["tablet_count"], 0);
let missing = TabletId::from_bytes([0xAB; 16]);
let err = handle
.transfer_leader(missing, identity.node_id)
.await
.unwrap_err();
assert!(
err.to_string().contains("hosts no live tablet"),
"transfer without tablet must fail closed: {err}"
);
let err = handle.split_tablet(missing, None).await.unwrap_err();
assert!(
err.to_string().contains("hosts no live tablet"),
"split without tablet must fail closed: {err}"
);
handle.shutdown().await.unwrap();
assert!(!handle.is_live().await);
}
#[tokio::test]
async fn start_fails_closed_without_cluster_init() {
let dir = tempdir().unwrap();
match ClusterRuntimeHandle::start(ClusterRuntimeOptions {
node_data: dir.path().to_path_buf(),
rpc_listen: "127.0.0.1:0".into(),
plaintext_test: true,
fast_timing: true,
})
.await
{
Ok(_) => panic!("expected NotInitialized without cluster init"),
Err(err) => assert!(
matches!(
err,
ClusterRuntimeError::Cluster(ClusterError::NotInitialized)
) || err
.to_string()
.to_ascii_lowercase()
.contains("not initialized")
|| err.to_string().contains("NotInitialized"),
"expected NotInitialized, got {err}"
),
}
}
#[test]
fn parse_key_hex_round_trip() {
let key = parse_key_hex("0a1b").unwrap();
assert_eq!(key.as_bytes(), &[0x0a, 0x1b]);
assert!(parse_key_hex("zz").is_err());
assert!(parse_key_hex("abc").is_err());
}
}