Use spawn_blocking for segment locking (#8069)

* Use spawn_blocking for segment locking

* Wrap blocking task in AbortOnDropHandle

* Use single blocking thread for fetching all segment size information

---------

Co-authored-by: timvisee <tim@visee.me>

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<CollectionSizeStats> {
+    ) -> CollectionResult<Option<CollectionSizeStats>> {
         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<Option<&CollectionSizeAtomicStats>> {
         self.collection_stats_cache
             .get_or_update_cache(|| Self::estimate_collection_size_stats(&self.shards_holder))
             .await

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<U>(
         &self,
         update_fn: impl FnOnce() -> U,
-    ) -> Option<&CollectionSizeAtomicStats>
+    ) -> CollectionResult<Option<&CollectionSizeAtomicStats>>
     where
-        U: Future<Output = Option<CollectionSizeStats>>,
+        U: Future<Output = CollectionResult<Option<CollectionSizeStats>>>,
     {
         // 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`.

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(());
     };
 

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<CollectionSizeStats> {
-        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<Option<CollectionSizeStats>> {
+        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::<Vec<_>>(),
+            ),
+            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<SaveOnDisk<PayloadIndexSchema>> {

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<CollectionSizeStats> {
+    pub async fn estimate_collection_size_stats(
+        &self,
+    ) -> CollectionResult<Option<CollectionSizeStats>> {
         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.