Repository: qdrant/qdrant
File: lib/collection/src/shards/replica_set/mod.rs (Rust)

Code before:
    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);
        };

        let mut total_vector_size = 0;
        let mut total_payload_size = 0;
        let mut total_points = 0;

        for segment in segments {
            let size_info = spawn_blocking(move || segment.get().read().size_info()).await?;
            total_vector_size += size_info.vectors_size_bytes;
            total_payload_size += size_info.payloads_size_bytes;
            total_points += size_info.num_points;
        }

        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>> {
        self.payload_index_schema.clone()
    }

    pub async fn optimizers_log(&self) -> Option<Arc<ParkingMutex<TrackerLog>>> {
        let local = self.local.read().await;
        local.as_ref().and_then(|shard| shard.optimizers_log())
    }

    pub async fn optimizations(&self) -> Option<LocalShardOptimizations> {
        let local = self.local.read().await;
        local.as_ref().and_then(|shard| shard.optimizations())
    }

Review comment (suggestion):
We should wrap this spawn in [`AbortOnDropHandle`](https://docs.rs/tokio-util/latest/tokio_util/task/struct.AbortOnDropHandle.html)

It's unlikely to have an effect in the majority of cases, but it could be useful if we've overloaded the thread pool. Adding it is basically free.

Code after:
    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);
        };

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

            (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>> {
        self.payload_index_schema.clone()
    }

    pub async fn optimizers_log(&self) -> Option<Arc<ParkingMutex<TrackerLog>>> {
        let local = self.local.read().await;
        local.as_ref().and_then(|shard| shard.optimizers_log())