Skip to main content

ringdb/payload/
pod.rs

1use memmap2::Mmap;
2use std::{
3    fs::File,
4    io::{BufWriter, Write},
5    marker::PhantomData,
6    path::Path,
7};
8use tempfile::TempPath;
9
10use super::traits::{PayloadBuilderOps, open_mmap};
11use super::{OwnedPayloadStore, RefPayloadStore};
12use crate::error::Result;
13use crate::persist::move_file;
14
15// ─── PodStoreBuilder ──────────────────────────────────────────────────────────
16//
17// Payloads are stored as raw bytes (`bytemuck::bytes_of`), back-to-back,
18// with no offset table. `fetch_ref(id)` returns a zero-copy `&T` in O(1):
19//
20//   offset = id * size_of::<T>()
21//   &T     = bytemuck::from_bytes(&mmap[offset..offset + size_of::<T>()])
22
23pub struct PodStoreBuilder<T> {
24    writer: BufWriter<File>,
25    temp_path: TempPath,
26    n_pushed: usize,
27    _marker: PhantomData<T>,
28}
29
30impl<T: bytemuck::Pod> PodStoreBuilder<T> {
31    pub fn new() -> Result<Self> {
32        let named = tempfile::NamedTempFile::new()?;
33        let (file, temp_path) = named.into_parts();
34        Ok(Self {
35            writer: BufWriter::new(file),
36            temp_path,
37            n_pushed: 0,
38            _marker: PhantomData,
39        })
40    }
41
42    fn push_inner(&mut self, payload: T) -> Result<()> {
43        self.writer.write_all(bytemuck::bytes_of(&payload))?;
44        self.n_pushed += 1;
45        Ok(())
46    }
47
48    fn finish_inner(self) -> Result<PodStore<T>> {
49        let Self {
50            writer,
51            temp_path,
52            n_pushed,
53            _marker,
54        } = self;
55        let total = (n_pushed * size_of::<T>()) as u64;
56        writer.into_inner().map_err(|e| e.into_error())?;
57        let mmap = open_mmap(temp_path.as_ref(), total)?;
58        Ok(PodStore {
59            mmap,
60            _marker: PhantomData,
61        })
62    }
63
64    fn finish_persisted_inner(self, payloads_path: &Path) -> Result<PodStore<T>> {
65        let Self {
66            writer,
67            temp_path,
68            n_pushed,
69            _marker,
70        } = self;
71        let total = (n_pushed * size_of::<T>()) as u64;
72        writer.into_inner().map_err(|e| e.into_error())?;
73        // move_file handles cross-filesystem moves with a copy fallback.
74        // If it fails, temp_path is still alive and its Drop will clean up.
75        move_file(temp_path.as_ref(), payloads_path)?;
76        let mmap = open_mmap(payloads_path, total)?;
77        Ok(PodStore {
78            mmap,
79            _marker: PhantomData,
80        })
81    }
82}
83
84impl<T: bytemuck::Pod> PayloadBuilderOps<T> for PodStoreBuilder<T> {
85    type Store = PodStore<T>;
86
87    fn push(&mut self, payload: T) -> Result<()> {
88        self.push_inner(payload)
89    }
90
91    fn finish(self) -> Result<PodStore<T>> {
92        self.finish_inner()
93    }
94
95    /// Pod storage has no offset table; `offsets_path` is ignored.
96    fn finish_persisted(self, payloads_path: &Path, _offsets_path: &Path) -> Result<PodStore<T>> {
97        self.finish_persisted_inner(payloads_path)
98    }
99}
100
101// ─── PodStore ─────────────────────────────────────────────────────────────────
102
103pub struct PodStore<T> {
104    mmap: Option<Mmap>,
105    // Declared after `mmap` so it drops after the mmap is released.
106    _marker: PhantomData<T>,
107}
108
109impl<T: bytemuck::Pod> PodStore<T> {
110    pub fn load(payloads_path: &Path) -> Result<Self> {
111        let total_bytes = std::fs::metadata(payloads_path)?.len();
112        let mmap = open_mmap(payloads_path, total_bytes)?;
113        Ok(PodStore {
114            mmap,
115            _marker: PhantomData,
116        })
117    }
118}
119
120impl<T: bytemuck::Pod> OwnedPayloadStore<T> for PodStore<T> {
121    /// Deserializes by copying `size_of::<T>()` bytes — no bincode, no heap alloc.
122    fn fetch_owned(&self, id: u32) -> crate::error::Result<T> {
123        Ok(*self.fetch_ref(id))
124    }
125}
126
127impl<T: bytemuck::Pod> RefPayloadStore<T> for PodStore<T> {
128    /// Zero-copy reference into the mmap — O(1), no allocation.
129    fn fetch_ref(&self, id: u32) -> &T {
130        let size = std::mem::size_of::<T>();
131        let offset = id as usize * size;
132        bytemuck::from_bytes(
133            &self.mmap.as_ref().expect("fetch_ref on empty store")[offset..offset + size],
134        )
135    }
136}