uqa_storage/key_value/diskann/
maintenance.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! A finite discovery snapshot and bounded current-state deletion share the owning provider.
8
9use std::sync::Arc;
10
11use uqa_core::memory::MemoryReservation;
12
13use crate::diskann_index::format::DiskANNGeneration;
14use crate::key_value::KeyValueRead;
15use crate::read_control::StorageReadControl;
16use crate::{KeyValueStore, StorageBackendResult};
17
18use super::{
19    invalid,
20    keys::{generation_prefix, state_generation, Keys, ROOT},
21    read_data_identity,
22    reclamation::{Reclamation, ReclamationMode},
23    source::KEY_PAGE_LIMIT,
24    validate_session, versioned_database, KeyValueDiskANNStore,
25};
26
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub enum DiskANNMaintenanceStatus {
29    /// A live build, retained source or selected generation prevents deletion in this pass.
30    Retained,
31    /// One bounded deletion committed; this same generation needs another step.
32    MoreRecords,
33    /// No current records remain for the generation.
34    Reclaimed,
35}
36
37#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub struct DiskANNMaintenanceStep {
39    pub generation: DiskANNGeneration,
40    pub status: DiskANNMaintenanceStatus,
41}
42
43/// Discover only generations visible at startup, without loading artifact payloads or building an identity directory. Each step checks one current generation and deletes at most 64 payload keys. New and earlier generations are reconsidered by the next pass. Errors preserve the current cursor and original physical attempt.
44pub struct KeyValueDiskANNMaintenance {
45    repository: Option<KeyValueDiskANNStore>,
46    read: Option<Arc<dyn KeyValueRead + Send + Sync>>,
47    database: Option<[u8; 16]>,
48    after: Option<DiskANNGeneration>,
49    current: Option<DiskANNGeneration>,
50    control: StorageReadControl,
51    _memory: MemoryReservation,
52}
53
54impl KeyValueDiskANNMaintenance {
55    pub fn start(
56        store: &Arc<dyn KeyValueStore>,
57        control: &StorageReadControl,
58    ) -> StorageBackendResult<Self> {
59        control.check()?;
60        versioned_database(&**store)?;
61        let memory = control.memory().reserve(std::mem::size_of::<Self>())?;
62        let session = store.open_session_with_cancellation(control.cancellation())?;
63        validate_session(&**store, &*session)?;
64        if session.in_transaction() {
65            return Err(invalid(
66                "maintenance requires an independent inactive session",
67            ));
68        }
69        let retained = session.open_retained_read_session(control.cancellation())?;
70        validate_session(&*session, &*retained)?;
71        let mut read = None;
72        retained.with_read_view(&mut |view| {
73            read = Some(view.retain(&[ROOT])?);
74            Ok(())
75        })?;
76        let read = read.ok_or_else(|| invalid("maintenance did not retain a read boundary"))?;
77        let database = read_data_identity(&*read, control)?;
78        if database.is_none() && read.contains_prefix_budgeted(ROOT, control)? {
79            return Err(invalid("generation records have no data identity"));
80        }
81        let repository = database
82            .map(|_| KeyValueDiskANNStore::connect(&session, control))
83            .transpose()?;
84        Ok(Self {
85            repository,
86            read: database.map(|_| read),
87            database,
88            after: None,
89            current: None,
90            control: control.clone(),
91            _memory: memory,
92        })
93    }
94
95    /// `None` completes this finite pass and releases its discovery snapshot. Retained generations do not prevent reaching later generations; incomplete deletion stays on its current generation.
96    pub fn step(&mut self) -> StorageBackendResult<Option<DiskANNMaintenanceStep>> {
97        self.control.check()?;
98        if self.current.is_none() {
99            self.current = self.next_generation()?;
100        }
101        let Some(generation) = self.current else {
102            self.read = None;
103            return Ok(None);
104        };
105        let repository = self
106            .repository
107            .as_ref()
108            .expect("discovered data has a repository");
109        let reclaimed = repository.reclaim_generation(
110            generation,
111            KEY_PAGE_LIMIT,
112            ReclamationMode::Maintenance,
113            &self.control,
114        )?;
115        let status = match reclaimed {
116            Reclamation::Complete => DiskANNMaintenanceStatus::Reclaimed,
117            Reclamation::More => DiskANNMaintenanceStatus::MoreRecords,
118            Reclamation::Retained => DiskANNMaintenanceStatus::Retained,
119        };
120        if reclaimed != Reclamation::More {
121            self.after = Some(generation);
122            self.current = None;
123        }
124        Ok(Some(DiskANNMaintenanceStep { generation, status }))
125    }
126
127    /// Resolve the exact original attempt before resuming a failed step. No discovery or deletion callback is replayed by this operation.
128    pub fn commit_pending(&self) -> StorageBackendResult<()> {
129        if let Some(repository) = &self.repository {
130            repository.commit_pending()?;
131        }
132        Ok(())
133    }
134
135    pub fn rollback_pending(&self) -> StorageBackendResult<()> {
136        if let Some(repository) = &self.repository {
137            repository.rollback_pending()?;
138        }
139        Ok(())
140    }
141
142    /// Drain one finite pass for an existing storage-maintenance invocation. Every step keeps the same per-batch bound; dropping the pass releases discovery history before ordinary provider version reclamation.
143    pub fn run(
144        store: &Arc<dyn KeyValueStore>,
145        control: &StorageReadControl,
146    ) -> StorageBackendResult<()> {
147        let mut pass = Self::start(store, control)?;
148        while pass.step()?.is_some() {}
149        super::KeyValueDiskANNMappingMaintenance::run(store, control)
150    }
151
152    fn next_generation(&self) -> StorageBackendResult<Option<DiskANNGeneration>> {
153        let Some(read) = &self.read else {
154            return Ok(None);
155        };
156        let database = self.database.expect("retained discovery data identity");
157        let after = self
158            .after
159            .map(|generation| Keys::new(generation).after_generation());
160        let mut found = None;
161        let mut failure = None;
162        read.control().check()?;
163        let source = read.visit_keys_after(
164            &generation_prefix(database),
165            after.as_ref().map(AsRef::as_ref),
166            1,
167            &self.control,
168            &mut |key| {
169                if failure.is_some() {
170                    return Err(invalid("generation discovery already failed"));
171                }
172                let outcome = (|| {
173                    self.control.check()?;
174                    read.control().check()?;
175                    if found.is_some() {
176                        return Err(invalid("generation discovery exceeded its key limit"));
177                    }
178                    let generation = state_generation(key)?;
179                    if generation.database() != database
180                        || self.after.is_some_and(|previous| {
181                            Keys::new(previous).prefix() >= Keys::new(generation).prefix()
182                        })
183                    {
184                        return Err(invalid(
185                            "generation discovery returned an unordered or foreign identity",
186                        ));
187                    }
188                    found = Some(generation);
189                    Ok(())
190                })();
191                if let Err(error) = outcome {
192                    failure = Some(error);
193                    return Err(invalid("generation discovery rejected a key"));
194                }
195                Ok(())
196            },
197        );
198        if let Some(error) = failure {
199            return Err(error);
200        }
201        source?;
202        read.control().check()?;
203        self.control.check()?;
204        Ok(found)
205    }
206}
207
208#[cfg(test)]
209mod tests;
uqa_storage/key_value/diskann/maintenance.rs

uqa_storage/key_value/diskann/
maintenance.rs
