spark_storage/
layout.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// On-disk layout for `--high-speed-swap`. One file per layer under
4// `--high-speed-swap-dir`, pre-allocated so the filesystem reserves the bytes
5// up-front (no surprise ENOSPC mid-decode).
6//
7// File names: `layer_{:05}.kv`. File contents are an opaque
8// `GroupLayout`-defined stripe; the `Layout` type owns the open `File`s.
9//
10// PLATFORMS. On Linux the files are opened `O_DIRECT` (the io_uring / cuFile
11// path needs it) and reserved with `posix_fallocate`. On Windows they are
12// opened buffered and reserved with `set_len`:
13//   * `FILE_FLAG_NO_BUFFERING` is the nearest O_DIRECT analogue but imposes
14//     sector alignment on every buffer, offset and length; the tier's records
15//     are 4 KiB-aligned but its bounce buffers are pinned host allocations
16//     whose alignment is CUDA's to choose, so the flag would be unsafe to
17//     assume. Buffered is correct, just not zero-copy.
18//   * `SetFileValidData` is the true fallocate analogue but needs
19//     SE_MANAGE_VOLUME_NAME; `set_len` reserves the range without it.
20
21use anyhow::{Context, Result};
22use std::fs::{File, OpenOptions};
23#[cfg(unix)]
24use std::os::fd::{AsRawFd, RawFd};
25use std::path::{Path, PathBuf};
26
27use crate::group::{GroupKey, GroupLayout};
28
29pub struct Layout {
30    pub dir: PathBuf,
31    pub spec: GroupLayout,
32    /// One `File` per layer. O_DIRECT on Linux for the io_uring / cuFile path;
33    /// buffered elsewhere. Held as `File` rather than `OwnedFd` so the portable
34    /// positional-I/O backend can use it on every platform.
35    files: Vec<File>,
36}
37
38impl Layout {
39    pub fn create(dir: &Path, spec: GroupLayout) -> Result<Self> {
40        std::fs::create_dir_all(dir).with_context(|| format!("mkdir {}", dir.display()))?;
41        let mut files = Vec::with_capacity(spec.num_layers as usize);
42        for layer in 0..spec.num_layers {
43            let p = dir.join(format!("layer_{layer:05}.kv"));
44            let mut opts = OpenOptions::new();
45            opts.read(true).write(true).create(true).truncate(false);
46            set_direct_flag(&mut opts);
47            let f = opts
48                .open(&p)
49                .with_context(|| format!("open {}", p.display()))?;
50            preallocate(&f, spec.bytes_per_layer())
51                .with_context(|| format!("preallocate {}", p.display()))?;
52            files.push(f);
53        }
54        Ok(Self {
55            dir: dir.to_path_buf(),
56            spec,
57            files,
58        })
59    }
60
61    /// Open an existing layout (panics if a file is missing or undersized).
62    pub fn open(dir: &Path, spec: GroupLayout) -> Result<Self> {
63        let mut files = Vec::with_capacity(spec.num_layers as usize);
64        for layer in 0..spec.num_layers {
65            let p = dir.join(format!("layer_{layer:05}.kv"));
66            let mut opts = OpenOptions::new();
67            opts.read(true).write(true);
68            set_direct_flag(&mut opts);
69            let f = opts
70                .open(&p)
71                .with_context(|| format!("open {}", p.display()))?;
72            let len = f.metadata()?.len();
73            if len < spec.bytes_per_layer() {
74                anyhow::bail!(
75                    "layer file {} is undersized: {} < {}",
76                    p.display(),
77                    len,
78                    spec.bytes_per_layer()
79                );
80            }
81            files.push(f);
82        }
83        Ok(Self {
84            dir: dir.to_path_buf(),
85            spec,
86            files,
87        })
88    }
89
90    /// Raw fd for the io_uring / cuFile paths, which are Linux-only.
91    #[cfg(unix)]
92    pub fn fd(&self, layer: u32) -> RawFd {
93        self.files[layer as usize].as_raw_fd()
94    }
95
96    /// The layer file itself — what the portable positional-I/O backend uses,
97    /// and the only accessor available on Windows.
98    pub fn file(&self, layer: u32) -> &File {
99        &self.files[layer as usize]
100    }
101
102    pub fn offset(&self, key: GroupKey) -> u64 {
103        self.spec.file_offset(key)
104    }
105
106    pub fn group_bytes(&self) -> u64 {
107        self.spec.group_bytes()
108    }
109}
110
111/// O_DIRECT on Linux; nothing to set elsewhere (see the header note).
112#[cfg(target_os = "linux")]
113fn set_direct_flag(opts: &mut OpenOptions) {
114    use std::os::unix::fs::OpenOptionsExt;
115    opts.custom_flags(libc::O_DIRECT);
116}
117
118#[cfg(not(target_os = "linux"))]
119fn set_direct_flag(_opts: &mut OpenOptions) {}
120
121#[cfg(unix)]
122fn preallocate(file: &File, size: u64) -> Result<()> {
123    // posix_fallocate is portable across ext4/xfs and reserves space without
124    // writing zeros; FALLOC_FL_KEEP_SIZE would be wrong here because we *do*
125    // want the file size to grow.
126    let fd = file.as_raw_fd();
127    let res = unsafe { libc::posix_fallocate(fd, 0, size as libc::off_t) };
128    if res != 0 {
129        anyhow::bail!("posix_fallocate({size}) failed: {res}");
130    }
131    Ok(())
132}
133
134// Windows: `set_len` extends the file to the requested size. NTFS keeps the
135// tail sparse until written, so this reserves the RANGE rather than the blocks
136// -- weaker than posix_fallocate, and the honest trade for not requiring the
137// SE_MANAGE_VOLUME privilege that SetFileValidData needs.
138#[cfg(windows)]
139fn preallocate(file: &File, size: u64) -> Result<()> {
140    file.set_len(size)
141        .with_context(|| format!("set_len({size})"))?;
142    Ok(())
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use crate::group::{GroupKey, KvKind};
149
150    #[test]
151    fn create_open_round_trip() {
152        let tmp = tempdir();
153        let spec = GroupLayout::new(2, 4, 2, 16, 128, 2, 4096);
154        {
155            let l = Layout::create(&tmp, spec).unwrap();
156            assert_eq!(l.spec.num_layers, 2);
157            // File should be size bytes_per_layer.
158            let p = tmp.join("layer_00000.kv");
159            let len = std::fs::metadata(&p).unwrap().len();
160            assert_eq!(len, spec.bytes_per_layer());
161        }
162        {
163            let l = Layout::open(&tmp, spec).unwrap();
164            let off = l.offset(GroupKey::new(0, 1, 1, KvKind::V));
165            assert_eq!(off, spec.file_offset(GroupKey::new(0, 1, 1, KvKind::V)));
166        }
167        std::fs::remove_dir_all(&tmp).ok();
168    }
169
170    fn tempdir() -> PathBuf {
171        let p = std::env::temp_dir().join(format!("atlas-storage-test-{}", std::process::id()));
172        std::fs::create_dir_all(&p).unwrap();
173        p
174    }
175}
spark_storage/layout.rs

spark_storage/
layout.rs
