Skip to main content

mongreldb_core/
replication.rs

1//! Replication bootstrap image and follower metadata.
2
3use crate::{MongrelError, Result};
4use serde::{Deserialize, Serialize};
5use std::collections::HashSet;
6use std::io::Write;
7use std::path::{Component, Path, PathBuf};
8
9const FORMAT_VERSION: u16 = 1;
10const REPLICA_MARKER: &str = "replica";
11const REPLICA_EPOCH: &str = "repl_epoch";
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub(crate) struct ReplicationFile {
15    path: PathBuf,
16    data: Vec<u8>,
17}
18
19impl ReplicationFile {
20    pub(crate) fn new(path: PathBuf, data: Vec<u8>) -> Self {
21        Self { path, data }
22    }
23}
24
25pub(crate) fn capture_files(root: &Path) -> Result<Vec<ReplicationFile>> {
26    let mut files = Vec::new();
27    capture_dir(root, root, &mut files)?;
28    files.sort_by(|a, b| a.path.cmp(&b.path));
29    Ok(files)
30}
31
32fn capture_dir(root: &Path, dir: &Path, files: &mut Vec<ReplicationFile>) -> Result<()> {
33    let mut entries = std::fs::read_dir(dir)?.collect::<std::io::Result<Vec<_>>>()?;
34    entries.sort_by_key(std::fs::DirEntry::file_name);
35    for entry in entries {
36        let path = entry.path();
37        let relative = path
38            .strip_prefix(root)
39            .map_err(|error| MongrelError::Other(error.to_string()))?;
40        if relative == Path::new("_meta/.lock")
41            || relative == Path::new("_meta/replica")
42            || relative == Path::new("_meta/repl_epoch")
43            || relative
44                .components()
45                .any(|component| matches!(component, Component::Normal(name) if name == "_cache"))
46        {
47            continue;
48        }
49        let file_type = entry.file_type()?;
50        if file_type.is_symlink() {
51            return Err(MongrelError::InvalidArgument(format!(
52                "replication snapshot refuses symlink {}",
53                path.display()
54            )));
55        }
56        if file_type.is_dir() {
57            capture_dir(root, &path, files)?;
58        } else if file_type.is_file() {
59            files.push(ReplicationFile::new(
60                relative.to_path_buf(),
61                std::fs::read(&path)?,
62            ));
63        }
64    }
65    Ok(())
66}
67
68/// A consistent database-directory image plus the leader commit epoch it
69/// covers. The image is opaque to HTTP; encode/decode use the core's versioned
70/// bincode envelope so server and client share one format.
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct ReplicationSnapshot {
73    version: u16,
74    epoch: u64,
75    files: Vec<ReplicationFile>,
76}
77
78/// Complete committed WAL transactions available after a follower epoch.
79#[derive(Debug, Clone)]
80pub struct ReplicationBatch {
81    pub current_epoch: u64,
82    pub earliest_epoch: Option<u64>,
83    pub requires_snapshot: bool,
84    pub records: Vec<crate::wal::Record>,
85}
86
87impl ReplicationSnapshot {
88    pub(crate) fn new(epoch: u64, files: Vec<ReplicationFile>) -> Self {
89        Self {
90            version: FORMAT_VERSION,
91            epoch,
92            files,
93        }
94    }
95
96    pub fn epoch(&self) -> u64 {
97        self.epoch
98    }
99
100    pub fn encode(&self) -> Result<Vec<u8>> {
101        Ok(bincode::serialize(self)?)
102    }
103
104    pub fn decode(bytes: &[u8]) -> Result<Self> {
105        let snapshot: Self = bincode::deserialize(bytes)?;
106        if snapshot.version != FORMAT_VERSION {
107            return Err(MongrelError::InvalidArgument(format!(
108                "unsupported replication snapshot version {}",
109                snapshot.version
110            )));
111        }
112        Ok(snapshot)
113    }
114
115    /// Atomically replace `destination` with this snapshot and mark it as a
116    /// read-only replica. Files are first written and fsynced in a sibling
117    /// staging directory; an existing destination is retained until install.
118    pub fn install(&self, destination: impl AsRef<Path>) -> Result<()> {
119        let destination = destination.as_ref();
120        let parent = destination
121            .parent()
122            .filter(|path| !path.as_os_str().is_empty())
123            .unwrap_or_else(|| Path::new("."));
124        std::fs::create_dir_all(parent)?;
125        let name = destination
126            .file_name()
127            .and_then(|name| name.to_str())
128            .ok_or_else(|| MongrelError::InvalidArgument("invalid replica destination".into()))?;
129        let nonce = std::time::SystemTime::now()
130            .duration_since(std::time::UNIX_EPOCH)
131            .unwrap_or_default()
132            .as_nanos();
133        let stage = parent.join(format!(
134            ".{name}.replica-stage-{}-{nonce}",
135            std::process::id()
136        ));
137        let backup = parent.join(format!(
138            ".{name}.replica-old-{}-{nonce}",
139            std::process::id()
140        ));
141
142        if stage.exists() || backup.exists() {
143            return Err(MongrelError::Conflict(
144                "replication staging path already exists".into(),
145            ));
146        }
147        std::fs::create_dir(&stage)?;
148        if let Err(error) = self.write_into(&stage) {
149            let _ = std::fs::remove_dir_all(&stage);
150            return Err(error);
151        }
152
153        let had_destination = destination.exists();
154        if had_destination {
155            std::fs::rename(destination, &backup)?;
156        }
157        if let Err(error) = std::fs::rename(&stage, destination) {
158            if had_destination {
159                let _ = std::fs::rename(&backup, destination);
160            }
161            let _ = std::fs::remove_dir_all(&stage);
162            return Err(error.into());
163        }
164        sync_dir(parent);
165        if had_destination {
166            std::fs::remove_dir_all(&backup)?;
167        }
168        Ok(())
169    }
170
171    fn write_into(&self, root: &Path) -> Result<()> {
172        let mut seen = HashSet::new();
173        for file in &self.files {
174            validate_relative_path(&file.path)?;
175            if !seen.insert(file.path.clone()) {
176                return Err(MongrelError::InvalidArgument(format!(
177                    "duplicate replication snapshot path {:?}",
178                    file.path
179                )));
180            }
181            let path = root.join(&file.path);
182            let parent = path.parent().expect("validated file has parent");
183            std::fs::create_dir_all(parent)?;
184            let mut output = std::fs::OpenOptions::new()
185                .create_new(true)
186                .write(true)
187                .open(&path)?;
188            output.write_all(&file.data)?;
189            output.sync_all()?;
190        }
191        if !root.join(crate::catalog::CATALOG_FILENAME).is_file() {
192            return Err(MongrelError::InvalidArgument(
193                "replication snapshot has no CATALOG".into(),
194            ));
195        }
196        let meta = root.join("_meta");
197        std::fs::create_dir_all(&meta)?;
198        write_synced(&meta.join(REPLICA_MARKER), b"read-only replica\n")?;
199        write_replica_epoch(root, self.epoch)?;
200        sync_dir(&meta);
201        sync_dir(root);
202        Ok(())
203    }
204}
205
206pub fn is_replica(root: impl AsRef<Path>) -> bool {
207    root.as_ref().join("_meta").join(REPLICA_MARKER).is_file()
208}
209
210pub fn replica_epoch(root: impl AsRef<Path>) -> Result<u64> {
211    let path = root.as_ref().join("_meta").join(REPLICA_EPOCH);
212    let value = std::fs::read_to_string(&path)
213        .map_err(|error| MongrelError::NotFound(format!("{}: {error}", path.display())))?;
214    value.trim().parse().map_err(|error| {
215        MongrelError::InvalidArgument(format!(
216            "invalid replica epoch in {}: {error}",
217            path.display()
218        ))
219    })
220}
221
222pub fn write_replica_epoch(root: impl AsRef<Path>, epoch: u64) -> Result<()> {
223    let meta = root.as_ref().join("_meta");
224    std::fs::create_dir_all(&meta)?;
225    let path = meta.join(REPLICA_EPOCH);
226    let temp = meta.join(format!("{REPLICA_EPOCH}.tmp"));
227    write_synced(&temp, epoch.to_string().as_bytes())?;
228    std::fs::rename(&temp, &path)?;
229    sync_dir(&meta);
230    Ok(())
231}
232
233fn validate_relative_path(path: &Path) -> Result<()> {
234    if path.as_os_str().is_empty()
235        || path
236            .components()
237            .any(|component| !matches!(component, Component::Normal(_)))
238    {
239        return Err(MongrelError::InvalidArgument(format!(
240            "unsafe replication snapshot path {:?}",
241            path
242        )));
243    }
244    Ok(())
245}
246
247fn write_synced(path: &Path, bytes: &[u8]) -> Result<()> {
248    let mut file = std::fs::File::create(path)?;
249    file.write_all(bytes)?;
250    file.sync_all()?;
251    Ok(())
252}
253
254fn sync_dir(path: &Path) {
255    if let Ok(dir) = std::fs::File::open(path) {
256        let _ = dir.sync_all();
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263
264    #[test]
265    fn snapshot_install_rejects_path_traversal() {
266        let dir = tempfile::tempdir().unwrap();
267        let snapshot = ReplicationSnapshot {
268            version: FORMAT_VERSION,
269            epoch: 1,
270            files: vec![ReplicationFile::new("../escape".into(), vec![1])],
271        };
272        assert!(snapshot.install(dir.path().join("replica")).is_err());
273        assert!(!dir.path().join("escape").exists());
274    }
275}