Repository: qdrant/qdrant
File: lib/collection/src/shards/local_shard/wal_ops.rs (Rust)

Code before:
use tokio::sync::{Mutex, oneshot};

use crate::operations::types::CollectionResult;
use crate::shards::local_shard::LocalShard;
use crate::update_handler::UpdateSignal;
use crate::update_workers::applied_seq::APPLIED_SEQ_SAVE_INTERVAL;

impl LocalShard {
    /// WAL is keeping more data, even if truncated.
    /// Useful if we expect to read old WAL records soon.
    pub async fn set_extended_wal_retention(&self) {
        self.wal.set_extended_retention().await;
    }

    /// WAL is keeping normal amount of data after truncation.
    pub async fn set_normal_wal_retention(&self) {
        self.wal.set_normal_retention().await;
    }

    /// Drop WAL from the last applied sequence number.
    /// Returns amount of removed records.
    pub async fn drop_wal(&self) -> CollectionResult<usize> {
        // First, lock the WAL to prevent new writes.
        let mut wal_lock = Mutex::lock_owned(self.wal.wal.clone()).await;

        // Next, lock the update handler to prevent changes in update handler state.
        let update_handler = self.update_handler.lock().await;

        // If there are no `op_num` in the `applied_seq`, it's unsafe to drop WAL.
        // Because there is no good save sequence number to drop from.
        let applied_seq = update_handler.applied_seq();
        if applied_seq.op_num().is_none() {
            return Ok(0);
        }

        // Send the plunger signal to the update handler.
        // It a marker that all previous updates are processed or skipped.
        let (tx, rx) = oneshot::channel();
        let plunger = UpdateSignal::Plunger(tx);
        self.update_sender.load().send(plunger).await?;

        // Set the skip updates flag to true to prevent processing new updates
        // while cleaning the update worker channel.
        update_handler
            .skip_updates_flag
            .store(true, std::sync::atomic::Ordering::Relaxed);

        // Wait until plunger is processed. Don't propagate errors from update worker because
        // it's important to switch flag back even if plunger failed.

Review comment (bug):
I'm afraid this can deadlock if the update worker just started processing a new operation, and tries to lock the WAL to get the operation for a pointer here:

https://github.com/qdrant/qdrant/blob/4c917ea155e6be8130481656dd7ef1da2b298b7c/lib/collection/src/update_workers/update_worker.rs#L64

Code after:
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use tokio::sync::{Mutex, oneshot};

use crate::operations::types::CollectionResult;
use crate::shards::local_shard::LocalShard;
use crate::update_handler::UpdateSignal;

/// Guard that sets an atomic bool to `true` on creation and back to `false` on drop.
/// Ensures the flag is always reset even on early return or panic.
struct SkipUpdatesGuard {
    flag: Arc<AtomicBool>,
}

impl SkipUpdatesGuard {
    fn new(flag: Arc<AtomicBool>) -> Self {
        flag.store(true, Ordering::Relaxed);
        Self { flag }
    }
}

impl Drop for SkipUpdatesGuard {
    fn drop(&mut self) {
        self.flag.store(false, Ordering::Relaxed);
    }
}

impl LocalShard {
    /// WAL is keeping more data, even if truncated.
    /// Useful if we expect to read old WAL records soon.
    pub async fn set_extended_wal_retention(&self) {
        self.wal.set_extended_retention().await;
    }

    /// WAL is keeping normal amount of data after truncation.
    pub async fn set_normal_wal_retention(&self) {
        self.wal.set_normal_retention().await;
    }

    /// Truncate unapplied WAL records.
    /// Returns amount of removed records.
    pub async fn truncate_unapplied_wal(&self) -> CollectionResult<usize> {
        // First, lock the `update_lock` to prevent new updates while dropping WAL.
        // Lock also `update_handler` to access its state.
        let _update_lock = self.update_lock.write().await;
        let update_handler = self.update_handler.lock().await;

        // Next, clear the update worker channel.