diskann-inmem 0.60.0

DiskANN3 is a composable library for bringing scalable, accurate and cost-effective vector indexing to multiple databases.
/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

//! # EBR lifecycle hooks for [`super::Store`]
//!
//! Please read this section carefully - the protocol is not difficult, but it *is* subtle.
//!
//! The transitions are a simplified version of the protocol described in [`crate::tag`]
//! that storage slots need to implement to be compatible. A state diagram is shown below:
//!
//! ```text
//!         +--------------- `reclaim` ----------------+
//!         |                                          |
//!         V                                          |
//!   +-----------+                               +----------+
//!   | Available |<---+                          | Retiring |
//!   +-----------+    |                          +----------+
//!         |          |                               ^
//!         |          |                               |
//!         |       `abort`                         `retire`
//!     `acquire`      |                               |
//!         |          |                               |
//!         |    +---------------+                   +-----------+
//!         +--->| Exclusive<'_> |---- `publish` --->| Published |
//!              +---------------+                   +-----------+
//!                    |
//!                 `freeze`
//!                    |
//!        +-----------+
//!        |
//!        V
//!    +--------+
//!    | Frozen |
//!    +--------+
//! ```
//!
//! ## Readable States
//!
//! * `published`: **New** references to slots may be given out in the "published" state. It is
//!   possible for a transition to go from "published" to "retiring" while references are lent
//!   out. This is fine as long as the lifetime of these references is bounded by a
//!   [`crate::epoch::Guard`]. Using [`super::Store::guard`] will provide such a guard.
//!
//! * `frozen`: Since "frozen" is a terminal state, it is safe to give out references to
//!   frozen slots.
//!
//!   Note that "frozen" points differ from normal points in that they are protected from
//!   being retired by the [`super::Store`]. This enables potential optimizations like
//!   pre-loading and then freezing all data to be inserted, allowing safe, concurrency-free
//!   read-only index builds. But for now, it mainly helps with algorithm correctness.
//!
//! ## Writable States
//!
//! * [`Exclusive`]: The exclusive stats is a little spooky. Slots can assume that an
//!   [`Exclusive`] for an index `i` is exclusive for its duration. This means that
//!   [`Exclusive`] implementations can lend out mutable references to its contents (for
//!   example, [`super::intrusive::Exclusive::as_mut_slice`]).
//!
//!   Code in [`super`] is very careful to maintain this invariant and all users of
//!   [`Exclusive`] must carefully maintain this as well.
//!
//! * `reclaim`: On a call to [`Slots::reclaim`], implementations may assume exclusive access
//!   to the indicated slot for the duration of the function call.
//!
//! ## Contracts
//!
//! Users of [`Slots`] must ensure that the lifecycle shown above is strictly observed.
//! Furthermore, for [`Exclusive`]s, exactly one of the terminal methods **must** be called.
//!
//! State transitions are driven by the authoritative [`super::Store`]. Before invoking a
//! transition, the store ensures the slot is not externally available in its previous
//! state. Further, the store commits the destination state only after the lifecycle API call
//! completes.
//!
//! ## Safety Considerations
//!
//! [`SlotsConfig::build`] is expected to receive two additional arguments:
//!
//! * [`epoch::RegistryHandle`]: A handle into the [`epoch::Registry`] used to protect the
//!   constructed [`Slots`]. This allows implementations to check the validity of
//!   [`epoch::Guard`]s via [`epoch::RegistryHandle::assert_guard_belongs`].
//!
//!   Unsafe code may rely on this check and all reader construction paths should use it
//!   as it is cheap.
//!
//! * [`tag::Authoritative`]: A reference to the authoritative tag source for the constructed
//!   [`Slots`]. Implementations can use [`tag::Authoritative::read_only`] to obtain a
//!   read-only, synchronizing view into the authoritative tag collection.

use std::fmt::Debug;

use crate::{epoch, num::IdLimit, tag};

use super::Lifecycle;

/// A configuration for a [`Slots`].
pub(crate) trait SlotsConfig: Debug {
    /// The type of the resulting [`Slots`].
    type Slots: Slots;

    /// Construction errors.
    type Error: std::error::Error + Send + Sync + 'static;

    /// Construct [`Self::Slots`].
    ///
    /// The [`IdLimit`] for the resulting store can be obtained from `tags.id_limit()`.
    ///
    /// # Safety
    ///
    /// This assumes that the resulting [`Slots`] is embedded in a [`crate::store::Store`]
    /// and that the following hold:
    ///
    /// * `handle` is an [`epoch::RegistryHandle`] into the [`epoch::Registry`] in the
    ///   containing `Store`.
    ///
    /// * `tags` is the [`tag::Authoritative`] collection in the containing `Store`.
    unsafe fn build(
        self,
        handle: epoch::RegistryHandle,
        tags: &tag::Authoritative,
    ) -> Result<Self::Slots, Self::Error>;
}

/// A lifecycle backend for [`super::Store`]'s EBR scheme.
///
/// See the [module level documentation](self) for details.
pub(crate) trait Slots: Debug + 'static {
    /// The writable [`Exclusive`] slot.
    type Exclusive<'a>: Exclusive;

    /// Return the exclusive upper bound for indices provided to this API.
    ///
    /// Callers should ensure that indices are in the range `[0..self.id_limit())`.
    fn id_limit(&self) -> IdLimit;

    /// Immediately transition slot `i` from the "available" state to the "slot" state.
    ///
    /// Implementations may panic when `i` is out-of-bounds, but must not rely on
    /// `i < Self::id_limit` for memory safety.
    ///
    /// # Safety
    ///
    /// Callers must ensure **all** of the following:
    ///
    /// 1. The slot is in the implicit "available" state according to the [module docs](self).
    ///
    /// 2. Access to slot `i` is exclusive before invoking this method and that exclusivity
    ///    is maintained until the returned [`Exclusive`] is consumed by a terminal method.
    ///
    /// 3. Exactly one of the [`Exclusive`] terminal methods is called. The [`Exclusive`]
    ///    **may not** be dropped or forgotten without one of these methods being called.
    unsafe fn acquire(&self, i: u32, _: Lifecycle) -> Self::Exclusive<'_>;

    /// Transition slot `i` from the "published" state to the "retiring" state.
    ///
    /// Implementations may panic when `i` is out-of-bounds, but must not rely on
    /// `i < Self::id_limit` for memory safety.
    ///
    /// # Safety
    ///
    /// The slot is in the implicit "published" state.
    unsafe fn retire(&self, i: u32, _: Lifecycle);

    /// Transition slot `i` from the "retiring" state to the "available" state.
    ///
    /// Implementations may panic when `i` is out-of-bounds, but must not rely on
    /// `i < Self::id_limit` for memory safety.
    ///
    /// # Safety
    ///
    /// Callers must ensure **all** of the following:
    ///
    /// 1. The slot is in the implicit "retiring" state.
    ///
    /// 2. All [`crate::epoch::Guard`]s for this [`Slots`] that could have obtained a
    ///    reference while this slot was in the "published" state have been dropped.
    unsafe fn reclaim(&self, i: u32, _: Lifecycle);
}

/// A writable slot for [`Slots`].
///
/// [`Exclusive`]s may assume that they have exclusive ownership of their slots for their
/// duration in accordance with [`Slots::acquire`].
pub(crate) trait Exclusive: Debug {
    /// Mark this slot as readable, transition it to the "published" state.
    fn publish(self, _: Lifecycle);

    /// Mark this slot as "frozen".
    fn freeze(self, _: Lifecycle);

    /// Abort any action, returning the slot to "available".
    fn abort(self, _: Lifecycle);
}