uqa_storage/key_value/diskann/identity/
reclamation.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Finite mapping cleanup requires absent generations and fences every new preparation.
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::super::{
19    invalid, keys::Keys as GenerationKeys, read_data_identity, validate_session,
20    KeyValueDiskANNStore,
21};
22use super::{
23    allocation::rejected_mapping_conflict, database_key, fixed, load, INDEX_GUARD_TAG, PREFIX,
24    ROOT, TABLE_GUARD_TAG,
25};
26
27const TABLE_BYTES: usize = PREFIX + 32;
28const INDEX_BYTES: usize = PREFIX + 48;
29
30#[derive(Clone, Copy)]
31struct MappingKey {
32    bytes: [u8; INDEX_BYTES],
33    len: usize,
34}
35
36impl MappingKey {
37    fn decode(key: &[u8], database: [u8; 16]) -> StorageBackendResult<Self> {
38        let valid = (key.len() == TABLE_BYTES && key[ROOT.len()] == 2)
39            || (key.len() == INDEX_BYTES && key[ROOT.len()] == 3);
40        if !valid || !key.starts_with(ROOT) || key[ROOT.len() + 1..PREFIX] != database {
41            return Err(invalid(
42                "mapping discovery requires a complete catalog identity",
43            ));
44        }
45        if key[PREFIX..].chunks_exact(16).any(|id| id == [0; 16]) {
46            return Err(invalid("mapping catalog identity must be nonzero"));
47        }
48        let mut bytes = [0; INDEX_BYTES];
49        bytes[..key.len()].copy_from_slice(key);
50        Ok(Self {
51            bytes,
52            len: key.len(),
53        })
54    }
55
56    fn table(self) -> Self {
57        let mut key = self;
58        key.bytes[ROOT.len()] = 2;
59        key.len = TABLE_BYTES;
60        key
61    }
62
63    fn guard(self) -> Self {
64        let mut key = self;
65        key.bytes[ROOT.len()] = if key.len == TABLE_BYTES {
66            TABLE_GUARD_TAG
67        } else {
68            INDEX_GUARD_TAG
69        };
70        key
71    }
72
73    fn head_prefix(self) -> Self {
74        let mut bytes = [0; INDEX_BYTES];
75        let prefix = super::super::publication::HEAD_PREFIX;
76        bytes[..prefix.len()].copy_from_slice(prefix);
77        let identity = &self.bytes[PREFIX..self.len];
78        bytes[prefix.len()..prefix.len() + identity.len()].copy_from_slice(identity);
79        Self {
80            bytes,
81            len: prefix.len() + identity.len(),
82        }
83    }
84
85    fn index_prefix(self) -> Self {
86        let mut key = self.table();
87        key.bytes[ROOT.len()] = 3;
88        key
89    }
90}
91
92impl AsRef<[u8]> for MappingKey {
93    fn as_ref(&self) -> &[u8] {
94        &self.bytes[..self.len]
95    }
96}
97
98/// A finite retained discovery pass over index mappings followed by table mappings. Each step considers one mapping and removes at most that record and its preparation guard. Generation payloads and retained history are separate maintenance owners.
99pub struct KeyValueDiskANNMappingMaintenance {
100    repository: KeyValueDiskANNStore,
101    read: Option<Arc<dyn KeyValueRead + Send + Sync>>,
102    database: Option<[u8; 16]>,
103    index: bool,
104    after: Option<MappingKey>,
105    current: Option<MappingKey>,
106    control: StorageReadControl,
107    _memory: MemoryReservation,
108}
109
110impl KeyValueDiskANNMappingMaintenance {
111    pub fn start(
112        store: &Arc<dyn KeyValueStore>,
113        control: &StorageReadControl,
114    ) -> StorageBackendResult<Self> {
115        control.check()?;
116        let memory = control.memory().reserve(std::mem::size_of::<Self>())?;
117        let repository = KeyValueDiskANNStore::connect(store, control)?;
118        let retained = repository
119            .owner
120            .store
121            .open_retained_read_session(control.cancellation())?;
122        validate_session(&*repository.owner.store, &*retained)?;
123        let mut read = None;
124        retained.with_read_view(&mut |view| {
125            read = Some(view.retain(&[ROOT])?);
126            Ok(())
127        })?;
128        let read =
129            read.ok_or_else(|| invalid("mapping maintenance did not retain a read boundary"))?;
130        let database = read_data_identity(&*read, control)?;
131        if database.is_none() && read.contains_prefix_budgeted(ROOT, control)? {
132            return Err(invalid("mapping records have no data identity"));
133        }
134        Ok(Self {
135            repository,
136            read: database.map(|_| read),
137            database,
138            index: true,
139            after: None,
140            current: None,
141            control: control.clone(),
142            _memory: memory,
143        })
144    }
145
146    /// `Some(true)` reclaimed a mapping (or observed a peer's deletion); `Some(false)` retained one. `None` finishes the pass. An uncertain outcome preserves this same key until original completion is resolved.
147    pub fn step(&mut self) -> StorageBackendResult<Option<bool>> {
148        self.control.check()?;
149        if self.current.is_none() {
150            self.current = self.next_key()?;
151        }
152        let Some(key) = self.current else {
153            self.read = None;
154            return Ok(None);
155        };
156        let reclaimed = self.repository.reclaim_mapping(
157            key,
158            self.database.expect("discovered data identity"),
159            &self.control,
160        )?;
161        self.after = Some(key);
162        self.current = None;
163        Ok(Some(reclaimed))
164    }
165
166    pub fn commit_pending(&self) -> StorageBackendResult<()> {
167        self.repository.commit_pending()
168    }
169    pub fn rollback_pending(&self) -> StorageBackendResult<()> {
170        self.repository.rollback_pending()
171    }
172
173    pub fn run(
174        store: &Arc<dyn KeyValueStore>,
175        control: &StorageReadControl,
176    ) -> StorageBackendResult<()> {
177        let mut pass = Self::start(store, control)?;
178        while pass.step()?.is_some() {}
179        Ok(())
180    }
181
182    fn next_key(&mut self) -> StorageBackendResult<Option<MappingKey>> {
183        let Some(read) = &self.read else {
184            return Ok(None);
185        };
186        let database = self.database.expect("retained mapping data identity");
187        loop {
188            self.control.check()?;
189            read.control().check()?;
190            let mut prefix = [0; PREFIX];
191            prefix[..ROOT.len()].copy_from_slice(ROOT);
192            prefix[ROOT.len()] = if self.index { 3 } else { 2 };
193            prefix[ROOT.len() + 1..].copy_from_slice(&database);
194            let mut found = None;
195            let mut failed = false;
196            let result = read.visit_keys_after(
197                &prefix,
198                self.after.as_ref().map(AsRef::as_ref),
199                1,
200                &self.control,
201                &mut |bytes| {
202                    if failed || found.is_some() {
203                        failed = true;
204                        return Err(invalid("mapping discovery exceeded its key limit"));
205                    }
206                    let decoded = MappingKey::decode(bytes, database).and_then(|key| {
207                        if !bytes.starts_with(&prefix)
208                            || self
209                                .after
210                                .as_ref()
211                                .is_some_and(|after| bytes <= after.as_ref())
212                        {
213                            return Err(invalid("mapping discovery returned an unordered key"));
214                        }
215                        Ok(key)
216                    });
217                    match decoded {
218                        Ok(key) => {
219                            found = Some(key);
220                            Ok(())
221                        }
222                        Err(error) => {
223                            failed = true;
224                            Err(error)
225                        }
226                    }
227                },
228            );
229            if failed {
230                return Err(invalid(
231                    "mapping discovery suppressed an invalid completion",
232                ));
233            }
234            result?;
235            self.control.check()?;
236            read.control().check()?;
237            if found.is_some() || !self.index {
238                return Ok(found);
239            }
240            self.index = false;
241            self.after = None;
242        }
243    }
244}
245
246impl KeyValueDiskANNStore {
247    fn reclaim_mapping(
248        &self,
249        key: MappingKey,
250        database: [u8; 16],
251        control: &StorageReadControl,
252    ) -> StorageBackendResult<bool> {
253        let mut reclaimed = false;
254        let result = self.mutate(control, &mut |read, batch| {
255            if read_data_identity(read, control)? != Some(database) {
256                return Err(invalid(
257                    "mapping reclamation belongs to another data identity",
258                ));
259            }
260            let Some(handle) = load(read, key.as_ref(), control)? else {
261                reclaimed = true;
262                return Ok(());
263            };
264            let table_key = key.table();
265            let table = if key.len == TABLE_BYTES {
266                handle
267            } else {
268                load(read, table_key.as_ref(), control)?
269                    .ok_or_else(|| invalid("index mapping has no table mapping"))?
270            };
271            let generation = DiskANNGeneration::new(
272                database,
273                table,
274                if key.len == TABLE_BYTES { 1 } else { handle },
275                1,
276            )?;
277            let physical = GenerationKeys::new(generation);
278            let suffix = if key.len == TABLE_BYTES { 16 } else { 8 };
279            if read.contains_prefix_budgeted(key.head_prefix().as_ref(), control)?
280                || read.contains_prefix_budgeted(
281                    &physical.prefix()[..physical.prefix().len() - suffix],
282                    control,
283                )?
284                || (key.len == TABLE_BYTES
285                    && read.contains_prefix_budgeted(key.index_prefix().as_ref(), control)?)
286            {
287                return Ok(());
288            }
289            let guard_key = key.guard();
290            let guard = fixed::<8>(control, |visit| {
291                read.visit_value_bounded(guard_key.as_ref(), 8, control, visit)
292            })?;
293            if guard.is_some_and(|bytes| u64::from_be_bytes(bytes) == 0) {
294                return Err(invalid("mapping preparation revision must be nonzero"));
295            }
296            batch.require_unchanged(&database_key())?;
297            batch.require_unchanged(key.as_ref())?;
298            batch.require_unchanged(guard_key.as_ref())?;
299            batch.delete(key.as_ref())?;
300            if guard.is_some() {
301                batch.delete(guard_key.as_ref())?;
302            }
303            reclaimed = true;
304            Ok(())
305        });
306        match result {
307            Err(error) if rejected_mapping_conflict(&error) => {
308                self.rollback_pending()?;
309                Ok(false)
310            }
311            Err(error) => Err(error),
312            Ok(()) => Ok(reclaimed),
313        }
314    }
315}
uqa_storage/key_value/diskann/identity/reclamation.rs

uqa_storage/key_value/diskann/identity/
reclamation.rs
