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

Code before:
                ));
            }
            yield_now().await;
        }
    }

    /// # Cancel safety
    ///
    /// This method is *not* cancel safe.
    async fn update(
        &self,
        operation: CollectionUpdateOperations,
        wait: bool,
        timeout: Option<Duration>,
        update_only_existing: bool,
        hw_measurement_acc: HwMeasurementAcc,
    ) -> CollectionResult<UpdateResult> {
        // `ShardRepilcaSet::update_impl` is not cancel safe, so this method is not cancel safe.

        // TODO: Optimize `remotes`/`local`/`clock` locking for the "happy path"?
        //
        // E.g., refactor `update`/`update_impl`, so that it would be possible to:
        // - lock `remotes`, `local`, `clock` (in specified order!) on the *first* iteration of the loop
        // - then release and lock `remotes` and `local` *only* for all next iterations
        // - but keep initial `clock` for the whole duration of `update`
        let mut clock = self.get_clock_with_timeout(timeout).await?;

        for attempt in 1..=UPDATE_MAX_CLOCK_REJECTED_RETRIES {
            let is_non_zero_tick = clock.current_tick().is_some();

            let res = self
                .update_impl(
                    operation.clone(),
                    wait,
                    timeout,
                    &mut clock,
                    update_only_existing,
                    hw_measurement_acc.clone(),
                )
                .await?;

            if let Some(res) = res {
                return Ok(res);
            }

            // Log a warning, if operation was rejected... but only if operation had a non-0 tick,
            // because operations with tick 0 should *always* be rejected and rejection is *expected*.
            if is_non_zero_tick && log::log_enabled!(log::Level::Warn) {
                if let Some(ids) = operation.point_ids() {
                    log::warn!(
                        "Operation affecting point IDs {ids:?} was rejected by some node(s), retrying... \

Review comment (suggestion):
And then here, you could do something like this to add a timeout:

```suggestion
        let mut clock = tokio::time::timeout(timeout.unwrap_or(Durtaion::MAX), self.get_clock()).await?;
```

Code after:
        operation: CollectionUpdateOperations,
        wait: bool,
        timeout: Option<Duration>,
        update_only_existing: bool,
        hw_measurement_acc: HwMeasurementAcc,
    ) -> CollectionResult<UpdateResult> {
        // `ShardRepilcaSet::update_impl` is not cancel safe, so this method is not cancel safe.

        // TODO: Optimize `remotes`/`local`/`clock` locking for the "happy path"?
        //
        // E.g., refactor `update`/`update_impl`, so that it would be possible to:
        // - lock `remotes`, `local`, `clock` (in specified order!) on the *first* iteration of the loop
        // - then release and lock `remotes` and `local` *only* for all next iterations
        // - but keep initial `clock` for the whole duration of `update`
        let clock_timeout = timeout.unwrap_or(Duration::MAX);
        let mut clock = tokio::time::timeout(clock_timeout, self.get_clock())
            .await
            .map_err(|_| {
                CollectionError::timeout(
                    clock_timeout,
                    format!("Failed to acquire clock for update operation within {timeout:?}"),
                )
            })?;

        for attempt in 1..=UPDATE_MAX_CLOCK_REJECTED_RETRIES {
            let is_non_zero_tick = clock.current_tick().is_some();

            let res = self
                .update_impl(
                    operation.clone(),
                    wait,
                    timeout,
                    &mut clock,
                    update_only_existing,
                    hw_measurement_acc.clone(),
                )
                .await?;

            if let Some(res) = res {
                return Ok(res);
            }

            // Log a warning, if operation was rejected... but only if operation had a non-0 tick,
            // because operations with tick 0 should *always* be rejected and rejection is *expected*.
            if is_non_zero_tick && log::log_enabled!(log::Level::Warn) {
                if let Some(ids) = operation.point_ids() {
                    log::warn!(
                        "Operation affecting point IDs {ids:?} was rejected by some node(s), retrying... \
                         (attempt {attempt}/{UPDATE_MAX_CLOCK_REJECTED_RETRIES})"
                    );
                } else {