link_r/
remote.rs

1// Author: support@lituus.io
2// Copyright: lituus-io, all rights reserved
3//
4//! Object storage for indexes (feature `remote`).
5//!
6//! Keep an index on an object store — Google Cloud Storage, a
7//! local directory, or memory — through objstore-r's [`ObjectStore`] trait.
8//!
9//! ```text
10//! {prefix}HEAD                          which generation is current (compare-and-swap only)
11//! {prefix}index/{generation:016x}.lnkr  the .lnkr image, byte for byte (create-only)
12//! ```
13//!
14//! - [`Remote::load`] fetches the current image and opens it with
15//!   [`LinkIndex::from_bytes`]: searches then run zero-copy over the fetched
16//!   buffer (no temp file, no second copy).
17//! - [`Remote::publish`] uploads a new generation only when the image changed
18//!   and swaps `HEAD` conditionally: of two publishers starting from the same
19//!   state exactly one wins and the other gets [`Error::Fenced`].
20//! - [`Remote::load_if_changed`] costs one metadata request when nothing
21//!   changed — the replica refresh loop.
22//!
23//! Futures are unboxed (`impl Future` over borrowed state); no `Arc`, no
24//! `dyn`; the internal `Mutex` is never held across an `.await`.
25
26pub use objstore_r::ObjectStore;
27
28use crate::facade::DefaultEmbedder;
29use crate::{Embedder, LinkIndex};
30use graph_r::hash::xxh3_64;
31use objstore_r::{fenced, Bytes};
32use std::future::Future;
33use std::sync::{Mutex, PoisonError};
34
35const MAGIC: u32 = 0x4D52_4E4C; // "LNRM"
36const VERSION: u16 = 1;
37/// `HEAD`'s sealed body: generation, length and checksum of the image.
38const HEAD_BODY: usize = 24;
39
40/// Result alias.
41pub type Result<T, E = Error> = std::result::Result<T, E>;
42
43/// What can go wrong.
44#[derive(Debug, thiserror::Error)]
45#[non_exhaustive]
46pub enum Error {
47    /// The index failed to load or serialize.
48    #[error(transparent)]
49    Link(#[from] crate::Error),
50    /// The object store failed.
51    #[error("object store: {0}")]
52    Store(objstore_r::Error),
53    /// Another publisher swapped `HEAD` first. Reload before publishing.
54    #[error("fenced: {reason}")]
55    Fenced {
56        /// Why.
57        reason: &'static str,
58    },
59    /// Published data failed verification.
60    #[error("corrupt remote index: {message}")]
61    Corrupt {
62        /// What was wrong.
63        message: String,
64    },
65}
66
67/// The fenced-publish outcomes keep their own variants.
68impl From<objstore_r::Error> for Error {
69    fn from(e: objstore_r::Error) -> Self {
70        match e {
71            objstore_r::Error::Fenced { reason } => Self::Fenced { reason },
72            objstore_r::Error::Corrupt { message } => Self::Corrupt { message },
73            other => Self::Store(other),
74        }
75    }
76}
77
78/// The published state.
79#[derive(Clone, Copy, Debug, PartialEq, Eq)]
80pub struct Published {
81    /// Monotonic publish generation (1 for the first).
82    pub generation: u64,
83    /// Image length in bytes.
84    pub len: u64,
85    /// xxh3 of the image.
86    pub sum: u64,
87    /// The publish uploaded a new image (false: unchanged, nothing written).
88    pub uploaded: bool,
89}
90
91#[derive(Clone, Copy, Debug, PartialEq, Eq)]
92struct Head {
93    generation: u64,
94    len: u64,
95    sum: u64,
96}
97
98impl Head {
99    fn encode(self) -> Vec<u8> {
100        let mut body = Vec::with_capacity(HEAD_BODY);
101        body.extend_from_slice(&self.generation.to_le_bytes());
102        body.extend_from_slice(&self.len.to_le_bytes());
103        body.extend_from_slice(&self.sum.to_le_bytes());
104        fenced::seal(MAGIC, VERSION, &body)
105    }
106
107    fn decode(b: &[u8]) -> Result<Self> {
108        let body = fenced::unseal(b, MAGIC, VERSION)?;
109        if body.len() != HEAD_BODY {
110            return Err(Error::Corrupt {
111                message: "HEAD has the wrong length".into(),
112            });
113        }
114        let u64_at = |i: usize| {
115            let mut a = [0u8; 8];
116            a.copy_from_slice(&body[i..i + 8]);
117            u64::from_le_bytes(a)
118        };
119        let head = Self {
120            generation: u64_at(0),
121            len: u64_at(8),
122            sum: u64_at(16),
123        };
124        if head.generation == 0 {
125            return Err(Error::Corrupt {
126                message: "HEAD names generation 0".into(),
127            });
128        }
129        Ok(head)
130    }
131}
132
133/// The CAS token plus what it named.
134#[derive(Clone, Copy, Debug)]
135struct Seen {
136    head_gen: u64,
137    head: Head,
138}
139
140/// An index location on an object store. See the crate docs.
141#[derive(Debug)]
142pub struct Remote<'o, O: ObjectStore> {
143    objects: &'o O,
144    prefix: String,
145    seen: Mutex<Option<Seen>>,
146}
147
148impl<'o, O: ObjectStore> Remote<'o, O> {
149    /// An index under `prefix` (empty, or a key prefix; `/` implied).
150    pub fn new(objects: &'o O, prefix: &str) -> Result<Self> {
151        let prefix = fenced::prefix(prefix)?;
152        Ok(Self {
153            objects,
154            prefix,
155            seen: Mutex::new(None),
156        })
157    }
158
159    fn head_key(&self) -> String {
160        format!("{}HEAD", self.prefix)
161    }
162    fn index_key(&self, generation: u64) -> String {
163        format!("{}index/{generation:016x}.lnkr", self.prefix)
164    }
165    fn seen(&self) -> Option<Seen> {
166        *self.seen.lock().unwrap_or_else(PoisonError::into_inner)
167    }
168    fn remember(&self, s: Option<Seen>) {
169        *self.seen.lock().unwrap_or_else(PoisonError::into_inner) = s;
170    }
171
172    async fn read_head(&self) -> Result<Option<Seen>> {
173        let Some(obj) = self.objects.get(&self.head_key()).await? else {
174            self.remember(None);
175            return Ok(None);
176        };
177        let seen = Seen {
178            head_gen: obj.generation,
179            head: Head::decode(&obj.bytes)?,
180        };
181        Ok(Some(seen))
182    }
183
184    async fn fetch_image(&self, head: Head) -> Result<Bytes> {
185        let key = self.index_key(head.generation);
186        Ok(fenced::fetch_verified(self.objects, &key, head.len, head.sum).await?)
187    }
188
189    /// The published generation, if any (reads `HEAD`).
190    pub fn generation(&self) -> impl Future<Output = Result<Option<u64>>> + '_ {
191        async move { Ok(self.read_head().await?.map(|s| s.head.generation)) }
192    }
193
194    /// Load the current index with a caller-supplied embedder (`None` when
195    /// nothing is published). The image is verified, then opened zero-copy.
196    pub fn load_with<'a, E: Embedder + 'a>(
197        &'a self,
198        embedder: E,
199    ) -> impl Future<Output = Result<Option<LinkIndex<E>>>> + 'a {
200        async move {
201            let Some(seen) = self.read_head().await? else {
202                return Ok(None);
203            };
204            let image = self.fetch_image(seen.head).await?;
205            let index = LinkIndex::from_bytes_with(image, embedder)?;
206            self.remember(Some(seen));
207            Ok(Some(index))
208        }
209    }
210
211    /// [`Remote::load_with`], or an empty in-memory index with the same
212    /// embedder when nothing is published yet (a writer's first start).
213    /// The flag says whether an image was loaded.
214    pub fn load_or_empty_with<'a, E: Embedder + 'a>(
215        &'a self,
216        embedder: E,
217    ) -> impl Future<Output = Result<(LinkIndex<E>, bool)>> + 'a {
218        async move {
219            let Some(seen) = self.read_head().await? else {
220                return Ok((LinkIndex::in_memory_with(embedder), false));
221            };
222            let image = self.fetch_image(seen.head).await?;
223            let index = LinkIndex::from_bytes_with(image, embedder)?;
224            self.remember(Some(seen));
225            Ok((index, true))
226        }
227    }
228
229    /// Like [`Remote::load_with`], but only when `HEAD` moved since this
230    /// `Remote` last loaded or published — one metadata request otherwise.
231    pub fn load_if_changed_with<'a, E: Embedder + 'a>(
232        &'a self,
233        embedder: E,
234    ) -> impl Future<Output = Result<Option<LinkIndex<E>>>> + 'a {
235        async move {
236            let meta = self.objects.head(&self.head_key()).await?;
237            match (meta, self.seen()) {
238                (None, _) => Ok(None),
239                (Some(m), Some(s)) if m.generation == s.head_gen => Ok(None),
240                _ => self.load_with(embedder).await,
241            }
242        }
243    }
244
245    /// Replica refresh in place: when `HEAD` moved since this `Remote` last
246    /// loaded or published, swap the new image into `index` (keeping its
247    /// embedder; see [`LinkIndex::reload_from_bytes`]). One metadata request
248    /// when nothing changed. Returns whether `index` changed.
249    pub fn load_if_changed_into<'a, E: Embedder>(
250        &'a self,
251        index: &'a mut LinkIndex<E>,
252    ) -> impl Future<Output = Result<bool>> + 'a {
253        async move {
254            let meta = self.objects.head(&self.head_key()).await?;
255            match (meta, self.seen()) {
256                (None, _) => return Ok(false),
257                (Some(m), Some(s)) if m.generation == s.head_gen => return Ok(false),
258                _ => {}
259            }
260            let Some(seen) = self.read_head().await? else {
261                return Ok(false);
262            };
263            let image = self.fetch_image(seen.head).await?;
264            index.reload_from_bytes(image)?;
265            self.remember(Some(seen));
266            Ok(true)
267        }
268    }
269
270    /// Publish `index` if its image differs from the published one. Uploads a
271    /// new generation create-only, then swaps `HEAD` conditionally; the
272    /// previous image is removed afterwards (best effort).
273    pub fn publish<'a, E: Embedder>(
274        &'a self,
275        index: &'a LinkIndex<E>,
276    ) -> impl Future<Output = Result<Published>> + 'a {
277        async move {
278            let image = index.to_bytes()?;
279            let sum = xxh3_64(&image);
280            let len = image.len() as u64;
281            let prev = match self.seen() {
282                Some(s) => Some(s),
283                None => self.read_head().await?,
284            };
285            if let Some(p) = prev.filter(|p| p.head.sum == sum && p.head.len == len) {
286                self.remember(Some(p));
287                return Ok(Published {
288                    generation: p.head.generation,
289                    len,
290                    sum,
291                    uploaded: false,
292                });
293            }
294            let head = Head {
295                generation: prev.map_or(1, |p| p.head.generation + 1),
296                len,
297                sum,
298            };
299            let key = self.index_key(head.generation);
300            fenced::put_immutable(self.objects, &key, Bytes::from(image))
301                .await
302                .map_err(|e| self.forget(e))?;
303            let body = Bytes::from(head.encode());
304            let head_gen = fenced::cas(
305                self.objects,
306                &self.head_key(),
307                prev.map(|p| p.head_gen),
308                body,
309            )
310            .await
311            .map_err(|e| self.forget(e))?;
312            self.remember(Some(Seen { head_gen, head }));
313            if let Some(p) = prev {
314                let _ = self
315                    .objects
316                    .delete(&self.index_key(p.head.generation))
317                    .await;
318            }
319            Ok(Published {
320                generation: head.generation,
321                len,
322                sum,
323                uploaded: true,
324            })
325        }
326    }
327
328    /// A failed publish leaves `HEAD` unknown: re-read before the next one.
329    fn forget(&self, e: objstore_r::Error) -> Error {
330        self.remember(None);
331        e.into()
332    }
333}
334
335impl<O: ObjectStore> Remote<'_, O> {
336    /// [`Remote::load_with`] with link-r's default embedder.
337    pub fn load(&self) -> impl Future<Output = Result<Option<LinkIndex>>> + '_ {
338        async move {
339            let embedder = default_embedder()?;
340            self.load_with(embedder).await
341        }
342    }
343
344    /// [`Remote::load_if_changed_with`] with link-r's default embedder.
345    pub fn load_if_changed(&self) -> impl Future<Output = Result<Option<LinkIndex>>> + '_ {
346        async move {
347            let embedder = default_embedder()?;
348            self.load_if_changed_with(embedder).await
349        }
350    }
351}
352
353fn default_embedder() -> Result<DefaultEmbedder> {
354    Ok(crate::facade::default_embedder()?)
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360
361    #[test]
362    fn head_round_trips_and_rejects_corruption() {
363        let h = Head {
364            generation: 9,
365            len: 1234,
366            sum: 77,
367        };
368        let b = h.encode();
369        assert_eq!(Head::decode(&b).unwrap(), h);
370        for i in 0..b.len() {
371            let mut c = b.clone();
372            c[i] ^= 1;
373            assert!(Head::decode(&c).is_err(), "flip at {i}");
374        }
375        assert!(Head::decode(&b[..b.len() - 1]).is_err());
376        assert!(Head::decode(&Head { generation: 0, ..h }.encode()).is_err());
377    }
378}
link_r/remote.rs

link_r/
remote.rs
