From 8bdbd9937686035dce530567657d58d364f87c31 Mon Sep 17 00:00:00 2001 From: Austin Barrington Date: Fri, 4 Sep 2026 23:01:20 +0100 Subject: [PATCH] Install the committed shard map on join before region movement. A joiner's map_version must match the cluster's before any peer-set change, so staging cannot run against a lagging epoch. --- .../src/adapters/cluster/sync_client.rs | 45 +++++++++ .../adapters/sharding/rocksdb_shard_map.rs | 87 ++++++++++++++++ .../src/application/cluster/bootstrap.rs | 6 +- hyperbytedb/src/application/runtime/mod.rs | 3 + .../src/application/shard_scheduler.rs | 26 +++++ hyperbytedb/src/domain/sharding/types.rs | 14 +++ hyperbytedb/src/ports/sharding.rs | 10 ++ hyperbytedb/tests/common/sharding_cluster.rs | 31 +++++- .../tests/sharding_cluster_integration.rs | 99 +++++++++++++++++++ 9 files changed, 319 insertions(+), 2 deletions(-) diff --git a/hyperbytedb/src/adapters/cluster/sync_client.rs b/hyperbytedb/src/adapters/cluster/sync_client.rs index 2f92af1..48797dd 100644 --- a/hyperbytedb/src/adapters/cluster/sync_client.rs +++ b/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,9 @@ 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. + self.sync_shard_map(&peer_addr).await?; let mut applied = 0u64; if let Some(ref shard_map) = self.shard_map { @@ -328,6 +333,23 @@ 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; + 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 +509,29 @@ impl SyncClient { } } +/// GET `/internal/shard/map` from `peer_addr` and return the committed snapshot. +pub async fn fetch_shard_map( + client: reqwest::Client, + peer_addr: &str, +) -> Result { + 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( diff --git a/hyperbytedb/src/adapters/sharding/rocksdb_shard_map.rs b/hyperbytedb/src/adapters/sharding/rocksdb_shard_map.rs index 1b7d91d..5867444 100644 --- a/hyperbytedb/src/adapters/sharding/rocksdb_shard_map.rs +++ b/hyperbytedb/src/adapters/sharding/rocksdb_shard_map.rs @@ -207,6 +207,47 @@ fn assemble_map(db: &DB) -> Result { 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(); diff --git a/hyperbytedb/src/application/cluster/bootstrap.rs b/hyperbytedb/src/application/cluster/bootstrap.rs index 9f90593..60214d8 100644 --- a/hyperbytedb/src/application/cluster/bootstrap.rs +++ b/hyperbytedb/src/application/cluster/bootstrap.rs @@ -98,6 +98,7 @@ impl ClusterBootstrap { wal: &Arc, points_sink: Option>, max_points_per_request: usize, + shard_map: Option>, ) -> 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(); diff --git a/hyperbytedb/src/application/runtime/mod.rs b/hyperbytedb/src/application/runtime/mod.rs index e2b90e2..5d77157 100644 --- a/hyperbytedb/src/application/runtime/mod.rs +++ b/hyperbytedb/src/application/runtime/mod.rs @@ -111,6 +111,9 @@ pub async fn serve(config: HyperbytedbConfig) -> anyhow::Result<()> { &wal_port, Some(sink_port), config.server.max_points_per_request, + rocks_shard_map + .as_ref() + .map(|m| m.clone() as Arc), ) .await?; Some( diff --git a/hyperbytedb/src/application/shard_scheduler.rs b/hyperbytedb/src/application/shard_scheduler.rs index 4c92d7e..22432e8 100644 --- a/hyperbytedb/src/application/shard_scheduler.rs +++ b/hyperbytedb/src/application/shard_scheduler.rs @@ -1169,6 +1169,24 @@ fn failover_watermark_safe(candidate_watermark: u64, max_peer_watermark: u64) -> candidate_watermark > 0 || max_peer_watermark == 0 } +/// Region data movement onto a joiner starts only after its committed +/// `map_version` matches the cluster's. Staging rows against a lagging map +/// would apply under the wrong epoch / peer set. +#[must_use] +pub fn joiner_map_caught_up(cluster_map_version: u64, joiner_map_version: u64) -> bool { + joiner_map_version == cluster_map_version +} + +/// Read a peer's committed `map_version` via `/internal/shard/map`. +pub async fn fetch_peer_map_version( + client: &reqwest::Client, + peer_addr: &str, +) -> Result { + let map = + crate::adapters::cluster::sync_client::fetch_shard_map(client.clone(), peer_addr).await?; + Ok(map.map_version) +} + #[allow(clippy::too_many_arguments)] async fn push_and_drop_range( peer_client: &Arc, @@ -1888,6 +1906,14 @@ mod tests { assert!(failover_watermark_safe(0, 0)); } + #[test] + fn joiner_map_catchup_blocks_movement_until_versions_match() { + assert!(joiner_map_caught_up(3, 3)); + assert!(!joiner_map_caught_up(3, 0)); + assert!(!joiner_map_caught_up(3, 2)); + assert!(!joiner_map_caught_up(3, 4)); + } + #[tokio::test] #[serial_test::serial(chdb)] async fn try_failover_proposes_transfer_primary() { diff --git a/hyperbytedb/src/domain/sharding/types.rs b/hyperbytedb/src/domain/sharding/types.rs index 4fd903f..8be463a 100644 --- a/hyperbytedb/src/domain/sharding/types.rs +++ b/hyperbytedb/src/domain/sharding/types.rs @@ -217,6 +217,20 @@ impl From<&ShardMap> for ShardMapJson { } } +impl From for ShardMap { + fn from(json: ShardMapJson) -> Self { + let mut spaces = HashMap::new(); + for space in json.spaces { + spaces.insert(space.key.clone(), space); + } + Self { + map_version: json.map_version, + next_region_id: json.next_region_id, + spaces, + } + } +} + #[cfg(test)] mod locate_tests { use super::*; diff --git a/hyperbytedb/src/ports/sharding.rs b/hyperbytedb/src/ports/sharding.rs index 0a75808..dc11bd0 100644 --- a/hyperbytedb/src/ports/sharding.rs +++ b/hyperbytedb/src/ports/sharding.rs @@ -27,6 +27,12 @@ pub trait ShardMapPort: Send + Sync { async fn apply_op(&self, op: ShardMapOp) -> Result; + /// Replace the local map with a peer snapshot (join catch-up). + /// + /// Used so a joiner's `map_version` matches the cluster before region + /// data movement starts. Does not change membership or region peers. + async fn replace_map(&self, map: ShardMap) -> Result<(), HyperbytedbError>; + /// Returns true when this node holds a replica of any series for the measurement. async fn node_owns_measurement( &self, @@ -73,6 +79,10 @@ impl ShardMapPort for DisabledShardMap { Err(HyperbytedbError::Internal("sharding is disabled".into())) } + async fn replace_map(&self, _map: ShardMap) -> Result<(), HyperbytedbError> { + Err(HyperbytedbError::Internal("sharding is disabled".into())) + } + async fn node_owns_measurement( &self, _node_id: u64, diff --git a/hyperbytedb/tests/common/sharding_cluster.rs b/hyperbytedb/tests/common/sharding_cluster.rs index 02f537c..b5e5ca1 100644 --- a/hyperbytedb/tests/common/sharding_cluster.rs +++ b/hyperbytedb/tests/common/sharding_cluster.rs @@ -47,6 +47,8 @@ pub struct ShardedTestNode { pub location_cache: Arc, pub membership: SharedMembership, pub query_port: Arc, + /// Shared with other in-process peers — libchdb allows one session per process. + pub chdb: SharedSession, flush: Arc, handle: tokio::task::JoinHandle<()>, shutdown: Option>, @@ -123,7 +125,7 @@ pub async fn start_sharded_node( let wal = Arc::new(RocksDbWal::open(&wal_dir).unwrap()); let metadata = Arc::new(RocksDbMetadata::open(&meta_dir).unwrap()); let chdb_adapter = Arc::new(ChdbQueryAdapter::from_shared(chdb.clone(), 0)); - let sink: Arc = Arc::new(ChdbNativeAdapter::new(chdb)); + let sink: Arc = Arc::new(ChdbNativeAdapter::new(chdb.clone())); let flush: Arc = Arc::new(FlushServiceImpl::new(wal.clone(), 0, sink.clone())); @@ -284,12 +286,39 @@ pub async fn start_sharded_node( location_cache, membership: shared_membership, query_port: chdb_adapter, + chdb, flush, handle, shutdown: Some(shutdown_tx), } } +pub async fn fetch_shard_map( + client: &reqwest::Client, + url: &str, +) -> hyperbytedb::domain::sharding::ShardMap { + let resp = client + .get(format!("{url}/internal/shard/map")) + .send() + .await + .unwrap(); + assert!( + resp.status().is_success(), + "shard map fetch failed: {}", + resp.status() + ); + let json: hyperbytedb::domain::sharding::ShardMapJson = resp.json().await.unwrap(); + json.into() +} + +pub async fn install_shard_map_from_peer(from: &ShardedTestNode, onto: &ShardedTestNode) { + let client = reqwest::Client::new(); + let map = fetch_shard_map(&client, &from.url).await; + onto.shard_map.replace_map(map).await.unwrap(); + let snap = onto.shard_map.snapshot().await.unwrap(); + onto.location_cache.refresh_from_map(&snap); +} + pub async fn start_sharded_single_node(dir: &Path, opts: ShardedClusterOptions) -> ShardedTestNode { let chdb_dir = dir.join("chdb-shared"); std::fs::create_dir_all(&chdb_dir).unwrap(); diff --git a/hyperbytedb/tests/sharding_cluster_integration.rs b/hyperbytedb/tests/sharding_cluster_integration.rs index 93a37f3..8be489b 100644 --- a/hyperbytedb/tests/sharding_cluster_integration.rs +++ b/hyperbytedb/tests/sharding_cluster_integration.rs @@ -797,6 +797,105 @@ async fn one_member_cluster_owns_region_after_first_write() { assert_eq!(region.primary, 1, "n=1 primary must be self"); } +/// P1.2: after a joiner is Active, its map_version equals the cluster's +/// before any region data movement (peer-set change) starts. +#[tokio::test] +#[serial(chdb)] +async fn joiner_map_version_matches_before_region_movement() { + let dir = tempfile::tempdir().unwrap(); + let opts = ShardedClusterOptions { + sharding: hyperbytedb::config::ShardingConfig { + enabled: true, + replication_factor: 3, + scatter_peer_timeout_ms: 500, + scatter_max_peer_attempts: 3, + ..Default::default() + }, + ..Default::default() + }; + let node1 = start_sharded_single_node(dir.path(), opts).await; + let client = reqwest::Client::new(); + create_db(&client, &node1.url, "p1db").await; + let resp = write_line( + &client, + &node1.url, + "p1db", + "cpu,host=solo value=1 1000000000", + ) + .await; + assert_eq!(resp.status(), reqwest::StatusCode::NO_CONTENT); + + let cluster_before = node1.shard_map.snapshot().await.unwrap(); + assert!( + cluster_before.map_version >= 1, + "first write must bump map_version" + ); + let region = cluster_before + .space("p1db", "autogen", "cpu") + .and_then(|s| s.regions.first()) + .expect("region after write"); + assert_eq!(region.peers, vec![1]); + + let l2 = bind_ephemeral().await; + let a2 = l2.local_addr().unwrap().to_string(); + { + let mut m = node1.membership.write().await; + m.add_node(hyperbytedb::domain::cluster::membership::NodeInfo { + node_id: 2, + addr: a2, + state: hyperbytedb::domain::cluster::membership::NodeState::Active, + joined_at: 0, + last_heartbeat: 0, + needs_sync: false, + }); + } + // libchdb is process-global; share node 1's session (pair-cluster harness). + let node2 = start_sharded_node( + dir.path(), + 2, + l2, + node1.membership.clone(), + &ShardedClusterOptions { + sharding: hyperbytedb::config::ShardingConfig { + enabled: true, + replication_factor: 3, + scatter_peer_timeout_ms: 500, + scatter_max_peer_attempts: 3, + ..Default::default() + }, + ..Default::default() + }, + node1.chdb.clone(), + ) + .await; + + let empty = node2.shard_map.snapshot().await.unwrap(); + assert_eq!(empty.map_version, 0, "joiner starts with an empty map"); + + install_shard_map_from_peer(&node1, &node2).await; + + let after = node2.shard_map.snapshot().await.unwrap(); + assert_eq!( + after.map_version, cluster_before.map_version, + "joiner map_version must equal the cluster's before movement" + ); + let joined_region = after + .space("p1db", "autogen", "cpu") + .and_then(|s| s.regions.first()) + .expect("installed region"); + assert_eq!( + joined_region.peers, + vec![1], + "catch-up must not add the joiner as a peer" + ); + assert!( + hyperbytedb::application::shard_scheduler::joiner_map_caught_up( + cluster_before.map_version, + after.map_version + ) + ); +} + async fn wait_for_database(client: &reqwest::Client, url: &str, db: &str) { for _ in 0..100 { let resp = query_sql(client, url, db, "SHOW DATABASES").await;