Repository: qdrant/qdrant
File: lib/collection/src/update_workers/update_worker.rs (Rust)

Code before:
        // Size of the unoptimized segment to be considered large enough for waiting.
        // If `None`, waiting is disabled.
        optimization_threshold: Option<usize>,
        segments: &LockedSegmentHolder,
        optimization_handles: Arc<TokioMutex<Vec<StoppableTaskHandle<bool>>>>,
        optimization_finished_receiver: &mut watch::Receiver<()>,
    ) -> CollectionResult<()> {
        let Some(optimization_threshold) = optimization_threshold else {
            // Waiting is disabled
            return Ok(());
        };

        loop {
            let locked_segments = segments.clone();
            let can_proceed = tokio::task::spawn_blocking(move || {
                let segments = locked_segments.read();
                let largest_unoptimized_segment_size_opt =
                    get_largest_unindexed_segment_vector_size(&segments);

                let largest_unoptimized_segment_size =
                    largest_unoptimized_segment_size_opt.unwrap_or(0);

                // True, if we can proceed with updates
                largest_unoptimized_segment_size <= optimization_threshold
            })
            .await
            .map_err(CollectionError::from)?;

            if can_proceed {
                return Ok(());
            }

            // Block only if there are running optimization that can terminate
            {
                let optimizations_guard = optimization_handles.lock().await;
                if optimizations_guard.iter().all(|h| h.is_finished()) {
                    return Ok(());
                }
            }

            // If unoptimized segments are too large, the only way it can be fixed is optimization
            // So we wait the notification of optimization completion to re-check the sizes
            log::debug!("waiting for optimization to allow updates");
            if let Err(err) = optimization_finished_receiver.changed().await {
                // this can be if optimization is cancelled, we don't need to wait anymore
                log::debug!("Optimization thread terminated with an error: {err}");
                return Ok(());
            }
        }
    }


Review comment (bug):
Isn't there an edge case where our current update will make it go over the threshold - in which case we have linking again?

Case:
- segment is 9kb
- update adds 2gb
- segment is 11kb and disappears in `indexed_only=true` search

Code after:
    /// Returns when all segments are smaller that the optimization_threshold.
    async fn wait_for_optimization(
        // Size of the unoptimized segment to be considered large enough for waiting.
        // If `None`, waiting is disabled.
        optimization_threshold_kb: Option<usize>,
        segments: &LockedSegmentHolder,
        optimization_handles: Arc<TokioMutex<Vec<StoppableTaskHandle<bool>>>>,
        optimization_finished_receiver: &mut watch::Receiver<()>,
    ) -> CollectionResult<()> {
        let Some(optimization_threshold_kb) = optimization_threshold_kb else {
            // Waiting is disabled
            return Ok(());
        };
        let optimization_threshold = optimization_threshold_kb.saturating_mul(BYTES_IN_KB);
        loop {
            let locked_segments = segments.clone();
            let can_proceed = tokio::task::spawn_blocking(move || {
                let segments = locked_segments.read();
                let largest_unoptimized_segment_size_opt =
                    get_largest_unindexed_segment_vector_size(&segments);

                let largest_unoptimized_segment_size =
                    largest_unoptimized_segment_size_opt.unwrap_or(0);

                // True, if we can proceed with updates
                largest_unoptimized_segment_size <= optimization_threshold
            })
            .await
            .map_err(CollectionError::from)?;

            if can_proceed {
                return Ok(());
            }

            // Block only if there are running optimization that can terminate
            {
                let optimizations_guard = optimization_handles.lock().await;
                if optimizations_guard.iter().all(|h| h.is_finished()) {
                    return Ok(());
                }
            }

            // If unoptimized segments are too large, the only way it can be fixed is optimization
            // So we wait the notification of optimization completion to re-check the sizes
            log::debug!("waiting for optimization to allow updates");
            if let Err(err) = optimization_finished_receiver.changed().await {
                // this can be if optimization is cancelled, we don't need to wait anymore
                log::debug!("Optimization thread terminated with an error: {err}");
                return Ok(());
            }
        }