feat(sharding): place regions and primaries onto joining nodes

hyperbytedb/src/adapters/cluster/sync_client.rs

@@ -7,6 +7,7 @@ use crate::domain::cluster::membership::{NodeState, SharedMembership};
 use crate::domain::cluster::sync::{
     JoinRequest, JoinResponse, MetadataSnapshot, SyncManifest, WalSyncResponse,
 };
+use crate::domain::sharding::{ShardMap, ShardMapJson};
 use crate::error::HyperbytedbError;
 use crate::ports::metadata::MetadataPort;
 use crate::ports::points_sink::PointsSinkPort;
@@ -133,6 +134,7 @@ impl SyncClient {
             );
 
             self.sync_metadata(&peer_addr).await?;
+            self.sync_shard_map(&peer_addr).await?;
 
             let updated_wal_seq = self.wal.last_sequence().await?;
             let applied = self.wal_catchup(&peer_addr, updated_wal_seq).await?;
@@ -178,6 +180,23 @@ impl SyncClient {
         );
 
         self.sync_metadata(&peer_addr).await?;
+        // Install the committed map before any region data movement so the
+        // joiner's map_version matches the cluster at Active.
+        //
+        // Best-effort: `/internal/shard/map` only exists on peers that have
+        // sharding enabled and run a build that serves it, so a 404 here says
+        // nothing about whether this node can catch up its WAL. Failing hard
+        // used to abort the join before `wal_catchup`, and after the retries
+        // were exhausted the node went Active having synced nothing at all.
+        // Skipping is safe: without a map there are no regions to move below,
+        // and the leader's own catch-up check gates placement onto this node.
+        if let Err(e) = self.sync_shard_map(&peer_addr).await {
+            tracing::warn!(
+                error = %e,
+                peer = %peer_addr,
+                "shard map install failed; continuing with WAL catch-up"
+            );
+        }
 
         let mut applied = 0u64;
         if let Some(ref shard_map) = self.shard_map {
@@ -328,6 +347,42 @@ impl SyncClient {
         Ok(())
     }
 
+    /// Copy the peer's committed shard map so `map_version` matches before
+    /// region rows are transferred onto this joiner.
+    pub async fn sync_shard_map(&self, peer_addr: &str) -> Result<(), HyperbytedbError> {
+        let Some(shard_map) = self.shard_map.as_ref() else {
+            return Ok(());
+        };
+        let remote = fetch_shard_map(self.client.clone(), peer_addr).await?;
+        let version = remote.map_version;
+
+        // Never move the map backwards. `map_version` counts local applies and
+        // the sync peer is whichever Active node came first out of a HashMap —
+        // not the leader, and not necessarily the most-applied node. Installing
+        // an older snapshot would drop ops Raft will never redeliver, because
+        // `replace_map` bypasses the state machine and leaves `last_applied`
+        // untouched; every later op would then fail StaleEpoch/UnknownRegion
+        // and the node would silently stop owning regions it holds data for.
+        let local = shard_map.snapshot().await?.map_version;
+        if !remote_map_is_newer(version, local) {
+            tracing::info!(
+                peer = %peer_addr,
+                remote_map_version = version,
+                local_map_version = local,
+                "peer shard map is not newer; keeping local map"
+            );
+            return Ok(());
+        }
+
+        shard_map.replace_map(remote).await?;
+        tracing::info!(
+            peer = %peer_addr,
+            map_version = version,
+            "installed peer shard map before region movement"
+        );
+        Ok(())
+    }
+
     async fn import_metadata_entry(
         &self,
         entry: &crate::domain::cluster::sync::MetadataEntry,
@@ -487,6 +542,39 @@ impl SyncClient {
     }
 }
 
+/// GET `/internal/shard/map` from `peer_addr` and return the committed snapshot.
+/// Whether a peer's shard map may replace the local one.
+///
+/// Strictly-newer only. Equal is a no-op, and older must be refused: a sync
+/// peer is whichever Active node came first out of a `HashMap`, so it is not
+/// necessarily the leader or the most-applied node.
+#[must_use]
+pub fn remote_map_is_newer(remote_version: u64, local_version: u64) -> bool {
+    remote_version > local_version
+}
+
+pub async fn fetch_shard_map(
+    client: reqwest::Client,
+    peer_addr: &str,
+) -> Result<ShardMap, HyperbytedbError> {
+    let url = format!("http://{peer_addr}/internal/shard/map");
+    let resp =
+        client.get(&url).send().await.map_err(|e| {
+            HyperbytedbError::PeerUnreachable(format!("shard map request failed: {e}"))
+        })?;
+    if !resp.status().is_success() {
+        return Err(HyperbytedbError::SyncFailed(format!(
+            "shard map request failed: {}",
+            resp.status()
+        )));
+    }
+    let json: ShardMapJson = resp
+        .json()
+        .await
+        .map_err(|e| HyperbytedbError::SyncFailed(format!("parse shard map: {e}")))?;
+    Ok(ShardMap::from(json))
+}
+
 /// Fail closed when the peer advertises a higher WAL watermark but returns no
 /// readable entries — usually because the leader truncated past the gap.
 fn verify_catchup_progress(
@@ -509,9 +597,19 @@ fn verify_catchup_progress(
 
 #[cfg(test)]
 mod tests {
-    use super::verify_catchup_progress;
+    use super::{remote_map_is_newer, verify_catchup_progress};
     use crate::error::HyperbytedbError;
 
+    #[test]
+    fn remote_map_installs_only_when_strictly_newer() {
+        assert!(remote_map_is_newer(10, 0), "fresh joiner takes the map");
+        assert!(remote_map_is_newer(10, 9));
+        assert!(!remote_map_is_newer(10, 10), "equal is a no-op");
+        // The case that silently strands a node: syncing from a less-applied
+        // peer would drop ops Raft never redelivers.
+        assert!(!remote_map_is_newer(8, 10));
+    }
+
     #[test]
     fn verify_catchup_progress_ok_when_caught_up() {
         verify_catchup_progress(5, 5, 0, 5).unwrap();

hyperbytedb/src/adapters/http/shard_handlers.rs

@@ -540,11 +540,15 @@ pub async fn handle_shard_transfer(
     Json(req): Json<ShardTransferPayload>,
 ) -> impl IntoResponse {
     if req.stage {
-        // Pre-commit staging (split optimization): the range is not owned by
-        // the destination in the committed map yet, so region/epoch checks are
-        // skipped. Guard rails: sender must be a known member and the local
-        // measurement must exist; rows outside `[start, end)` are filtered
-        // during apply.
+        // Pre-commit staging (split optimization, and replica placement onto a
+        // joiner): the range is not owned by the destination in the committed
+        // map yet, so region/epoch checks are skipped. The only guard rail is
+        // that the sender must be a known member; rows outside `[start, end)`
+        // are filtered during apply.
+        //
+        // There is deliberately no local-measurement check. A joiner receiving
+        // a measurement it has never seen has no catalog row yet, and
+        // `apply_transfer_push` creates one via `prepare_batch_metadata`.
         let Some(membership) = state.membership.as_ref() else {
             return (
                 StatusCode::SERVICE_UNAVAILABLE,
@@ -561,20 +565,6 @@ pub async fn handle_shard_transfer(
                 .into_response();
         }
         drop(m);
-        if state
-            .metadata
-            .get_measurement(&req.db, &req.rp, &req.measurement)
-            .await
-            .ok()
-            .flatten()
-            .is_none()
-        {
-            return (
-                StatusCode::NOT_FOUND,
-                Json(serde_json::json!({"error": "measurement not found for staging"})),
-            )
-                .into_response();
-        }
     } else {
         let Some(ctx) = state.shard_routing.as_ref() else {
             return sharding_disabled();
@@ -710,19 +700,33 @@ pub async fn handle_shard_rehome(
             .into_response();
     }
 
-    let outcome = match push_region_transfer_data(
-        peer_client,
-        &state.metadata,
-        &state.wal,
-        Some(&state.query_port),
-        state.node_id,
-        &ctx_key(&req),
-        region,
-        req.dest_primary,
-        state.max_points_per_request,
-    )
-    .await
-    {
+    let outcome = match if req.stage {
+        crate::application::shard_transfer::stage_region_transfer_data(
+            peer_client,
+            &state.metadata,
+            &state.wal,
+            Some(&state.query_port),
+            state.node_id,
+            &ctx_key(&req),
+            region,
+            req.dest_primary,
+            state.max_points_per_request,
+        )
+        .await
+    } else {
+        push_region_transfer_data(
+            peer_client,
+            &state.metadata,
+            &state.wal,
+            Some(&state.query_port),
+            state.node_id,
+            &ctx_key(&req),
+            region,
+            req.dest_primary,
+            state.max_points_per_request,
+        )
+        .await
+    } {
         Ok(o) => o,
         Err(e) => {
             counter!("hyperbytedb_shard_transfer_failures_total").increment(1);
@@ -734,18 +738,19 @@ pub async fn handle_shard_rehome(
         }
     };
 
-    if let Err(e) = complete_region_transfer(
-        peer_client,
-        &state.metadata,
-        Some(&state.points_sink),
-        state.node_id,
-        &ctx_key(&req),
-        region,
-        req.dest_primary,
-        outcome.transfer_id,
-        req.drop_source,
-    )
-    .await
+    if !req.stage
+        && let Err(e) = complete_region_transfer(
+            peer_client,
+            &state.metadata,
+            Some(&state.points_sink),
+            state.node_id,
+            &ctx_key(&req),
+            region,
+            req.dest_primary,
+            outcome.transfer_id,
+            req.drop_source,
+        )
+        .await
     {
         counter!("hyperbytedb_shard_transfer_failures_total").increment(1);
         return (

hyperbytedb/src/adapters/sharding/rocksdb_shard_map.rs

@@ -207,6 +207,47 @@ fn assemble_map(db: &DB) -> Result<ShardMap, HyperbytedbError> {
     Ok(map)
 }
 
+/// Atomically persist the full map (join catch-up / snapshot install).
+fn persist_full_map(db: &DB, map: &ShardMap) -> Result<(), HyperbytedbError> {
+    let existing = load_spaces(db)?;
+    let incoming: std::collections::HashSet<&MeasurementKey> = map.spaces.keys().collect();
+    let mut batch = WriteBatch::default();
+    batch.put(
+        META_KEY,
+        serde_json::to_vec(&PersistedShardMapMeta {
+            map_version: map.map_version,
+            next_region_id: map.next_region_id,
+        })
+        .map_err(|e| {
+            HyperbytedbError::ShardMap(crate::error::ChainedError::with_context(
+                "shard map meta serialize",
+                e,
+            ))
+        })?,
+    );
+    for space in &existing {
+        if !incoming.contains(&space.key) {
+            batch.delete(space_key(&space.key));
+        }
+    }
+    for space in map.spaces.values() {
+        batch.put(
+            space_key(&space.key),
+            serde_json::to_vec(&PersistedShardSpace {
+                space: space.clone(),
+            })
+            .map_err(|e| {
+                HyperbytedbError::ShardMap(crate::error::ChainedError::with_context(
+                    "shard space serialize",
+                    e,
+                ))
+            })?,
+        );
+    }
+    db.write(batch)
+        .map_err(|e| HyperbytedbError::Storage(e.to_string().into()))
+}
+
 /// Atomically persist the global counters plus the one space an op touched.
 fn persist_space(
     db: &DB,
@@ -326,6 +367,20 @@ impl ShardMapPort for RocksDbShardMap {
         Ok(map)
     }
 
+    async fn replace_map(&self, map: ShardMap) -> Result<(), HyperbytedbError> {
+        let _guard = self.apply_lock.lock().await;
+        for space in map.spaces.values() {
+            space
+                .validate()
+                .map_err(|e| HyperbytedbError::ShardMap(e.into()))?;
+        }
+        map.validate_global_region_ids()
+            .map_err(|e| HyperbytedbError::ShardMap(e.into()))?;
+        persist_full_map(&self.db, &map)?;
+        *self.cache.write() = Arc::new(map);
+        Ok(())
+    }
+
     async fn node_owns_measurement(
         &self,
         node_id: u64,
@@ -392,6 +447,38 @@ mod tests {
         assert_eq!(map.next_region_id, 2);
     }
 
+    #[tokio::test]
+    async fn replace_map_installs_peer_snapshot() {
+        let dir = tempfile::tempdir().unwrap();
+        let local = RocksDbShardMap::open(dir.path(), true).unwrap();
+        local
+            .apply_op(ShardMapOp::BootstrapMeasurement {
+                key: MeasurementKey::new("db", "rp", "old"),
+                region: region(1, 0, u64::MAX),
+            })
+            .await
+            .unwrap();
+
+        let incoming = crate::domain::sharding::ShardMap {
+            map_version: 4,
+            next_region_id: 3,
+            spaces: [(
+                MeasurementKey::new("db", "rp", "cpu"),
+                crate::domain::sharding::MeasurementShardSpace {
+                    key: MeasurementKey::new("db", "rp", "cpu"),
+                    regions: vec![region(2, 0, u64::MAX)],
+                },
+            )]
+            .into_iter()
+            .collect(),
+        };
+        local.replace_map(incoming).await.unwrap();
+        let snap = local.snapshot().await.unwrap();
+        assert_eq!(snap.map_version, 4);
+        assert!(snap.space("db", "rp", "old").is_none());
+        assert!(snap.space("db", "rp", "cpu").is_some());
+    }
+
     #[tokio::test]
     async fn legacy_single_key_format_migrates_on_open() {
         let dir = tempfile::tempdir().unwrap();

hyperbytedb/src/application/cluster/bootstrap.rs

@@ -98,6 +98,7 @@ impl ClusterBootstrap {
         wal: &Arc<dyn WalPort>,
         points_sink: Option<Arc<dyn PointsSinkPort>>,
         max_points_per_request: usize,
+        shard_map: Option<Arc<dyn crate::ports::sharding::ShardMapPort>>,
     ) -> anyhow::Result<()> {
         {
             let mut m = self.membership.write().await;
@@ -107,7 +108,7 @@ impl ClusterBootstrap {
 
         tracing::info!("startup phase: syncing with cluster before accepting traffic");
 
-        let sync_client = SyncClient::with_points_sink(
+        let mut sync_client = SyncClient::with_points_sink(
             config.node_id,
             config.cluster_addr.clone(),
             self.membership.clone(),
@@ -117,6 +118,9 @@ impl ClusterBootstrap {
             max_points_per_request,
             self.peer_addrs.clone(),
         );
+        if let Some(map) = shard_map {
+            sync_client = sync_client.with_shard_map(map);
+        }
 
         let dbs = metadata.list_databases().await?;
         let has_data = !dbs.is_empty();