diff --git a/lib/collection/src/collection/resharding.rs b/lib/collection/src/collection/resharding.rs index bd3704bb58..bb11a7a303 100644 --- a/lib/collection/src/collection/resharding.rs +++ b/lib/collection/src/collection/resharding.rs @@ -40,23 +40,19 @@ impl Collection { { let mut shard_holder = self.shards_holder.write().await; - // Idempotent: if the same resharding is already in progress, this is a - // re-application of an already-applied consensus entry (e.g. after crash - // recovery). Skip without error so we don't silently diverge from peers - // that applied it successfully on the first attempt. - if let Some(state) = shard_holder.resharding_state() - && state.matches(&resharding_key) - { - log::warn!( - "Resharding {resharding_key} is already in progress, skipping start (idempotent re-apply)", - ); - return Ok(()); - } - shard_holder.check_start_resharding(&resharding_key)?; - // If scaling up, create a new replica set - let replica_set = if resharding_key.direction == ReshardingDirection::Up { + // If scaling up, create a new replica set. + // + // Idempotent: only create the replica set if one for this shard id + // does not exist yet. A replica set may already be present as a + // leftover from a previous partial apply (e.g. crash between shard + // directory creation and persisting resharding state) or from an + // earlier successful apply we are now replaying. Recreating it + // would wipe the existing shard contents. + let replica_set = if resharding_key.direction == ReshardingDirection::Up + && !shard_holder.contains_shard(resharding_key.shard_id) + { let replica_set = self .create_replica_set( resharding_key.shard_id, @@ -78,19 +74,25 @@ impl Collection { if resharding_key.direction == ReshardingDirection::Up { let mut config = self.collection_config.write().await; match config.params.sharding_method.unwrap_or_default() { - // If adding a shard, increase persisted count so we load it on restart + // Idempotent: set the shard count to the new shard's id + // plus one (shard ids are contiguous from zero, so this is + // the count after adding the shard). A replay converges to + // the same value instead of incrementing again. ShardingMethod::Auto => { - debug_assert_eq!(config.params.shard_number.get(), resharding_key.shard_id); - - config.params.shard_number = config - .params - .shard_number + resharding_key + .debug_assert_targets_last_shard(config.params.shard_number.get()); + let new_shard_number = resharding_key + .shard_id .checked_add(1) + .and_then(NonZeroU32::new) .expect("cannot have more than u32::MAX shards after resharding"); - if let Err(err) = config.save(&self.path) { - log::error!( - "Failed to update and save collection config during resharding: {err}", - ); + if config.params.shard_number != new_shard_number { + config.params.shard_number = new_shard_number; + if let Err(err) = config.save(&self.path) { + log::error!( + "Failed to update and save collection config during resharding: {err}", + ); + } } } // Custom shards don't use the persisted count, we don't change it @@ -187,12 +189,10 @@ impl Collection { pub async fn finish_resharding(&self, resharding_key: ReshardKey) -> CollectionResult<()> { let mut shard_holder = self.shards_holder.write().await; - // Idempotent: if no resharding is in progress, finish was already applied - if shard_holder.resharding_state().is_none() { - log::warn!("finish_resharding: no resharding in progress, skipping"); - return Ok(()); - } - + // Each step below is individually idempotent, so on replay (e.g. after + // a crash that applied only part of the operation) we fall through + // every step to reconcile any partially-applied state instead of + // short-circuiting on the first condition that already holds. shard_holder.check_finish_resharding(&resharding_key)?; shard_holder.finish_resharding_unchecked(&resharding_key)?; @@ -209,21 +209,23 @@ impl Collection { { let mut config = self.collection_config.write().await; match config.params.sharding_method.unwrap_or_default() { - // If removing a shard, decrease persisted count so we don't load it on restart + // Idempotent: set the shard count to the id of the just + // removed shard (which equals the new count, since shard + // ids are contiguous and this was the highest one). A + // replay converges to the same value instead of + // decrementing again. ShardingMethod::Auto => { - debug_assert_eq!( - config.params.shard_number.get() - 1, - resharding_key.shard_id, - ); - - config.params.shard_number = - NonZeroU32::new(config.params.shard_number.get() - 1) - .expect("cannot have zero shards after finishing resharding"); - - if let Err(err) = config.save(&self.path) { - log::error!( - "Failed to update and save collection config during resharding: {err}" - ); + resharding_key + .debug_assert_targets_last_shard(config.params.shard_number.get()); + let new_shard_number = NonZeroU32::new(resharding_key.shard_id) + .expect("cannot have zero shards after finishing resharding down"); + if config.params.shard_number != new_shard_number { + config.params.shard_number = new_shard_number; + if let Err(err) = config.save(&self.path) { + log::error!( + "Failed to update and save collection config during resharding: {err}" + ); + } } } // Custom shards don't use the persisted count, we don't change it @@ -246,12 +248,10 @@ impl Collection { let shard_holder = self.shards_holder.read().await; - // Idempotent: if no resharding is in progress, abort was already applied - if shard_holder.resharding_state().is_none() { - log::warn!("abort_resharding: no resharding in progress, skipping"); - return Ok(()); - } - + // Each step below is individually idempotent, so on replay (e.g. after + // a crash that applied only part of the abort) we fall through every + // step to reconcile any partially-applied state instead of + // short-circuiting on the first condition that already holds. if !force { shard_holder.check_abort_resharding(&resharding_key)?; } else { @@ -266,21 +266,16 @@ impl Collection { self.invalidate_clean_local_shards([resharding_key.shard_id]) .await; } - // On resharding down: existing shards may have new points moved into them - ReshardingDirection::Down => match shard_holder.rings.get(&resharding_key.shard_key) { - Some(HashRingRouter::Resharding { old: _, new }) => { - self.invalidate_clean_local_shards(new.nodes().clone()) - .await; - } - Some(HashRingRouter::Single(ring)) => { - debug_assert!(false, "must have resharding hash ring during resharding"); + // On resharding down: existing shards may have new points moved + // into them. Use the router's node set, which works regardless of + // whether the ring has already been rolled back to `Single` on a + // replay. + ReshardingDirection::Down => { + if let Some(ring) = shard_holder.rings.get(&resharding_key.shard_key) { self.invalidate_clean_local_shards(ring.nodes().clone()) .await; } - None => { - debug_assert!(false, "must have hash ring for resharding key"); - } - }, + } } // Abort all resharding transfer related to this specific resharding operation @@ -301,21 +296,22 @@ impl Collection { if resharding_key.direction == ReshardingDirection::Up { let mut config = self.collection_config.write().await; match config.params.sharding_method.unwrap_or_default() { - // If removing a shard, decrease persisted count so we don't load it on restart + // Idempotent: set the shard count to the aborted shard's id + // (shard ids are contiguous from zero, so this is the count + // after removing the shard). A replay converges to the same + // value instead of decrementing again. ShardingMethod::Auto => { - debug_assert_eq!( - config.params.shard_number.get() - 1, - resharding_key.shard_id, - ); - - config.params.shard_number = - NonZeroU32::new(config.params.shard_number.get() - 1) - .expect("cannot have zero shards after aborting resharding"); - - if let Err(err) = config.save(&self.path) { - log::error!( - "Failed to update and save collection config during resharding: {err}" - ); + resharding_key + .debug_assert_targets_last_shard(config.params.shard_number.get()); + let new_shard_number = NonZeroU32::new(resharding_key.shard_id) + .expect("cannot have zero shards after aborting resharding up"); + if config.params.shard_number != new_shard_number { + config.params.shard_number = new_shard_number; + if let Err(err) = config.save(&self.path) { + log::error!( + "Failed to update and save collection config during resharding: {err}" + ); + } } } // Custom shards don't use the persisted count, we don't change it diff --git a/lib/collection/src/shards/resharding.rs b/lib/collection/src/shards/resharding.rs index a1979902e3..71bb89475d 100644 --- a/lib/collection/src/shards/resharding.rs +++ b/lib/collection/src/shards/resharding.rs @@ -98,6 +98,23 @@ pub struct ReshardKey { pub shard_key: Option, } +impl ReshardKey { + /// Pin the auto-sharding invariant that resharding always targets the + /// last shard id. `shard_number` must equal `shard_id` or `shard_id + 1`; + /// any other value would silently corrupt the count via the id-derived + /// `shard_number` update. + pub(crate) fn debug_assert_targets_last_shard(&self, shard_number: u32) { + let shard_id = self.shard_id; + debug_assert!( + shard_id + .checked_add(1) + .is_some_and(|next| shard_number == shard_id || shard_number == next), + "auto-sharding resharding must target the last shard id; \ + shard_number={shard_number} shard_id={shard_id}", + ); + } +} + impl fmt::Display for ReshardKey { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{}/{}/{:?}", self.peer_id, self.shard_id, self.shard_key) diff --git a/lib/collection/src/shards/shard_holder/mod.rs b/lib/collection/src/shards/shard_holder/mod.rs index 34e56c12fb..bd286030fb 100644 --- a/lib/collection/src/shards/shard_holder/mod.rs +++ b/lib/collection/src/shards/shard_holder/mod.rs @@ -204,8 +204,12 @@ impl ShardHolder { shard_id: ShardId, shard_key: &ShardKey, ) -> CollectionResult<()> { + // Idempotent: only rewrite the mapping if the shard id is actually + // present under this key. A replay after a successful removal finds + // nothing to remove and skips the unnecessary disk write. self.key_mapping.write_optional(|key_mapping| { - if !key_mapping.contains_key(shard_key) { + let shard_ids = key_mapping.get(shard_key)?; + if !shard_ids.contains(&shard_id) { return None; } diff --git a/lib/collection/src/shards/shard_holder/resharding.rs b/lib/collection/src/shards/shard_holder/resharding.rs index 0f85775672..afabf7ef81 100644 --- a/lib/collection/src/shards/shard_holder/resharding.rs +++ b/lib/collection/src/shards/shard_holder/resharding.rs @@ -61,10 +61,13 @@ impl ShardHolder { assert_resharding_state_consistency(&state, ring, shard_key); if let Some(state) = state.deref() { + // Idempotent: a matching state means start_resharding was + // already applied on this peer. Caller falls through each + // step to reconcile any partially-applied state, so we must + // not short-circuit with an error that would be silently + // swallowed by apply_entries and cause divergence. return if state.matches(resharding_key) { - Err(CollectionError::bad_request(format!( - "resharding {resharding_key} is already in progress:\n{state:#?}" - ))) + Ok(()) } else { Err(CollectionError::bad_request(format!( "another resharding is in progress:\n{state:#?}" @@ -93,15 +96,14 @@ impl ShardHolder { match resharding_key.direction { ReshardingDirection::Up => { if has_shard { - // Allow re-application: the shard may exist as a leftover from a - // previous incomplete start attempt (e.g. crash after key_mapping - // was persisted but before resharding_state was written). - // create_shard_dir will clean up the stale directory and add_shard - // will evict the old entry. + // Idempotent: the replica set may already exist as a + // leftover from a previous partial apply or as the result + // of an earlier successful apply we are replaying. The + // caller skips create_replica_set in that case and falls + // through; we must not block on this condition. log::warn!( "Shard {shard_id} already exists during resharding start, \ - likely a leftover from a previous incomplete attempt, \ - it will be recreated" + treating as idempotent re-apply and reusing existing replica set", ); } } @@ -157,15 +159,15 @@ impl ShardHolder { .await?; } - self.resharding_state.write(|state| { - debug_assert!( - state.is_none(), - "resharding is already in progress:\n{state:#?}", - ); - - *state = Some(ReshardState::new( - uuid, direction, peer_id, shard_id, shard_key, - )); + // Idempotent: if matching state is already persisted (replay after a + // successful apply), leave it alone. Only write if missing or stale. + self.resharding_state.write_optional(|state| { + let new_state = + ReshardState::new(uuid, direction, peer_id, shard_id, shard_key.clone()); + match state { + Some(existing) if *existing == new_state => None, + _ => Some(Some(new_state)), + } })?; Ok(()) @@ -269,10 +271,13 @@ impl ShardHolder { } pub fn finish_resharding_unchecked(&mut self, _: &ReshardKey) -> CollectionResult<()> { - self.resharding_state.write(|state| { - debug_assert!(state.is_some(), "resharding is not in progress"); - *state = None; - })?; + // Idempotent: if state is already cleared (replay after a successful + // finish/abort), leave it alone so we don't spuriously touch the file. + self.resharding_state.write_optional( + |state| { + if state.is_some() { Some(None) } else { None } + }, + )?; Ok(()) } @@ -442,8 +447,13 @@ impl ShardHolder { // Drop the shard if let Some(shard_key) = shard_key { + // Idempotent: only rewrite the mapping if the shard id is + // actually present under this key. A replay after a + // successful abort finds nothing to remove and skips the + // unnecessary disk write. self.key_mapping.write_optional(|key_mapping| { - if !key_mapping.contains_key(shard_key) { + let shard_ids = key_mapping.get(shard_key)?; + if !shard_ids.contains(&shard_id) { return None; } @@ -464,15 +474,12 @@ impl ShardHolder { } } - self.resharding_state.write(|state| { - debug_assert!( - state - .as_ref() - .is_some_and(|state| state.matches(&resharding_key)), - "resharding {resharding_key} is not in progress:\n{state:#?}" - ); - - state.take(); + // Idempotent: if state is already cleared, or holds a different + // resharding (already superseded), leave it alone. Only clear a state + // that matches the resharding_key we are aborting. + self.resharding_state.write_optional(|state| match state { + Some(state) if state.matches(&resharding_key) => Some(None), + _ => None, })?; Ok(()) @@ -904,6 +911,109 @@ mod tests { ); } + // ------------------------------------------------------------------ + // check_start_resharding idempotency + // ------------------------------------------------------------------ + + #[test] + fn test_start_check_idempotent_when_matching_resharding_active() { + let (_dir, mut holder) = make_holder(); + let key = make_reshard_key(1); + + // Set up ring and resharding state as they would look after a + // successful `start_resharding` had been applied on this peer. + holder + .rings + .get_mut(&None) + .unwrap() + .start_resharding(key.shard_id, key.direction); + holder + .resharding_state + .write(|state| { + *state = Some(ReshardState::new( + key.uuid, + key.direction, + key.peer_id, + key.shard_id, + key.shard_key.clone(), + )); + }) + .unwrap(); + + // Re-applying start for the same key must not fail. Otherwise the + // replay would be silently swallowed and the subsequent idempotent + // steps in `Collection::start_resharding` would be skipped — leaving + // half-applied state unreconciled on this peer. + let result = holder.check_start_resharding(&key); + assert!( + result.is_ok(), + "check_start_resharding must return Ok when the same resharding is already active (idempotent re-apply), got error: {result:?}", + ); + } + + // ------------------------------------------------------------------ + // finish_resharding_unchecked idempotency + // ------------------------------------------------------------------ + + #[test] + fn test_finish_unchecked_idempotent_when_no_resharding_active() { + let (_dir, mut holder) = make_holder(); + let key = make_reshard_key(1); + + // Replay `finish` after the resharding state has already been cleared + // (partial apply between the state write and a subsequent step). + // The unchecked helper must not panic or error in this case. + let result = holder.finish_resharding_unchecked(&key); + assert!( + result.is_ok(), + "finish_resharding_unchecked must return Ok when no resharding state is present, got error: {result:?}", + ); + assert!( + holder.resharding_state.read().is_none(), + "state must remain cleared after idempotent finish", + ); + } + + // ------------------------------------------------------------------ + // start_resharding_unchecked idempotency + // ------------------------------------------------------------------ + + #[tokio::test] + async fn test_start_unchecked_idempotent_when_matching_state_present() { + let (_dir, mut holder) = make_holder(); + let key = make_reshard_key(1); + + let expected_state = ReshardState::new( + key.uuid, + key.direction, + key.peer_id, + key.shard_id, + key.shard_key.clone(), + ); + + // Pre-seed matching state to simulate a replay after a successful + // start had already persisted the resharding state. + holder + .resharding_state + .write(|state| { + *state = Some(expected_state.clone()); + }) + .unwrap(); + + // Re-apply. `new_shard = None` matches the caller's behavior when it + // has detected that the replica set already exists. + holder + .start_resharding_unchecked(key.clone(), None) + .await + .expect("start_resharding_unchecked must be idempotent on replay"); + + assert_eq!( + holder.resharding_state.read().clone(), + Some(expected_state), + "matching state must be preserved verbatim across a replay", + ); + } + // ------------------------------------------------------------------ // Divergence scenario: proves the bug end-to-end // ------------------------------------------------------------------ diff --git a/tests/consensus_tests/test_resharding_replay.py b/tests/consensus_tests/test_resharding_replay.py new file mode 100644 index 0000000000..dd9072fd04 --- /dev/null +++ b/tests/consensus_tests/test_resharding_replay.py @@ -0,0 +1,88 @@ +"""Resharding start replay reconciles a partial-apply window where the +config.json save crashed after resharding_state.json was persisted.""" + +import json +import pathlib + +import requests + +from .assertions import assert_http_ok +from .fixtures import create_collection +from .utils import ( + get_cluster_info, + get_collection_cluster_info, + processes, + start_cluster, + start_peer, + wait_collection_exists_and_active_on_all_peers, + wait_for, + wait_for_collection_resharding_operations_count, + wait_for_peer_online, +) + +COLLECTION = "test_resharding_replay" + + +def test_start_resharding_replay_reconciles_pre_config_save_crash(tmp_path: pathlib.Path): + peer_uris, peer_dirs, bootstrap_uri = start_cluster(tmp_path, 3) + + target_idx = 2 + target_dir = peer_dirs[target_idx] + target_peer_id = get_cluster_info(peer_uris[target_idx])["peer_id"] + + create_collection(peer_uris[0], collection=COLLECTION, shard_number=2, replication_factor=1) + wait_collection_exists_and_active_on_all_peers(COLLECTION, peer_uris) + + raft_state = target_dir / "storage" / "raft_state.json" + config = target_dir / "storage" / "collections" / COLLECTION / "config.json" + new_shard_dir = target_dir / "storage" / "collections" / COLLECTION / "2" + + commit_before = json.loads(raft_state.read_text())["state"]["hard_state"]["commit"] + + resp = requests.post( + f"{peer_uris[0]}/collections/{COLLECTION}/cluster", + json={"start_resharding": {"direction": "up", "peer_id": target_peer_id, "shard_key": None}}, + ) + assert_http_ok(resp) + wait_for_collection_resharding_operations_count(peer_uris[target_idx], COLLECTION, 1) + wait_for(lambda: new_shard_dir.exists() + and json.loads(config.read_text())["params"]["shard_number"] == 3) + + p = processes.pop(target_idx) + restart_port = p.p2p_port + p.kill() + + # Simulate a crash between resharding_state.json and config.json saves. + cfg = json.loads(config.read_text()) + cfg["params"]["shard_number"] = 2 + config.write_text(json.dumps(cfg)) + + # Force replay of the start_resharding entry. + state = json.loads(raft_state.read_text()) + state["apply_progress_queue"] = [commit_before + 1, state["state"]["hard_state"]["commit"]] + raft_state.write_text(json.dumps(state)) + + peer_uris[target_idx] = start_peer(target_dir, f"peer_replay_{target_idx}.log", bootstrap_uri, port=restart_port) + new_uri = peer_uris[target_idx] + wait_for_peer_online(new_uri) + + # Pre-fix: matching-state early return short-circuits replay; shard 2 never + # gets re-registered on this peer. Post-fix: replay falls through and + # reconciles via contains_shard(2)=false → create_replica_set + config save. + def has_new_shard(): + try: + info = get_collection_cluster_info(new_uri, COLLECTION) + except requests.RequestException: + return False + return any(s["shard_id"] == 2 for s in info.get("local_shards", [])) + + wait_for(has_new_shard) + assert json.loads(config.read_text())["params"]["shard_number"] == 3 + + resp = requests.post( + f"{peer_uris[0]}/collections/{COLLECTION}/cluster", + json={"abort_resharding": {}}, + ) + assert_http_ok(resp) + for uri in peer_uris: + wait_for_collection_resharding_operations_count(uri, COLLECTION, 0)