diskann_record/backend/
memory.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! In-memory save/load contexts.
7//!
8//! [`MemorySaveContext`] and [`MemoryContext`] mirror the disk-backed contexts but
9//! keep the manifest value and every side-car artifact in memory. Saving through an
10//! [`MemorySaveContext`] yields an [`MemoryContext`] (its `SaveContext::Output`) that can
11//! be loaded directly via the `load` entry point.
12//!
13//! Unlike the disk path, this round trip never serializes through JSON: the manifest
14//! [`Value`] is deep-copied to `'static` via [`Value::into_owned`] and side-car artifacts
15//! are buffered in memory through a [`std::io::Cursor`]. As a result these contexts are
16//! available regardless of the `disk` / `serde` features.
17
18use std::{collections::HashMap, io::Cursor, sync::Mutex};
19
20use crate::{
21    Value,
22    load::{self, LoadContext, Reader},
23    save::{self, Handle, SaveContext, Writer, delegate_write_and_seek},
24};
25
26/// A save-side `SaveContext` that keeps every side-car artifact and the committed
27/// manifest value in memory.
28///
29/// `SaveContext::finish` consumes the context and returns an [`MemoryContext`] ready
30/// to be loaded with the `load` entry point.
31///
32/// # Cleanup on failure
33///
34/// Failures in the same process are automatically cleaned up.
35#[derive(Debug, Default)]
36pub struct MemorySaveContext {
37    files: Mutex<HashMap<String, Option<Vec<u8>>>>,
38}
39
40impl MemorySaveContext {
41    /// Create an empty in-memory save context.
42    pub fn new() -> Self {
43        Self::default()
44    }
45}
46
47impl SaveContext for MemorySaveContext {
48    type Output = MemoryContext;
49
50    fn write(&self, key: Option<&str>) -> save::Result<Writer<'_>> {
51        // Mirror the disk context: a human-readable hint must be a simple relative file
52        // name. Absolute paths, parent traversal, and multi-component paths cannot produce
53        // a single, well-formed key, so they are ignored and treated as if no hint had been
54        // supplied.
55        let key = key.filter(|key| {
56            let mut components = std::path::Path::new(key).components();
57            matches!(components.next(), Some(std::path::Component::Normal(_)))
58                && components.next().is_none()
59        });
60
61        let mut files = self
62            .files
63            .lock()
64            .unwrap_or_else(|poison| poison.into_inner());
65
66        // Prefix each artifact with the count of artifacts written so far, so that reusing
67        // the same `key` (or omitting it) still yields a unique name.
68        let name = match key {
69            Some(key) => format!("{:03}-{}", files.len(), key),
70            None => format!("{:03}", files.len()),
71        };
72
73        if files.contains_key(&name) {
74            return Err(save::Error::message(format!(
75                "generated artifact name {:?} collides with an existing artifact",
76                name,
77            )));
78        }
79        // Reserve the name with an empty slot so the count advances and concurrent writers
80        // cannot collide; the slot is filled with the real bytes by `Writer::finish`. A slot
81        // that is never filled is reported by `SaveContext::finish`.
82        files.insert(name.clone(), None);
83        drop(files);
84
85        Ok(Writer::new(
86            MemoryWriter {
87                cursor: Cursor::new(Vec::new()),
88                parent: self,
89            },
90            name,
91        ))
92    }
93
94    fn finish(self, value: Value<'_>) -> save::Result<MemoryContext> {
95        let files = self
96            .files
97            .into_inner()
98            .unwrap_or_else(|poison| poison.into_inner());
99        let files = files
100            .into_iter()
101            .map(|(name, bytes)| match bytes {
102                Some(bytes) => Ok((name, bytes)),
103                None => Err(save::Error::message(format!(
104                    "artifact {:?} was reserved but never finished",
105                    name,
106                ))),
107            })
108            .collect::<save::Result<HashMap<_, _>>>()?;
109        Ok(MemoryContext {
110            files,
111            value: value.into_owned(),
112        })
113    }
114}
115
116/// An in-memory [`WriterInner`](save::WriterInner) that buffers bytes in a [`Cursor`] and,
117/// on finish, deposits the completed buffer into its parent context's file store.
118#[derive(Debug)]
119struct MemoryWriter<'a> {
120    cursor: Cursor<Vec<u8>>,
121    parent: &'a MemorySaveContext,
122}
123
124impl save::WriterInner for MemoryWriter<'_> {
125    fn finish(self: Box<Self>, name: String) -> save::Result<Handle> {
126        self.parent
127            .files
128            .lock()
129            .unwrap_or_else(|poison| poison.into_inner())
130            .insert(name.clone(), Some(self.cursor.into_inner()));
131        Ok(Handle::new(name))
132    }
133}
134
135delegate_write_and_seek!(cursor, MemoryWriter<'_>);
136
137/// A load-side `LoadContext` backed entirely by in-memory buffers.
138///
139/// Produced by [`MemorySaveContext`] via `SaveContext::finish`. Holds the committed
140/// manifest [`Value`] and every side-car artifact as an in-memory byte buffer, so loading
141/// never serializes through JSON or touches the filesystem.
142#[derive(Debug)]
143pub struct MemoryContext {
144    files: HashMap<String, Vec<u8>>,
145    value: Value<'static>,
146}
147
148impl LoadContext for MemoryContext {
149    fn value(&self) -> load::Result<&Value<'_>> {
150        Ok(&self.value)
151    }
152
153    fn read(&self, key: &str) -> load::Result<Reader<'_>> {
154        match self.files.get(key) {
155            Some(bytes) => Ok(Reader::new(Cursor::new(bytes.as_slice()))),
156            None => Err(
157                load::Error::from(load::error::Kind::MissingFile).context(format!(
158                    "handle references artifact {:?} which is not registered in this context",
159                    key,
160                )),
161            ),
162        }
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use std::io::{Read, Write};
169
170    use super::*;
171    use crate::{Version, load, save};
172
173    #[derive(Debug, PartialEq)]
174    struct Doc {
175        name: String,
176        blob: Vec<u8>,
177    }
178
179    impl save::Save for Doc {
180        const VERSION: Version = Version::new(1, 0);
181        fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
182            let mut io = context.write(Some("blob.bin"))?;
183            io.write_all(&self.blob).map_err(save::Error::new)?;
184            let mut record = crate::save_fields!(self, context, [name]);
185            record.insert("blob", io.finish()?)?;
186            Ok(record)
187        }
188    }
189
190    impl load::Load<'_> for Doc {
191        const VERSION: Version = Version::new(1, 0);
192        fn load(object: load::Object<'_>) -> load::Result<Self> {
193            crate::load_fields!(object, [name: String, blob: save::Handle]);
194            let mut io = object.read(&blob)?;
195            let mut blob = Vec::new();
196            io.read_to_end(&mut blob).map_err(load::Error::new)?;
197            Ok(Self { name, blob })
198        }
199        fn load_legacy(_: load::Object<'_>) -> load::Result<Self> {
200            Err(load::error::Kind::UnknownVersion.into())
201        }
202    }
203
204    #[test]
205    fn round_trips_without_serde_or_disk() {
206        let doc = Doc {
207            name: "example".to_owned(),
208            blob: vec![1, 2, 3, 4, 5],
209        };
210
211        let context = save::save(&doc, MemorySaveContext::new()).unwrap();
212        let restored: Doc = load::load(&context).unwrap();
213
214        assert_eq!(doc, restored);
215    }
216
217    #[test]
218    fn read_rejects_unregistered_artifact() {
219        let context = MemoryContext {
220            files: HashMap::new(),
221            value: Value::Null,
222        };
223        let err = context
224            .read("missing.bin")
225            .err()
226            .expect("an unregistered artifact must be rejected");
227        assert!(format!("{err}").contains("not registered in this context"));
228    }
229
230    #[test]
231    fn write_rejects_name_collision() {
232        let ctx = MemorySaveContext::new();
233        // The count prefix normally makes a collision impossible, so seed the bookkeeping map
234        // directly with the exact name the next `write` will generate (one entry => count 1 =>
235        // `001-artifact.bin`).
236        ctx.files
237            .lock()
238            .unwrap()
239            .insert("001-artifact.bin".to_string(), None);
240        let err = SaveContext::write(&ctx, Some("artifact.bin"))
241            .expect_err("a generated name that is already registered must be rejected");
242        assert!(format!("{err}").contains("collides with an existing artifact"));
243    }
244
245    #[test]
246    fn finish_rejects_unfinished_artifact() {
247        let ctx = MemorySaveContext::new();
248        // Reserve an artifact slot but drop the writer without calling `finish`, leaving the
249        // slot empty.
250        let writer = SaveContext::write(&ctx, Some("artifact.bin")).unwrap();
251        drop(writer);
252        let err = ctx
253            .finish(Value::Null)
254            .expect_err("finish must fail when an artifact was reserved but never finished");
255        assert!(format!("{err}").contains("was reserved but never finished"));
256    }
257
258    #[test]
259    fn write_names_anonymous_artifact_with_count_prefix() {
260        let ctx = MemorySaveContext::new();
261        // Passing `None` as the hint exercises the count-only naming branch.
262        let handle = SaveContext::write(&ctx, None).unwrap().finish().unwrap();
263        assert_eq!(handle.as_str(), "000");
264    }
265
266    #[test]
267    fn load_dispatches_to_load_legacy_on_version_mismatch() {
268        // Build an object whose version does not match `Doc::VERSION` (1.0) so the `Loadable`
269        // blanket dispatches to `Doc::load_legacy`, which refuses with `UnknownVersion`.
270        let value = save::Record::empty()
271            .into_value(Version::new(2, 0))
272            .into_owned();
273        let context = MemoryContext {
274            files: HashMap::new(),
275            value,
276        };
277        let err = load::load::<Doc, _>(&context)
278            .expect_err("a version mismatch must dispatch to load_legacy, which refuses");
279        assert!(format!("{err}").contains("unknown version"));
280    }
281
282    // Regression for the NaN float-field bug, demonstrated end-to-end through the in-memory
283    // backend (no serde / disk required). A struct carrying NaN-valued `f64` / `f32` fields is
284    // saved and reloaded; the reload previously FAILED with `NumberOutOfRange` because
285    // `Number::as_f64` / `as_f32` rejected NaN.
286    #[derive(Debug)]
287    struct Floats {
288        finite: f64,
289        nan64: f64,
290        nan32: f32,
291    }
292
293    impl save::Save for Floats {
294        const VERSION: Version = Version::new(1, 0);
295        fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
296            Ok(crate::save_fields!(self, context, [finite, nan64, nan32]))
297        }
298    }
299
300    impl load::Load<'_> for Floats {
301        const VERSION: Version = Version::new(1, 0);
302        fn load(object: load::Object<'_>) -> load::Result<Self> {
303            crate::load_fields!(object, [finite: f64, nan64: f64, nan32: f32]);
304            Ok(Self {
305                finite,
306                nan64,
307                nan32,
308            })
309        }
310        fn load_legacy(_: load::Object<'_>) -> load::Result<Self> {
311            Err(load::error::Kind::UnknownVersion.into())
312        }
313    }
314
315    #[test]
316    fn nan_float_fields_round_trip_in_memory() {
317        let value = Floats {
318            finite: 1.5,
319            nan64: f64::NAN,
320            nan32: f32::NAN,
321        };
322
323        let context = save::save(&value, MemorySaveContext::new()).unwrap();
324        let restored: Floats = load::load(&context).expect("NaN float fields must reload");
325
326        assert_eq!(restored.finite, 1.5);
327        assert!(restored.nan64.is_nan(), "f64 NaN field lost on reload");
328        assert!(restored.nan32.is_nan(), "f32 NaN field lost on reload");
329    }
330}
diskann_record/backend/memory.rs

diskann_record/backend/
memory.rs
