diff --git a/lib/collection/src/collection/collection_ops.rs b/lib/collection/src/collection/collection_ops.rs index 138d2d9b92..433f062bf9 100644 --- a/lib/collection/src/collection/collection_ops.rs +++ b/lib/collection/src/collection/collection_ops.rs @@ -213,7 +213,7 @@ impl Collection { /// Handle replica changes /// - /// add and remove replicas from replica set + /// Remove replicas from replica set pub async fn handle_replica_changes( &self, replica_changes: Vec, @@ -223,17 +223,17 @@ impl Collection { } let shard_holder = self.shards_holder.read().await; + let mut to_remove = Vec::with_capacity(replica_changes.len()); for change in replica_changes { let (shard_id, peer_id) = match change { Change::Remove(shard_id, peer_id) => (shard_id, peer_id), }; - let Some(replica_set) = shard_holder.get_shard(shard_id) else { + let Some(replica_set) = shard_holder.get_shard(shard_id).cloned() else { return Err(CollectionError::bad_request(format!( - "Shard {} of {} not found", - shard_id, - self.name() + "Shard {shard_id} of {} not found", + self.name(), ))); }; @@ -268,9 +268,15 @@ impl Collection { .get_transfers(|transfer| transfer.from == peer_id || transfer.to == peer_id) }; - // ...and cancel transfer tasks and remove transfers from internal state + to_remove.push((replica_set, peer_id, transfers)); + } + + // Must release shard holder lock for abort_shard_transfer_and_resharding + drop(shard_holder); + + for (replica_set, peer_id, transfers) in to_remove { for transfer in transfers { - self.abort_shard_transfer_and_resharding(transfer.key(), Some(&shard_holder)) + self.abort_shard_transfer_and_resharding(transfer.key()) .await?; } @@ -283,6 +289,7 @@ impl Collection { // the transfer should be cancelled (see the block right above this comment), // so no special handling is needed. } + Ok(()) } diff --git a/lib/collection/src/collection/mod.rs b/lib/collection/src/collection/mod.rs index c5c1fee072..41ed88e447 100644 --- a/lib/collection/src/collection/mod.rs +++ b/lib/collection/src/collection/mod.rs @@ -402,8 +402,8 @@ impl Collection { new_state: ReplicaState, from_state: Option, ) -> CollectionResult<()> { - let shard_holder = self.shards_holder.read().await; - let replica_set = shard_holder + let mut shard_holder = self.shards_holder.read().await; + let mut replica_set = shard_holder .get_shard(shard_id) .ok_or_else(|| shard_not_found_error(shard_id))?; @@ -451,14 +451,46 @@ impl Collection { ))); } + // Abort resharding *before* persisting the new replica state. + // + // If we did this after the persist, a crash between the durable + // `replica_state.json` write and `abort_resharding` would leave + // `resharding_state.json` stuck `Some` on the replaying peer (a + // `current_state`-based gate reads `Dead` from disk on replay and + // skips the abort). Doing it first means a partial-apply crash + // either re-runs the whole entry (replica state still reads as a + // resharding state, gate fires, the idempotent abort re-runs) or no + // longer needs the abort to fire at all (resharding_state already + // cleared by the prior attempt). Either way, every peer converges. + // + // Covers both directions: up (`Resharding`) and down (`ReshardingScaleDown`). + // For up, `abort_resharding` drops the shard entirely; we detect the + // missing shard below and skip the persist. For down, the shard + // stays, abort just reverts the receiver back to `Active`, and the + // persist below overwrites that with `Dead`. + if new_state == ReplicaState::Dead && current_state.is_some_and(|s| s.is_resharding()) { + let resharding_state = shard_holder.resharding_state.read().clone(); + if let Some(state) = resharding_state { + // `abort_resharding` grabs the shard_holder write lock internally. + drop(shard_holder); + self.abort_resharding(state.key(), false).await?; + // For up direction, abort dropped the shard. If it's gone, the + // replica we'd be deactivating no longer exists on this peer + // and the rest of the function is a no-op. + shard_holder = self.shards_holder.read().await; + let Some(rs) = shard_holder.get_shard(shard_id) else { + return Ok(()); + }; + replica_set = rs; + } + } + // Update replica status replica_set .ensure_replica_with_state(peer_id, new_state) .await?; if new_state == ReplicaState::Dead { - let resharding_state = shard_holder.resharding_state.read().clone(); - let all_nodes_fixed_cancellation = self .channel_service .all_peers_at_version(&ABORT_TRANSFERS_ON_SHARD_DROP_FIX_FROM_VERSION); @@ -476,36 +508,11 @@ impl Collection { // Functions below lock `shard_holder`! drop(shard_holder); - let mut abort_resharding_result = CollectionResult::Ok(()); - - // Abort resharding, if resharding shard is marked as `Dead`. - // - // This branch should only be triggered, if resharding is currently at `MigratingPoints` - // stage, because target shard should be marked as `Active`, when all resharding transfers - // are successfully completed, and so the check *right above* this one would be triggered. - // - // So, if resharding reached `ReadHashRingCommitted`, this branch *won't* be triggered, - // and resharding *won't* be cancelled. The update request should *fail* with "failed to - // update all replicas of a shard" error. - // - // If resharding reached `ReadHashRingCommitted`, and this branch is triggered *somehow*, - // then `Collection::abort_resharding` call should return an error, so no special handling - // is needed. - let is_resharding = current_state - .as_ref() - .is_some_and(ReplicaState::is_resharding); - if is_resharding && let Some(state) = resharding_state { - abort_resharding_result = self.abort_resharding(state.key(), false).await; - } - // Terminate transfer if source or target replicas are now dead for transfer in related_transfers { - self.abort_shard_transfer_and_resharding(transfer.key(), None) + self.abort_shard_transfer_and_resharding(transfer.key()) .await?; } - - // Propagate resharding errors now - abort_resharding_result?; } // If not initialized yet, we need to check if it was initialized by this call @@ -668,7 +675,7 @@ impl Collection { } for transfer in self.get_related_transfers(peer_id).await { - self.abort_shard_transfer_and_resharding(transfer.key(), None) + self.abort_shard_transfer_and_resharding(transfer.key()) .await?; } diff --git a/lib/collection/src/collection/shard_transfer.rs b/lib/collection/src/collection/shard_transfer.rs index 9cef748a28..a187cd7dd7 100644 --- a/lib/collection/src/collection/shard_transfer.rs +++ b/lib/collection/src/collection/shard_transfer.rs @@ -436,34 +436,35 @@ impl Collection { pub async fn abort_shard_transfer_and_resharding( &self, transfer_key: ShardTransferKey, - shard_holder: Option<&ShardHolder>, ) -> CollectionResult<()> { - let mut shard_holder_guard = None; + // Look up transfer and any resharding state we need to abort + let resharding_state = { + let shard_holder = self.shards_holder.read().await; - let shard_holder = match shard_holder { - Some(shard_holder) => shard_holder, - None => shard_holder_guard.insert(self.shards_holder.read().await), + let Some(transfer) = shard_holder.get_transfer(&transfer_key) else { + return Ok(()); + }; + + if transfer.is_resharding() { + shard_holder.resharding_state.read().clone() + } else { + None + } }; + // Abort resharding before the transfer to be idempotent + if let Some(state) = resharding_state { + self.abort_resharding(state.key(), false).await?; + } + + // Resharding may already have aborted the transfer so we check it again + let shard_holder = self.shards_holder.read().await; + let Some(transfer) = shard_holder.get_transfer(&transfer_key) else { return Ok(()); }; - let is_resharding_transfer = transfer.is_resharding(); - self.abort_shard_transfer(transfer, shard_holder).await?; - - if is_resharding_transfer { - let resharding_state = shard_holder.resharding_state.read().clone(); - - // `abort_resharding` locks `shard_holder`! - drop(shard_holder_guard); - - if let Some(state) = resharding_state { - self.abort_resharding(state.key(), false).await?; - } - } - - Ok(()) + self.abort_shard_transfer(transfer, &shard_holder).await } /// Initiate local partial shard diff --git a/lib/collection/src/shards/replica_set/update.rs b/lib/collection/src/shards/replica_set/update.rs index 76753ea32f..beb10c6f1c 100644 --- a/lib/collection/src/shards/replica_set/update.rs +++ b/lib/collection/src/shards/replica_set/update.rs @@ -743,9 +743,21 @@ impl ShardReplicaSet { self.shard_id, ); - // Deactivate replica in consensus if it matches the state we expect - // Always deactivate the replica if its in a shard transfer related state - let from_state = Some(peer_state).filter(|state| !state.is_partial_or_recovery()); + // Deactivate replica in consensus if it matches the state we expect. + // We filter out transient transfer states (Partial, Recovery, etc.) + // because those can change rapidly between proposal and apply, and + // we want the deactivation to go through even if the state has + // moved on. Resharding/ReshardingScaleDown are stable through the + // resharding lifecycle (only flipped by finish_migrating_points + // → Active or by this same deactivation flow → Dead), so we keep + // them populated. Carrying the captured pre-state in the consensus + // log entry is what lets the local `abort_resharding` side-effect + // in `Collection::set_shard_replica_state` re-fire correctly on + // replay — without it, replay reads `current_state` as Dead from + // disk and silently skips the side-effect, leaving + // resharding_state stale. + let from_state = Some(peer_state) + .filter(|state| !state.is_partial_or_recovery() || state.is_resharding()); self.add_locally_disabled(Some(state), *peer_id, from_state); } diff --git a/lib/storage/src/content_manager/toc/collection_meta_ops.rs b/lib/storage/src/content_manager/toc/collection_meta_ops.rs index e88f128e62..ddc6e0ec1d 100644 --- a/lib/storage/src/content_manager/toc/collection_meta_ops.rs +++ b/lib/storage/src/content_manager/toc/collection_meta_ops.rs @@ -576,7 +576,7 @@ impl TableOfContent { } ShardTransferOperations::RecoveryToPartial(transfer) | ShardTransferOperations::SnapshotRecovered(transfer) => { - // Validate transfer exists to prevent double handling + // Validate transfer exists transfer::helpers::validate_transfer_exists( &transfer, &collection.state().await.transfers, @@ -635,7 +635,7 @@ impl TableOfContent { )?; log::warn!("Aborting shard transfer: {reason}"); collection - .abort_shard_transfer_and_resharding(transfer, None) + .abort_shard_transfer_and_resharding(transfer) .await?; } }; diff --git a/tests/consensus_tests/fixtures.py b/tests/consensus_tests/fixtures.py index 987276bbba..1929c85dc3 100644 --- a/tests/consensus_tests/fixtures.py +++ b/tests/consensus_tests/fixtures.py @@ -96,6 +96,7 @@ def upsert_random_points( num_cities=None, headers={}, extra_payload=None, + timeout=None, ): extra_payload = extra_payload or {} @@ -129,6 +130,7 @@ def upsert_random_points( "shard_key": shard_key, }, headers=headers, + timeout=timeout, ) if fail_on_error: assert_http_ok(r_batch) diff --git a/tests/consensus_tests/test_resharding_replay.py b/tests/consensus_tests/test_resharding_replay.py index dd9072fd04..d35756e106 100644 --- a/tests/consensus_tests/test_resharding_replay.py +++ b/tests/consensus_tests/test_resharding_replay.py @@ -8,17 +8,7 @@ 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, -) +from .utils import * COLLECTION = "test_resharding_replay" diff --git a/tests/consensus_tests/test_resharding_set_replica_dead_replay.py b/tests/consensus_tests/test_resharding_set_replica_dead_replay.py new file mode 100644 index 0000000000..c730ab300c --- /dev/null +++ b/tests/consensus_tests/test_resharding_set_replica_dead_replay.py @@ -0,0 +1,170 @@ +"""SetShardReplicaState(Resharding -> Dead) must clear resharding_state on every peer. + +Reorder fix in `Collection::set_shard_replica_state`: when this entry deactivates +a `Resharding` replica, the local `abort_resharding` runs *before* the durable +replica_state write. With the old order (persist first, abort after), a crash +between the two left `resharding_state` stuck `Some` on the replaying peer; the +replay's gate read `current_state` back as `Dead` from disk and silently +skipped the abort. With the reorder, that intermediate state is unreachable — +either the abort completes (and resharding_state is cleared) before the persist +runs, or the entry hasn't been "applied" from consensus' point of view and +re-runs from the top on replay. + +The test exercises the natural failure-detection path: kill a peer holding a +`Resharding` replica, drive updates that fail on it, and verify every peer's +`resharding_operations` clears — including target after it catches up. +""" + +import json +import pathlib + +import requests + +from .assertions import assert_http_ok +from .fixtures import create_collection, upsert_random_points +from .utils import * + +COLLECTION = "test_resharding_set_replica_dead_replay" + + +def _all_replicas(info): + for local in info["local_shards"]: + yield {**local, "peer_id": info["peer_id"]} + for remote in info["remote_shards"]: + yield remote + + +def test_set_replica_dead_clears_resharding_state(tmp_path: pathlib.Path): + peer_uris, peer_dirs, bootstrap_uri = start_cluster(tmp_path, 3) + + target_idx = 2 + target_dir = peer_dirs[target_idx] + peer_ids = [get_cluster_info(uri)["peer_id"] for uri in peer_uris] + target_peer_id = peer_ids[target_idx] + + # 2 shards (RF=1). Resharding up adds shard 2 on the target peer in + # `Resharding` state. We then drive a resharding-stream-records transfer + # from an existing shard to shard 2 on a non-target peer so shard 2 has + # *two* `Resharding` replicas — without the second replica, the + # failure-detection path's "last active peer" guard refuses to propose + # `SetShardReplicaState(Dead)` and the deactivation never happens. + create_collection(peer_uris[0], collection=COLLECTION, shard_number=2, replication_factor=1) + wait_collection_exists_and_active_on_all_peers(COLLECTION, peer_uris) + + new_shard_dir = target_dir / "storage" / "collections" / COLLECTION / "2" + new_replica_state = new_shard_dir / "replica_state.json" + + 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) + + for uri in peer_uris: + wait_for_collection_resharding_operations_count(uri, COLLECTION, 1) + wait_for(lambda: new_replica_state.exists()) + wait_for( + lambda: json.loads(new_replica_state.read_text())["peers"].get(str(target_peer_id)) + == "Resharding" + ) + + # Pick a transfer source on a non-target peer (auto-distribution may put + # shard 0 on the target peer). Stream from that peer's existing shard + # into shard 2 on the third peer, giving shard 2 a second Resharding + # replica there. + info = get_collection_cluster_info(peer_uris[0], COLLECTION) + source_idx = source_peer_id = source_shard = None + for r in _all_replicas(info): + replica_idx = peer_ids.index(r["peer_id"]) + if replica_idx == target_idx: + continue + source_idx = replica_idx + source_peer_id = r["peer_id"] + source_shard = r["shard_id"] + break + assert source_idx is not None, "no non-target peer holds an existing shard" + extra_idx = next(i for i in range(3) if i != target_idx and i != source_idx) + extra_peer_id = peer_ids[extra_idx] + driver_uri = peer_uris[source_idx] + + resp = requests.post( + f"{peer_uris[0]}/collections/{COLLECTION}/cluster", + json={ + "replicate_shard": { + "from_peer_id": source_peer_id, + "to_peer_id": extra_peer_id, + "shard_id": source_shard, + "to_shard_id": 2, + "method": "resharding_stream_records", + } + }, + ) + assert_http_ok(resp) + wait_for_collection_shard_transfers_count(peer_uris[0], COLLECTION, 0) + + # Kill target. Updates routed to shard 2 will fail on target's Resharding + # replica; the second peer's Resharding replica keeps the shard "live" + # in active_or_resharding_peers so the failure handler is permitted to + # propose `SetShardReplicaState(Resharding -> Dead)` for target. + p = processes.pop(target_idx) + restart_port = p.p2p_port + p.kill() + + # Running peers apply `SetShardReplicaState(Dead)`; the local + # `abort_resharding` runs first (per the reorder), clearing + # `resharding_state` before the replica state write. + # + # With weak ordering the upsert returns as soon as the local replica + # acks, leaving the fan-out failure to target's dead replica to land + # asynchronously — a single batch sometimes returns 200 without ever + # observing the dead target. Drive a fresh batch on each poll until + # every running peer has applied the deactivation and cleared its + # `resharding_state`. + # + # Cap each upsert at 5s so a stuck gRPC channel to the dead peer can't + # block a single iteration past the wait_for budget — without this, + # `requests.put` has no timeout and a hung fan-out leaves wait_for + # unable to iterate at all (test hangs until pytest's outer deadline + # rather than failing within 60s). + running_uris = [uri for idx, uri in enumerate(peer_uris) if idx != target_idx] + + def all_running_peers_aborted_resharding() -> bool: + try: + upsert_random_points( + driver_uri, + num=100, + collection_name=COLLECTION, + fail_on_error=False, + ordering="weak", + timeout=5, + ) + except requests.RequestException: + pass + return all( + not get_collection_cluster_info(uri, COLLECTION).get("resharding_operations") + for uri in running_uris + ) + + wait_for( + all_running_peers_aborted_resharding, + wait_for_timeout=60, + wait_for_interval=1, + ) + + # Restart target. It catches up via raft, applies the same entry, and + # the same reorder makes it clear `resharding_state` and drop shard 2 + # on this peer too. + peer_uris[target_idx] = start_peer( + target_dir, f"peer_catchup_{target_idx}.log", bootstrap_uri, port=restart_port + ) + catchup_uri = peer_uris[target_idx] + wait_for_peer_online(catchup_uri) + wait_for_collection_resharding_operations_count(catchup_uri, COLLECTION, 0) + + # All peers must agree — no divergence in `resharding_operations`. + for uri in peer_uris: + info = get_collection_cluster_info(uri, COLLECTION) + assert not info.get("resharding_operations"), ( + f"Resharding state divergence on peer {uri}: " + f"resharding_operations={info.get('resharding_operations')!r}" + )