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

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

    #[allow(clippy::too_many_arguments)]
    fn update_worker_internal(
        collection_name: CollectionId,
        operation: CollectionUpdateOperations,
        op_num: SeqNumberType,
        wait: bool,
        wal: LockedWal,
        segments: LockedSegmentHolder,
        update_operation_lock: Arc<tokio::sync::RwLock<()>>,
        update_tracker: UpdateTracker,
        hw_measurements: HwMeasurementAcc,
    ) -> CollectionResult<usize> {
        // If wait flag is set, explicitly flush WAL first
        if wait {
            wal.blocking_lock().flush().map_err(|err| {
                CollectionError::service_error(format!(
                    "Can't flush WAL before operation {op_num} - {err}"
                ))
            })?;
        }

Review comment (bug):
This has a data race!

Because we use `send()`, `changed()` does not trigger if the signal it sent between the above `can_proceed` and this `changed()` barrier.

There's to ways to solve it:
1. use [`send_modify()`](https://docs.rs/tokio/latest/tokio/sync/watch/struct.Sender.html#method.send_modify) to also trigger this barrier if it was sent before
2. move the `changed()` call before the `can_proceed` check and only await here

I believe the second solution is a bit better, potentially skipping some extra loops.

Relevant docs:
- https://docs.rs/tokio/latest/tokio/sync/watch/struct.Sender.html#method.send
- https://docs.rs/tokio/latest/tokio/sync/watch/struct.Sender.html#method.send_modify
- https://docs.rs/tokio/latest/tokio/sync/watch/struct.Receiver.html#method.changed
- https://docs.rs/tokio/latest/tokio/sync/watch/index.html#change-notifications

Code after:

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

    #[allow(clippy::too_many_arguments)]
    fn update_worker_internal(
        collection_name: CollectionId,
        operation: CollectionUpdateOperations,
        op_num: SeqNumberType,
        wait: bool,
        wal: LockedWal,
        segments: LockedSegmentHolder,
        update_operation_lock: Arc<tokio::sync::RwLock<()>>,
        update_tracker: UpdateTracker,
        hw_measurements: HwMeasurementAcc,
    ) -> CollectionResult<usize> {
        // If wait flag is set, explicitly flush WAL first
        if wait {
            wal.blocking_lock().flush().map_err(|err| {
                CollectionError::service_error(format!(
                    "Can't flush WAL before operation {op_num} - {err}"
                ))