upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-07-23 11:11:00 -05:00
Code Issues Packages Projects Releases Wiki Activity

Make resharding operations (shard holder) idempotent (#8789)

Browse Source 
* Make resharding state transitions idempotent on replay

Why: consensus entries may be re-applied after a crash (partial state
on disk) or during raft recovery. The unchecked state-transition
helpers used `debug_assert!` to require a specific starting state, so
a replay would panic in debug or silently overwrite in release.

How to apply: use write_optional so the state file is only touched
when the in-memory state needs to change. This also avoids unnecessary
fsyncs when a replay is a no-op.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Skip replica set creation on resharding start replay

Why: on replay, `create_replica_set` would call `create_shard_dir`,
which removes and recreates the existing shard directory -- wiping
any shard contents that had already been migrated or written since
the first apply.

How to apply: check `contains_shard(shard_id)` before creating the
replica set and pass `None` to `start_resharding_unchecked` when a
replica set with the target shard id is already present. Also
relaxes `check_start_resharding` to return `Ok` (instead of a
swallowed `bad_request`) when a matching resharding state is
persisted, so the caller can fall through each idempotent step.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Derive resharding shard count from shard id for idempotency

Why: the start/finish/abort paths used ++/-- on the persisted shard
number with a `debug_assert` pinning the expected starting value. On
replay (e.g. after a crash between the shard holder mutation and the
config save) this either panics or produces a wrong count.

Since resharding always targets the last shard (auto sharding assigns
contiguous ids from zero), the target count is a pure function of the
shard id: `shard_id + 1` for start up, `shard_id` for finish down and
abort up. Set it directly and skip the save when it already matches,
so replay converges to the same value without touching the file.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Fall through resharding finish/abort steps on replay

Why: the early return on `resharding_state.is_none()` assumed that a
missing state means the operation was fully applied. But a crash can
leave the state cleared while the shard drop, key mapping removal, or
shard count update are still pending. Short-circuiting then skips the
reconciling work those replays need to do.

How to apply: drop the early return so every step runs; each step is
already individually idempotent (check_*, drop_and_remove_shard,
remove_shard_from_key_mapping, the set-based shard count update).
The abort_resharding down-invalidation path now reads nodes from the
router regardless of its variant, so a replay over an already-rolled-
back ring doesn't trip the removed debug_assert!s.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Skip shard key mapping write when nothing to remove

Why: `write_optional` returns `Some` whenever the shard key is
present, even if the shard id we're removing is already gone. That
still triggers a JSON rewrite and a cache notification on every
replay of a finished finish/abort.

How to apply: check that the shard id is actually in the set before
returning `Some`. If the set is missing or already doesn't contain
the id, return `None` so the on-disk file and in-memory data are
left untouched.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Add tests for resharding replay idempotency

Covers the paths that must not error or panic on replay:
- `check_start_resharding` when matching state is already present
- `start_resharding_unchecked` preserves matching state verbatim
- `finish_resharding_unchecked` when state is already cleared

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Simplify finish_resharding_unchecked closure

Replace a manual `match` on `Option` with `as_ref().map(...)` to
satisfy clippy::manual_map.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* test: add a test for resharding replay (#8839)

* [ai] add a debug_assert on shard_number (#8846)

* Update lib/collection/src/shards/shard_holder/resharding.rs

Co-authored-by: Roman Titov <ffuugoo@users.noreply.github.com>

* Reformat

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: tellet-q <166374656+tellet-q@users.noreply.github.com>
Co-authored-by: Roman Titov <ffuugoo@users.noreply.github.com>
This commit is contained in:
Tim Visée
2026-04-30 14:20:40 +02:00
committed by timvisee
parent cd2b86491c
commit 963bf32b5d
5 changed files with 326 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

150
lib/collection/src/collection/resharding.rs
View File

@@ -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

17
lib/collection/src/shards/resharding.rs
View File

@@ -98,6 +98,23 @@ pub struct ReshardKey {
pub shard_key: Option<ShardKey>,
}
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)

6
lib/collection/src/shards/shard_holder/mod.rs
View File

@@ -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;
}

176
lib/collection/src/shards/shard_holder/resharding.rs
View File

@@ -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
// ------------------------------------------------------------------

88
tests/consensus_tests/test_resharding_replay.py Normal file
View File

@@ -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)