333 lines
10 KiB
Rust
333 lines
10 KiB
Rust
use futures::StreamExt;
|
|
use serde::{Deserialize, Serialize};
|
|
use thiserror::Error;
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
|
pub struct TenantPlacementConfig {
|
|
pub virtual_nodes_per_node: usize,
|
|
pub nodes: Vec<NodePlacement>,
|
|
pub tenants: std::collections::HashMap<String, String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
|
pub struct NodePlacement {
|
|
pub node_id: String,
|
|
pub tenant_range: String,
|
|
}
|
|
|
|
pub fn placement_constraint_for_tenant_range(tenant_range: &str) -> String {
|
|
format!("node.labels.tenant_range == {}", tenant_range)
|
|
}
|
|
|
|
pub fn placement_constraints_for_node(node: &NodePlacement) -> Vec<String> {
|
|
vec![placement_constraint_for_tenant_range(&node.tenant_range)]
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct MigrationPlan {
|
|
pub tenant_id: String,
|
|
pub from_node: String,
|
|
pub to_node: String,
|
|
pub actions: Vec<MigrationAction>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum MigrationAction {
|
|
DrainTenant { tenant_id: String },
|
|
UpdatePlacement { tenant_id: String, node_id: String },
|
|
ReloadConfig,
|
|
}
|
|
|
|
pub fn plan_graceful_tenant_migration(
|
|
tenant_id: impl Into<String>,
|
|
from_node: impl Into<String>,
|
|
to_node: impl Into<String>,
|
|
) -> MigrationPlan {
|
|
let tenant_id = tenant_id.into();
|
|
let from_node = from_node.into();
|
|
let to_node = to_node.into();
|
|
|
|
MigrationPlan {
|
|
tenant_id: tenant_id.clone(),
|
|
from_node,
|
|
to_node: to_node.clone(),
|
|
actions: vec![
|
|
MigrationAction::DrainTenant {
|
|
tenant_id: tenant_id.clone(),
|
|
},
|
|
MigrationAction::UpdatePlacement {
|
|
tenant_id,
|
|
node_id: to_node,
|
|
},
|
|
MigrationAction::ReloadConfig,
|
|
],
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Error)]
|
|
pub enum TenantPlacementKvError {
|
|
#[error("NATS connection error: {0}")]
|
|
Connection(String),
|
|
#[error("KV error: {0}")]
|
|
Kv(String),
|
|
#[error("Config parse error: {0}")]
|
|
Parse(String),
|
|
#[error("Unsupported key operation")]
|
|
UnsupportedOperation,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct TenantPlacementKvClient {
|
|
kv: async_nats::jetstream::kv::Store,
|
|
}
|
|
|
|
impl TenantPlacementKvClient {
|
|
pub async fn connect(
|
|
nats_url: impl Into<String>,
|
|
bucket: impl Into<String>,
|
|
) -> Result<Self, TenantPlacementKvError> {
|
|
Self::connect_with_timeout(nats_url, bucket, std::time::Duration::from_secs(2)).await
|
|
}
|
|
|
|
pub async fn connect_with_timeout(
|
|
nats_url: impl Into<String>,
|
|
bucket: impl Into<String>,
|
|
timeout: std::time::Duration,
|
|
) -> Result<Self, TenantPlacementKvError> {
|
|
let nats_url = nats_url.into();
|
|
let bucket = bucket.into();
|
|
|
|
let client = tokio::time::timeout(timeout, async_nats::connect(nats_url))
|
|
.await
|
|
.map_err(|_| TenantPlacementKvError::Connection("connect timeout".to_string()))?
|
|
.map_err(|e| TenantPlacementKvError::Connection(e.to_string()))?;
|
|
|
|
let jetstream = async_nats::jetstream::new(client);
|
|
|
|
let kv = match jetstream.get_key_value(&bucket).await {
|
|
Ok(kv) => kv,
|
|
Err(_) => jetstream
|
|
.create_key_value(async_nats::jetstream::kv::Config {
|
|
bucket: bucket.clone(),
|
|
..Default::default()
|
|
})
|
|
.await
|
|
.map_err(|e| TenantPlacementKvError::Kv(e.to_string()))?,
|
|
};
|
|
|
|
Ok(Self { kv })
|
|
}
|
|
|
|
pub async fn get_json(
|
|
&self,
|
|
key: &str,
|
|
) -> Result<Option<serde_json::Value>, TenantPlacementKvError> {
|
|
let entry = self
|
|
.kv
|
|
.entry(key)
|
|
.await
|
|
.map_err(|e| TenantPlacementKvError::Kv(e.to_string()))?;
|
|
|
|
match entry {
|
|
Some(entry) => serde_json::from_slice::<serde_json::Value>(&entry.value)
|
|
.map(Some)
|
|
.map_err(|e| TenantPlacementKvError::Parse(e.to_string())),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub async fn put_json(
|
|
&self,
|
|
key: &str,
|
|
value: &serde_json::Value,
|
|
) -> Result<(), TenantPlacementKvError> {
|
|
let bytes =
|
|
serde_json::to_vec(value).map_err(|e| TenantPlacementKvError::Parse(e.to_string()))?;
|
|
self.kv
|
|
.put(key, bytes.into())
|
|
.await
|
|
.map_err(|e| TenantPlacementKvError::Kv(e.to_string()))?;
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn watch_json(
|
|
&self,
|
|
pattern: &str,
|
|
) -> Result<
|
|
std::pin::Pin<
|
|
Box<
|
|
dyn futures::Stream<Item = Result<serde_json::Value, TenantPlacementKvError>>
|
|
+ Send,
|
|
>,
|
|
>,
|
|
TenantPlacementKvError,
|
|
> {
|
|
let watch = self
|
|
.kv
|
|
.watch(pattern)
|
|
.await
|
|
.map_err(|e| TenantPlacementKvError::Kv(e.to_string()))?;
|
|
|
|
Ok(Box::pin(watch.filter_map(|entry| async move {
|
|
match entry {
|
|
Ok(entry) => match entry.operation {
|
|
async_nats::jetstream::kv::Operation::Put => {
|
|
match serde_json::from_slice::<serde_json::Value>(&entry.value) {
|
|
Ok(v) => Some(Ok(v)),
|
|
Err(e) => Some(Err(TenantPlacementKvError::Parse(e.to_string()))),
|
|
}
|
|
}
|
|
async_nats::jetstream::kv::Operation::Delete
|
|
| async_nats::jetstream::kv::Operation::Purge => None,
|
|
},
|
|
Err(e) => Some(Err(TenantPlacementKvError::Kv(e.to_string()))),
|
|
}
|
|
})))
|
|
}
|
|
|
|
pub async fn load_config_with_fallback(
|
|
nats_url: impl Into<String>,
|
|
bucket: impl Into<String>,
|
|
key: &str,
|
|
fallback_path: &str,
|
|
) -> Result<serde_json::Value, TenantPlacementKvError> {
|
|
let try_kv = match Self::connect_with_timeout(
|
|
nats_url,
|
|
bucket,
|
|
std::time::Duration::from_millis(300),
|
|
)
|
|
.await
|
|
{
|
|
Ok(client) => match client.get_json(key).await {
|
|
Ok(Some(v)) => Ok(v),
|
|
Ok(None) => Err(TenantPlacementKvError::Kv("missing key".to_string())),
|
|
Err(e) => Err(e),
|
|
},
|
|
Err(e) => Err(e),
|
|
};
|
|
|
|
match try_kv {
|
|
Ok(v) => Ok(v),
|
|
Err(_) => {
|
|
let raw = std::fs::read_to_string(fallback_path)
|
|
.map_err(|e| TenantPlacementKvError::Kv(e.to_string()))?;
|
|
if fallback_path.ends_with(".json") {
|
|
serde_json::from_str(&raw)
|
|
.map_err(|e| TenantPlacementKvError::Parse(e.to_string()))
|
|
} else {
|
|
let yaml: serde_yaml::Value = serde_yaml::from_str(&raw)
|
|
.map_err(|e| TenantPlacementKvError::Parse(e.to_string()))?;
|
|
let json = serde_json::to_value(yaml)
|
|
.map_err(|e| TenantPlacementKvError::Parse(e.to_string()))?;
|
|
Ok(json)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use futures::StreamExt;
|
|
|
|
#[test]
|
|
fn stack_file_is_valid_yaml() {
|
|
let raw = std::fs::read_to_string("../swarm/stacks/platform.yml").unwrap();
|
|
let _: serde_yaml::Value = serde_yaml::from_str(&raw).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn stack_services_count() {
|
|
let raw = std::fs::read_to_string("../swarm/stacks/platform.yml").unwrap();
|
|
let doc: serde_yaml::Value = serde_yaml::from_str(&raw).unwrap();
|
|
let services = doc.get("services").and_then(|v| v.as_mapping()).unwrap();
|
|
assert!(services.contains_key(serde_yaml::Value::String("nats".to_string())));
|
|
assert!(services.contains_key(serde_yaml::Value::String("gateway".to_string())));
|
|
assert!(services.contains_key(serde_yaml::Value::String("aggregate".to_string())));
|
|
}
|
|
|
|
#[test]
|
|
fn tenant_placement_config_loads() {
|
|
let raw = std::fs::read_to_string("../swarm/tenant-placement.yaml").unwrap();
|
|
let cfg: TenantPlacementConfig = serde_yaml::from_str(&raw).unwrap();
|
|
assert_eq!(cfg.virtual_nodes_per_node, 200);
|
|
assert!(cfg.nodes.iter().any(|n| n.node_id == "node-a"));
|
|
assert_eq!(cfg.tenants.get("tenant-a").unwrap(), "node-a");
|
|
}
|
|
|
|
#[test]
|
|
fn placement_constraint_generated_correctly() {
|
|
let node = NodePlacement {
|
|
node_id: "node-a".to_string(),
|
|
tenant_range: "00-3f".to_string(),
|
|
};
|
|
let constraints = placement_constraints_for_node(&node);
|
|
assert_eq!(constraints, vec!["node.labels.tenant_range == 00-3f"]);
|
|
}
|
|
|
|
#[test]
|
|
fn graceful_tenant_migration_plan_is_ordered() {
|
|
let plan = plan_graceful_tenant_migration("tenant-a", "node-a", "node-b");
|
|
assert_eq!(plan.tenant_id, "tenant-a");
|
|
assert_eq!(
|
|
plan.actions,
|
|
vec![
|
|
MigrationAction::DrainTenant {
|
|
tenant_id: "tenant-a".to_string(),
|
|
},
|
|
MigrationAction::UpdatePlacement {
|
|
tenant_id: "tenant-a".to_string(),
|
|
node_id: "node-b".to_string(),
|
|
},
|
|
MigrationAction::ReloadConfig,
|
|
]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn tenant_placement_kv_falls_back_to_local_file() {
|
|
let tmp = tempfile::tempdir().unwrap();
|
|
let path = tmp.path().join("placement.yaml");
|
|
std::fs::write(
|
|
&path,
|
|
r#"
|
|
virtual_nodes_per_node: 100
|
|
nodes:
|
|
- node_id: "node-a"
|
|
tenant_range: "00-ff"
|
|
tenants:
|
|
tenant-a: "node-a"
|
|
"#,
|
|
)
|
|
.unwrap();
|
|
|
|
let cfg = TenantPlacementKvClient::load_config_with_fallback(
|
|
"nats://127.0.0.1:1",
|
|
"TENANT_PLACEMENT",
|
|
"placement",
|
|
path.to_string_lossy().as_ref(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(cfg["virtual_nodes_per_node"], 100);
|
|
assert_eq!(cfg["tenants"]["tenant-a"], "node-a");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn tenant_placement_kv_watch_returns_stream() {
|
|
let result = TenantPlacementKvClient::connect_with_timeout(
|
|
"nats://127.0.0.1:1",
|
|
"TENANT_PLACEMENT",
|
|
std::time::Duration::from_millis(50),
|
|
)
|
|
.await;
|
|
assert!(result.is_err());
|
|
|
|
let mut stream =
|
|
futures::stream::empty::<Result<serde_json::Value, TenantPlacementKvError>>();
|
|
assert!(stream.next().await.is_none());
|
|
}
|
|
}
|