ronantakizawa/codereview-bench huggingface dataset:
Repository: qdrant/qdrant
File: lib/common/common/src/is_alive_lock.rs (Rust)

Code before:
/// Structure which ensures that the lock is alive at the time of locking,
/// and will prevent dropping while guarded.
///
/// This structure explicitly doesn't implement Clone, so that `handle` is used instead.
#[derive(Debug)]
pub struct IsAliveLock {
    inner: Arc<Mutex<bool>>,
}

impl IsAliveLock {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(true)),
        }
    }

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

    /// Waits for lock and poisons without dropping.
    /// Lock will no longer be usable after this.
    pub fn poison(&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.inner.lock() = false;
    }
}

/// 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 {

Review comment (suggestion):
I am a bit confused by the term `poison`.

For me it refers to handling panics in threads holding locks so I feel the term is ambiguous here.

Code after:
/// Structure which ensures that the lock is alive at the time of locking,
/// and will prevent dropping while guarded.
///
/// Dropping this structure will also mark as dead, preventing future access through any dangling
/// handles.
///
/// This structure explicitly doesn't implement Clone, so that `handle` is used instead.
#[derive(Debug)]
pub struct IsAliveLock {
    inner: Arc<Mutex<bool>>,
}

impl IsAliveLock {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(true)),
        }
    }

    /// 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>>,
Repository: qdrant/qdrant
File: lib/common/common/src/is_alive_lock.rs (Rust)

Code before:
    }
}

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();

        assert!(handle.lock_if_alive().is_some());

        lock.mark_dead();

        assert!(handle.lock_if_alive().is_none());
    }


Review comment (suggestion):
```suggestion
    /// 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<'_>> {
```

Code after:
    /// 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
pub struct IsAliveGuard(#[expect(dead_code)] ArcMutexGuard<RawMutex, bool>);

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

    #[test]
    fn test_marking_dead() {
Repository: qdrant/qdrant
File: lib/common/common/src/is_alive_lock.rs (Rust)

Code before:
    }

    /// Get a handle for this lock.
    pub fn handle(&self) -> IsAliveHandle {
        IsAliveHandle {
            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::*;

Review comment (bug):
Also we must not forget that implicitly blocking on a drop here may be a potential footgun deadlocking Tokio runtimes with a low runtime. Currently the only thing that protects us against it is that the lock holders (flusher tasks) don't use threads in the same pool.

We can specialize this structure to flushers, and debug assert that all handles never call in [any runtime context](https://docs.rs/tokio/latest/tokio/runtime/struct.Handle.html#method.try_current).

Code after:
        Self {
            inner: Arc::new(Mutex::new(true)),
        }
    }

    /// 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
Repository: qdrant/qdrant
File: lib/common/common/src/is_alive_lock.rs (Rust)

Code before:

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> {
        if let Some(mutex) = self.inner.upgrade() {
            // Since the Weak<Mutex> was upgraded, this is the safe way of keeping the mutex
            // alive while referencing it with the guard, and returning both.
            IsAliveGuardCell::try_new(mutex, |mutex| {
                let guard = mutex.lock();
                if *guard {
                    Ok(IsAliveGuardInner(guard))
                } else {
                    Err(())
                }
            })
            .ok()
            // hide implementation details with this newtype
            .map(IsAliveGuard)
        } else {
            None
        }
    }
}

/// Guards a `true` boolean
///
/// This is an opaque wrapper that hides the self_cell implementation details.
pub struct IsAliveGuard(#[expect(dead_code)] IsAliveGuardCell);

// Private self_cell implementation

Review comment (suggestion):
Nit, lets remove nesting:

```rust
let Some(mutex) = self.inner.upgrade() else { return None; };
```

Code after:
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
pub struct IsAliveGuard(#[expect(dead_code)] ArcMutexGuard<RawMutex, bool>);

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

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

        assert!(handle.lock_if_alive().is_some());

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

Code before:
                }
            })
            .ok()
            // hide implementation details with this newtype
            .map(IsAliveGuard)
        } else {
            None
        }
    }
}

/// Guards a `true` boolean
///
/// This is an opaque wrapper that hides the self_cell implementation details.
pub struct IsAliveGuard(#[expect(dead_code)] IsAliveGuardCell);

// Private self_cell implementation
self_cell!(
    struct IsAliveGuardCell {
        owner: Arc<Mutex<bool>>,

        #[covariant]
        dependent: IsAliveGuardInner,
    }
);
struct IsAliveGuardInner<'a>(#[expect(dead_code)] MutexGuard<'a, bool>);

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

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

        assert!(handle.lock_if_alive().is_some());

        lock.blocking_mark_dead();

        assert!(handle.lock_if_alive().is_none());
    }

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

        assert!(handle.lock_if_alive().is_some());

        // dropping the handle does not poison the lock

Review comment (suggestion):
I didn't realize it'd become so ugly just because we have to hold the Arc alive.

Can't we use the ArcMutexGuard type here?

https://docs.rs/parking_lot/latest/parking_lot/struct.ArcMutexGuard.html ([via](https://docs.rs/parking_lot/latest/parking_lot/type.Mutex.html#method.lock_arc))

I believe it only needs the Arc when obtaining the lock, after which we can drop the upgraded reference.

If possible, I strongly prefer that over this self cell hack.

Code after:
}

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

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

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

        assert!(handle.lock_if_alive().is_some());

        lock.blocking_mark_dead();

        assert!(handle.lock_if_alive().is_none());
    }

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

        assert!(handle.lock_if_alive().is_some());

        // dropping the handle does not poison the lock
        drop(handle);
        let handle = lock.handle();
        assert!(handle.lock_if_alive().is_some());

        // dropping the parent poisons the lock
        drop(lock);
        assert!(handle.lock_if_alive().is_none());
    }

    #[test]
    fn test_parent_waits_for_guard() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let lock = IsAliveLock::new();
        let handle = lock.handle();

        // Test following sequence
        // | tick | lock           | handle        |
        // | ---- | -------------  | ------        |
        // |   0  |                | guarded       |
        // |   1  | drop (waiting) | guarded       |
        // |   2  |      (waiting) | guarded       |
Repository: qdrant/qdrant
File: lib/edge/python/src/query.rs (Rust)

Code before:
pub struct PyPrefetch(ShardPrefetch);

#[pyclass_repr]
#[pymethods]
impl PyPrefetch {
    #[new]
    pub fn new(
        prefetches: Vec<PyPrefetch>,
        query: Option<PyScoringQuery>,
        limit: usize,
        params: Option<PySearchParams>,
        filter: Option<PyFilter>,
        score_threshold: Option<f32>,
    ) -> Self {
        Self(ShardPrefetch {
            prefetches: PyPrefetch::peel_vec(prefetches),
            query: query.map(ScoringQuery::from),
            limit,
            params: params.map(SearchParams::from),
            filter: filter.map(Filter::from),
            score_threshold: score_threshold.map(OrderedFloat),
        })
    }

    #[getter]
    pub fn prefetches(&self) -> &[PyPrefetch] {
        PyPrefetch::wrap_slice(&self.0.prefetches)
    }

    #[getter]
    pub fn query(&self) -> Option<PyScoringQuery> {
        self.0.query.clone().map(PyScoringQuery)
    }

    #[getter]
    pub fn limit(&self) -> usize {
        self.0.limit
    }

    #[getter]
    pub fn params(&self) -> Option<PySearchParams> {
        self.0.params.map(PySearchParams)
    }

    #[getter]
    pub fn filter(&self) -> Option<PyFilter> {
        self.0.filter.clone().map(PyFilter)
    }

    #[getter]
    pub fn score_threshold(&self) -> Option<f32> {

Review comment (suggestion):
`params` after `score_threshold`?

```suggestion
        filter = None,
        score_threshold = None,
        params = None,
```

Code after:
            score_threshold: _,
            limit: _,
            offset: _,
            params: _,
            with_vector: _,
            with_payload: _,
        } = self.0;
    }
}

#[pyclass(name = "Prefetch")]
#[derive(Clone, Debug, Into, TransparentWrapper)]
#[repr(transparent)]
pub struct PyPrefetch(ShardPrefetch);

#[pyclass_repr]
#[pymethods]
impl PyPrefetch {
    #[new]
    #[pyo3(signature = (
        limit,
        query = None,
        prefetches = None,
        params = None,
        filter = None,
        score_threshold = None,
    ))]
    pub fn new(
        limit: usize,
        query: Option<PyScoringQuery>,
        prefetches: Option<Vec<PyPrefetch>>,
        params: Option<PySearchParams>,
        filter: Option<PyFilter>,
        score_threshold: Option<f32>,
    ) -> Self {
        Self(ShardPrefetch {
            prefetches: PyPrefetch::peel_vec(prefetches.unwrap_or_default()),
            limit,
            query: query.map(ScoringQuery::from),
            params: params.map(SearchParams::from),
            filter: filter.map(Filter::from),
            score_threshold: score_threshold.map(OrderedFloat),
        })
    }

    #[getter]
    pub fn prefetches(&self) -> &[PyPrefetch] {
        PyPrefetch::wrap_slice(&self.0.prefetches)
    }

    #[getter]
Repository: qdrant/qdrant
File: lib/edge/python/src/search.rs (Rust)

Code before:
use derive_more::Into;
use ordered_float::OrderedFloat;
use pyo3::prelude::*;
use shard::query::query_enum::QueryEnum;
use shard::search::CoreSearchRequest;

use crate::repr::*;
use crate::*;

#[pyclass(name = "SearchRequest")]
#[derive(Clone, Debug, Into)]
pub struct PySearchRequest(CoreSearchRequest);

#[pyclass_repr]
#[pymethods]
impl PySearchRequest {
    #[new]
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        query: PyQuery,
        filter: Option<PyFilter>,
        params: Option<PySearchParams>,
        limit: usize,
        offset: usize,
        with_vector: Option<PyWithVector>,
        with_payload: Option<PyWithPayload>,
        score_threshold: Option<f32>,
    ) -> Self {
        Self(CoreSearchRequest {
            query: QueryEnum::from(query),
            filter: filter.map(Filter::from),
            params: params.map(SearchParams::from),
            limit,
            offset,
            with_vector: with_vector.map(WithVector::from),
            with_payload: with_payload.map(WithPayloadInterface::from),
            score_threshold,
        })
    }

    #[getter]
    pub fn query(&self) -> &PyQuery {
        PyQuery::wrap_ref(&self.0.query)
    }

    #[getter]
    pub fn filter(&self) -> Option<&PyFilter> {
        self.0.filter.as_ref().map(PyFilter::wrap_ref)
    }

    #[getter]

Review comment (suggestion):
`params` after `score_threshold`?

```suggestion
        with_vector = None,
        with_payload = None,
        score_threshold = None,
        params = None,
```

Code after:
use derive_more::Into;
use ordered_float::OrderedFloat;
use pyo3::prelude::*;
use shard::query::query_enum::QueryEnum;
use shard::search::CoreSearchRequest;

use crate::repr::*;
use crate::*;

#[pyclass(name = "SearchRequest")]
#[derive(Clone, Debug, Into)]
pub struct PySearchRequest(CoreSearchRequest);

#[pyclass_repr]
#[pymethods]
impl PySearchRequest {
    #[new]
    #[pyo3(signature = (
        query,
        limit,
        offset = None,
        filter = None,
        params = None,
        with_vector = None,
        with_payload = None,
        score_threshold = None,
    ))]
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        query: PyQuery,
        limit: usize,
        offset: Option<usize>,
        filter: Option<PyFilter>,
        params: Option<PySearchParams>,
        with_vector: Option<PyWithVector>,
        with_payload: Option<PyWithPayload>,
        score_threshold: Option<f32>,
    ) -> Self {
        Self(CoreSearchRequest {
            query: QueryEnum::from(query),
            limit,
            offset: offset.unwrap_or(0),
            filter: filter.map(Filter::from),
            params: params.map(SearchParams::from),
            with_vector: with_vector.map(WithVector::from),
            with_payload: with_payload.map(WithPayloadInterface::from),
            score_threshold,
        })
    }

    #[getter]
Repository: qdrant/qdrant
File: lib/shard/src/wal.rs (Rust)

Code before:
    /// Last index that is still available in logical WAL.
    pub fn last_index(&self) -> u64 {
        self.wal.last_index()
    }

    pub fn segment_capacity(&self) -> usize {
        self.options.segment_capacity
    }

    pub fn set_extended_retention(&mut self) {
        let normal_retention = self.options.retain_closed.get();
        self.wal
            .set_retention(normal_retention * INCREASED_RETENTION_FACTOR);
    }

    pub fn set_normal_retention(&mut self) {
        let normal_retention = self.options.retain_closed.get();
        self.wal.set_retention(normal_retention);
    }

    pub fn drop_from(&mut self, from_index: u64) -> Result<()> {
        self.wal
            .truncate(from_index)
            .map_err(|err| WalError::TruncateWalError(format!("{err:?}")))
    }
}

#[derive(Debug, Deserialize, Serialize)]
struct WalState {
    pub ack_index: u64,
}

impl WalState {
    pub fn new(ack_index: u64) -> Self {
        Self { ack_index }
    }
}

pub type Result<T, E = WalError> = result::Result<T, E>;

#[derive(Debug, Error)]
#[error("{0}")]
pub enum WalError {
    #[error("Can't init WAL: {0}")]
    InitWalError(String),
    #[error("Can't write WAL: {0}")]
    WriteWalError(String),
    #[error("Can't truncate WAL: {0}")]
    TruncateWalError(String),
    #[error("Operation rejected by WAL for old clock")]
    ClockRejected,

Review comment (bug):
Let's add a (debug) assertion here:

`from_index >= self.first_index()`

to prevent truncating entries that we already acknowledged. It won't prevent truncating items that have just been applied but it's better than nothing.

Code after:
    /// Last index that is still available in logical WAL.
    pub fn last_index(&self) -> u64 {
        self.wal.last_index()
    }

    pub fn segment_capacity(&self) -> usize {
        self.options.segment_capacity
    }

    pub fn set_extended_retention(&mut self) {
        let normal_retention = self.options.retain_closed.get();
        self.wal
            .set_retention(normal_retention * INCREASED_RETENTION_FACTOR);
    }

    pub fn set_normal_retention(&mut self) {
        let normal_retention = self.options.retain_closed.get();
        self.wal.set_retention(normal_retention);
    }

    pub fn drop_from(&mut self, from_index: u64) -> Result<()> {
        debug_assert!(from_index >= self.first_index());
        self.wal
            .truncate(from_index)
            .map_err(|err| WalError::TruncateWalError(format!("{err:?}")))
    }
}

#[derive(Debug, Deserialize, Serialize)]
struct WalState {
    pub ack_index: u64,
}

impl WalState {
    pub fn new(ack_index: u64) -> Self {
        Self { ack_index }
    }
}

pub type Result<T, E = WalError> = result::Result<T, E>;

#[derive(Debug, Error)]
#[error("{0}")]
pub enum WalError {
    #[error("Can't init WAL: {0}")]
    InitWalError(String),
    #[error("Can't write WAL: {0}")]
    WriteWalError(String),
    #[error("Can't truncate WAL: {0}")]
    TruncateWalError(String),
    #[error("Operation rejected by WAL for old clock")]