Clear shard data before snapshot recovery transfer (#8782)

* [ai] On shard snapshot transfer recovery, drop existing shard before recovery

* [ai] Add integration test to assert clearing behavior

* [ai] Debug assert our replica is not active when we clear it

* [ai] Tweak assertion

* Fix flaky test, replica may temporarily not be visible

* Replace debug assertion with runtime error
This commit is contained in:
Tim Visée
2026-04-30 11:26:01 +02:00
committed by timvisee
parent 4bf6f4d183
commit 33baf70038
6 changed files with 283 additions and 4 deletions

View File

@@ -374,6 +374,23 @@ impl Collection {
.assert_shard_exists(shard_id)
}
/// Drop the local shard and clear its on-disk data, before a shard snapshot
/// transfer downloads a replacement snapshot. See
/// [`ShardReplicaSet::clear_local_for_snapshot_recovery`] for details and safety
/// constraints.
pub async fn clear_local_shard_for_snapshot_recovery(
&self,
shard_id: ShardId,
) -> CollectionResult<()> {
self.shards_holder
.read()
.await
.get_shard(shard_id)
.ok_or_else(|| shard_not_found_error(shard_id))?
.clear_local_for_snapshot_recovery(&self.path)
.await
}
pub async fn try_take_partial_snapshot_recovery_lock(
&self,
shard_id: ShardId,

View File

@@ -330,6 +330,57 @@ impl ShardReplicaSet {
}
}
/// Drop the in-memory local shard (if any) and remove its on-disk data files.
///
/// Used by shard snapshot transfers to free disk space before the receiving node
/// downloads the new snapshot, avoiding having both the old data and the incoming
/// snapshot on disk at the same time. The configuration files and replica state
/// are preserved, so the shard directory remains a valid (empty) shard.
///
/// Only safe to call while the shard is in a state that prevents user requests
/// (`PartialSnapshot` during a shard transfer). Do NOT call this from a
/// user-triggered URL recovery path, where the shard may still be serving queries.
///
/// Writes the shard initializing flag before clearing, so that a crash between
/// here and the end of the subsequent `restore_local_replica_from` call causes
/// the shard to be loaded as a dummy on startup and re-recovered.
pub async fn clear_local_for_snapshot_recovery(
&self,
collection_path: &Path,
) -> CollectionResult<()> {
// Callers must only invoke this while the shard is in a state that cannot
// be a source of truth (e.g. `PartialSnapshot` during a shard transfer).
// Clearing a source-of-truth replica would silently drop data that may
// still be serving queries.
if self
.peer_state(self.this_peer_id())
.is_some_and(|s| s.can_be_source_of_truth())
{
return Err(CollectionError::service_error(format!(
"clear_local_for_snapshot_recovery called on a peer that can be source-of-truth {}:{}",
self.collection_id, self.shard_id,
)));
}
let mut local = self.local.write().await;
// Mark the shard as initializing before touching disk, so a crash during or
// after clearing is detected on next startup and the shard is reloaded as a
// dummy that triggers recovery.
let shard_flag = shard_initializing_flag_path(collection_path, self.shard_id);
let flag_file = tokio_fs::File::create(&shard_flag).await?;
flag_file.sync_all().await?;
sync_parent_dir_async(&shard_flag).await?;
if let Some(shard) = local.take() {
shard.stop_gracefully().await;
}
LocalShard::clear(&self.shard_path).await?;
Ok(())
}
pub async fn get_partial_snapshot_manifest(&self) -> CollectionResult<SnapshotManifest> {
self.local
.read()

View File

@@ -194,6 +194,17 @@ pub async fn recover_shard_snapshot(
.await
.start_shard_recovery(shard_id);
// For shard transfers, drop the existing shard and clear its on-disk data
// before downloading the new snapshot so we don't need space for both copies.
// Safe because the shard is in `PartialSnapshot` state for the duration of
// the transfer and will not serve user requests. Not done for user-triggered
// URL recovery, where the shard may still be active.
if matches!(snapshot_priority, SnapshotPriority::ShardTransfer) {
collection
.clear_local_shard_for_snapshot_recovery(shard_id)
.await?;
}
let download_task = async {
let DownloadResult {
snapshot,

View File

@@ -344,7 +344,8 @@ def test_dirty_shard_handling_with_active_replicas(tmp_path: pathlib.Path, trans
wait_for_collection_shard_transfers_count(peer_api_uris[-1], COLLECTION_NAME, 0)
# Wait for all replicas to be active on the receiving peer
wait_for_all_replicas_active(peer_api_uris[-1], COLLECTION_NAME)
# Must have at least one replica on the last peer (replica may temporarily disappear when peer is restarted during snapshot recovery)
wait_for_all_replicas_active(peer_api_uris[-1], COLLECTION_NAME, min_local_replicas=1)
# Assert that the local shard is active and not empty
try:

View File

@@ -0,0 +1,197 @@
import pathlib
from time import sleep
from .fixtures import create_collection, upsert_random_points
from .utils import * # includes `time` module, `requests`, `processes`, etc.
# The receiver now pre-clears its shard data on disk before downloading the
# snapshot (for SnapshotPriority::ShardTransfer only). These tests pin down the
# externally-observable behavior of that flow: the cluster keeps serving
# queries while the receiver is in `PartialSnapshot`, and a mid-transfer crash
# of the receiver heals via the dummy-shard → dead → retransfer path.
N_PEERS = 3
N_SHARDS = 3
N_REPLICA = 2
COLLECTION_NAME = "test_collection"
def _count_points(peer_url):
r = requests.post(
f"{peer_url}/collections/{COLLECTION_NAME}/points/count",
json={"exact": True},
)
assert_http_ok(r)
return r.json()["result"]["count"]
def _pick_transfer_endpoints(peer_api_uris):
src = get_collection_cluster_info(peer_api_uris[0], COLLECTION_NAME)
dst = get_collection_cluster_info(peer_api_uris[2], COLLECTION_NAME)
dst_shard_ids = {s["shard_id"] for s in dst["local_shards"]}
# Pick a shard that lives on peer 0 but NOT on peer 2, so replicate_shard
# actually sends a snapshot across rather than being a no-op.
src_shard = next(
s for s in src["local_shards"] if s["shard_id"] not in dst_shard_ids
)
return src["peer_id"], dst["peer_id"], src_shard["shard_id"]
def _start_snapshot_replicate(peer_api_uris, from_peer_id, to_peer_id, shard_id):
r = requests.post(
f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster",
json={
"replicate_shard": {
"shard_id": shard_id,
"from_peer_id": from_peer_id,
"to_peer_id": to_peer_id,
"method": "snapshot",
}
},
)
assert_http_ok(r)
# While a shard snapshot transfer is in progress, the other replicas must keep
# serving read traffic — the receiver's shard is being cleared + refilled on
# disk, but consensus has it in `PartialSnapshot` so requests should route to
# healthy replicas. Verifies the happy path end-to-end with consistency.
def test_cluster_queryable_during_shard_snapshot_transfer(tmp_path: pathlib.Path):
assert_project_root()
peer_api_uris, _peer_dirs, _bootstrap_uri = start_cluster(
tmp_path, N_PEERS, 25000
)
create_collection(
peer_api_uris[0], shard_number=N_SHARDS, replication_factor=N_REPLICA
)
wait_collection_exists_and_active_on_all_peers(
collection_name=COLLECTION_NAME, peer_api_uris=peer_api_uris
)
upsert_random_points(peer_api_uris[0], 3000)
baseline = _count_points(peer_api_uris[0])
assert baseline == 3000
from_peer_id, to_peer_id, shard_id = _pick_transfer_endpoints(peer_api_uris)
_start_snapshot_replicate(peer_api_uris, from_peer_id, to_peer_id, shard_id)
# Poll until the transfer is finished, continuously querying every peer.
# Every count request must succeed and return the full count — with 2
# replicas per shard, the receiver going through PartialSnapshot must not
# impact availability.
saw_transfer_in_progress = False
deadline = time.time() + 30
while time.time() < deadline:
info = get_collection_cluster_info(peer_api_uris[0], COLLECTION_NAME)
transfers = info.get("shard_transfers", [])
if transfers:
saw_transfer_in_progress = True
for uri in peer_api_uris:
assert _count_points(uri) == baseline
if not transfers and saw_transfer_in_progress:
break
sleep(RETRY_INTERVAL_SEC)
assert saw_transfer_in_progress, (
"Expected to observe at least one in-progress transfer"
)
wait_for_collection_shard_transfers_count(peer_api_uris[0], COLLECTION_NAME, 0)
wait_for_all_replicas_active(peer_api_uris[0], COLLECTION_NAME)
for uri in peer_api_uris:
assert _count_points(uri) == baseline
# The destination now carries the extra replica — 2 original shards + the
# replicated one.
dst_info = get_collection_cluster_info(peer_api_uris[2], COLLECTION_NAME)
assert len(dst_info["local_shards"]) == 3
assert any(
s["shard_id"] == shard_id and s["state"] == "Active"
for s in dst_info["local_shards"]
)
# Kill the receiving node while a snapshot transfer is in progress — deliberately
# after the pre-clear has had time to run. On restart the receiver must surface
# the shard as dirty (via the init flag written before clearing), the consensus
# layer must abort the interrupted transfer and recover the replica, and the
# final cluster must be consistent.
def test_receiver_restart_during_shard_snapshot_transfer(tmp_path: pathlib.Path):
assert_project_root()
peer_api_uris, peer_dirs, bootstrap_uri = start_cluster(
tmp_path, N_PEERS, 25500
)
create_collection(
peer_api_uris[0], shard_number=N_SHARDS, replication_factor=N_REPLICA
)
wait_collection_exists_and_active_on_all_peers(
collection_name=COLLECTION_NAME, peer_api_uris=peer_api_uris
)
# Enough points that the snapshot transfer is not instantaneous, so we
# have a window to crash the receiver after pre-clear has started.
upsert_random_points(peer_api_uris[0], 20000, batch_size=500)
baseline = _count_points(peer_api_uris[0])
assert baseline == 20000
from_peer_id, to_peer_id, shard_id = _pick_transfer_endpoints(peer_api_uris)
# Capture ports/dir of the receiver so we can restart at the same address
# and keep the same peer id. `start_peer` lays ports out as port+0 (p2p),
# port+1 (grpc), port+2 (http), so we pass the p2p_port as the seed.
receiver_proc = processes[2]
receiver_port = receiver_proc.p2p_port
receiver_dir = pathlib.Path(peer_dirs[2])
receiver_url = peer_api_uris[2]
_start_snapshot_replicate(peer_api_uris, from_peer_id, to_peer_id, shard_id)
# Wait until the transfer is actually active from the sender's POV, then
# crash the receiver. This ensures the clear-and-download flow has been
# entered on the receiver (or is about to be).
wait_for_collection_shard_transfers_count(peer_api_uris[0], COLLECTION_NAME, 1)
sleep(1.0)
receiver_proc.kill()
processes.remove(receiver_proc)
# Restart the receiver at the same HTTP/gRPC/p2p ports so it rejoins the
# cluster with the same peer id.
restarted_url = start_peer(
receiver_dir,
"peer_2_restarted.log",
bootstrap_uri,
port=receiver_port,
)
assert restarted_url == receiver_url
wait_for_peer_online(restarted_url)
# The sender's transfer attempt either completed before the kill or was
# aborted by it. Either way, the cluster should converge: any stray
# transfer drains, the dummy-shard-on-restart triggers recovery of the
# cleared replica, and eventually all replicas end up Active again.
# Give the cluster a generous budget for this — snapshot recovery on
# a dirty shard goes through a full retransfer.
wait_for(
lambda: get_collection_cluster_info(
peer_api_uris[0], COLLECTION_NAME
).get("shard_transfers", []) == [],
wait_for_timeout=90,
)
wait_for_all_replicas_active(
peer_api_uris[0], COLLECTION_NAME
)
# Data on every live peer (including the restarted one) must match the
# baseline. The point count check implicitly verifies each peer can route
# to a healthy replica of every shard.
for uri in [peer_api_uris[0], peer_api_uris[1], restarted_url]:
assert _count_points(uri) == baseline, (
f"Point count mismatch at {uri}"
)

View File

@@ -631,9 +631,11 @@ def check_collection_shard_transfer_progress(peer_api_uri: str, collection_name:
return False
def check_all_replicas_active(peer_api_uri: str, collection_name: str, headers={}) -> bool:
def check_all_replicas_active(peer_api_uri: str, collection_name: str, headers={}, min_local_replicas=0) -> bool:
try:
collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
if len(collection_cluster_info["local_shards"]) < min_local_replicas:
return False
for shard in collection_cluster_info["local_shards"]:
if shard['state'] != 'Active':
return False
@@ -703,9 +705,9 @@ def wait_for_some_replicas_not_active(peer_api_uri: str, collection_name: str):
raise e
def wait_for_all_replicas_active(peer_api_uri: str, collection_name: str, headers={}):
def wait_for_all_replicas_active(peer_api_uri: str, collection_name: str, headers={}, min_local_replicas=0):
try:
wait_for(check_all_replicas_active, peer_api_uri, collection_name, headers=headers)
wait_for(check_all_replicas_active, peer_api_uri, collection_name, headers=headers, min_local_replicas=min_local_replicas)
except Exception as e:
print_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
raise e