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

Code before:
    }

    fn highest_replica_peer_id(&self) -> Option<PeerId> {
        self.replica_state.read().peers().keys().max().cloned()
    }

    pub async fn get_clock_with_timeout(
        &self,
        timeout: Option<Duration>,
    ) -> CollectionResult<ClockGuard> {
        let start_time = std::time::Instant::now();
        loop {
            if let Some(clock) = self.clock_set.lock().await.get_clock() {
                break Ok(clock);
            }
            if let Some(timeout) = timeout
                && start_time.elapsed() >= timeout
            {
                break Err(CollectionError::timeout(
                    timeout,
                    format!("Failed to acquire clock for update operation within {timeout:?}"),
                ));
            }
            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();


Review comment (suggestion):
What I meant is, instead of `get_clock_with_timeout` do:

```suggestion
    pub async fn get_clock(&self) -> ClockGuard {
        loop {
            match self.clock_set.lock().await.get_clock() {
                Some(clock) => return clock,
                None => yield_now().await,
            }
        }
    }
```

Code after:
    }

    fn highest_replica_peer_id(&self) -> Option<PeerId> {
        self.replica_state.read().peers().keys().max().cloned()
    }

    pub async fn get_clock(&self) -> ClockGuard {
        loop {
            match self.clock_set.lock().await.get_clock() {
                Some(clock) => return clock,
                // Prevent blocking async runtime with spinlock
                None => 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 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(),