upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-08-06 10:00:58 -05:00
Code Issues Packages Projects Releases Wiki Activity

Prefer owned versus referenced usage of PeerId/ShardId, they are Copy (#5344)

Browse Source
This commit is contained in:
Tim Visée
2024-10-31 15:44:49 +01:00
committed by timvisee
parent 4904e623a5
commit 1128ac8ff8
23 changed files with 153 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
lib/collection/src/collection/collection_ops.rs
View File

@@ -198,7 +198,7 @@ impl Collection {
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) else {
return Err(CollectionError::BadRequest {
description: format!("Shard {} of {} not found", shard_id, self.name()),
});
@@ -333,7 +333,6 @@ impl Collection {
// extract shards info
for (shard_id, replica_set) in shards_holder.get_shards() {
let shard_id = *shard_id;
let peers = replica_set.peers();
if replica_set.has_local_shard().await {

40
lib/collection/src/collection/mod.rs
View File

@@ -329,7 +329,7 @@ impl Collection {
}
pub async fn contains_shard(&self, shard_id: ShardId) -> bool {
self.shards_holder.read().await.contains_shard(&shard_id)
self.shards_holder.read().await.contains_shard(shard_id)
}
pub async fn wait_local_shard_replica_state(
@@ -340,7 +340,7 @@ impl Collection {
) -> CollectionResult<()> {
let shard_holder_read = self.shards_holder.read().await;
let shard = shard_holder_read.get_shard(&shard_id);
let shard = shard_holder_read.get_shard(shard_id);
let Some(replica_set) = shard else {
return Err(CollectionError::NotFound {
what: "Shard {shard_id}".into(),
@@ -359,16 +359,16 @@ impl Collection {
) -> CollectionResult<()> {
let shard_holder = self.shards_holder.read().await;
let replica_set = shard_holder
.get_shard(&shard_id)
.get_shard(shard_id)
.ok_or_else(|| shard_not_found_error(shard_id))?;
log::debug!(
"Changing shard {}:{shard_id} replica state from {:?} to {state:?}",
self.id,
replica_set.peer_state(&peer_id),
replica_set.peer_state(peer_id),
);
let current_state = replica_set.peer_state(&peer_id);
let current_state = replica_set.peer_state(peer_id);
// Validation:
//
@@ -380,7 +380,7 @@ impl Collection {
.read()
.contains_key(&peer_id);
let replica_exists = replica_set.peer_state(&peer_id).is_some();
let replica_exists = replica_set.peer_state(peer_id).is_some();
if !peer_exists && !replica_exists {
return Err(CollectionError::bad_input(format!(
@@ -443,14 +443,14 @@ impl Collection {
}
replica_set
.ensure_replica_with_state(&peer_id, state)
.ensure_replica_with_state(peer_id, state)
.await?;
if state == ReplicaState::Dead {
// TODO(resharding): Abort all resharding transfers!?
// Terminate transfer if source or target replicas are now dead
let related_transfers = shard_holder.get_related_transfers(&shard_id, &peer_id);
let related_transfers = shard_holder.get_related_transfers(shard_id, peer_id);
// `abort_shard_transfer` locks `shard_holder`!
drop(shard_holder);
@@ -485,7 +485,7 @@ impl Collection {
pub async fn shard_recovery_point(&self, shard_id: ShardId) -> CollectionResult<RecoveryPoint> {
let shard_holder_read = self.shards_holder.read().await;
let shard = shard_holder_read.get_shard(&shard_id);
let shard = shard_holder_read.get_shard(shard_id);
let Some(replica_set) = shard else {
return Err(CollectionError::NotFound {
what: format!("Shard {shard_id}"),
@@ -502,7 +502,7 @@ impl Collection {
) -> CollectionResult<()> {
let shard_holder_read = self.shards_holder.read().await;
let shard = shard_holder_read.get_shard(&shard_id);
let shard = shard_holder_read.get_shard(shard_id);
let Some(replica_set) = shard else {
return Err(CollectionError::NotFound {
what: "Shard {shard_id}".into(),
@@ -524,7 +524,7 @@ impl Collection {
let shard_info = ShardInfo {
replicas: replicas.peers(),
};
(*shard_id, shard_info)
(shard_id, shard_info)
})
.collect(),
resharding,
@@ -579,7 +579,7 @@ impl Collection {
}
// Check for un-reported finished transfers
let outgoing_transfers = shard_holder.get_outgoing_transfers(&self.this_peer_id);
let outgoing_transfers = shard_holder.get_outgoing_transfers(self.this_peer_id);
let tasks_lock = self.transfer_tasks.lock().await;
for transfer in outgoing_transfers {
match tasks_lock
@@ -619,29 +619,29 @@ impl Collection {
// Check for proper replica states
for replica_set in shard_holder.all_shards() {
let this_peer_id = &replica_set.this_peer_id();
let this_peer_id = replica_set.this_peer_id();
let shard_id = replica_set.shard_id;
let peers = replica_set.peers();
let this_peer_state = peers.get(this_peer_id).copied();
let this_peer_state = peers.get(&this_peer_id).copied();
let is_last_active = peers.values().filter(|state| **state == Active).count() == 1;
if this_peer_state == Some(Initializing) {
// It is possible, that collection creation didn't report
// Try to activate shard, as the collection clearly exists
on_finish_init(*this_peer_id, shard_id);
on_finish_init(this_peer_id, shard_id);
continue;
}
if self.shared_storage_config.node_type == NodeType::Listener {
if this_peer_state == Some(Active) && !is_last_active {
// Convert active node from active to listener
on_convert_to_listener(*this_peer_id, shard_id);
on_convert_to_listener(this_peer_id, shard_id);
continue;
}
} else if this_peer_state == Some(Listener) {
// Convert listener node to active
on_convert_from_listener(*this_peer_id, shard_id);
on_convert_from_listener(this_peer_id, shard_id);
continue;
}
@@ -654,7 +654,7 @@ impl Collection {
// Respect shard transfer limit, consider already proposed transfers in our counts
let (mut incoming, outgoing) = shard_holder.count_shard_transfer_io(this_peer_id);
incoming += proposed.get(this_peer_id).copied().unwrap_or(0);
incoming += proposed.get(&this_peer_id).copied().unwrap_or(0);
if self.check_auto_shard_transfer_limit(incoming, outgoing) {
log::trace!("Postponing automatic shard {shard_id} transfer to stay below limit on this node (incoming: {incoming}, outgoing: {outgoing})");
continue;
@@ -680,7 +680,7 @@ impl Collection {
for replica_id in replica_set.active_remote_shards().await {
let transfer = ShardTransfer {
from: replica_id,
to: *this_peer_id,
to: this_peer_id,
shard_id,
to_shard_id: None,
sync: true,
@@ -693,7 +693,7 @@ impl Collection {
}
// Respect shard transfer limit, consider already proposed transfers in our counts
let (incoming, mut outgoing) = shard_holder.count_shard_transfer_io(&replica_id);
let (incoming, mut outgoing) = shard_holder.count_shard_transfer_io(replica_id);
outgoing += proposed.get(&replica_id).copied().unwrap_or(0);
if self.check_auto_shard_transfer_limit(incoming, outgoing) {
log::trace!("Postponing automatic shard {shard_id} transfer to stay below limit on peer {replica_id} (incoming: {incoming}, outgoing: {outgoing})");

2
lib/collection/src/collection/point_ops.rs
View File

@@ -92,7 +92,7 @@ impl Collection {
let result = tokio::task::spawn(async move {
let _update_lock = update_lock;
let Some(shard) = shard_holder.get_shard(&shard_selection) else {
let Some(shard) = shard_holder.get_shard(shard_selection) else {
return Ok(None);
};

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

@@ -194,7 +194,7 @@ impl Collection {
if resharding_key.direction == ReshardingDirection::Down {
// Remove the shard we've now migrated all points out of
if let Some(shard_key) = &resharding_key.shard_key {
shard_holder.remove_shard_from_key_mapping(&resharding_key.shard_id, shard_key)?;
shard_holder.remove_shard_from_key_mapping(resharding_key.shard_id, shard_key)?;
}
shard_holder
.drop_and_remove_shard(resharding_key.shard_id)

28
lib/collection/src/collection/shard_transfer.rs
View File

@@ -19,7 +19,7 @@ use crate::shards::transfer::{
};
impl Collection {
pub async fn get_outgoing_transfers(&self, current_peer_id: &PeerId) -> Vec<ShardTransfer> {
pub async fn get_outgoing_transfers(&self, current_peer_id: PeerId) -> Vec<ShardTransfer> {
self.shards_holder
.read()
.await
@@ -67,10 +67,10 @@ impl Collection {
.unwrap_or(shard_transfer.shard_id);
let shards_holder = self.shards_holder.read().await;
let from_replica_set = shards_holder.get_shard(&from_shard_id).ok_or_else(|| {
let from_replica_set = shards_holder.get_shard(from_shard_id).ok_or_else(|| {
CollectionError::service_error(format!("Shard {from_shard_id} doesn't exist"))
})?;
let to_replica_set = shards_holder.get_shard(&to_shard_id).ok_or_else(|| {
let to_replica_set = shards_holder.get_shard(to_shard_id).ok_or_else(|| {
CollectionError::service_error(format!("Shard {to_shard_id} doesn't exist"))
})?;
let _was_not_transferred =
@@ -112,7 +112,7 @@ impl Collection {
debug_assert!(old_shard.is_none(), "We should not have a local shard yet");
} else {
to_replica_set
.ensure_replica_with_state(&shard_transfer.to, initial_state)
.ensure_replica_with_state(shard_transfer.to, initial_state)
.await?;
}
@@ -206,10 +206,10 @@ impl Collection {
let mut is_dest_replica_active = false;
let dest_replica_set =
shard_holder.get_shard(&transfer.to_shard_id.unwrap_or(transfer.shard_id));
shard_holder.get_shard(transfer.to_shard_id.unwrap_or(transfer.shard_id));
if let Some(replica_set) = dest_replica_set {
if replica_set.peer_state(&transfer.to).is_some() {
if replica_set.peer_state(transfer.to).is_some() {
// Promote *destination* replica/shard to `Active` if:
//
// - replica *exists*
@@ -227,7 +227,7 @@ impl Collection {
};
if transfer.to == self.this_peer_id {
replica_set.set_replica_state(&transfer.to, state)?;
replica_set.set_replica_state(transfer.to, state)?;
} else {
replica_set.add_remote(transfer.to, state).await?;
}
@@ -237,7 +237,7 @@ impl Collection {
}
// Handle *source* replica
let src_replica_set = shard_holder.get_shard(&transfer.shard_id);
let src_replica_set = shard_holder.get_shard(transfer.shard_id);
if let Some(replica_set) = src_replica_set {
if transfer.sync || is_resharding_transfer {
@@ -303,8 +303,8 @@ impl Collection {
let shard_id = transfer_key.to_shard_id.unwrap_or(transfer_key.shard_id);
if let Some(replica_set) = shard_holder.get_shard(&shard_id) {
if replica_set.peer_state(&transfer.to).is_some() {
if let Some(replica_set) = shard_holder.get_shard(shard_id) {
if replica_set.peer_state(transfer.to).is_some() {
if is_resharding_transfer {
// If *resharding* shard transfer failed, we don't need/want to change replica state:
// - on transfer failure, the whole resharding would be aborted (see below),
@@ -319,7 +319,7 @@ impl Collection {
// and so failed transfer does not introduce any inconsistencies to points
// that are not affected by resharding in all other shards
} else if transfer.sync {
replica_set.set_replica_state(&transfer.to, ReplicaState::Dead)?;
replica_set.set_replica_state(transfer.to, ReplicaState::Dead)?;
} else {
replica_set.remove_peer(transfer.to).await?;
}
@@ -362,7 +362,7 @@ impl Collection {
async move {
let shards_holder = shards_holder.read_owned().await;
let Some(replica_set) = shards_holder.get_shard(&shard_id) else {
let Some(replica_set) = shards_holder.get_shard(shard_id) else {
return Err(CollectionError::service_error(format!(
"Shard {shard_id} doesn't exist, repartition is not supported yet"
)));
@@ -390,7 +390,7 @@ impl Collection {
// We can guarantee that replica_set is not None, cause we checked it before
// and `shards_holder` is holding the lock.
// This is a workaround for lifetime checker.
let replica_set = shards_holder.get_shard(&shard_id).unwrap();
let replica_set = shards_holder.get_shard(shard_id).unwrap();
let shard_transfer_registered = shards_holder.shard_transfers.wait_for(
|shard_transfers| {
shard_transfers.iter().any(|shard_transfer| {
@@ -409,7 +409,7 @@ impl Collection {
&& replica_set.wait_for_state_condition_sync(
|state| {
state
.get_peer_state(&this_peer_id)
.get_peer_state(this_peer_id)
.map_or(false, |peer_state| peer_state.is_partial_or_recovery())
},
defaults::CONSENSUS_META_OP_WAIT,

4
lib/collection/src/collection/snapshots.rs
View File

@@ -100,7 +100,7 @@ impl Collection {
// Create snapshot of each shard
for (shard_id, replica_set) in shards_holder.get_shards() {
let shard_snapshot_path =
shard_versioning::versioned_shard_path(Path::new(""), *shard_id, 0);
shard_versioning::versioned_shard_path(Path::new(""), shard_id, 0);
// If node is listener, we can save whatever currently is in the storage
let save_wal = self.shared_storage_config.node_type != NodeType::Listener;
@@ -284,7 +284,7 @@ impl Collection {
) -> CollectionResult<SnapshotStream> {
let shard = OwnedRwLockReadGuard::try_map(
Arc::clone(&self.shards_holder).read_owned().await,
|x| x.get_shard(&shard_id),
|x| x.get_shard(shard_id),
)
.map_err(|_| shard_not_found_error(shard_id))?;

2
lib/collection/src/collection/state_management.rs
View File

@@ -96,7 +96,7 @@ impl Collection {
// and create new shards if needed
for (shard_id, shard_info) in shards {
match self.shards_holder.read().await.get_shard(&shard_id) {
match self.shards_holder.read().await.get_shard(shard_id) {
Some(replica_set) => replica_set.apply_state(shard_info.replicas).await?,
None => {
let shard_replicas: Vec<_> = shard_info.replicas.keys().copied().collect();

10
lib/collection/src/shards/replica_set/execute_read_operation.rs
View File

@@ -55,8 +55,8 @@ impl ShardReplicaSet {
let read_consistency = read_consistency.unwrap_or_default();
let local_count = usize::from(self.peer_state(&self.this_peer_id()).is_some());
let active_local_count = usize::from(self.peer_is_active(&self.this_peer_id()));
let local_count = usize::from(self.peer_state(self.this_peer_id()).is_some());
let active_local_count = usize::from(self.peer_is_active(self.this_peer_id()));
let remotes = self.remotes.read().await;
@@ -65,7 +65,7 @@ impl ShardReplicaSet {
// TODO(resharding): Handle resharded shard?
let active_remotes_count = remotes
.iter()
.filter(|remote| self.peer_is_active(&remote.peer_id))
.filter(|remote| self.peer_is_active(remote.peer_id))
.count();
let total_count = local_count + remotes_count;
@@ -159,7 +159,7 @@ impl ShardReplicaSet {
Err(_) => (self.local.read().right_future(), false, None),
};
let local_is_active = self.peer_is_active(&self.this_peer_id());
let local_is_active = self.peer_is_active(self.this_peer_id());
let local_operation = if local_is_active {
let local_operation = async {
@@ -183,7 +183,7 @@ impl ShardReplicaSet {
// TODO(resharding): Handle resharded shard?
let mut active_remotes: Vec<_> = remotes
.iter()
.filter(|remote| self.peer_is_active(&remote.peer_id))
.filter(|remote| self.peer_is_active(remote.peer_id))
.collect();
active_remotes.shuffle(&mut rand::thread_rng());

46
lib/collection/src/shards/replica_set/mod.rs
View File

@@ -236,7 +236,7 @@ impl ShardReplicaSet {
replica_state
.write(|rs| {
let this_peer_id = rs.this_peer_id;
let local_state = rs.remove_peer_state(&this_peer_id);
let local_state = rs.remove_peer_state(this_peer_id);
if let Some(state) = local_state {
rs.set_peer_state(this_peer_id, state);
}
@@ -366,7 +366,7 @@ impl ShardReplicaSet {
active_peers.len() == 1 && active_peers.contains(&peer_id)
}
pub fn peer_state(&self, peer_id: &PeerId) -> Option<ReplicaState> {
pub fn peer_state(&self, peer_id: PeerId) -> Option<ReplicaState> {
self.replica_state.read().get_peer_state(peer_id).copied()
}
@@ -376,7 +376,7 @@ impl ShardReplicaSet {
replica_state
.active_peers()
.into_iter()
.filter(|peer_id| !self.is_locally_disabled(peer_id))
.filter(|&peer_id| !self.is_locally_disabled(peer_id))
.collect()
}
@@ -387,7 +387,7 @@ impl ShardReplicaSet {
replica_state
.active_peers()
.into_iter()
.filter(|peer_id| !self.is_locally_disabled(peer_id) && *peer_id != this_peer_id)
.filter(|&peer_id| !self.is_locally_disabled(peer_id) && peer_id != this_peer_id)
.collect()
}
@@ -417,7 +417,7 @@ impl ShardReplicaSet {
) -> CollectionResult<()> {
self.wait_for(
move |replica_set_state| {
replica_set_state.get_peer_state(&replica_set_state.this_peer_id) == Some(&state)
replica_set_state.get_peer_state(replica_set_state.this_peer_id) == Some(&state)
},
timeout,
)
@@ -438,7 +438,7 @@ impl ShardReplicaSet {
timeout: Duration,
) -> CollectionResult<()> {
self.wait_for(
move |replica_set_state| replica_set_state.get_peer_state(&peer_id) == Some(&state),
move |replica_set_state| replica_set_state.get_peer_state(peer_id) == Some(&state),
timeout,
)
.await
@@ -537,7 +537,7 @@ impl ShardReplicaSet {
self.replica_state.write(|rs| {
rs.is_local = false;
let this_peer_id = rs.this_peer_id;
rs.remove_peer_state(&this_peer_id);
rs.remove_peer_state(this_peer_id);
})?;
self.update_locally_disabled(self.this_peer_id());
@@ -583,7 +583,7 @@ impl ShardReplicaSet {
pub async fn remove_remote(&self, peer_id: PeerId) -> CollectionResult<()> {
self.replica_state.write(|rs| {
rs.remove_peer_state(&peer_id);
rs.remove_peer_state(peer_id);
})?;
self.update_locally_disabled(peer_id);
@@ -597,19 +597,19 @@ impl ShardReplicaSet {
/// Ensure that remote shard is initialized.
pub async fn ensure_replica_with_state(
&self,
peer_id: &PeerId,
peer_id: PeerId,
state: ReplicaState,
) -> CollectionResult<()> {
if *peer_id == self.this_peer_id() {
if peer_id == self.this_peer_id() {
self.set_replica_state(peer_id, state)?;
} else {
// Create remote shard if necessary
self.add_remote(*peer_id, state).await?;
self.add_remote(peer_id, state).await?;
}
Ok(())
}
pub fn set_replica_state(&self, peer_id: &PeerId, state: ReplicaState) -> CollectionResult<()> {
pub fn set_replica_state(&self, peer_id: PeerId, state: ReplicaState) -> CollectionResult<()> {
log::debug!(
"Changing local shard {}:{} state from {:?} to {state:?}",
self.collection_id,
@@ -618,12 +618,12 @@ impl ShardReplicaSet {
);
self.replica_state.write(|rs| {
if rs.this_peer_id == *peer_id {
if rs.this_peer_id == peer_id {
rs.is_local = true;
}
rs.set_peer_state(*peer_id, state);
rs.set_peer_state(peer_id, state);
})?;
self.update_locally_disabled(*peer_id);
self.update_locally_disabled(peer_id);
Ok(())
}
@@ -877,11 +877,11 @@ impl ShardReplicaSet {
/// Check whether a peer is registered as `active`.
/// Unknown peers are not active.
fn peer_is_active(&self, peer_id: &PeerId) -> bool {
fn peer_is_active(&self, peer_id: PeerId) -> bool {
self.peer_state(peer_id) == Some(ReplicaState::Active) && !self.is_locally_disabled(peer_id)
}
fn peer_is_active_or_resharding(&self, peer_id: &PeerId) -> bool {
fn peer_is_active_or_resharding(&self, peer_id: PeerId) -> bool {
let is_active_or_resharding = matches!(
self.peer_state(peer_id),
Some(ReplicaState::Active | ReplicaState::Resharding)
@@ -892,8 +892,8 @@ impl ShardReplicaSet {
is_active_or_resharding && !is_locally_disabled
}
fn is_locally_disabled(&self, peer_id: &PeerId) -> bool {
self.locally_disabled_peers.read().is_disabled(*peer_id)
fn is_locally_disabled(&self, peer_id: PeerId) -> bool {
self.locally_disabled_peers.read().is_disabled(peer_id)
}
/// Locally disable given peer
@@ -1007,16 +1007,16 @@ pub struct ReplicaSetState {
}
impl ReplicaSetState {
pub fn get_peer_state(&self, peer_id: &PeerId) -> Option<&ReplicaState> {
self.peers.get(peer_id)
pub fn get_peer_state(&self, peer_id: PeerId) -> Option<&ReplicaState> {
self.peers.get(&peer_id)
}
pub fn set_peer_state(&mut self, peer_id: PeerId, state: ReplicaState) {
self.peers.insert(peer_id, state);
}
pub fn remove_peer_state(&mut self, peer_id: &PeerId) -> Option<ReplicaState> {
self.peers.remove(peer_id)
pub fn remove_peer_state(&mut self, peer_id: PeerId) -> Option<ReplicaState> {
self.peers.remove(&peer_id)
}
pub fn peers(&self) -> HashMap<PeerId, ReplicaState> {

30
lib/collection/src/shards/replica_set/update.rs
View File

@@ -37,7 +37,7 @@ impl ShardReplicaSet {
let local = self.local.read().await;
if let Some(local_shard) = local.deref() {
match self.peer_state(&self.this_peer_id()) {
match self.peer_state(self.this_peer_id()) {
Some(
ReplicaState::Active
| ReplicaState::Partial
@@ -129,7 +129,7 @@ impl ShardReplicaSet {
peer_ids
.into_iter()
.filter(|peer_id| self.peer_is_active_or_resharding(peer_id)) // re-acquire replica_state read lock
.filter(|&peer_id| self.peer_is_active_or_resharding(peer_id)) // re-acquire replica_state read lock
.max()
}
@@ -207,11 +207,11 @@ impl ShardReplicaSet {
// Target all remote peers that can receive updates
let updatable_remote_shards: Vec<_> = remotes
.iter()
.filter(|rs| self.is_peer_updatable(&rs.peer_id))
.filter(|rs| self.is_peer_updatable(rs.peer_id))
.collect();
// Local is defined and can receive updates
let local_is_updatable = local.is_some() && self.is_peer_updatable(&this_peer_id);
let local_is_updatable = local.is_some() && self.is_peer_updatable(this_peer_id);
if updatable_remote_shards.is_empty() && !local_is_updatable {
return Err(CollectionError::service_error(format!(
@@ -227,8 +227,8 @@ impl ShardReplicaSet {
let mut update_futures = Vec::with_capacity(updatable_remote_shards.len() + 1);
if let Some(local) = local.deref() {
if self.is_peer_updatable(&this_peer_id) {
let local_wait = if self.peer_state(&this_peer_id) == Some(ReplicaState::Listener) {
if self.is_peer_updatable(this_peer_id) {
let local_wait = if self.peer_state(this_peer_id) == Some(ReplicaState::Listener) {
false
} else {
wait
@@ -376,7 +376,7 @@ impl ShardReplicaSet {
self.handle_failed_replicas(
failures
.iter()
.filter(|(peer_id, _)| self.peer_is_resharding(peer_id)),
.filter(|(peer_id, _)| self.peer_is_resharding(*peer_id)),
&self.replica_state.read(),
update_only_existing,
);
@@ -388,7 +388,7 @@ impl ShardReplicaSet {
if !successes
.iter()
.any(|(peer_id, _)| self.peer_is_active_or_resharding(peer_id))
.any(|&(peer_id, _)| self.peer_is_active_or_resharding(peer_id))
{
return Err(CollectionError::service_error(format!(
"Failed to apply operation to at least one `Active` replica. \
@@ -416,7 +416,7 @@ impl ShardReplicaSet {
/// Whether to send updates to the given peer
///
/// A peer in dead state, or a locally disabled peer, will not accept updates.
fn is_peer_updatable(&self, peer_id: &PeerId) -> bool {
fn is_peer_updatable(&self, peer_id: PeerId) -> bool {
let res = match self.peer_state(peer_id) {
Some(ReplicaState::Active) => true,
Some(ReplicaState::Partial) => true,
@@ -434,7 +434,7 @@ impl ShardReplicaSet {
res && !self.is_locally_disabled(peer_id)
}
fn peer_is_resharding(&self, peer_id: &PeerId) -> bool {
fn peer_is_resharding(&self, peer_id: PeerId) -> bool {
self.peer_state(peer_id) == Some(ReplicaState::Resharding)
&& !self.is_locally_disabled(peer_id)
}
@@ -454,7 +454,7 @@ impl ShardReplicaSet {
self.shard_id,
);
let Some(&peer_state) = state.get_peer_state(peer_id) else {
let Some(&peer_state) = state.get_peer_state(*peer_id) else {
continue;
};
@@ -557,10 +557,10 @@ mod tests {
// at build time the replicas are all dead, they need to be activated
assert_eq!(rs.highest_alive_replica_peer_id(), None);
rs.set_replica_state(&1, ReplicaState::Active).unwrap();
rs.set_replica_state(&3, ReplicaState::Active).unwrap();
rs.set_replica_state(&4, ReplicaState::Active).unwrap();
rs.set_replica_state(&5, ReplicaState::Partial).unwrap();
rs.set_replica_state(1, ReplicaState::Active).unwrap();
rs.set_replica_state(3, ReplicaState::Active).unwrap();
rs.set_replica_state(4, ReplicaState::Active).unwrap();
rs.set_replica_state(5, ReplicaState::Partial).unwrap();
assert_eq!(rs.highest_replica_peer_id(), Some(5));
assert_eq!(rs.highest_alive_replica_peer_id(), Some(4));

23
lib/collection/src/shards/resharding/stage_migrate_points.rs
View File

@@ -132,12 +132,11 @@ async fn drive_up(
let source_peer_ids = {
let shard_holder = shard_holder.read().await;
let replica_set =
shard_holder.get_shard(&source_shard_id).ok_or_else(|| {
CollectionError::service_error(format!(
"Shard {source_shard_id} not found in the shard holder for resharding",
))
})?;
let replica_set = shard_holder.get_shard(source_shard_id).ok_or_else(|| {
CollectionError::service_error(format!(
"Shard {source_shard_id} not found in the shard holder for resharding",
))
})?;
let active_peer_ids = replica_set.active_shards().await;
if active_peer_ids.is_empty() {
@@ -147,13 +146,13 @@ async fn drive_up(
}
// Respect shard transfer limits, always allow local transfers
let (incoming, _) = shard_holder.count_shard_transfer_io(&this_peer_id);
let (incoming, _) = shard_holder.count_shard_transfer_io(this_peer_id);
if incoming < incoming_limit {
active_peer_ids
.into_iter()
.filter(|peer_id| {
.filter(|&peer_id| {
let (_, outgoing) = shard_holder.count_shard_transfer_io(peer_id);
outgoing < outgoing_limit || peer_id == &this_peer_id
outgoing < outgoing_limit || peer_id == this_peer_id
})
.collect()
} else if active_peer_ids.contains(&this_peer_id) {
@@ -299,14 +298,14 @@ async fn drive_down(
let source_replica_set =
shard_holder
.get_shard(&reshard_key.shard_id)
.get_shard(reshard_key.shard_id)
.ok_or_else(|| {
CollectionError::service_error(format!(
"Shard {} not found in the shard holder for resharding",
reshard_key.shard_id,
))
})?;
let target_replica_set = shard_holder.get_shard(&target_shard_id).ok_or_else(|| {
let target_replica_set = shard_holder.get_shard(target_shard_id).ok_or_else(|| {
CollectionError::service_error(format!(
"Shard {target_shard_id} not found in the shard holder for resharding",
))
@@ -415,7 +414,7 @@ async fn migrate_local(
// Normally consensus takes care of this, but we don't use consensus here
{
let shard_holder = shard_holder.read().await;
let replica_set = shard_holder.get_shard(&source_shard_id).ok_or_else(|| {
let replica_set = shard_holder.get_shard(source_shard_id).ok_or_else(|| {
CollectionError::service_error(format!(
"Shard {source_shard_id} not found in the shard holder for resharding",
))

2
lib/collection/src/shards/resharding/stage_propagate_deletes.rs
View File

@@ -55,7 +55,7 @@ pub(super) async fn drive(
loop {
let shard_holder = shard_holder.read().await;
let replica_set = shard_holder.get_shard(&source_shard_id).ok_or_else(|| {
let replica_set = shard_holder.get_shard(source_shard_id).ok_or_else(|| {
CollectionError::service_error(format!(
"Shard {source_shard_id} not found in the shard holder for resharding",
))

10
lib/collection/src/shards/resharding/stage_replicate.rs
View File

@@ -65,7 +65,7 @@ pub(super) async fn drive(
// Ensure we don't exceed the outgoing transfer limits
let shard_holder = shard_holder.read().await;
let (_, outgoing) = shard_holder.count_shard_transfer_io(&this_peer_id);
let (_, outgoing) = shard_holder.count_shard_transfer_io(this_peer_id);
if outgoing >= outgoing_limit {
log::trace!("Postponing resharding replication transfer to stay below transfer limit (outgoing: {outgoing})");
sleep(SHARD_TRANSFER_IO_LIMIT_RETRY_INTERVAL).await;
@@ -73,7 +73,7 @@ pub(super) async fn drive(
}
// Select peers that don't have this replica yet
let Some(replica_set) = shard_holder.get_shard(&reshard_key.shard_id) else {
let Some(replica_set) = shard_holder.get_shard(reshard_key.shard_id) else {
return Err(CollectionError::service_error(format!(
"Shard {} not found in the shard holder for resharding",
reshard_key.shard_id,
@@ -99,9 +99,9 @@ pub(super) async fn drive(
.unwrap();
let candidate_peers: Vec<_> = candidate_peers
.into_iter()
.filter(|(peer_id, shard_count)| {
.filter(|&(peer_id, shard_count)| {
let (incoming, _) = shard_holder.count_shard_transfer_io(peer_id);
lowest_shard_count == *shard_count && incoming < incoming_limit
lowest_shard_count == shard_count && incoming < incoming_limit
})
.map(|(peer_id, _)| peer_id)
.collect();
@@ -188,7 +188,7 @@ async fn has_enough_replicas(
.get();
let current_replication_factor = {
let shard_holder_read = shard_holder.read().await;
let Some(replica_set) = shard_holder_read.get_shard(&reshard_key.shard_id) else {
let Some(replica_set) = shard_holder_read.get_shard(reshard_key.shard_id) else {
return Err(CollectionError::service_error(format!(
"Shard {} not found in the shard holder for resharding",
reshard_key.shard_id,

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

@@ -146,7 +146,7 @@ impl ShardHolder {
pub fn remove_shard_from_key_mapping(
&mut self,
shard_id: &ShardId,
shard_id: ShardId,
shard_key: &ShardKey,
) -> Result<(), CollectionError> {
self.key_mapping.write_optional(|key_mapping| {
@@ -155,10 +155,10 @@ impl ShardHolder {
}
let mut key_mapping = key_mapping.clone();
key_mapping.get_mut(shard_key).unwrap().remove(shard_id);
key_mapping.get_mut(shard_key).unwrap().remove(&shard_id);
Some(key_mapping)
})?;
self.shard_id_to_key_mapping.remove(shard_id);
self.shard_id_to_key_mapping.remove(&shard_id);
Ok(())
}
@@ -272,16 +272,16 @@ impl ShardHolder {
Ok(())
}
pub fn contains_shard(&self, shard_id: &ShardId) -> bool {
self.shards.contains_key(shard_id)
pub fn contains_shard(&self, shard_id: ShardId) -> bool {
self.shards.contains_key(&shard_id)
}
pub fn get_shard(&self, shard_id: &ShardId) -> Option<&ShardReplicaSet> {
self.shards.get(shard_id)
pub fn get_shard(&self, shard_id: ShardId) -> Option<&ShardReplicaSet> {
self.shards.get(&shard_id)
}
pub fn get_shards(&self) -> impl Iterator<Item = (&ShardId, &ShardReplicaSet)> {
self.shards.iter()
pub fn get_shards(&self) -> impl Iterator<Item = (ShardId, &ShardReplicaSet)> {
self.shards.iter().map(|(id, shard)| (*id, shard))
}
pub fn all_shards(&self) -> impl Iterator<Item = &ShardReplicaSet> {
@@ -414,12 +414,12 @@ impl ShardHolder {
///
/// This only includes shard transfers that are in consensus for the current collection. A
/// shard transfer that has just been proposed may not be included yet.
pub fn count_shard_transfer_io(&self, peer_id: &PeerId) -> (usize, usize) {
pub fn count_shard_transfer_io(&self, peer_id: PeerId) -> (usize, usize) {
let (mut incoming, mut outgoing) = (0, 0);
for transfer in self.shard_transfers.read().iter() {
incoming += usize::from(transfer.to == *peer_id);
outgoing += usize::from(transfer.from == *peer_id);
incoming += usize::from(transfer.to == peer_id);
outgoing += usize::from(transfer.from == peer_id);
}
(incoming, outgoing)
@@ -477,13 +477,9 @@ impl ShardHolder {
Some(resharding_operations)
}
pub fn get_related_transfers(
&self,
shard_id: &ShardId,
peer_id: &PeerId,
) -> Vec<ShardTransfer> {
pub fn get_related_transfers(&self, shard_id: ShardId, peer_id: PeerId) -> Vec<ShardTransfer> {
self.get_transfers(|transfer| {
transfer.shard_id == *shard_id && (transfer.from == *peer_id || transfer.to == *peer_id)
transfer.shard_id == shard_id && (transfer.from == peer_id || transfer.to == peer_id)
})
}
@@ -700,11 +696,11 @@ impl ShardHolder {
let is_local =
replica_set.this_peer_id() == local_peer_id && replica_set.is_local().await;
let is_initializing =
replica_set.peer_state(&local_peer_id) == Some(ReplicaState::Initializing);
replica_set.peer_state(local_peer_id) == Some(ReplicaState::Initializing);
if not_distributed && is_local && is_initializing {
log::warn!("Local shard {collection_id}:{} stuck in Initializing state, changing to Active", replica_set.shard_id);
replica_set
.set_replica_state(&local_peer_id, ReplicaState::Active)
.set_replica_state(local_peer_id, ReplicaState::Active)
.expect("Failed to set local shard state");
}
let shard_key = shard_id_to_key_mapping.get(&shard_id).cloned();
@@ -719,7 +715,7 @@ impl ShardHolder {
}
pub async fn assert_shard_exists(&self, shard_id: ShardId) -> CollectionResult<()> {
match self.get_shard(&shard_id) {
match self.get_shard(shard_id) {
Some(_) => Ok(()),
None => Err(shard_not_found_error(shard_id)),
}
@@ -727,7 +723,7 @@ impl ShardHolder {
async fn assert_shard_is_local(&self, shard_id: ShardId) -> CollectionResult<()> {
let is_local_shard = self
.is_shard_local(&shard_id)
.is_shard_local(shard_id)
.await
.ok_or_else(|| shard_not_found_error(shard_id))?;
@@ -745,7 +741,7 @@ impl ShardHolder {
shard_id: ShardId,
) -> CollectionResult<()> {
let is_local_shard = self
.is_shard_local_or_queue_proxy(&shard_id)
.is_shard_local_or_queue_proxy(shard_id)
.await
.ok_or_else(|| shard_not_found_error(shard_id))?;
@@ -759,7 +755,7 @@ impl ShardHolder {
}
/// Returns true if shard is explicitly local, false otherwise.
pub async fn is_shard_local(&self, shard_id: &ShardId) -> Option<bool> {
pub async fn is_shard_local(&self, shard_id: ShardId) -> Option<bool> {
match self.get_shard(shard_id) {
Some(shard) => Some(shard.is_local().await),
None => None,
@@ -767,7 +763,7 @@ impl ShardHolder {
}
/// Returns true if shard is explicitly local or is queue proxy shard, false otherwise.
pub async fn is_shard_local_or_queue_proxy(&self, shard_id: &ShardId) -> Option<bool> {
pub async fn is_shard_local_or_queue_proxy(&self, shard_id: ShardId) -> Option<bool> {
match self.get_shard(shard_id) {
Some(shard) => Some(shard.is_local().await || shard.is_queue_proxy().await),
None => None,
@@ -779,7 +775,7 @@ impl ShardHolder {
let mut res = Vec::with_capacity(1);
for (shard_id, replica_set) in self.get_shards() {
if replica_set.has_local_shard().await {
res.push(*shard_id);
res.push(shard_id);
}
}
res
@@ -788,7 +784,7 @@ impl ShardHolder {
/// Count how many shard replicas are on the given peer.
pub fn count_peer_shards(&self, peer_id: PeerId) -> usize {
self.get_shards()
.filter(|(_, replica_set)| replica_set.peer_state(&peer_id).is_some())
.filter(|(_, replica_set)| replica_set.peer_state(peer_id).is_some())
.count()
}
@@ -819,8 +815,8 @@ impl ShardHolder {
.collect()
}
pub fn get_outgoing_transfers(&self, current_peer_id: &PeerId) -> Vec<ShardTransfer> {
self.get_transfers(|transfer| transfer.from == *current_peer_id)
pub fn get_outgoing_transfers(&self, current_peer_id: PeerId) -> Vec<ShardTransfer> {
self.get_transfers(|transfer| transfer.from == current_peer_id)
}
/// # Cancel safety
@@ -836,7 +832,7 @@ impl ShardHolder {
let snapshots_path = Self::snapshots_path_for_shard_unchecked(snapshots_path, shard_id);
let shard = self
.get_shard(&shard_id)
.get_shard(shard_id)
.ok_or_else(|| shard_not_found_error(shard_id))?;
let snapshot_manager = shard.get_snapshots_storage_manager()?;
snapshot_manager.list_snapshots(&snapshots_path).await
@@ -856,7 +852,7 @@ impl ShardHolder {
// and would be deleted, if future is cancelled
let shard = self
.get_shard(&shard_id)
.get_shard(shard_id)
.ok_or_else(|| shard_not_found_error(shard_id))?;
if !shard.is_local().await && !shard.is_queue_proxy().await {
@@ -987,7 +983,7 @@ impl ShardHolder {
temp_dir: &Path,
cancel: cancel::CancellationToken,
) -> CollectionResult<()> {
if !self.contains_shard(&shard_id) {
if !self.contains_shard(shard_id) {
return Err(shard_not_found_error(shard_id));
}
@@ -1065,7 +1061,7 @@ impl ShardHolder {
// (see `VectorsConfig::check_compatible_with_segment_config`)
let replica_set = self
.get_shard(&shard_id)
.get_shard(shard_id)
.ok_or_else(|| shard_not_found_error(shard_id))?;
// `ShardReplicaSet::restore_local_replica_from` is *not* cancel safe

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

@@ -306,7 +306,7 @@ impl ShardHolder {
// Revert replicas in `Resharding` state back into `Active` state
for (peer, state) in shard.peers() {
if state == ReplicaState::Resharding {
shard.set_replica_state(&peer, ReplicaState::Active)?;
shard.set_replica_state(peer, ReplicaState::Active)?;
}
}
@@ -334,8 +334,8 @@ impl ShardHolder {
// Remove new shard if resharding up
if direction == ReshardingDirection::Up {
if let Some(shard) = self.get_shard(&shard_id) {
match shard.peer_state(&peer_id) {
if let Some(shard) = self.get_shard(shard_id) {
match shard.peer_state(peer_id) {
Some(ReplicaState::Resharding) => {
log::debug!("removing peer {peer_id} from {shard_id} replica set");
shard.remove_peer(peer_id).await?;
@@ -514,7 +514,7 @@ impl ShardHolder {
}
pub async fn cleanup_local_shard(&self, shard_id: ShardId) -> CollectionResult<UpdateResult> {
let shard = self.get_shard(&shard_id).ok_or_else(|| {
let shard = self.get_shard(shard_id).ok_or_else(|| {
CollectionError::not_found(format!("shard {shard_id} does not exist"))
})?;
@@ -535,7 +535,7 @@ impl ShardHolder {
}
pub fn hash_ring_filter(&self, shard_id: ShardId) -> Option<hash_ring::HashRingFilter> {
if !self.contains_shard(&shard_id) {
if !self.contains_shard(shard_id) {
return None;
}

4
lib/collection/src/shards/transfer/driver.rs
View File

@@ -166,7 +166,7 @@ pub async fn revert_proxy_shard_to_local(
shard_holder: &ShardHolder,
shard_id: ShardId,
) -> CollectionResult<bool> {
let replica_set = match shard_holder.get_shard(&shard_id) {
let replica_set = match shard_holder.get_shard(shard_id) {
None => return Ok(false),
Some(replica_set) => replica_set,
};
@@ -241,7 +241,7 @@ where
if is_err || is_cancelled {
// Revert queue proxy if we still have any to prepare for the next attempt
if let Some(shard) = shards_holder.read().await.get_shard(&transfer.shard_id) {
if let Some(shard) = shards_holder.read().await.get_shard(transfer.shard_id) {
shard.revert_queue_proxy_local().await;
}
}

6
lib/collection/src/shards/transfer/resharding_stream_records.rs
View File

@@ -41,7 +41,7 @@ pub(crate) async fn transfer_resharding_stream_records(
{
let shard_holder = shard_holder.read().await;
let Some(replica_set) = shard_holder.get_shard(&shard_id) else {
let Some(replica_set) = shard_holder.get_shard(shard_id) else {
return Err(CollectionError::service_error(format!(
"Shard {shard_id} cannot be proxied because it does not exist"
)));
@@ -91,7 +91,7 @@ pub(crate) async fn transfer_resharding_stream_records(
loop {
let shard_holder = shard_holder.read().await;
let Some(replica_set) = shard_holder.get_shard(&shard_id) else {
let Some(replica_set) = shard_holder.get_shard(shard_id) else {
// Forward proxy gone?!
// That would be a programming error.
return Err(CollectionError::service_error(format!(
@@ -120,7 +120,7 @@ pub(crate) async fn transfer_resharding_stream_records(
// Update cutoff point on remote shard, disallow recovery before our current last seen
{
let shard_holder = shard_holder.read().await;
let Some(replica_set) = shard_holder.get_shard(&shard_id) else {
let Some(replica_set) = shard_holder.get_shard(shard_id) else {
// Forward proxy gone?!
// That would be a programming error.
return Err(CollectionError::service_error(format!(

2
lib/collection/src/shards/transfer/snapshot.rs
View File

@@ -175,7 +175,7 @@ pub(super) async fn transfer_snapshot(
let shard_holder_read = shard_holder.read().await;
let local_rest_address = channel_service.current_rest_address(transfer_config.from)?;
let transferring_shard = shard_holder_read.get_shard(&shard_id);
let transferring_shard = shard_holder_read.get_shard(shard_id);
let Some(replica_set) = transferring_shard else {
return Err(CollectionError::service_error(format!(
"Shard {shard_id} cannot be queue proxied because it does not exist"

6
lib/collection/src/shards/transfer/stream_records.rs
View File

@@ -38,7 +38,7 @@ pub(super) async fn transfer_stream_records(
{
let shard_holder = shard_holder.read().await;
let Some(replica_set) = shard_holder.get_shard(&shard_id) else {
let Some(replica_set) = shard_holder.get_shard(shard_id) else {
return Err(CollectionError::service_error(format!(
"Shard {shard_id} cannot be proxied because it does not exist"
)));
@@ -76,7 +76,7 @@ pub(super) async fn transfer_stream_records(
loop {
let shard_holder = shard_holder.read().await;
let Some(replica_set) = shard_holder.get_shard(&shard_id) else {
let Some(replica_set) = shard_holder.get_shard(shard_id) else {
// Forward proxy gone?!
// That would be a programming error.
return Err(CollectionError::service_error(format!(
@@ -105,7 +105,7 @@ pub(super) async fn transfer_stream_records(
// Update cutoff point on remote shard, disallow recovery before our current last seen
{
let shard_holder = shard_holder.read().await;
let Some(replica_set) = shard_holder.get_shard(&shard_id) else {
let Some(replica_set) = shard_holder.get_shard(shard_id) else {
// Forward proxy gone?!
// That would be a programming error.
return Err(CollectionError::service_error(format!(

2
lib/collection/src/shards/transfer/wal_delta.rs
View File

@@ -98,7 +98,7 @@ pub(super) async fn transfer_wal_delta(
let shard_holder_read = shard_holder.read().await;
let transferring_shard = shard_holder_read.get_shard(&shard_id);
let transferring_shard = shard_holder_read.get_shard(shard_id);
let Some(replica_set) = transferring_shard else {
return Err(CollectionError::service_error(format!(
"Shard {shard_id} cannot be queue proxied because it does not exist"

8
lib/collection/src/tests/points_dedup.rs
View File

@@ -109,13 +109,13 @@ async fn fixture() -> Collection {
let op = OperationWithClockTag::from(CollectionUpdateOperations::PointOperation(
PointOperations::UpsertPoints(PointInsertOperationsInternal::PointsList(vec![
PointStructPersisted {
id: u64::from(*shard_id).into(),
id: u64::from(shard_id).into(),
vector: VectorStructPersisted::Single(
(0..DIM).map(|_| rng.gen_range(0.0..1.0)).collect(),
),
payload: Some(Payload(Map::from_iter([(
"num".to_string(),
Value::from(-(*shard_id as i32)),
Value::from(-(shard_id as i32)),
)]))),
},
PointStructPersisted {
@@ -125,7 +125,7 @@ async fn fixture() -> Collection {
),
payload: Some(Payload(Map::from_iter([(
"num".to_string(),
Value::from(100 - *shard_id as i32),
Value::from(100 - shard_id as i32),
)]))),
},
])),
@@ -139,7 +139,7 @@ async fn fixture() -> Collection {
// Activate all shards
for shard_id in 0..SHARD_COUNT {
collection
.set_shard_replica_state(shard_id as ShardId, PEER_ID, ReplicaState::Active, None)
.set_shard_replica_state(shard_id, PEER_ID, ReplicaState::Active, None)
.await
.expect("failed to active shard");
}

8
lib/collection/src/tests/snapshot_test.rs
View File

@@ -148,15 +148,15 @@ async fn _test_snapshot_collection(node_type: NodeType) {
{
let shards_holder = &recovered_collection.shards_holder.read().await;
let replica_ser_0 = shards_holder.get_shard(&0).unwrap();
let replica_ser_0 = shards_holder.get_shard(0).unwrap();
assert!(replica_ser_0.is_local().await);
let replica_ser_1 = shards_holder.get_shard(&1).unwrap();
let replica_ser_1 = shards_holder.get_shard(1).unwrap();
assert!(replica_ser_1.is_local().await);
let replica_ser_2 = shards_holder.get_shard(&2).unwrap();
let replica_ser_2 = shards_holder.get_shard(2).unwrap();
assert!(!replica_ser_2.is_local().await);
assert_eq!(replica_ser_2.peers().len(), 1);
let replica_ser_3 = shards_holder.get_shard(&3).unwrap();
let replica_ser_3 = shards_holder.get_shard(3).unwrap();
assert!(replica_ser_3.is_local().await);
assert_eq!(replica_ser_3.peers().len(), 3); // 2 remotes + 1 local

2
lib/storage/src/content_manager/toc/mod.rs
View File

@@ -461,7 +461,7 @@ impl TableOfContent {
let collections = self.collections.read().await;
if let Some(proposal_sender) = &self.consensus_proposal_sender {
for collection in collections.values() {
for transfer in collection.get_outgoing_transfers(&self.this_peer_id).await {
for transfer in collection.get_outgoing_transfers(self.this_peer_id).await {
let cancel_transfer =
ConsensusOperations::abort_transfer(collection.name(), transfer, reason);
proposal_sender.send(cancel_transfer)?;