use std::sync::atomic::{AtomicI64, Ordering};
use super::reembed::SearchState;
enum ResPhase {
Reserved,
Committed,
Done,
}
pub(crate) struct ReservationGuard<'a> {
state: &'a SearchState,
pending_op_count: Option<&'a AtomicI64>,
rid: &'a str,
seq: u64,
phase: ResPhase,
}
impl<'a> ReservationGuard<'a> {
pub(crate) fn with_pending_op(
state: &'a SearchState,
pending_op_count: &'a AtomicI64,
rid: &'a str,
seq: u64,
) -> Self {
Self {
state,
pending_op_count: Some(pending_op_count),
rid,
seq,
phase: ResPhase::Reserved,
}
}
pub(crate) fn publish_only(state: &'a SearchState, rid: &'a str, seq: u64) -> Self {
Self {
state,
pending_op_count: None,
rid,
seq,
phase: ResPhase::Reserved,
}
}
pub(crate) fn mark_committed(&mut self) {
self.phase = ResPhase::Committed;
}
pub(crate) fn count_pending_op_on_completion(&mut self, counter: &'a AtomicI64) {
self.pending_op_count = Some(counter);
}
pub(crate) fn complete(&mut self) -> bool {
let published = self.discharge_committed();
self.phase = ResPhase::Done;
published
}
fn discharge_committed(&self) -> bool {
let published = self.state.vec_index.publish(self.rid, self.seq);
if let Some(counter) = self.pending_op_count {
counter.fetch_add(1, Ordering::Relaxed);
}
published
}
}
impl Drop for ReservationGuard<'_> {
fn drop(&mut self) {
match self.phase {
ResPhase::Reserved => {
self.state.vec_index.remove_appended(self.rid, self.seq);
}
ResPhase::Committed => {
self.discharge_committed();
}
ResPhase::Done => {}
}
}
}
pub(crate) struct BatchReservationGuard<'a> {
state: &'a SearchState,
entries: Vec<(String, u64)>,
phase: ResPhase,
}
impl<'a> BatchReservationGuard<'a> {
pub(crate) fn new(state: &'a SearchState, capacity: usize) -> Self {
Self {
state,
entries: Vec::with_capacity(capacity),
phase: ResPhase::Reserved,
}
}
pub(crate) fn reserve(
&mut self,
rid: String,
embedding: Vec<f32>,
seq: u64,
) -> crate::error::Result<()> {
if self
.state
.vec_index
.append_reserved(rid.clone(), embedding, seq)?
== crate::vector::delta_index::ReservedAppend::AlreadyPresent
{
return Err(crate::error::YantrikDbError::InvalidInput(format!(
"freshly minted rid {rid} already present in the delta at seq \
{seq} — engine invariant violation"
)));
}
self.entries.push((rid, seq));
Ok(())
}
pub(crate) fn mark_committed(&mut self) {
self.phase = ResPhase::Committed;
}
pub(crate) fn complete(&mut self) -> bool {
let all_published = self.discharge_committed();
self.phase = ResPhase::Done;
all_published
}
fn discharge_committed(&self) -> bool {
let mut all_published = true;
for (rid, seq) in &self.entries {
all_published &= self.state.vec_index.publish(rid, *seq);
}
all_published
}
}
impl Drop for BatchReservationGuard<'_> {
fn drop(&mut self) {
match self.phase {
ResPhase::Reserved => {
for (rid, seq) in &self.entries {
self.state.vec_index.remove_appended(rid, *seq);
}
}
ResPhase::Committed => {
self.discharge_committed();
}
ResPhase::Done => {}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::YantrikDB;
fn db() -> YantrikDB {
YantrikDB::new(":memory:", 8).unwrap()
}
#[test]
fn drop_before_commit_removes_the_reservation() {
let db = db();
let state = db.search_state.load_full();
let _ = state
.vec_index
.append_reserved("r1".into(), vec![0.1; 8], 7)
.unwrap();
{
let _g = ReservationGuard::publish_only(&state, "r1", 7);
}
assert!(
!state.vec_index.remove_appended("r1", 7),
"guard's Drop must have removed the reservation (nothing left to remove)"
);
}
#[test]
fn unwind_before_commit_removes_the_reservation() {
let db = db();
let state = db.search_state.load_full();
let _ = state
.vec_index
.append_reserved("r2".into(), vec![0.1; 8], 8)
.unwrap();
let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let _g = ReservationGuard::publish_only(&state, "r2", 8);
panic!("simulated panic between reserve and commit");
}));
assert!(res.is_err(), "panic must propagate");
assert!(
!state.vec_index.remove_appended("r2", 8),
"unwind must not leak the reservation"
);
}
#[test]
fn unwind_after_commit_publishes_rather_than_stranding() {
let db = db();
let state = db.search_state.load_full();
let _ = state
.vec_index
.append_reserved("r3".into(), vec![0.1; 8], 9)
.unwrap();
let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let mut g = ReservationGuard::publish_only(&state, "r3", 9);
g.mark_committed();
panic!("simulated panic between commit and publish");
}));
assert!(res.is_err(), "panic must propagate");
assert!(
!state.vec_index.publish("r3", 9),
"Drop must have published it already (a second publish finds nothing unpublished)"
);
assert!(
state.vec_index.remove_appended("r3", 9),
"the entry must still EXIST — post-commit it is published, never removed"
);
}
#[test]
fn complete_then_drop_discharges_exactly_once() {
let db = db();
let state = db.search_state.load_full();
let _ = state
.vec_index
.append_reserved("r4".into(), vec![0.1; 8], 10)
.unwrap();
let before = db.pending_op_count.load(Ordering::Relaxed);
{
let mut g = ReservationGuard::with_pending_op(&state, &db.pending_op_count, "r4", 10);
g.mark_committed();
assert!(g.complete(), "complete() publishes and reports it");
}
assert_eq!(
db.pending_op_count.load(Ordering::Relaxed),
before + 1,
"exactly one increment — complete() then Drop must not double-count"
);
}
#[test]
fn publish_only_never_touches_the_pending_counter() {
let db = db();
let state = db.search_state.load_full();
let before = db.pending_op_count.load(Ordering::Relaxed);
let _ = state
.vec_index
.append_reserved("r5".into(), vec![0.1; 8], 11)
.unwrap();
{
let mut g = ReservationGuard::publish_only(&state, "r5", 11);
g.mark_committed();
g.complete();
}
let _ = state
.vec_index
.append_reserved("r6".into(), vec![0.1; 8], 12)
.unwrap();
{
let mut g = ReservationGuard::publish_only(&state, "r6", 12);
g.mark_committed();
}
assert_eq!(
db.pending_op_count.load(Ordering::Relaxed),
before,
"publish_only must never move pending_op_count, on either discharge path"
);
}
#[test]
fn batch_drop_before_commit_removes_every_reservation() {
let db = db();
let state = db.search_state.load_full();
{
let mut g = BatchReservationGuard::new(&state, 3);
g.reserve("b1".into(), vec![0.1; 8], 21).unwrap();
g.reserve("b2".into(), vec![0.2; 8], 22).unwrap();
}
assert!(
!state.vec_index.remove_appended("b1", 21),
"first reservation must have been removed by Drop"
);
assert!(
!state.vec_index.remove_appended("b2", 22),
"second reservation must have been removed by Drop"
);
}
#[test]
fn batch_unwind_after_commit_publishes_all_rather_than_stranding() {
let db = db();
let state = db.search_state.load_full();
let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let mut g = BatchReservationGuard::new(&state, 2);
g.reserve("b3".into(), vec![0.1; 8], 23).unwrap();
g.reserve("b4".into(), vec![0.2; 8], 24).unwrap();
g.mark_committed();
panic!("simulated panic between the batch RELEASE and publish");
}));
assert!(res.is_err(), "panic must propagate");
for (rid, seq) in [("b3", 23u64), ("b4", 24u64)] {
assert!(
!state.vec_index.publish(rid, seq),
"{rid} must already be published by Drop"
);
assert!(
state.vec_index.remove_appended(rid, seq),
"{rid} must still EXIST — post-commit it is published, never removed"
);
}
}
#[test]
fn batch_complete_then_drop_discharges_exactly_once_and_never_counts() {
let db = db();
let state = db.search_state.load_full();
let before = db.pending_op_count.load(Ordering::Relaxed);
{
let mut g = BatchReservationGuard::new(&state, 2);
g.reserve("b5".into(), vec![0.1; 8], 25).unwrap();
g.reserve("b6".into(), vec![0.2; 8], 26).unwrap();
g.mark_committed();
assert!(g.complete(), "complete() publishes all and reports it");
}
assert!(
state.vec_index.remove_appended("b5", 25),
"published entry must survive the Done-phase Drop"
);
assert_eq!(
db.pending_op_count.load(Ordering::Relaxed),
before,
"the batch guard must never move pending_op_count"
);
}
}