use std::collections::HashMap;
use std::sync::Mutex;
pub const BACKFILLING_AFTER_SECS: f64 = 5.0;
pub const ACTIVE_AFTER_SECS: f64 = 15.0;
pub const SEARCHABLE_AFTER_SECS: f64 = 30.0;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum VectorIndexPhase {
Allocating,
Backfilling,
ActiveNotSearchable,
Searchable,
}
impl VectorIndexPhase {
pub(crate) fn index_status(self) -> &'static str {
match self {
Self::Allocating | Self::Backfilling => "CREATING",
Self::ActiveNotSearchable | Self::Searchable => "ACTIVE",
}
}
pub(crate) fn backfilling(self) -> Option<bool> {
match self {
Self::Allocating => Some(false),
Self::Backfilling => Some(true),
Self::ActiveNotSearchable | Self::Searchable => None,
}
}
pub(crate) fn is_searchable(self) -> bool {
matches!(self, Self::Searchable)
}
pub(crate) fn is_creating(self) -> bool {
matches!(self, Self::Allocating | Self::Backfilling)
}
pub(crate) fn holds_table_updating(self) -> bool {
matches!(self, Self::Allocating)
}
pub(crate) fn accepts_delete(self) -> bool {
!matches!(self, Self::Allocating)
}
}
pub(crate) fn phase_of(armed_at: Option<f64>, now: f64) -> VectorIndexPhase {
let Some(armed_at) = armed_at else {
return VectorIndexPhase::Searchable;
};
let elapsed = now - armed_at;
if elapsed < BACKFILLING_AFTER_SECS {
VectorIndexPhase::Allocating
} else if elapsed < ACTIVE_AFTER_SECS {
VectorIndexPhase::Backfilling
} else if elapsed < SEARCHABLE_AFTER_SECS {
VectorIndexPhase::ActiveNotSearchable
} else {
VectorIndexPhase::Searchable
}
}
#[derive(Debug, Clone, Default)]
pub(crate) struct VectorIndexPhases(HashMap<String, VectorIndexPhase>);
impl VectorIndexPhases {
pub(crate) fn get(&self, index_name: &str) -> VectorIndexPhase {
self.0
.get(index_name)
.copied()
.unwrap_or(VectorIndexPhase::Searchable)
}
pub(crate) fn any_holding_table_updating(&self) -> bool {
self.0.values().any(|p| p.holds_table_updating())
}
pub(crate) fn any_creating(&self) -> bool {
self.0.values().any(|p| p.is_creating())
}
}
#[derive(Default)]
pub struct VectorIndexLifecycle {
armed: Mutex<HashMap<String, HashMap<String, f64>>>,
}
impl VectorIndexLifecycle {
pub fn new() -> Self {
Self::default()
}
fn armed(&self) -> std::sync::MutexGuard<'_, HashMap<String, HashMap<String, f64>>> {
self.armed.lock().unwrap_or_else(|e| e.into_inner())
}
pub(crate) fn arm<S: crate::storage_backend::StorageBackend>(
&self,
storage: &S,
table_name: &str,
index_name: &str,
) {
self.arm_at(table_name, index_name, now(storage));
}
pub(crate) fn disarm(&self, table_name: &str, index_name: &str) {
let mut armed = self.armed();
if let Some(table) = armed.get_mut(table_name) {
table.remove(index_name);
if table.is_empty() {
armed.remove(table_name);
}
}
}
pub(crate) fn forget_table(&self, table_name: &str) {
self.armed().remove(table_name);
}
#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
pub(crate) fn forget_all(&self) {
self.armed().clear();
}
#[cfg(test)]
pub(crate) fn is_empty(&self) -> bool {
self.armed().is_empty()
}
pub(crate) fn phases_armed_on<S: crate::storage_backend::StorageBackend>(
&self,
storage: &S,
table_name: &str,
) -> VectorIndexPhases {
self.phases_armed_on_at(table_name, now(storage))
}
pub(crate) fn phases_of<'a, S: crate::storage_backend::StorageBackend>(
&self,
storage: &S,
table_name: &str,
index_names: impl Iterator<Item = &'a str>,
) -> VectorIndexPhases {
self.phases_of_at(table_name, index_names, now(storage))
}
pub(crate) fn phase_of_index<S: crate::storage_backend::StorageBackend>(
&self,
storage: &S,
table_name: &str,
index_name: &str,
) -> VectorIndexPhase {
phase_of(self.armed_at(table_name, index_name), now(storage))
}
fn arm_at(&self, table_name: &str, index_name: &str, now: f64) {
self.armed()
.entry(table_name.to_string())
.or_default()
.insert(index_name.to_string(), now);
}
fn armed_at(&self, table_name: &str, index_name: &str) -> Option<f64> {
self.armed().get(table_name)?.get(index_name).copied()
}
fn phases_armed_on_at(&self, table_name: &str, now: f64) -> VectorIndexPhases {
VectorIndexPhases(
self.armed()
.get(table_name)
.into_iter()
.flatten()
.map(|(name, armed_at)| (name.clone(), phase_of(Some(*armed_at), now)))
.collect(),
)
}
fn phases_of_at<'a>(
&self,
table_name: &str,
index_names: impl Iterator<Item = &'a str>,
now: f64,
) -> VectorIndexPhases {
let armed = self.armed();
let table = armed.get(table_name);
VectorIndexPhases(
index_names
.map(|name| {
let armed_at = table.and_then(|t| t.get(name)).copied();
(name.to_string(), phase_of(armed_at, now))
})
.collect(),
)
}
}
fn now<S: crate::storage_backend::StorageBackend>(storage: &S) -> f64 {
storage.clock().now_unix_secs_f64()
}
#[cfg(test)]
mod tests {
use super::*;
const ARMED: f64 = 1_700_000_000.0;
#[test]
fn while_allocating_the_table_updates_and_the_index_refuses_everything() {
let phase = phase_of(Some(ARMED), ARMED + BACKFILLING_AFTER_SECS - 0.001);
assert_eq!(phase, VectorIndexPhase::Allocating);
assert_eq!(phase.index_status(), "CREATING");
assert_eq!(phase.backfilling(), Some(false));
assert!(!phase.is_searchable());
assert!(phase.is_creating());
assert!(phase.holds_table_updating());
assert!(!phase.accepts_delete());
}
#[test]
fn while_backfilling_the_table_is_active_and_the_index_takes_a_delete() {
let phase = phase_of(Some(ARMED), ARMED + ACTIVE_AFTER_SECS - 0.001);
assert_eq!(phase, VectorIndexPhase::Backfilling);
assert_eq!(phase.index_status(), "CREATING");
assert_eq!(phase.backfilling(), Some(true));
assert!(!phase.is_searchable());
assert!(phase.is_creating());
assert!(!phase.holds_table_updating());
assert!(phase.accepts_delete());
}
#[test]
fn between_the_thresholds_reports_active_and_still_refuses_a_search() {
let phase = phase_of(Some(ARMED), ARMED + ACTIVE_AFTER_SECS + 0.001);
assert_eq!(phase, VectorIndexPhase::ActiveNotSearchable);
assert_eq!(phase.index_status(), "ACTIVE");
assert_eq!(phase.backfilling(), None);
assert!(!phase.is_searchable());
assert!(!phase.is_creating());
assert!(!phase.holds_table_updating());
assert!(phase.accepts_delete());
}
#[test]
fn beyond_the_searchable_threshold_serves_the_search() {
let phase = phase_of(Some(ARMED), ARMED + SEARCHABLE_AFTER_SECS + 0.001);
assert_eq!(phase, VectorIndexPhase::Searchable);
assert_eq!(phase.index_status(), "ACTIVE");
assert_eq!(phase.backfilling(), None);
assert!(phase.is_searchable());
}
#[test]
fn exactly_on_a_threshold_resolves_to_the_later_phase() {
assert_eq!(
phase_of(Some(ARMED), ARMED + BACKFILLING_AFTER_SECS),
VectorIndexPhase::Backfilling
);
assert_eq!(
phase_of(Some(ARMED), ARMED + ACTIVE_AFTER_SECS),
VectorIndexPhase::ActiveNotSearchable
);
assert_eq!(
phase_of(Some(ARMED), ARMED + SEARCHABLE_AFTER_SECS),
VectorIndexPhase::Searchable
);
}
#[test]
fn an_index_with_no_entry_is_active_and_searchable() {
let phase = phase_of(None, ARMED);
assert_eq!(phase, VectorIndexPhase::Searchable);
assert_eq!(phase.index_status(), "ACTIVE");
assert_eq!(phase.backfilling(), None);
assert!(phase.is_searchable());
}
#[test]
fn a_clock_that_went_backwards_reads_as_just_armed() {
assert_eq!(
phase_of(Some(ARMED), ARMED - 60.0),
VectorIndexPhase::Allocating
);
}
#[test]
fn reading_does_not_change_what_a_later_read_returns() {
let lifecycle = VectorIndexLifecycle::new();
lifecycle.arm_at("t", "vix", ARMED);
for _ in 0..50 {
let phases = lifecycle.phases_of_at("t", ["vix"].into_iter(), ARMED + 1.0);
assert_eq!(phases.get("vix"), VectorIndexPhase::Allocating);
assert!(phases.any_creating());
}
let later = lifecycle.phases_of_at("t", ["vix"].into_iter(), ARMED + SEARCHABLE_AFTER_SECS);
assert_eq!(later.get("vix"), VectorIndexPhase::Searchable);
assert!(!later.any_creating());
}
#[test]
fn the_unfiltered_set_carries_every_armed_index_on_that_table_alone() {
let lifecycle = VectorIndexLifecycle::new();
lifecycle.arm_at("t", "a", ARMED);
lifecycle.arm_at("t", "b", ARMED - SEARCHABLE_AFTER_SECS);
lifecycle.arm_at("other", "c", ARMED);
let phases = lifecycle.phases_armed_on_at("t", ARMED);
assert_eq!(phases.get("a"), VectorIndexPhase::Allocating);
assert_eq!(phases.get("b"), VectorIndexPhase::Searchable);
assert_eq!(phases.get("c"), VectorIndexPhase::Searchable);
assert!(phases.any_creating());
let none = lifecycle.phases_armed_on_at("nothing-armed-here", ARMED);
assert_eq!(none.get("a"), VectorIndexPhase::Searchable);
assert!(!none.any_creating());
}
#[test]
fn an_index_the_table_no_longer_defines_is_not_reported() {
let lifecycle = VectorIndexLifecycle::new();
lifecycle.arm_at("t", "gone", ARMED);
let phases = lifecycle.phases_of_at("t", ["still-here"].into_iter(), ARMED + 1.0);
assert_eq!(phases.get("still-here"), VectorIndexPhase::Searchable);
assert!(!phases.any_creating());
}
#[test]
fn disarming_forgets_one_index_and_leaves_its_neighbour() {
let lifecycle = VectorIndexLifecycle::new();
lifecycle.arm_at("t", "a", ARMED);
lifecycle.arm_at("t", "b", ARMED);
lifecycle.disarm("t", "a");
let phases = lifecycle.phases_of_at("t", ["a", "b"].into_iter(), ARMED + 1.0);
assert_eq!(phases.get("a"), VectorIndexPhase::Searchable);
assert_eq!(phases.get("b"), VectorIndexPhase::Allocating);
}
#[test]
fn forgetting_a_table_clears_every_index_on_it() {
let lifecycle = VectorIndexLifecycle::new();
lifecycle.arm_at("t", "a", ARMED);
lifecycle.arm_at("other", "a", ARMED);
lifecycle.forget_table("t");
let gone = lifecycle.phases_of_at("t", ["a"].into_iter(), ARMED + 1.0);
assert_eq!(gone.get("a"), VectorIndexPhase::Searchable);
let kept = lifecycle.phases_of_at("other", ["a"].into_iter(), ARMED + 1.0);
assert_eq!(kept.get("a"), VectorIndexPhase::Allocating);
}
#[test]
fn re_arming_restarts_the_window() {
let lifecycle = VectorIndexLifecycle::new();
lifecycle.arm_at("t", "vix", ARMED);
lifecycle.arm_at("t", "vix", ARMED + SEARCHABLE_AFTER_SECS);
let phases = lifecycle.phases_of_at(
"t",
["vix"].into_iter(),
ARMED + SEARCHABLE_AFTER_SECS + 1.0,
);
assert_eq!(phases.get("vix"), VectorIndexPhase::Allocating);
}
}