diff --git a/minikv-core/src/lib.rs b/minikv-core/src/lib.rs index ce2a03a..ac90eed 100644 --- a/minikv-core/src/lib.rs +++ b/minikv-core/src/lib.rs @@ -1,6 +1,7 @@ pub mod error; pub mod hashing; pub mod locking; +pub mod rebalance; pub mod rebuild; pub mod record; pub mod replication; diff --git a/minikv-core/src/rebalance.rs b/minikv-core/src/rebalance.rs new file mode 100644 index 0000000..f2079e4 --- /dev/null +++ b/minikv-core/src/rebalance.rs @@ -0,0 +1,236 @@ +//! Rebalance objects to their deterministic ideal volume set. +//! +//! Rebalancing is required when: +//! - Volume servers are added or removed. +//! - The configured replica count changes. +//! - Objects were written to non-ideal volumes (e.g. due to temporary outages). +//! +//! For each key, the ideal replica set is computed using `key_to_volume`. +//! The object is then migrated so that the set of volumes physically +//! storing it matches this ideal set. +//! +//! # Algorithm (per key) +//! +//! 1. Issue HEAD requests to all volumes recorded in the DB to determine +//! which volumes actually contain the object. +//! 2. If no volume is reachable → return `Ok(false)` (object unavailable). +//! 3. If the reachable set already matches the ideal set → no-op. +//! 4. GET the object from the first reachable volume. +//! 5. PUT the object to any ideal volume that does not already contain it. +//! 6. Update the DB record to the ideal volume list. +//! 7. DELETE the object from volumes no longer in the ideal set. +//! +//! # Metadata Handling +//! +//! - `content_type` is preserved from the existing record. +//! - `hash` is cleared during rebalance. The object body is copied between +//! volumes and is not re-verified here. +//! - `deleted` is always written as `Deleted::No`. +//! +//! # Concurrency +//! +//! `rebalance_all` limits concurrency using a semaphore with 16 +//! concurrent tasks. + +use std::sync::Arc; +use std::time::Duration; + +use tokio::sync::Semaphore; +use tracing::{error, info, warn}; + +use crate::error::Error; +use crate::hashing::key_to_path; +use crate::record::{Deleted, Record}; +use crate::replication::{remote_delete, remote_get, remote_head, remote_put}; +use crate::state::AppState; +use crate::volumes::key_to_volume; + +/// Rebalance a single key to its ideal volume set. +/// +/// `volumes` are the volumes recorded in the DB. +/// `kvolumes` are the ideal volumes computed by `key_to_volume`. +/// +/// Returns: +/// - `Ok(true)` if already balanced or successfully migrated. +/// - `Ok(false)` if the object is missing or a migration step failed. +/// - `Err(_)` for unexpected internal failures. +pub async fn rebalance_key( + state: &AppState, + key: &[u8], + volumes: &[String], + kvolumes: &[String], +) -> Result { + let kp = key_to_path(key); + + // Step 1: find volumes that actually have the data. + let mut reachable: Vec = Vec::new(); + for vol in volumes { + let url = format!("http://{vol}{kp}"); + match remote_head(&state.http_client, &url, Duration::from_secs(60)).await { + Ok(true) => reachable.push(vol.clone()), + Ok(false) => {} + Err(e) => { + warn!(?e, url, "rebalance HEAD error"); + return Ok(false); + } + } + } + + if reachable.is_empty() { + warn!( + key = ?String::from_utf8_lossy(key), + "rebalance impossible. Object missing from all volumes" + ); + return Ok(false); + } + + // Step 2: check if already in ideal position. + if !crate::volumes::needs_rebalance(&reachable, kvolumes) { + return Ok(true); + } + + info!( + key = ?String::from_utf8_lossy(key), + from = ?reachable, + to = ?kvolumes, + "rebalancing key" + ); + + // Step 3: read object from first reachable volume. + let mut body = None; + for vol in &reachable { + let url = format!("http://{vol}{kp}"); + match remote_get(&state.http_client, &url).await { + Ok(b) => { + body = Some(b); + break; + } + Err(e) => warn!(?e, url, "rebalance GET error"), + } + } + let body = match body { + Some(b) => b, + None => { + error!(key = ?String::from_utf8_lossy(key), "rebalance: could not read from any reachable volume"); + return Ok(false); + } + }; + + // Step 4: PUT to volumes that need it (not already in reachable set). + for vol in kvolumes { + if reachable.contains(vol) { + continue; // already there + } + let url = format!("http://{vol}{kp}"); + if let Err(e) = remote_put(&state.http_client, &url, body.clone()).await { + warn!(?e, url, "rebalance PUT error"); + return Ok(false); + } + } + + // Step 5: update DB + // preserve content_type from the existing record. + // Hash is intentionally cleared during rebalance: the body may have been + // copied across volumes and we cannot re-verify it here without re-reading. + // This matches the original Go behaviour. + let existing = state.get_record(key).await; + if !state + .put_record( + key, + Record { + volumes: kvolumes.to_vec(), + deleted: Deleted::No, + hash: None, + content_type: existing.content_type, + }, + ) + .await + { + error!("rebalance: DB put failed"); + return Ok(false); + } + + // Step 6: DELETE from volumes no longer needed. + for vol in &reachable { + if kvolumes.contains(vol) { + continue; // still needed + } + let url = format!("http://{vol}{kp}"); + if let Err(e) = remote_delete(&state.http_client, &url).await { + warn!(?e, url, "rebalance DELETE error"); + return Ok(false); + } + } + + Ok(true) +} + +/// Rebalance all records currently stored in the DB. +/// +/// Performs a full DB scan and spawns bounded async tasks to +/// rebalance each key independently. +/// +/// Corrupt records are skipped. +pub async fn rebalance_all(state: Arc) { + info!("starting full rebalance to {:?}", state.volumes); + + // Collect all entries (blocking LevelDB scan, run off the async executor). + let db = Arc::clone(&state.db); + let entries: Vec<(Vec, Vec)> = + tokio::task::spawn_blocking(move || db.scan_all().unwrap_or_default()) + .await + .unwrap_or_default(); + + info!("rebalance: scanning {} keys", entries.len()); + + let sem = Arc::new(Semaphore::new(16)); + let mut handles = Vec::new(); + + for (raw_key, raw_val) in entries { + let rec = match Record::decode(&raw_val) { + Ok(r) => r, + Err(e) => { + warn!(?e, "skipping corrupt record during rebalance"); + continue; + } + }; + + let kvolumes = key_to_volume(&raw_key, &state.volumes, state.replicas, state.subvolumes); + let state = Arc::clone(&state); + let sem = Arc::clone(&sem); + + let handle = tokio::spawn(async move { + let _permit = sem.acquire().await.expect("semaphore closed"); + let result = rebalance_key(&state, &raw_key, &rec.volumes, &kvolumes).await; + if let Err(e) = result { + warn!(?e, key = ?String::from_utf8_lossy(&raw_key), "rebalance error"); + } + }); + handles.push(handle); + } + + for h in handles { + let _ = h.await; + } + + info!("rebalance complete"); +} + +#[cfg(test)] +mod tests { + use crate::volumes::needs_rebalance; + + #[test] + fn no_rebalance_when_already_ideal() { + let current = vec!["a".into(), "b".into()]; + let ideal = vec!["a".into(), "b".into()]; + assert!(!needs_rebalance(¤t, &ideal)); + } + + #[test] + fn rebalance_needed_when_different() { + let current = vec!["a".into(), "b".into()]; + let ideal = vec!["a".into(), "c".into()]; + assert!(needs_rebalance(¤t, &ideal)); + } +} diff --git a/minikv-core/src/rebuild.rs b/minikv-core/src/rebuild.rs index d4ece79..4c0246f 100644 --- a/minikv-core/src/rebuild.rs +++ b/minikv-core/src/rebuild.rs @@ -6,7 +6,7 @@ //! //! The rebuild process is fully deterministic with respect to: //! - key decoding -//! - volume selection (`key2volume`) +//! - volume selection (`key_to_volume`) //! - replica ordering //! //! It does not recover hash or content-type metadata, as that information @@ -21,7 +21,7 @@ //! 3. For each file entry: //! - Base64-decode the filename to obtain the raw key. //! - Merge the current volume into the key's record. -//! - Reorder volumes according to `key2volume`, preserving unknown +//! - Reorder volumes according to `key_to_volume`, preserving unknown //! volumes at the end. //! 4. Write the reconstructed record with: //! - `deleted = No` @@ -99,7 +99,7 @@ fn is_subvolume_dir(entry: &AutoindexEntry) -> bool { /// /// - Decodes the filename into raw key bytes. /// - Acquires a per-key lock to prevent concurrent modification. -/// - Computes the ideal replica ordering via `key2volume`. +/// - Computes the ideal replica ordering via `key_to_volume`. /// - Merges the current volume into the existing record. /// - Reorders volumes deterministically. /// - Writes a reconstructed `Record`. @@ -255,7 +255,7 @@ pub async fn rebuild_all(state: Arc) { }; // Check if volume uses subvolume directories. - let has_subvolumes = top_listing.iter().any(|e| is_subvolume_dir(e)); + let has_subvolumes = top_listing.iter().any(is_subvolume_dir); if has_subvolumes { for sv in top_listing.iter().filter(|e| is_subvolume_dir(e)) { diff --git a/minikv-core/tests/rebalance.rs b/minikv-core/tests/rebalance.rs new file mode 100644 index 0000000..8fe05e1 --- /dev/null +++ b/minikv-core/tests/rebalance.rs @@ -0,0 +1,181 @@ +//! Rebalance unit tests. +//! +//! These tests validate: +//! - Pure `needs_rebalance` semantics (order and membership sensitivity). +//! - `rebalance_key` behaviour using HTTP mocks instead of real volume servers. +//! - Correct handling of no-op, migration, and missing-object scenarios. +//! +//! An in-memory `MemStore` is used as the metadata backend. + +use dashmap::DashMap; +use minikv_core::locking::KeyLock; +use minikv_core::rebalance::rebalance_key; +use minikv_core::replication::build_volume_client; +use minikv_core::state::AppState; +use minikv_core::storage::mem::MemStore; +use minikv_core::volumes::needs_rebalance; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; +use wiremock::matchers::method; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +/// Build a test `AppState` backed by `MemStore` and pointing at `volumes`. +fn make_state(volumes: Vec) -> Arc { + let replicas = volumes.len().min(2); + Arc::new(AppState { + db: Arc::new(MemStore::default()), + key_lock: KeyLock::new(), + upload_ids: DashMap::new(), + volumes, + vol_rewrite: HashMap::new(), + fallback: None, + replicas, + subvolumes: 1, + protect: false, + checksum: false, + accel_redirect: false, + vol_timeout: Duration::from_secs(5), + http_client: build_volume_client(), + }) +} + +// ── needs_rebalance unit tests ──────────────────────────────────────────── + +#[test] +fn needs_rebalance_same_order_false() { + let v = vec!["a".into(), "b".into()]; + assert!(!needs_rebalance(&v, &v)); +} + +#[test] +fn needs_rebalance_different_order_true() { + let a = vec!["a".into(), "b".into()]; + let b = vec!["b".into(), "a".into()]; + assert!(needs_rebalance(&a, &b)); +} + +#[test] +fn needs_rebalance_different_members_true() { + let a = vec!["a".into(), "b".into()]; + let b = vec!["a".into(), "c".into()]; + assert!(needs_rebalance(&a, &b)); +} + +#[test] +fn needs_rebalance_different_lengths_true() { + let a = vec!["a".into(), "b".into()]; + let b = vec!["a".into()]; + assert!(needs_rebalance(&a, &b)); +} + +// ── rebalance_key integration tests with mock HTTP ──────────────────────── + +/// Returns the key2path for "testkey". We use this to build mock URL paths. +fn test_key_path() -> String { + minikv_core::hashing::key_to_path(b"testkey") +} + +#[tokio::test] +async fn rebalance_key_already_ideal_is_noop() { + let server_a = MockServer::start().await; + let server_b = MockServer::start().await; + + let host_a = server_a.uri().strip_prefix("http://").unwrap().to_string(); + let host_b = server_b.uri().strip_prefix("http://").unwrap().to_string(); + + let _kp = test_key_path(); + + // Both volumes respond 200 to HEAD. + Mock::given(method("HEAD")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server_a) + .await; + Mock::given(method("HEAD")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server_b) + .await; + + let state = make_state(vec![host_a.clone(), host_b.clone()]); + let current = vec![host_a.clone(), host_b.clone()]; + let ideal = vec![host_a.clone(), host_b.clone()]; + + let result = rebalance_key(&state, b"testkey", ¤t, &ideal).await; + assert!(result.unwrap()); +} + +#[tokio::test] +async fn rebalance_key_missing_from_all_volumes_returns_false() { + let server = MockServer::start().await; + + Mock::given(method("HEAD")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + + let host = server.uri().strip_prefix("http://").unwrap().to_string(); + let state = make_state(vec![host.clone()]); + + let result = rebalance_key( + &state, + b"testkey", + std::slice::from_ref(&host), + std::slice::from_ref(&host), + ) + .await; + assert!(!result.unwrap()); +} + +#[tokio::test] +async fn rebalance_key_moves_to_new_volume() { + let server_old = MockServer::start().await; + let server_new = MockServer::start().await; + + let _kp = test_key_path(); + + // Old volume: responds 200 to HEAD and serves content. + Mock::given(method("HEAD")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server_old) + .await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(b"hello".to_vec())) + .mount(&server_old) + .await; + + // New volume: accepts PUT, old volume serves DELETE. + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(201)) + .mount(&server_new) + .await; + Mock::given(method("DELETE")) + .respond_with(ResponseTemplate::new(204)) + .mount(&server_old) + .await; + + let host_old = server_old + .uri() + .strip_prefix("http://") + .unwrap() + .to_string(); + let host_new = server_new + .uri() + .strip_prefix("http://") + .unwrap() + .to_string(); + + let state = make_state(vec![host_old.clone(), host_new.clone()]); + + // Write the record in DB as currently on old volume. + state.db.put(b"testkey", b"hello,world").unwrap(); + + let result = rebalance_key( + &state, + b"testkey", + std::slice::from_ref(&host_old), + std::slice::from_ref(&host_new), + ) + .await; + + assert!(result.is_ok(), "rebalance should not error: {result:?}"); +}