metrale_storage/
cascade_backend.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `CascadeBackend`: a write-back cache of groups in pinned host
4//! memory (T1) in front of another `StorageBackend`. Writes land in T1; when
5//! T1 is full, the least recently used group is written down to the backing.
6//! Reads copy T1 hits to the device and forward misses to the backing. No
7//! tier transforms bytes.
8//!
9//! Owner: metrale-storage KV tier.
10//! Invariants:
11//! - A read of a group resident in T1 is served from T1; misses are not
12//!   copied into T1.
13//! - A write waits for the event recorded after the last async hit copies
14//!   before it overwrites a slot.
15
16use std::ffi::c_void;
17
18use anyhow::{Context, Result};
19
20use crate::backend::{ReadRequest, StorageBackend};
21use crate::cascade_policy::SlotCache;
22use crate::cuda_min::{CudaEvent, PinnedBuffer, copy_h_to_d_async, stream_sync};
23use crate::group::{GroupKey, GroupLayout};
24
25/// 2026-09-25: The T1 bytes: one pinned buffer of `cap_slots * group_bytes`,
26/// slot `i` at `ptr + i * group_bytes`.
27struct PinnedStore {
28    buf: PinnedBuffer,
29    group_bytes: usize,
30}
31
32impl PinnedStore {
33    fn new(cap_slots: u32, group_bytes: usize) -> Result<Self> {
34        let bytes = (cap_slots as usize)
35            .checked_mul(group_bytes)
36            .context("CascadeBackend T1 size overflow")?;
37        Ok(Self {
38            buf: PinnedBuffer::new(bytes).context("alloc T1 pinned store")?,
39            group_bytes,
40        })
41    }
42    #[inline]
43    fn slot_host_ptr(&self, slot: u32) -> *const c_void {
44        // 2026-09-25: SAFETY: `SlotCache` hands out slots below `cap_slots`, so
45        // the offset is inside the buffer.
46        unsafe {
47            (self.buf.ptr as *const u8).add(slot as usize * self.group_bytes) as *const c_void
48        }
49    }
50    /// 2026-09-25: Copy `slot`'s bytes into a new Vec, so the caller can pass
51    /// them to `&mut backing` without borrowing the store.
52    fn slot_bytes(&self, slot: u32) -> Vec<u8> {
53        // 2026-09-25: SAFETY: as above; a slot spans `group_bytes` bytes.
54        unsafe {
55            std::slice::from_raw_parts(
56                (self.buf.ptr as *const u8).add(slot as usize * self.group_bytes),
57                self.group_bytes,
58            )
59            .to_vec()
60        }
61    }
62    fn write_slot(&mut self, slot: u32, src: &[u8]) {
63        debug_assert_eq!(src.len(), self.group_bytes);
64        // 2026-09-25: SAFETY: the slot is in range. `src` must be `group_bytes`
65        // long, which only the debug assertion above checks.
66        unsafe {
67            std::ptr::copy_nonoverlapping(
68                src.as_ptr(),
69                (self.buf.ptr as *mut u8).add(slot as usize * self.group_bytes),
70                self.group_bytes,
71            );
72        }
73    }
74}
75
76pub struct CascadeBackend {
77    hot: SlotCache,
78    store: PinnedStore,
79    backing: Box<dyn StorageBackend>,
80    group_bytes: usize,
81    /// 2026-09-25: The event `read_async` records after its hit copies. The
82    /// next `write_from_host` waits on it before it overwrites a slot, so the
83    /// host does not write a slot that a queued copy still reads.
84    last_read_event: Option<CudaEvent>,
85}
86
87// 2026-09-25: SAFETY: the pinned store's pointer is used only by `&mut self`
88// methods; the one `&self` method, `group_layout`, calls the backing, which the
89// trait bound makes `Sync`.
90unsafe impl Sync for CascadeBackend {}
91
92impl CascadeBackend {
93    pub fn new(
94        backing: Box<dyn StorageBackend>,
95        layout: GroupLayout,
96        cap_slots: u32,
97    ) -> Result<Self> {
98        let group_bytes = layout.group_bytes() as usize;
99        tracing::info!(
100            "high-speed-swap: T1 cascade cache = {cap_slots} slots × {group_bytes} B = {:.1} GiB local pinned RAM, backing below",
101            (cap_slots as f64 * group_bytes as f64) / (1024.0 * 1024.0 * 1024.0),
102        );
103        Ok(Self {
104            hot: SlotCache::new(cap_slots),
105            store: PinnedStore::new(cap_slots, group_bytes)?,
106            backing,
107            group_bytes,
108            last_read_event: None,
109        })
110    }
111
112    /// 2026-09-25: Write every resident T1 group down to the backing.
113    fn flush_all(&mut self) -> Result<()> {
114        for (key, slot) in self.hot.residents() {
115            let bytes = self.store.slot_bytes(slot);
116            self.backing.write_from_host(key, &bytes)?;
117        }
118        Ok(())
119    }
120
121    /// 2026-09-25: Queue a copy of each T1 hit on `stream` and forward the
122    /// misses to the backing. Sync: misses go through `backing.read`, then the
123    /// stream is synced. Async: misses go through `backing.read_async`, an event
124    /// is recorded after the hit copies, and the stream is not synced.
125    fn read_common(&mut self, requests: &[ReadRequest], stream: u64, is_async: bool) -> Result<()> {
126        let keys: Vec<GroupKey> = requests.iter().map(|r| r.group).collect();
127        let (hits, misses) = self.hot.plan_read(&keys);
128        for (i, slot) in &hits {
129            let src = self.store.slot_host_ptr(*slot);
130            copy_h_to_d_async(requests[*i].dst_dev_ptr, src, self.group_bytes, stream)?;
131            self.hot.touch(*slot);
132        }
133        if is_async && !hits.is_empty() {
134            let ev = CudaEvent::new()?;
135            ev.record(stream)?;
136            self.last_read_event = Some(ev);
137        }
138        // 2026-09-25: Misses are not copied into T1; only writes fill it. The
139        // final sync also covers the hit copies.
140        if !misses.is_empty() {
141            let miss_reqs: Vec<ReadRequest> = misses.iter().map(|&i| requests[i]).collect();
142            if is_async {
143                self.backing.read_async(&miss_reqs, stream)?;
144            } else {
145                self.backing.read(&miss_reqs, stream)?;
146            }
147        }
148        if !is_async {
149            stream_sync(stream)?;
150        }
151        Ok(())
152    }
153}
154
155impl StorageBackend for CascadeBackend {
156    fn write_from_host(&mut self, key: GroupKey, src: &[u8]) -> Result<()> {
157        let plan = self.hot.plan_write(key);
158        // 2026-09-25: The victim's bytes are still in the slot: write them down
159        // first.
160        if let Some((victim_key, victim_slot)) = plan.flush_victim {
161            let victim_bytes = self.store.slot_bytes(victim_slot);
162            self.backing
163                .write_from_host(victim_key, &victim_bytes)
164                .context("cascade: flush T1 victim to backing")?;
165        }
166        // 2026-09-25: An async hit copy may still read this slot.
167        if let Some(ev) = self.last_read_event.take() {
168            ev.sync()?;
169        }
170        self.store.write_slot(plan.slot, src);
171        Ok(())
172    }
173
174    fn read(&mut self, requests: &[ReadRequest], stream: u64) -> Result<()> {
175        self.read_common(requests, stream, false)
176    }
177
178    fn read_async(&mut self, requests: &[ReadRequest], stream: u64) -> Result<()> {
179        self.read_common(requests, stream, true)
180    }
181
182    fn register_landing_region(&mut self, base: u64, len: usize) -> Result<()> {
183        // 2026-09-25: Forwarded, so the backing's zero-copy reads of misses can
184        // land in the region; T1 hits are plain copies.
185        self.backing.register_landing_region(base, len)
186    }
187
188    fn group_layout(&self) -> GroupLayout {
189        // 2026-09-25: The backing's geometry. The block methods keep the trait
190        // defaults, which fan out through this backend's per-group methods.
191        self.backing.group_layout()
192    }
193}
194
195impl Drop for CascadeBackend {
196    fn drop(&mut self) {
197        // 2026-09-25: Write T1 down before the pinned store is freed; errors
198        // are ignored.
199        let _ = self.flush_all();
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use crate::backend::PosixBackend;
207    use crate::cuda_min::{CudaCtx, DeviceBuffer, copy_d_to_h_async, stream_sync};
208    use crate::group::KvKind;
209    use crate::layout::Layout;
210
211    fn tempdir(name: &str) -> std::path::PathBuf {
212        let p = std::env::temp_dir().join(format!("metrale-cascade-{name}-{}", std::process::id()));
213        let _ = std::fs::remove_dir_all(&p);
214        std::fs::create_dir_all(&p).unwrap();
215        p
216    }
217
218    /// 2026-09-25: With a 2-slot T1 over 4 groups: write groups 0 and 1,
219    /// `read_async` both (two T1 hits), then write groups 2 and 3, which evict
220    /// 0 and 1 from the slots the async copies read. The async results, and a
221    /// sync read of all four groups (hits 2 and 3, misses 0 and 1 from the
222    /// backing), must match what was written.
223    #[test]
224    #[ignore = "requires GPU"]
225    fn read_async_eviction_no_corruption() {
226        let _ctx = CudaCtx::new(0).expect("cuda init");
227        let dir = tempdir("evict");
228        let spec = GroupLayout::new(1, 4, 1, 16, 128, 2, 4096);
229        let layout = Layout::create(&dir, spec).unwrap();
230        let bytes = spec.group_bytes() as usize;
231        let backing = Box::new(PosixBackend::new(layout).unwrap());
232        let mut cascade = CascadeBackend::new(backing, spec, 2).unwrap();
233
234        let keys: Vec<GroupKey> = (0..4u32)
235            .map(|b| GroupKey::new(0, b, 0, KvKind::K))
236            .collect();
237        let pat = |b: usize| -> Vec<u8> {
238            (0..bytes)
239                .map(|i| ((i * 7 + b * 11) & 0xFF) as u8)
240                .collect()
241        };
242
243        cascade.write_from_host(keys[0], &pat(0)).unwrap();
244        cascade.write_from_host(keys[1], &pat(1)).unwrap();
245
246        let d0 = DeviceBuffer::new(bytes).unwrap();
247        let d1 = DeviceBuffer::new(bytes).unwrap();
248        cascade
249            .read_async(
250                &[
251                    ReadRequest {
252                        group: keys[0],
253                        dst_dev_ptr: d0.ptr,
254                    },
255                    ReadRequest {
256                        group: keys[1],
257                        dst_dev_ptr: d1.ptr,
258                    },
259                ],
260                _ctx.stream,
261            )
262            .unwrap();
263        cascade.write_from_host(keys[2], &pat(2)).unwrap();
264        cascade.write_from_host(keys[3], &pat(3)).unwrap();
265        stream_sync(_ctx.stream).unwrap();
266        let readback = |d: &DeviceBuffer, want: &[u8]| {
267            let mut got = vec![0u8; bytes];
268            copy_d_to_h_async(got.as_mut_ptr() as *mut c_void, d.ptr, bytes, _ctx.stream).unwrap();
269            stream_sync(_ctx.stream).unwrap();
270            assert_eq!(&got, want, "cascade async restore corrupted");
271        };
272        readback(&d0, &pat(0));
273        readback(&d1, &pat(1));
274
275        let devs: Vec<DeviceBuffer> = (0..4).map(|_| DeviceBuffer::new(bytes).unwrap()).collect();
276        let reqs: Vec<ReadRequest> = keys
277            .iter()
278            .zip(&devs)
279            .map(|(k, d)| ReadRequest {
280                group: *k,
281                dst_dev_ptr: d.ptr,
282            })
283            .collect();
284        cascade.read(&reqs, _ctx.stream).unwrap();
285        for (b, d) in devs.iter().enumerate() {
286            readback(d, &pat(b));
287        }
288        drop(cascade);
289        std::fs::remove_dir_all(&dir).ok();
290    }
291}
metrale_storage/cascade_backend.rs

metrale_storage/
cascade_backend.rs
