diskann_record/save/context.rs
1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Save-side context and side-car writer.
7//!
8//! [`Context`] is the cheap handle handed to every [`super::Save::save`] impl. It owns
9//! nothing visible to the user; cloning it is free, and it can be passed to children to
10//! propagate the same artifact-tracking state.
11//!
12//! [`Writer`] is the borrowed side-car artifact handle returned by [`Context::write`].
13//! It implements [`std::io::Write`] and [`std::io::Seek`]; calling
14//! [`Writer::finish`] flushes the buffer and yields a [`Handle`] that can be inserted
15//! into a [`super::Record`].
16
17use std::io::BufWriter;
18
19use crate::save::{Error, Handle, Result, Value};
20
21/// Generate forwarding [`std::io::Write`] and [`std::io::Seek`] impls for `$T` that
22/// delegate every method to its `$field` member.
23macro_rules! delegate_write_and_seek {
24 ($field:ident, $T:ty) => {
25 impl std::io::Write for $T {
26 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
27 self.$field.write(buf)
28 }
29 fn flush(&mut self) -> std::io::Result<()> {
30 self.$field.flush()
31 }
32 fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
33 self.$field.write_vectored(bufs)
34 }
35 fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
36 self.$field.write_all(buf)
37 }
38 fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::io::Result<()> {
39 self.$field.write_fmt(args)
40 }
41 }
42
43 impl std::io::Seek for $T {
44 fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
45 self.$field.seek(pos)
46 }
47 fn rewind(&mut self) -> std::io::Result<()> {
48 self.$field.rewind()
49 }
50 fn stream_position(&mut self) -> std::io::Result<u64> {
51 self.$field.stream_position()
52 }
53 fn seek_relative(&mut self, offset: i64) -> std::io::Result<()> {
54 self.$field.seek_relative(offset)
55 }
56 }
57 };
58}
59
60pub(crate) use delegate_write_and_seek;
61
62/// The backing store for a save operation.
63///
64/// A `SaveContext` decides where side-car artifacts are written ([`SaveContext::write`])
65/// and how the final manifest is committed ([`SaveContext::finish`]). The concrete
66/// implementations live under `crate::backend`: a disk-backed context (under the `disk`
67/// feature) and an in-memory context. Alternative implementations (e.g. a virtual
68/// filesystem) can be supplied for testing or to avoid touching the filesystem.
69///
70/// The generic [`save`](super::save) entry point is parameterized over this trait so
71/// that the base crate carries no hard dependency on any particular implementation.
72pub(crate) trait SaveContext {
73 /// The value produced once the manifest has been committed by
74 /// [`SaveContext::finish`]. For the disk-backed context this is `()`.
75 type Output;
76
77 /// Allocate a new side-car artifact, optionally tagged with a human-readable `key`.
78 ///
79 /// The artifact's on-disk name (and the value stored in its [`Handle`]) is prefixed
80 /// with the count of artifacts written so far; when `key` is `Some`, it is appended as
81 /// `{count}-{key}` purely to aid debugging. Reusing the same `key` across calls is
82 /// allowed — the count prefix disambiguates the artifacts.
83 ///
84 /// # Errors
85 ///
86 /// Returns [`Error`] if `key` is `Some` but not a simple relative file name, or if the
87 /// underlying artifact cannot be created.
88 fn write(&self, key: Option<&str>) -> Result<Writer<'_>>;
89
90 /// Commit the manifest `value`, consuming the context.
91 ///
92 /// # Errors
93 ///
94 /// Returns [`Error`] if the manifest cannot be serialized or committed.
95 fn finish(self, value: Value<'_>) -> Result<Self::Output>;
96}
97
98/// Object-safe view of the artifact-allocating half of a [`SaveContext`].
99///
100/// [`Context`] holds a `&dyn GetWrite` so that the same handle can be threaded through
101/// every [`Save::save`](super::Save) impl regardless of the concrete context type.
102/// [`SaveContext::finish`] (which consumes `self` and names an associated type) is not
103/// object safe, so it is deliberately excluded from this trait.
104pub(super) trait GetWrite {
105 fn write(&self, key: Option<&str>) -> Result<Writer<'_>>;
106}
107
108impl<T> GetWrite for T
109where
110 T: SaveContext,
111{
112 fn write(&self, key: Option<&str>) -> Result<Writer<'_>> {
113 <T as SaveContext>::write(self, key)
114 }
115}
116
117/// A cheap, clonable handle threaded through every [`Save::save`](super::Save) impl.
118///
119/// `Context` exposes one operation — [`Context::write`] — for allocating a side-car
120/// artifact. The same context is passed to nested [`Save`](super::Save) impls (typically
121/// via the [`save_fields!`](crate::save_fields) macro), so a single save tree shares
122/// artifact-name bookkeeping. It borrows the backing `SaveContext` as an object-safe
123/// `GetWrite` so that the save tree is agnostic to the concrete context type.
124#[derive(Clone)]
125pub struct Context<'a> {
126 inner: &'a dyn GetWrite,
127}
128
129impl<'a> Context<'a> {
130 pub(super) fn new(inner: &'a dyn GetWrite) -> Self {
131 Self { inner }
132 }
133
134 /// Allocate a new side-car artifact in the manifest directory, optionally tagging it
135 /// with a human-readable `key`.
136 ///
137 /// The artifact is named with the count of artifacts written so far (with `key`, when
138 /// `Some`, appended as a readability hint), so the same `key` may be passed to
139 /// multiple calls. The returned [`Writer`] is positioned at offset 0 and implements
140 /// [`std::io::Write`] / [`std::io::Seek`]. Call [`Writer::finish`] to obtain a
141 /// [`Handle`] that may be inserted into a [`Record`](super::Record).
142 ///
143 /// # Errors
144 ///
145 /// Returns [`Error`] if `key` is `Some` but not a simple relative file name, or if the
146 /// underlying file cannot be created.
147 pub fn write(&self, key: Option<&str>) -> Result<Writer<'_>> {
148 self.inner.write(key)
149 }
150}
151
152/// The backend-specific half of a [`Writer`].
153///
154/// Each [`SaveContext`](super::SaveContext) implementation supplies its own
155/// `WriterInner` (e.g. an in-memory cursor or an on-disk file). [`Writer`] wraps it in a
156/// [`BufWriter`] and forwards [`std::io::Write`] / [`std::io::Seek`] to it; on
157/// [`Writer::finish`] the (flushed) inner is consumed to commit the artifact and yield a
158/// [`Handle`].
159pub(crate) trait WriterInner: std::io::Write + std::io::Seek + std::fmt::Debug {
160 /// Commit the completed artifact under `name`, returning its [`Handle`].
161 fn finish(self: Box<Self>, name: String) -> Result<Handle>;
162}
163
164/// A borrowed side-car artifact writer produced by [`Context::write`].
165///
166/// Implements [`std::io::Write`] and [`std::io::Seek`]. Writes are buffered; calling
167/// [`Writer::finish`] flushes the buffer, commits the artifact through the backing
168/// writer, and returns a [`Handle`].
169#[derive(Debug)]
170pub struct Writer<'a> {
171 inner: BufWriter<Box<dyn WriterInner + 'a>>,
172 name: String,
173}
174
175impl<'a> Writer<'a> {
176 /// Wrap a backend-specific [`WriterInner`] into a buffered [`Writer`] named `name`.
177 pub(crate) fn new<T>(inner: T, name: String) -> Self
178 where
179 T: WriterInner + 'a,
180 {
181 Self {
182 inner: BufWriter::new(Box::new(inner)),
183 name,
184 }
185 }
186
187 /// Flush and close the writer, returning a [`Handle`] for the artifact.
188 ///
189 /// Insert the returned handle into a [`Record`](super::Record) (typically via
190 /// [`Record::insert`](super::Record::insert)) so that load-side code can locate the
191 /// artifact through the manifest.
192 pub fn finish(self) -> Result<Handle> {
193 // into_inner() flushes the buffered bytes into the backend before we commit it.
194 let inner = self
195 .inner
196 .into_inner()
197 .map_err(|err| Error::new(err.into_error()))?;
198 inner.finish(self.name)
199 }
200}
201
202delegate_write_and_seek!(inner, Writer<'_>);