From 1fa1c613b0a0f1210ef0df5ffbdda28fc92df936 Mon Sep 17 00:00:00 2001 From: Danil Shaymurzin Date: Fri, 16 Jan 2026 20:51:29 +0500 Subject: [PATCH] feat: implement cluster health monitoring and chunk replication - Add new cluster module with health check and replication loops - Track node status and chunk locations in AppState - Implement chunk replication when under-replicated - Add new API endpoint for registering chunk locations - Update cluster info endpoint to include storage nodes --- Cargo.toml | 1 + src/bin/gateway.rs | 20 +++- src/cluster.rs | 132 +++++++++++++++++++++++++++ src/handlers.rs | 223 ++++++++++++++++++++++++++++++++++++++++----- src/http.rs | 1 + src/lib.rs | 1 + 6 files changed, 356 insertions(+), 22 deletions(-) create mode 100644 src/cluster.rs diff --git a/Cargo.toml b/Cargo.toml index bf6be47..26b4531 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ async-stream = "0.3" tokio-stream = "0.1" uuid.workspace = true hex.workspace = true +reqwest = { version = "0.12", features = ["json"] } [dev-dependencies] tower = "0.5" diff --git a/src/bin/gateway.rs b/src/bin/gateway.rs index 47c887c..2e5dce3 100644 --- a/src/bin/gateway.rs +++ b/src/bin/gateway.rs @@ -1,4 +1,4 @@ -use monoce_gateway::{HttpGateway, AppState}; +use monoce_gateway::{HttpGateway, AppState, cluster::{health_check_loop, replication_loop}}; use std::net::SocketAddr; use tracing_subscriber; @@ -27,6 +27,24 @@ async fn main() -> anyhow::Result<()> { // Create app state let state = AppState::new(node_id, cluster_id); + // Spawn health check background task + let health_state = state.clone(); + tokio::spawn(async move { + loop { + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + health_check_loop(&health_state).await; + } + }); + + // Spawn replication reconciliation background task + let replication_state = state.clone(); + tokio::spawn(async move { + loop { + tokio::time::sleep(std::time::Duration::from_secs(10)).await; + replication_loop(&replication_state).await; + } + }); + // Create and start gateway let gateway = HttpGateway::new(addr, state); diff --git a/src/cluster.rs b/src/cluster.rs new file mode 100644 index 0000000..f40aeb0 --- /dev/null +++ b/src/cluster.rs @@ -0,0 +1,132 @@ +use crate::handlers::AppState; +use std::time::{SystemTime, UNIX_EPOCH}; + +pub async fn health_check_loop(state: &AppState) { + let mut registry = state.node_registry.write().await; + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + + for (node_id, record) in registry.nodes.iter_mut() { + let is_healthy = check_node_health(&record.address).await; + let was_online = record.is_online; + record.is_online = is_healthy; + if is_healthy { + record.last_seen_ms = now; + } + if was_online != is_healthy { + if is_healthy { + tracing::info!("Node {} is now ONLINE", node_id); + } else { + tracing::warn!("Node {} is now OFFLINE", node_id); + } + } + } +} + +async fn check_node_health(address: &str) -> bool { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(2)) + .build() + .ok(); + if let Some(client) = client { + client.get(format!("{}/health", address)) + .send() + .await + .map(|r| r.status().is_success()) + .unwrap_or(false) + } else { + false + } +} + +pub async fn replication_loop(state: &AppState) { + // Get current chunk locations and node states + let registry = state.node_registry.read().await; + let locations = state.chunk_locations.read().await; + + let online_nodes: Vec = registry.nodes.values() + .filter(|n| n.is_online) + .map(|n| n.address.clone()) + .collect(); + + if online_nodes.is_empty() { + return; + } + + // Collect chunks that need replication + let mut replication_tasks: Vec<(String, String, String)> = Vec::new(); + + for (chunk_cid, node_addrs) in locations.iter() { + let online_replicas: Vec<&String> = node_addrs.iter() + .filter(|addr| online_nodes.contains(addr)) + .collect(); + + // If under-replicated (less than 2 replicas online), try to replicate + if online_replicas.len() < 2 && !online_replicas.is_empty() { + // Find a node that doesn't have this chunk + let missing_nodes: Vec<&String> = online_nodes.iter() + .filter(|n| !node_addrs.contains(n)) + .collect(); + + if let Some(dest_node) = missing_nodes.first() { + if let Some(src_node) = online_replicas.first() { + replication_tasks.push(( + chunk_cid.clone(), + (*src_node).clone(), + (*dest_node).clone(), + )); + } + } + } + } + + // Release read locks before performing replications + drop(locations); + drop(registry); + + // Perform replications and update locations + for (chunk_cid, src_node, dest_node) in replication_tasks { + tracing::info!( + "Replicating chunk {} from {} to {}", + chunk_cid, src_node, dest_node + ); + + match replicate_chunk(&chunk_cid, &src_node, &dest_node).await { + Ok(_) => { + // Update chunk_locations after successful replication + let mut locations = state.chunk_locations.write().await; + if let Some(addrs) = locations.get_mut(&chunk_cid) { + if !addrs.contains(&dest_node) { + addrs.push(dest_node.clone()); + } + } + tracing::info!("Successfully replicated chunk {} to {}", chunk_cid, dest_node); + } + Err(e) => { + tracing::error!("Failed to replicate chunk {}: {}", chunk_cid, e); + } + } + } +} + +async fn replicate_chunk(cid: &str, src: &str, dest: &str) -> Result<(), Box> { + let client = reqwest::Client::new(); + + // Fetch from source + let data = client.get(format!("{}/api/v1/chunks/{}", src, cid)) + .send() + .await? + .bytes() + .await?; + + // Upload to destination + client.put(format!("{}/api/v1/chunks/{}", dest, cid)) + .body(data) + .send() + .await? + .error_for_status()?; + + Ok(()) +} diff --git a/src/handlers.rs b/src/handlers.rs index b7b8ac1..48ed91e 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -6,6 +6,7 @@ use axum::{ }; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; +use std::collections::HashMap; use std::sync::Arc; use tokio::sync::RwLock; @@ -13,6 +14,23 @@ use monoce_common::{Cid, cid::Codec}; use monoce_meta::{MetaStateMachine, namespace::NamespaceRecord}; use monoce_node::store::{ChunkStore, MemoryChunkStore}; +// ============================================================================ +// Node Registry +// ============================================================================ + +#[derive(Debug, Clone, Serialize)] +pub struct NodeRecord { + pub node_id: String, + pub address: String, + pub is_online: bool, + pub last_seen_ms: u64, +} + +#[derive(Debug, Default, Clone)] +pub struct NodeRegistry { + pub nodes: HashMap, +} + // ============================================================================ // Application State // ============================================================================ @@ -23,19 +41,73 @@ pub struct AppState { pub store: Arc>, pub node_id: String, pub cluster_id: String, + pub node_registry: Arc>, + pub chunk_locations: Arc>>>, } impl AppState { pub fn new(node_id: String, cluster_id: String) -> Self { + let mut registry = NodeRegistry::default(); + + // Parse MONOCE_STORAGE_NODES env var + // Format: "node-1=http://127.0.0.1:9001,node-2=http://127.0.0.1:9002,node-3=http://127.0.0.1:9003" + if let Ok(nodes_env) = std::env::var("MONOCE_STORAGE_NODES") { + tracing::info!("Parsing storage nodes: {}", nodes_env); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + + for node_spec in nodes_env.split(',') { + let node_spec = node_spec.trim(); + if node_spec.is_empty() { + continue; + } + if let Some((storage_node_id, address)) = node_spec.split_once('=') { + tracing::info!("Adding storage node: {} -> {}", storage_node_id, address); + registry.nodes.insert( + storage_node_id.to_string(), + NodeRecord { + node_id: storage_node_id.to_string(), + address: address.to_string(), + is_online: true, + last_seen_ms: now, + }, + ); + } + } + tracing::info!("Loaded {} storage nodes", registry.nodes.len()); + } else { + tracing::warn!("MONOCE_STORAGE_NODES not set, running in single-node mode"); + } + Self { meta: Arc::new(RwLock::new(MetaStateMachine::default())), store: Arc::new(RwLock::new(MemoryChunkStore::default())), node_id, cluster_id, + node_registry: Arc::new(RwLock::new(registry)), + chunk_locations: Arc::new(RwLock::new(HashMap::new())), } } } +pub async fn check_node_health(address: &str) -> bool { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(2)) + .build() + .ok(); + if let Some(client) = client { + client.get(format!("{}/health", address)) + .send() + .await + .map(|r| r.status().is_success()) + .unwrap_or(false) + } else { + false + } +} + // ============================================================================ // Request/Response types // ============================================================================ @@ -218,22 +290,52 @@ pub async fn delete_object( // ============================================================================ pub async fn create_session( - State(_state): State, + State(state): State, Json(req): Json, ) -> (StatusCode, Json) { let session_id = uuid::Uuid::new_v4().to_string(); - let _rf = req.replication_factor.unwrap_or(1); // Default to 1 for local demo + let rf = req.replication_factor.unwrap_or(1) as usize; let ttl = req.ttl_ms.unwrap_or(300_000); let expires_at = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_millis() as u64 + ttl; - let placements: Vec = req.chunks.iter().map(|c| { + // Get list of online nodes from registry + let registry = state.node_registry.read().await; + let online_nodes: Vec = registry + .nodes + .values() + .filter(|n| n.is_online) + .map(|n| n.address.clone()) + .collect(); + drop(registry); + + let placements: Vec = req.chunks.iter().enumerate().map(|(idx, c)| { + let node_addresses = if online_nodes.is_empty() { + // Fall back to gateway address (single-node mode) + vec!["http://localhost:8080".to_string()] + } else { + // Round-robin selection of rf distinct nodes + let mut selected = Vec::with_capacity(rf); + for i in 0..rf { + let node_idx = (idx + i) % online_nodes.len(); + let addr = &online_nodes[node_idx]; + if !selected.contains(addr) { + selected.push(addr.clone()); + } + } + // If we couldn't get enough distinct nodes, just use what we have + if selected.is_empty() { + vec!["http://localhost:8080".to_string()] + } else { + selected + } + }; + ChunkPlacement { chunk_cid: c.cid.clone(), - // Use localhost:8080 (this gateway) as the storage node for demo - node_addresses: vec!["http://localhost:8080".to_string()], + node_addresses, lease_expires_at: expires_at, } }).collect(); @@ -407,9 +509,15 @@ pub async fn delete_namespace( // Chunk Handlers // ============================================================================ +#[derive(Debug, Deserialize)] +pub struct PutChunkQuery { + pub node_id: Option, +} + pub async fn put_chunk( State(state): State, Path(cid_hex): Path, + Query(query): Query, body: Bytes, ) -> impl IntoResponse { let computed_cid = Cid::from_bytes(&body, Codec::Chunk); @@ -431,11 +539,34 @@ pub async fn put_chunk( let mut store = state.store.write().await; match store.put(chunk) { - Ok(_) => (StatusCode::CREATED, Json(json!({ - "cid": cid_hex, - "size": size, - "status": "stored" - }))).into_response(), + Ok(_) => { + drop(store); + + // Record chunk location + let node_address = if let Some(ref node_id) = query.node_id { + // Look up node address from registry + let registry = state.node_registry.read().await; + registry + .nodes + .get(node_id) + .map(|n| n.address.clone()) + .unwrap_or_else(|| "http://localhost:8080".to_string()) + } else { + "http://localhost:8080".to_string() + }; + + let mut locations = state.chunk_locations.write().await; + locations + .entry(cid_hex.clone()) + .or_insert_with(Vec::new) + .push(node_address); + + (StatusCode::CREATED, Json(json!({ + "cid": cid_hex, + "size": size, + "status": "stored" + }))).into_response() + } Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({ "error": e.to_string() }))).into_response(), @@ -460,11 +591,45 @@ pub async fn get_chunk( } pub async fn get_chunk_locations( - State(_state): State, - Path(_cid_hex): Path, + State(state): State, + Path(cid_hex): Path, ) -> Json> { - // For this demo, all chunks are stored locally on this gateway - Json(vec!["http://localhost:8080".to_string()]) + let locations = state.chunk_locations.read().await; + let mut result = locations.get(&cid_hex).cloned().unwrap_or_default(); + + // Check if chunk is also stored locally + let store = state.store.read().await; + let is_stored_locally = store.list().iter().any(|cid| cid.to_hex() == cid_hex); + drop(store); + + if is_stored_locally && !result.contains(&"http://localhost:8080".to_string()) { + result.push("http://localhost:8080".to_string()); + } + + // If no locations found, return empty + Json(result) +} + +#[derive(Debug, Deserialize)] +pub struct RegisterChunkLocationRequest { + pub node_address: String, +} + +pub async fn register_chunk_location( + State(state): State, + Path(cid_hex): Path, + Json(req): Json, +) -> impl IntoResponse { + let mut locations = state.chunk_locations.write().await; + let addrs = locations.entry(cid_hex.clone()).or_insert_with(Vec::new); + if !addrs.contains(&req.node_address) { + addrs.push(req.node_address.clone()); + } + (StatusCode::OK, Json(json!({ + "cid": cid_hex, + "node_address": req.node_address, + "status": "registered" + }))) } // ============================================================================ @@ -474,17 +639,33 @@ pub async fn get_chunk_locations( pub async fn get_cluster( State(state): State, ) -> Json { + let registry = state.node_registry.read().await; + + // Start with this gateway node as leader + let mut members = vec![NodeInfo { + node_id: state.node_id.clone(), + address: "127.0.0.1:8080".to_string(), + role: "leader".to_string(), + state: "healthy".to_string(), + is_leader: true, + }]; + + // Add storage nodes from registry + for record in registry.nodes.values() { + members.push(NodeInfo { + node_id: record.node_id.clone(), + address: record.address.clone(), + role: "storage".to_string(), + state: if record.is_online { "healthy" } else { "offline" }.to_string(), + is_leader: false, + }); + } + Json(ClusterInfo { cluster_id: state.cluster_id.clone(), term: 1, leader_id: state.node_id.clone(), - members: vec![NodeInfo { - node_id: state.node_id.clone(), - address: "127.0.0.1:8080".to_string(), - role: "leader".to_string(), - state: "healthy".to_string(), - is_leader: true, - }], + members, }) } diff --git a/src/http.rs b/src/http.rs index 23d70f2..8c37932 100644 --- a/src/http.rs +++ b/src/http.rs @@ -55,6 +55,7 @@ fn api_routes() -> Router { .route("/chunks/{cid}", get(handlers::get_chunk)) .route("/chunks/{cid}", put(handlers::put_chunk)) .route("/chunks/{cid}/locations", get(handlers::get_chunk_locations)) + .route("/chunks/{cid}/locations", post(handlers::register_chunk_location)) // Cluster .route("/cluster", get(handlers::get_cluster)) diff --git a/src/lib.rs b/src/lib.rs index fd46aa2..01e0efe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ pub mod http; pub mod grpc; pub mod handlers; pub mod streaming; +pub mod cluster; pub use http::HttpGateway; pub use handlers::AppState;