diff --git a/lib/collection/src/collection/mod.rs b/lib/collection/src/collection/mod.rs index 037be4653d..499656329c 100644 --- a/lib/collection/src/collection/mod.rs +++ b/lib/collection/src/collection/mod.rs @@ -165,7 +165,7 @@ impl Collection { let shared_shard_holder = SharedShardHolder::new(shard_holder); let collection_stats_cache = CollectionSizeStatsCache::new_with_values( - Self::estimate_collection_size_stats(&shared_shard_holder).await, + Self::estimate_collection_size_stats(&shared_shard_holder).await?, ); // Once the config is persisted - the collection is considered to be successfully created. @@ -284,7 +284,9 @@ impl Collection { let shared_shard_holder = SharedShardHolder::new(shard_holder); let collection_stats_cache = CollectionSizeStatsCache::new_with_values( - Self::estimate_collection_size_stats(&shared_shard_holder).await, + Self::estimate_collection_size_stats(&shared_shard_holder) + .await + .expect("Failed to load collection size stats"), ); Self { @@ -867,14 +869,16 @@ impl Collection { async fn estimate_collection_size_stats( shards_holder: &SharedShardHolder, - ) -> Option { + ) -> CollectionResult> { let shard_lock = shards_holder.read().await; shard_lock.estimate_collection_size_stats().await } /// Returns estimations of collection sizes. This values are cached and might be not 100% up to date. /// The cache gets updated every 32 calls. - pub(crate) async fn estimated_collection_stats(&self) -> Option<&CollectionSizeAtomicStats> { + pub(crate) async fn estimated_collection_stats( + &self, + ) -> CollectionResult> { self.collection_stats_cache .get_or_update_cache(|| Self::estimate_collection_size_stats(&self.shards_holder)) .await diff --git a/lib/collection/src/common/collection_size_stats.rs b/lib/collection/src/common/collection_size_stats.rs index 9556e91425..e13601fe07 100644 --- a/lib/collection/src/common/collection_size_stats.rs +++ b/lib/collection/src/common/collection_size_stats.rs @@ -1,6 +1,8 @@ use std::future::Future; use std::sync::atomic::{AtomicUsize, Ordering}; +use crate::operations::types::CollectionResult; + /// Amount of requests that have to be done until the cached data gets updated. const UPDATE_INTERVAL: usize = 32; @@ -33,18 +35,20 @@ impl CollectionSizeStatsCache { pub async fn get_or_update_cache( &self, update_fn: impl FnOnce() -> U, - ) -> Option<&CollectionSizeAtomicStats> + ) -> CollectionResult> where - U: Future>, + U: Future>>, { // Update if necessary if self.check_need_update_and_increment() { - let updated = update_fn().await?; + let Some(updated) = update_fn().await? else { + return Ok(None); + }; self.update(updated); } // Give caller access to cached (inner) values which are always updated if required - self.stats.as_ref() + Ok(self.stats.as_ref()) } /// Sets all cache values to `new_stats`. diff --git a/lib/collection/src/operations/verification/update.rs b/lib/collection/src/operations/verification/update.rs index f2a2437d0f..7fff44dd5b 100644 --- a/lib/collection/src/operations/verification/update.rs +++ b/lib/collection/src/operations/verification/update.rs @@ -71,7 +71,7 @@ impl StrictModeVerification for SetPayload { strict_mode_config: &StrictModeConfig, ) -> CollectionResult<()> { if let Some(payload_size_limit_bytes) = strict_mode_config.max_collection_payload_size_bytes - && let Some(local_stats) = collection.estimated_collection_stats().await + && let Some(local_stats) = collection.estimated_collection_stats().await? { check_collection_payload_size_limit(payload_size_limit_bytes, local_stats)?; } @@ -252,7 +252,7 @@ async fn check_collection_size_limit( return Ok(()); } - let Some(stats) = collection.estimated_collection_stats().await else { + let Some(stats) = collection.estimated_collection_stats().await? else { return Ok(()); }; diff --git a/lib/collection/src/shards/replica_set/mod.rs b/lib/collection/src/shards/replica_set/mod.rs index f8ae0edf5a..f89afda301 100644 --- a/lib/collection/src/shards/replica_set/mod.rs +++ b/lib/collection/src/shards/replica_set/mod.rs @@ -25,6 +25,7 @@ use segment::types::{ExtendedPointId, Filter, ShardKey}; use serde::{Deserialize, Serialize}; use tokio::runtime::Handle; use tokio::sync::{Mutex, RwLock}; +use tokio::task::spawn_blocking; use tokio_util::task::AbortOnDropHandle; use self::partial_snapshot_meta::PartialSnapshotMeta; @@ -1306,40 +1307,48 @@ impl ShardReplicaSet { /// Returns the estimated size of all local segments. /// Since this locks all segments you should cache this value in performance critical scenarios! - pub(crate) async fn calculate_local_shard_stats(&self) -> Option { - self.local - .read() - .await - .as_ref() - .map(|i| match i { - Shard::Local(local) => { - let mut total_vector_size = 0; - let mut total_payload_size = 0; - let mut total_points = 0; + pub(crate) async fn calculate_local_shard_stats( + &self, + ) -> CollectionResult> { + let Some(segments) = self.local.read().await.as_ref().and_then(|i| match i { + Shard::Local(local) => Some( + // Collect the segments first so we don't have the segment holder locked for the entire duration of the loop. + local + .segments + .read() + .iter() + .map(|i| i.1.clone()) + .collect::>(), + ), + Shard::Proxy(_) | Shard::ForwardProxy(_) | Shard::QueueProxy(_) | Shard::Dummy(_) => { + None + } + }) else { + return Ok(None); + }; - // Collect the segments first so we don't have the segment holder locked for the entire duration of the loop. - let segments: Vec<_> = - local.segments.read().iter().map(|i| i.1).cloned().collect(); + let handle = spawn_blocking(move || { + let mut total_vector_size = 0; + let mut total_payload_size = 0; + let mut total_points = 0; - for segment in segments { - let size_info = segment.get().read().size_info(); - total_vector_size += size_info.vectors_size_bytes; - total_payload_size += size_info.payloads_size_bytes; - total_points += size_info.num_points; - } + for segment in segments { + let size_info = segment.get().read().size_info(); + total_vector_size += size_info.vectors_size_bytes; + total_payload_size += size_info.payloads_size_bytes; + total_points += size_info.num_points; + } - Some(CollectionSizeStats { - vector_storage_size: total_vector_size, - payload_storage_size: total_payload_size, - points_count: total_points, - }) - } - Shard::Proxy(_) - | Shard::ForwardProxy(_) - | Shard::QueueProxy(_) - | Shard::Dummy(_) => None, - }) - .unwrap_or_default() + (total_vector_size, total_payload_size, total_points) + }); + let (total_vector_size, total_payload_size, total_points) = + AbortOnDropHandle::new(handle).await?; + + Ok(Some(CollectionSizeStats { + vector_storage_size: total_vector_size, + payload_storage_size: total_payload_size, + points_count: total_points, + })) } pub(crate) fn payload_index_schema(&self) -> Arc> { diff --git a/lib/collection/src/shards/shard_holder/mod.rs b/lib/collection/src/shards/shard_holder/mod.rs index 3e21dd7226..ee014e3ea9 100644 --- a/lib/collection/src/shards/shard_holder/mod.rs +++ b/lib/collection/src/shards/shard_holder/mod.rs @@ -1520,30 +1520,32 @@ impl ShardHolder { } /// Estimates the collections size based on local shard data. Returns `None` if no shard for the collection was found locally. - pub async fn estimate_collection_size_stats(&self) -> Option { + pub async fn estimate_collection_size_stats( + &self, + ) -> CollectionResult> { if self.is_distributed().await { // In distributed, we estimate the whole collection size by using a single local shard and multiply by amount of shards in the collection. for shard in self.shards.iter() { - if let Some(shard_stats) = shard.1.calculate_local_shard_stats().await { + if let Some(shard_stats) = shard.1.calculate_local_shard_stats().await? { // TODO(resharding) take into account the ongoing resharding and exclude shards that are being filled from multiplication. // Project the single shards size to the full collection. let collection_estimate = shard_stats.multiplied_with(self.shards.len()); - return Some(collection_estimate); + return Ok(Some(collection_estimate)); } } - return None; + return Ok(None); } // Local mode: return collection size estimations using all shards. let mut stats = CollectionSizeStats::default(); for shard in self.shards.iter() { - if let Some(shard_stats) = shard.1.calculate_local_shard_stats().await { + if let Some(shard_stats) = shard.1.calculate_local_shard_stats().await? { stats.accumulate_metrics_from(&shard_stats); } } - Some(stats) + Ok(Some(stats)) } /// Returns `true` if the collection is distributed across multiple nodes.