Repository: qdrant/qdrant
File: lib/common/common/src/is_alive_lock.rs (Rust)

Code before:
            inner: self.inner.clone(),
        }
    }

    /// Waits for lock and marks as dead without dropping.
    /// Lock will no longer be usable after this.
    pub fn mark_dead(&self) {
        *self.inner.lock() = false
    }
}

impl Default for IsAliveLock {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for IsAliveLock {
    fn drop(&mut self) {
        // prevent dangling handles from accessing the lock
        self.mark_dead();
    }
}

/// Handle for `IsAliveLock` which can access the lock at a later time.
///
/// This is a separate structure so it does not change the boolean on drop.
pub struct IsAliveHandle {
    inner: Arc<Mutex<bool>>,
}

impl IsAliveHandle {
    /// Get a guard of this lock if the parent hasn't been dropped
    pub fn lock_if_alive(&self) -> Option<IsAliveGuard<'_>> {
        let guard = self.inner.lock();
        guard.then_some(IsAliveGuard(guard))
    }
}

/// Guards a `true` boolean
#[expect(dead_code)]
pub struct IsAliveGuard<'a>(MutexGuard<'a, bool>);

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_marking_dead() {
        let lock = IsAliveLock::new();
        let handle = lock.handle();

Review comment (suggestion):
I'm afraid this may introduce a huge problem! A problem that breaks flushing _once again_. If so, it would have existed before and it not introduced in this PR. Though I do suggest to fix it in here.

Specifically, I'm afraid about graceful shutdown, undefined cancellation/abort order and WAL acknowledgement.

We have a flush worker in each local shard to flush all segments. It writes the latest acknowledged version into the WAL, so that we only replay newer versions on restart.

That logic you can find here: https://github.com/qdrant/qdrant/blob/6fe46d0d6097ba772f2cb654f8e02b99f6db2181/lib/collection/src/update_workers/flush_workers.rs#L65-L97

I'm afraid that on regular (graceful) shutdown we might:
1. shortcut flushing in some structures (due to is alive lock)
2. segment holder will still assume successful flush
3. update worker will still acknowledge newest version even though it isn't fully persisted

I'm not exactly sure this **can** happen. The problem is that we currently don't have strict control over the order of dropping/cancellation/aborting of flush tasks and futures. I don't suggest to make it well ordered as that'll be a huge effort.

My guess is that the only right way of solving this is to not shortcut flushers with `Ok(())`, but with `Err(Cancelled)`. That way we can explicitly handle cancellation down the line.

Code after:
    /// Get a handle for this lock.
    pub fn handle(&self) -> IsAliveHandle {
        IsAliveHandle {
            inner: Arc::downgrade(&self.inner),
        }
    }

    /// Waits for lock and marks as dead without dropping.
    /// Lock will no longer be usable after this.
    pub fn blocking_mark_dead(&self) {
        *self.inner.lock() = false
    }
}

impl Default for IsAliveLock {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for IsAliveLock {
    fn drop(&mut self) {
        // prevent dangling handles from accessing the lock
        self.blocking_mark_dead();
    }
}

/// Handle for `IsAliveLock` which can access the lock at a later time.
///
/// This is a separate structure so it does not change the boolean on drop.
pub struct IsAliveHandle {
    inner: Weak<Mutex<bool>>,
}

impl IsAliveHandle {
    /// Get a guard of this lock if the parent hasn't been dropped
    #[must_use = "Guard must be held for lifetime of operation, abort if None is returned"]
    pub fn lock_if_alive(&self) -> Option<IsAliveGuard> {
        let mutex = self.inner.upgrade()?;

        let is_alive = Mutex::lock_arc(&mutex);

        if *is_alive {
            Some(IsAliveGuard(is_alive))
        } else {
            None
        }
    }
}

/// Guards a `true` boolean