Support rebalancing objects to their ideal volume set (#9)

* Fix cargo clippy lint warnings

* Support rebalancing objects to their ideal volume set

Author: ekkolon

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;

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<bool, Error> {
+    let kp = key_to_path(key);
+
+    // Step 1: find volumes that actually have the data.
+    let mut reachable: Vec<String> = 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<AppState>) {
+    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<u8>, Vec<u8>)> =
+        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(&current, &ideal));
+    }
+
+    #[test]
+    fn rebalance_needed_when_different() {
+        let current = vec!["a".into(), "b".into()];
+        let ideal = vec!["a".into(), "c".into()];
+        assert!(needs_rebalance(&current, &ideal));
+    }
+}

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<AppState>) {
         };
 
         // 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)) {