uqa_storage/diskann_index/build/
temporary.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use std::io::{Seek, SeekFrom, Write};
8use std::path::Path;
9use std::sync::Arc;
10
11use parking_lot::Mutex;
12
13use super::invalid;
14use crate::read_control::StorageReadControl;
15use crate::temporary_file::BlockTemporaryFile;
16use crate::{StorageBackendError, StorageBackendResult};
17
18pub(super) type File = BlockTemporaryFile<4096>;
19
20#[derive(Debug, thiserror::Error, PartialEq, Eq)]
21pub enum DiskANNTemporaryError {
22    #[error("temporary storage requires {required} bytes, exceeding limit {limit}")]
23    Limit { required: u64, limit: u64 },
24    #[error("temporary storage size overflow")]
25    SizeOverflow,
26}
27
28impl From<DiskANNTemporaryError> for StorageBackendError {
29    fn from(error: DiskANNTemporaryError) -> Self {
30        StorageBackendError::backend("DiskANN build", error)
31    }
32}
33
34#[derive(Debug)]
35struct State {
36    limit: u64,
37    used: u64,
38    peak: u64,
39}
40
41/// Shared allowance for encrypted file lengths, including framing and rounded blocks. Filesystem allocation metadata is separate; clones share the same live limit.
42#[derive(Debug, Clone)]
43pub struct DiskANNTemporaryBudget(Arc<Mutex<State>>);
44
45impl DiskANNTemporaryBudget {
46    pub fn new(limit: u64) -> Self {
47        Self(Arc::new(Mutex::new(State {
48            limit,
49            used: 0,
50            peak: 0,
51        })))
52    }
53
54    pub fn limit(&self) -> u64 {
55        self.0.lock().limit
56    }
57    pub fn used(&self) -> u64 {
58        self.0.lock().used
59    }
60    pub fn peak(&self) -> u64 {
61        self.0.lock().peak
62    }
63
64    pub(super) fn shares_allowance(&self, other: &Self) -> bool {
65        Arc::ptr_eq(&self.0, &other.0)
66    }
67}
68
69struct Reservation {
70    budget: DiskANNTemporaryBudget,
71    bytes: u64,
72}
73
74impl Reservation {
75    fn grow_to(&mut self, bytes: u64) -> StorageBackendResult<()> {
76        let additional = bytes.saturating_sub(self.bytes);
77        let mut state = self.budget.0.lock();
78        let required = state
79            .used
80            .checked_add(additional)
81            .ok_or(DiskANNTemporaryError::SizeOverflow)?;
82        if required > state.limit {
83            return Err(DiskANNTemporaryError::Limit {
84                required,
85                limit: state.limit,
86            }
87            .into());
88        }
89        state.used = required;
90        state.peak = state.peak.max(required);
91        self.bytes += additional;
92        Ok(())
93    }
94}
95
96impl Drop for Reservation {
97    fn drop(&mut self) {
98        self.budget.0.lock().used -= self.bytes;
99    }
100}
101
102pub(super) struct TemporaryRun {
103    // The file closes before its physical-byte reservation is released.
104    file: File,
105    reservation: Reservation,
106    length: u64,
107    failed: bool,
108}
109
110impl TemporaryRun {
111    pub(super) fn uses_allowance(&self, budget: &DiskANNTemporaryBudget) -> bool {
112        self.reservation.budget.shares_allowance(budget)
113    }
114
115    pub(super) fn new(
116        directory: &Path,
117        budget: &DiskANNTemporaryBudget,
118        control: &StorageReadControl,
119    ) -> StorageBackendResult<Self> {
120        control.check()?;
121        Ok(Self {
122            file: File::new_in(directory).map_err(io_error)?,
123            reservation: Reservation {
124                budget: budget.clone(),
125                bytes: 0,
126            },
127            length: 0,
128            failed: false,
129        })
130    }
131
132    /// A failed write poisons this unpublished run. Its owner must discard it, retaining its entire reservation until the file closes.
133    pub(super) fn append(
134        &mut self,
135        bytes: &[u8],
136        control: &StorageReadControl,
137    ) -> StorageBackendResult<()> {
138        control.check()?;
139        if self.failed {
140            return Err(invalid("temporary run previously failed"));
141        }
142        let end = self
143            .length
144            .checked_add(bytes.len() as u64)
145            .ok_or(DiskANNTemporaryError::SizeOverflow)?;
146        self.reservation
147            .grow_to(File::physical_len_for(end).map_err(io_error)?)?;
148        self.failed = true;
149        self.file
150            .seek(SeekFrom::Start(self.length))
151            .map_err(io_error)?;
152        for chunk in bytes.chunks(4096) {
153            control.check()?;
154            self.file.write_all(chunk).map_err(io_error)?;
155        }
156        control.check()?;
157        self.length = end;
158        self.failed = false;
159        Ok(())
160    }
161
162    pub(super) fn read<T>(
163        &self,
164        control: &StorageReadControl,
165        read: impl FnOnce(&mut File) -> StorageBackendResult<T>,
166    ) -> StorageBackendResult<T> {
167        control.check()?;
168        if self.failed {
169            return Err(invalid("temporary run previously failed"));
170        }
171        let mut file = self.file.reopen().map_err(io_error)?;
172        let result = read(&mut file)?;
173        control.check()?;
174        Ok(result)
175    }
176
177    #[cfg(test)]
178    pub(super) fn path(&self) -> &Path {
179        self.file.path()
180    }
181
182    #[cfg(test)]
183    pub(super) fn block_io_counts(&self) -> (usize, u64) {
184        self.file.block_io_counts()
185    }
186}
187
188pub(super) fn io_error(error: std::io::Error) -> StorageBackendError {
189    StorageBackendError::backend("DiskANN temporary storage", error)
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195
196    #[test]
197    fn failed_partial_run_is_unreadable_and_retains_its_charge_until_discarded() {
198        let directory = tempfile::tempdir().unwrap();
199        let control = StorageReadControl::with_limit(4096);
200        let temporary = DiskANNTemporaryBudget::new(1 << 20);
201        let mut run = TemporaryRun::new(directory.path(), &temporary, &control).unwrap();
202        let path = run.path().to_owned();
203        run.file.fail_write_after(4096 + 24 + 2 + 16 + 2 + 30);
204        assert!(matches!(
205            run.append(&[1; 9000], &control),
206            Err(StorageBackendError::Backend { .. })
207        ));
208        assert_eq!(temporary.used(), File::physical_len_for(9000).unwrap());
209        assert!(std::fs::metadata(&path).unwrap().len() < temporary.used());
210        assert!(run.read(&control, |_| Ok(())).is_err());
211        assert!(run.append(&[2], &control).is_err());
212        drop(run);
213        assert_eq!(temporary.used(), 0);
214        assert!(!path.exists());
215    }
216}
uqa_storage/diskann_index/build/temporary.rs

uqa_storage/diskann_index/build/
temporary.rs
