uqa_storage/key_value/
diskann.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Database-owned immutable `DiskANN` generations above the common versioned byte store.
8
9use std::sync::Arc;
10
11use parking_lot::Mutex;
12use uqa_core::memory::{Budgeted, BudgetedVec};
13
14use crate::catalog::new_nonzero_catalog_identity;
15use crate::diskann_index::format::DiskANNGeneration;
16use crate::mvcc::{DatabaseId, IdentifierRequest, ResourceLease, WeakResourceLease};
17use crate::read_control::StorageReadControl;
18use crate::{
19    KeyValueStore, StorageBackendError, StorageBackendResult, StorageSessionAffinity,
20    StorageTransactionModel,
21};
22
23use super::{KeyValueMutation, KeyValueRead};
24
25mod build;
26pub(super) mod conformance;
27pub mod identifiers;
28mod identity;
29mod keys;
30mod maintenance;
31mod ownership;
32mod pruning;
33pub mod publication;
34mod reclamation;
35mod source;
36mod staging;
37pub(in crate::key_value) mod state;
38
39pub(super) const READ_PREFIX: &[u8] = keys::ROOT;
40
41use keys::{database_key, Keys};
42use state::{data_identity, fixed};
43
44pub use identity::KeyValueDiskANNMappingMaintenance;
45pub use maintenance::{
46    DiskANNMaintenanceStatus, DiskANNMaintenanceStep, KeyValueDiskANNMaintenance,
47};
48pub use pruning::KeyValueDiskANNPruner;
49pub use source::KeyValueDiskANNSource;
50pub use staging::KeyValueDiskANNStage;
51pub use state::DiskANNStageStatus;
52
53struct Owner {
54    store: Arc<dyn KeyValueStore>,
55    database: DatabaseId,
56    affinity: StorageSessionAffinity,
57    writer: Mutex<Option<ResourceLease>>,
58    leases: Mutex<BudgetedVec<WeakResourceLease>>,
59}
60
61/// A dedicated physical staging session. Its completion methods resolve existing evaluated attempts; they never replay a mutation or complete the caller's SQL transaction.
62#[derive(Clone)]
63pub struct KeyValueDiskANNStore {
64    owner: Arc<Budgeted<Owner>>,
65}
66
67impl KeyValueDiskANNStore {
68    /// Atomically establish catalog mappings and an owned Writing generation. Publication still requires the retained source's current-definition guard and complete build coverage; `start` remains idempotent.
69    pub fn allocate_bound_stage(
70        &self,
71        scope: &crate::diskann_index::catalog::DiskANNIndexScope,
72        control: &StorageReadControl,
73    ) -> StorageBackendResult<KeyValueDiskANNStage> {
74        self.allocate_catalog_stage(scope, control)
75    }
76
77    pub fn connect(
78        source: &Arc<dyn KeyValueStore>,
79        control: &StorageReadControl,
80    ) -> StorageBackendResult<Self> {
81        control.check()?;
82        let database = versioned_database(&**source)?;
83        let store = source.open_session_with_cancellation(control.cancellation())?;
84        let affinity = validate_session(&**source, &*store)?;
85        if store.in_transaction()
86            || store.identifier_allocator().is_none()
87            || store.resource_leases().is_none()
88        {
89            return Err(invalid(
90                "staging requires an inactive session, durable identifiers and resource leases",
91            ));
92        }
93        // Verify retained reads before any initialization writes. The temporary lease does not enumerate records.
94        let retained = store.open_retained_read_session(control.cancellation())?;
95        validate_session(&*store, &*retained)?;
96        drop(retained);
97        let owner = Budgeted::new(
98            Owner {
99                store,
100                database,
101                affinity,
102                writer: Mutex::new(None),
103                leases: Mutex::new(BudgetedVec::new(control.memory())),
104            },
105            control.memory().empty_reservation(),
106        )
107        .into_shared()?;
108        Ok(Self { owner })
109    }
110
111    /// Establish a persistent data identity once. It remains separate from transaction-history identity across backup restoration.
112    pub fn initialize(&self, control: &StorageReadControl) -> StorageBackendResult<[u8; 16]> {
113        let mut selected = None;
114        self.mutate(control, &mut |read, batch| {
115            selected = read_data_identity(read, control)?;
116            if selected.is_none() {
117                let identity = new_nonzero_catalog_identity("DiskANN", "data")?;
118                let mut bytes = [1; 17];
119                bytes[1..].copy_from_slice(&identity);
120                batch.put(&database_key(), &bytes)?;
121                selected = Some(identity);
122            }
123            Ok(())
124        })?;
125        selected.ok_or_else(|| invalid("initialization did not select a data identity"))
126    }
127
128    pub fn data_identity(&self, control: &StorageReadControl) -> StorageBackendResult<[u8; 16]> {
129        let _writer = self.owner.writer.lock();
130        self.idle(control)?;
131        stored_data_identity(&*self.owner.store, control)
132    }
133
134    /// Reserve a never-reused generation number. Starting the returned handle is a separate fallible operation, so the caller retains its identity if creation loses its commit reply.
135    pub fn allocate_stage(
136        &self,
137        table: u64,
138        index: u64,
139        control: &StorageReadControl,
140    ) -> StorageBackendResult<KeyValueDiskANNStage> {
141        let _writer = self.owner.writer.lock();
142        self.idle(control)?;
143        let database = stored_data_identity(&*self.owner.store, control)?;
144        let provisional = DiskANNGeneration::new(database, table, index, 1)?;
145        let keys = Keys::new(provisional);
146        let allocator = self
147            .owner
148            .store
149            .identifier_allocator()
150            .ok_or_else(|| invalid("durable generation identifiers are unavailable"))?;
151        // Custom versioned stores may retain predecessor reservations. Preserve
152        // their per-index floor even when no provider format migration ran.
153        let minimum = match allocator.identifier_watermark(keys.legacy_allocation_namespace())? {
154            Some(previous) => previous.checked_add(1).ok_or_else(|| {
155                crate::mvcc::VersionError::IdentifiersExhausted.into_storage_error()
156            })?,
157            None => 1,
158        };
159        let allocation = allocator.allocate_identifiers(
160            identifiers::NAMESPACE,
161            IdentifierRequest::Reserve {
162                minimum,
163                maximum: u64::MAX,
164                count: std::num::NonZeroU64::MIN,
165            },
166        )?;
167        let generation = DiskANNGeneration::new(database, table, index, allocation.watermark())?;
168        Ok(KeyValueDiskANNStage::reserved(
169            self.clone(),
170            generation,
171            self.reserve_owner(control)?,
172        ))
173    }
174
175    /// Reopen an existing staging owner; an absent or discarded generation cannot be recreated by this path.
176    pub fn resume_stage(
177        &self,
178        generation: DiskANNGeneration,
179        control: &StorageReadControl,
180    ) -> StorageBackendResult<KeyValueDiskANNStage> {
181        KeyValueDiskANNStage::resume(self.clone(), generation, control)
182    }
183
184    pub fn open_source(
185        &self,
186        generation: DiskANNGeneration,
187        control: &StorageReadControl,
188    ) -> StorageBackendResult<Arc<KeyValueDiskANNSource>> {
189        KeyValueDiskANNSource::capture(self, generation, DiskANNStageStatus::Sealed, control)
190    }
191
192    /// Finish the exact retained attempt, including acknowledgement after a durable commit with a lost reply.
193    pub fn commit_pending(&self) -> StorageBackendResult<()> {
194        let mut writer = self.owner.writer.lock();
195        let result = if self.owner.store.in_transaction() {
196            self.owner.store.commit_transaction()
197        } else {
198            Ok(())
199        };
200        if !self.owner.store.in_transaction() {
201            *writer = None;
202        }
203        result
204    }
205
206    /// Resolve or abort the exact retained attempt. A confirmed durable commit remains a committed outcome from the underlying session.
207    pub fn rollback_pending(&self) -> StorageBackendResult<()> {
208        let mut writer = self.owner.writer.lock();
209        let result = if self.owner.store.in_transaction() {
210            self.owner.store.rollback_transaction()
211        } else {
212            Ok(())
213        };
214        if !self.owner.store.in_transaction() {
215            *writer = None;
216        }
217        result
218    }
219
220    fn idle(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
221        control.check()?;
222        if versioned_database(&*self.owner.store)? != self.owner.database
223            || self.owner.store.transaction_affinity().as_ref() != Some(&self.owner.affinity)
224        {
225            return Err(invalid("staging session changed its database or affinity"));
226        }
227        if self.owner.store.in_transaction() {
228            return Err(invalid(
229                "resolve the retained staging attempt before another operation",
230            ));
231        }
232        Ok(())
233    }
234
235    fn mutate(
236        &self,
237        control: &StorageReadControl,
238        operation: &mut KeyValueMutation<'_>,
239    ) -> StorageBackendResult<()> {
240        self.mutate_owned(None, control, operation)
241    }
242
243    fn mutate_owned(
244        &self,
245        lease: Option<&ResourceLease>,
246        control: &StorageReadControl,
247        operation: &mut KeyValueMutation<'_>,
248    ) -> StorageBackendResult<()> {
249        let mut writer = self.owner.writer.lock();
250        self.idle(control)?;
251        *writer = lease.cloned();
252        let result = self.owner.store.with_mutation(&mut |read, batch| {
253            operation(read, batch)?;
254            control.check()
255        });
256        if !self.owner.store.in_transaction() {
257            *writer = None;
258        }
259        result
260    }
261}
262
263fn versioned_database(store: &dyn KeyValueStore) -> StorageBackendResult<DatabaseId> {
264    match store.transaction_model() {
265        StorageTransactionModel::VersionedConcurrent { database }
266            if store.transaction_affinity().is_some() =>
267        {
268            Ok(database)
269        }
270        _ => Err(invalid(
271            "versioned session and retention capabilities are required",
272        )),
273    }
274}
275
276fn validate_session(
277    source: &dyn KeyValueStore,
278    retained: &dyn KeyValueStore,
279) -> StorageBackendResult<StorageSessionAffinity> {
280    if versioned_database(source)? != versioned_database(retained)?
281        || source.storage_identity()? != retained.storage_identity()?
282        || source.transaction_affinity() == retained.transaction_affinity()
283    {
284        return Err(invalid(
285            "independent session has a different provider or shared affinity",
286        ));
287    }
288    retained
289        .transaction_affinity()
290        .ok_or_else(|| invalid("session affinity is missing"))
291}
292
293fn stored_data_identity(
294    store: &dyn KeyValueStore,
295    control: &StorageReadControl,
296) -> StorageBackendResult<[u8; 16]> {
297    fixed(control, |visit| {
298        store.visit_value_bounded(&database_key(), 17, control, visit)
299    })?
300    .map(data_identity)
301    .transpose()?
302    .ok_or_else(|| invalid("persistent data identity is missing"))
303}
304
305fn read_data_identity(
306    read: &dyn KeyValueRead,
307    control: &StorageReadControl,
308) -> StorageBackendResult<Option<[u8; 16]>> {
309    fixed(control, |visit| {
310        read.visit_value_bounded(&database_key(), 17, control, visit)
311    })?
312    .map(data_identity)
313    .transpose()
314}
315
316fn invalid(message: &str) -> StorageBackendError {
317    StorageBackendError::Other(format!("invalid DiskANN Key/Value generation: {message}"))
318}
uqa_storage/key_value/diskann.rs

uqa_storage/key_value/
diskann.rs
