spark_storage/backend/
posix.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Phase-2 reference backend. Single pinned bounce buffer, positional read +
4// `cuMemcpyHtoDAsync`, stream-sync after every memcpy to avoid the next read
5// overwriting in-flight DMA. Slow-but-deterministic; used by tests as the
6// oracle the io_uring backend is compared against.
7//
8// Named "posix" for its history; it is now PORTABLE. The positional read/write
9// go through `atlas_tier::pio`, which is `pread`/`pwrite` on unix and
10// `seek_read`/`seek_write` on Windows, so this is the backend the Windows
11// build uses (io_uring has no Windows analogue).
12
13use anyhow::{Context, Result, bail};
14use std::ffi::c_void;
15
16use super::{ReadRequest, StorageBackend};
17use crate::cuda_min::{PinnedBuffer, copy_h_to_d_async, stream_sync};
18use crate::group::{GroupKey, GroupLayout};
19use crate::layout::Layout;
20
21pub struct PosixBackend {
22    layout: Layout,
23    bounce: PinnedBuffer,
24}
25
26impl PosixBackend {
27    pub fn new(layout: Layout) -> Result<Self> {
28        let bounce = PinnedBuffer::new(layout.group_bytes() as usize)
29            .context("alloc pinned bounce buffer")?;
30        Ok(Self { layout, bounce })
31    }
32    pub fn layout(&self) -> &Layout {
33        &self.layout
34    }
35}
36
37impl StorageBackend for PosixBackend {
38    fn read(&mut self, requests: &[ReadRequest], stream: u64) -> Result<()> {
39        let bytes = self.layout.group_bytes() as usize;
40        let bounce_ptr = self.bounce.ptr;
41        for req in requests {
42            let off = self.layout.offset(req.group);
43            // SAFETY: `bounce_ptr` is a pinned host allocation of at least
44            // `group_bytes()`, owned by `self.bounce` for the lifetime of this
45            // call and not aliased -- the loop serialises on `stream_sync`.
46            let buf = unsafe { std::slice::from_raw_parts_mut(bounce_ptr as *mut u8, bytes) };
47            atlas_tier::pio::read_exact_at(self.layout.file(req.group.layer), buf, off)
48                .with_context(|| format!("read {bytes}@{off}"))?;
49            // The pinned bounce buffer is shared across all requests in this
50            // call; we must let the H→D DMA complete before the next pread
51            // overwrites the buffer, otherwise the second cuMemcpyHtoDAsync
52            // will read partial / stale bytes. Phase-3 io_uring backend uses
53            // multiple registered buffers and avoids this serialization.
54            copy_h_to_d_async(req.dst_dev_ptr, bounce_ptr as *const c_void, bytes, stream)?;
55            stream_sync(stream)?;
56        }
57        Ok(())
58    }
59
60    fn write_from_host(&mut self, key: GroupKey, src: &[u8]) -> Result<()> {
61        let bytes = self.layout.group_bytes() as usize;
62        if src.len() != bytes {
63            bail!(
64                "write_from_host: src len {} != group bytes {bytes}",
65                src.len()
66            );
67        }
68        // O_DIRECT requires page-aligned source. Stage through the pinned
69        // bounce buffer (which is page-aligned per cuMemAllocHost contract).
70        unsafe {
71            std::ptr::copy_nonoverlapping(src.as_ptr(), self.bounce.ptr as *mut u8, bytes);
72        }
73        let off = self.layout.offset(key);
74        // SAFETY: as above -- pinned, group-sized, exclusively owned here.
75        let buf = unsafe { std::slice::from_raw_parts(self.bounce.ptr as *const u8, bytes) };
76        atlas_tier::pio::write_all_at(self.layout.file(key.layer), buf, off)
77            .with_context(|| format!("write {bytes}@{off}"))?;
78        // fsync would be needed for crash durability; skipped for the test
79        // path where the file is single-process / single-run.
80        Ok(())
81    }
82
83    fn group_layout(&self) -> GroupLayout {
84        self.layout.spec
85    }
86}
87
88impl PosixBackend {
89    /// Test helper: drop the page cache for the layer files so subsequent
90    /// reads actually hit NVMe (otherwise small tests trivially read from RAM).
91    ///
92    /// `posix_fadvise(DONTNEED)` has no portable equivalent; on Windows the
93    /// cache is dropped by reopening with `FILE_FLAG_NO_BUFFERING`, which this
94    /// tier deliberately does not use. A no-op there means a Windows run of
95    /// the timing tests may read from RAM -- which is why this is a test
96    /// helper and not a correctness primitive.
97    #[cfg(unix)]
98    pub fn drop_pagecache(&self) {
99        for layer in 0..self.layout.spec.num_layers {
100            let fd = self.layout.fd(layer);
101            unsafe { libc::posix_fadvise(fd, 0, 0, libc::POSIX_FADV_DONTNEED) };
102        }
103    }
104
105    #[cfg(not(unix))]
106    pub fn drop_pagecache(&self) {}
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use crate::group::{GroupLayout, KvKind};
113
114    fn tempdir(name: &str) -> std::path::PathBuf {
115        let p = std::env::temp_dir().join(format!("atlas-storage-{}-{}", name, std::process::id()));
116        let _ = std::fs::remove_dir_all(&p);
117        std::fs::create_dir_all(&p).unwrap();
118        p
119    }
120
121    #[test]
122    #[ignore = "requires GPU"]
123    fn write_then_read_round_trip() {
124        // CUDA must be initialised before any pinned-host allocation.
125        let _ctx = crate::cuda_min::CudaCtx::new(0).expect("cuda init");
126        let dir = tempdir("rt");
127        let spec = GroupLayout::new(1, 2, 1, 16, 128, 2, 4096);
128        let layout = Layout::create(&dir, spec).unwrap();
129        let mut backend = PosixBackend::new(layout).unwrap();
130        let bytes = backend.layout().group_bytes() as usize;
131        let pat: Vec<u8> = (0..bytes).map(|i| (i & 0xFF) as u8).collect();
132        let key = GroupKey::new(0, 1, 0, KvKind::V);
133        backend.write_from_host(key, &pat).unwrap();
134        backend.drop_pagecache();
135
136        let dev = crate::cuda_min::DeviceBuffer::new(bytes).unwrap();
137        let req = ReadRequest {
138            group: key,
139            dst_dev_ptr: dev.ptr,
140        };
141        // Construct a stream from the (already-existing) ctx to satisfy the
142        // backend signature.
143        backend.read(&[req], _ctx.stream).unwrap();
144        let mut host_back = vec![0_u8; bytes];
145        crate::cuda_min::copy_d_to_h_async(
146            host_back.as_mut_ptr() as *mut c_void,
147            dev.ptr,
148            bytes,
149            _ctx.stream,
150        )
151        .unwrap();
152        crate::cuda_min::stream_sync(_ctx.stream).unwrap();
153        assert_eq!(host_back, pat);
154        std::fs::remove_dir_all(&dir).ok();
155    }
156}
spark_storage/backend/posix.rs

spark_storage/backend/
posix.rs
