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

Code before:
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.
        let plunger_result = rx.await;

        // Update worker channel is now clean.
        // Reset the flag and check the plunger result.
        update_handler
            .skip_updates_flag
            .store(false, std::sync::atomic::Ordering::Relaxed);

Review comment (suggestion):
I'm afraid this might not match up with the queue/WAL. Because this number is only updated **after** application, and not at the moment we popped it from the channel.

Would it be possible to:
- set skip updates flag
- store `op_num` of first operation that is popped of the channel
- drain whole channel
- truncate from that `op_num` onwards

because then we can **ensure** that `op_num` never hit the update worker and everything before it is (eventually) applied.

Code after:
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.
        // To do this, set the skip updates flag. The guard resets it when dropped.
        // New updates won't be skipped because the update_lock is held until we're done.
        let _skip_updates_guard = SkipUpdatesGuard::new(update_handler.skip_updates.clone());

        // Then, Send the plunger signal to the update handler.
        // It's a marker that all previous updates are processed or skipped.
        let (tx, rx) = oneshot::channel();