Skip to main content

mongreldb_core/
pitr.rs

1//! Point-in-time recovery archives built from an online base backup plus
2//! checksummed, transaction-complete logical WAL chunks.
3
4use crate::backup::verify_backup;
5use crate::catalog::TableState;
6use crate::epoch::Epoch;
7use crate::wal::{Op, Record};
8use crate::{Database, MongrelError, Result};
9use fs2::FileExt;
10use serde::{Deserialize, Serialize};
11use sha2::{Digest, Sha256};
12use std::collections::{HashMap, HashSet};
13use std::io::Write;
14use std::path::{Path, PathBuf};
15
16const FORMAT_VERSION: u16 = 1;
17const MANIFEST_FILE: &str = "pitr.json";
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum PitrTarget {
21    Latest,
22    Epoch(u64),
23    TimestampNanos(u64),
24}
25
26#[derive(Clone, Copy)]
27pub enum PitrCredentials<'a> {
28    None,
29    Encryption(&'a str),
30    User {
31        username: &'a str,
32        password: &'a str,
33    },
34    EncryptionAndUser {
35        passphrase: &'a str,
36        username: &'a str,
37        password: &'a str,
38    },
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
42pub struct PitrCommitPoint {
43    pub epoch: u64,
44    pub unix_nanos: u64,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
48pub struct PitrChunkRef {
49    pub file: String,
50    pub from_epoch: u64,
51    pub through_epoch: u64,
52    pub records: usize,
53    pub bytes: u64,
54    pub sha256: String,
55    pub commits: Vec<PitrCommitPoint>,
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
59pub struct PitrArchiveManifest {
60    pub format_version: u16,
61    pub base_epoch: u64,
62    pub base_unix_nanos: u64,
63    pub archived_through_epoch: u64,
64    pub last_commit_unix_nanos: u64,
65    pub chunks: Vec<PitrChunkRef>,
66}
67
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct PitrArchiveReport {
70    pub archive: PathBuf,
71    pub from_epoch: u64,
72    pub through_epoch: u64,
73    pub records: usize,
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize)]
77struct PitrChunk {
78    format_version: u16,
79    from_epoch: u64,
80    through_epoch: u64,
81    records: Vec<Record>,
82    commits: Vec<PitrCommitPoint>,
83}
84
85impl Database {
86    /// Initialize a PITR archive with a consistent online base backup.
87    pub fn create_pitr_archive(&self, destination: impl AsRef<Path>) -> Result<PitrArchiveReport> {
88        let destination = destination.as_ref();
89        let (destination, parent, stage) = prepare_destination(destination, "pitr-stage")?;
90        std::fs::create_dir(&stage)?;
91        let outcome = (|| {
92            let backup = self.hot_backup(stage.join("base"))?;
93            let now = unix_nanos();
94            let manifest = PitrArchiveManifest {
95                format_version: FORMAT_VERSION,
96                base_epoch: backup.epoch,
97                base_unix_nanos: now,
98                archived_through_epoch: backup.epoch,
99                last_commit_unix_nanos: now,
100                chunks: Vec::new(),
101            };
102            write_manifest(&stage, &manifest)?;
103            crate::backup::sync_directories(&stage)?;
104            if destination.exists() {
105                return Err(MongrelError::Conflict(format!(
106                    "PITR archive already exists: {}",
107                    destination.display()
108                )));
109            }
110            std::fs::rename(&stage, &destination)?;
111            sync_dir(&parent)?;
112            // Keep a practical WAL window for callers that archive periodically.
113            self.set_replication_wal_retention_segments(64);
114            Ok(PitrArchiveReport {
115                archive: destination,
116                from_epoch: backup.epoch,
117                through_epoch: backup.epoch,
118                records: 0,
119            })
120        })();
121        if outcome.is_err() && stage.exists() {
122            let _ = std::fs::remove_dir_all(stage);
123        }
124        outcome
125    }
126
127    /// Append all complete commits since the archive watermark. Spilled-run
128    /// commits are converted to ordinary logical Put records while their run
129    /// payload remains available. A retention gap fails closed.
130    pub fn archive_pitr(&self, archive: impl AsRef<Path>) -> Result<PitrArchiveReport> {
131        let archive = archive.as_ref().canonicalize()?;
132        let lock_path = archive.join(".archive.lock");
133        let lock = std::fs::OpenOptions::new()
134            .create(true)
135            .truncate(false)
136            .read(true)
137            .write(true)
138            .open(lock_path)?;
139        lock.lock_exclusive()?;
140        let mut manifest = read_pitr_manifest(&archive)?;
141        let from_epoch = manifest.archived_through_epoch;
142        let batch = self.replication_batch_since(from_epoch)?;
143        if batch.current_epoch == from_epoch {
144            return Ok(PitrArchiveReport {
145                archive,
146                from_epoch,
147                through_epoch: from_epoch,
148                records: 0,
149            });
150        }
151        let has_spilled = batch.records.iter().any(|record| {
152            matches!(&record.op, Op::TxnCommit { added_runs, .. } if !added_runs.is_empty())
153        });
154        if batch.requires_snapshot && !has_spilled {
155            return Err(MongrelError::Conflict(format!(
156                "PITR WAL retention gap after epoch {from_epoch}; create a new base archive"
157            )));
158        }
159        if batch
160            .earliest_epoch
161            .is_some_and(|earliest| earliest > from_epoch.saturating_add(1))
162        {
163            return Err(MongrelError::Conflict(format!(
164                "PITR WAL retention gap: earliest retained epoch is {}",
165                batch.earliest_epoch.unwrap()
166            )));
167        }
168
169        let records = materialize_spilled_records(self, batch.records)?;
170        if records.is_empty() {
171            return Err(MongrelError::Conflict(
172                "PITR source advanced but no complete WAL transactions remain".into(),
173            ));
174        }
175        let mut timestamps = HashMap::new();
176        let mut commit_epochs = Vec::new();
177        for record in &records {
178            match record.op {
179                Op::CommitTimestamp { unix_nanos } => {
180                    timestamps.insert(record.txn_id, unix_nanos);
181                }
182                Op::TxnCommit { epoch, .. } => commit_epochs.push((record.txn_id, epoch)),
183                _ => {}
184            }
185        }
186        commit_epochs.sort_by_key(|(_, epoch)| *epoch);
187        let archive_time = unix_nanos();
188        let mut last_timestamp = manifest.last_commit_unix_nanos;
189        let commits = commit_epochs
190            .into_iter()
191            .map(|(txn_id, epoch)| {
192                let timestamp = timestamps
193                    .get(&txn_id)
194                    .copied()
195                    .unwrap_or(archive_time)
196                    .max(last_timestamp);
197                last_timestamp = timestamp;
198                PitrCommitPoint {
199                    epoch,
200                    unix_nanos: timestamp,
201                }
202            })
203            .collect::<Vec<_>>();
204        let through_epoch = commits
205            .last()
206            .map(|commit| commit.epoch)
207            .ok_or_else(|| MongrelError::Conflict("PITR batch has no commit marker".into()))?;
208        let chunk = PitrChunk {
209            format_version: FORMAT_VERSION,
210            from_epoch,
211            through_epoch,
212            records,
213            commits: commits.clone(),
214        };
215        let bytes = bincode::serialize(&chunk)?;
216        let file = format!("wal-{from_epoch:020}-{through_epoch:020}.bin");
217        let chunk_path = archive.join(&file);
218        let mut output = std::fs::OpenOptions::new()
219            .create_new(true)
220            .write(true)
221            .open(&chunk_path)?;
222        output.write_all(&bytes)?;
223        output.sync_all()?;
224        manifest.chunks.push(PitrChunkRef {
225            file,
226            from_epoch,
227            through_epoch,
228            records: chunk.records.len(),
229            bytes: bytes.len() as u64,
230            sha256: sha256_bytes(&bytes),
231            commits,
232        });
233        manifest.archived_through_epoch = through_epoch;
234        manifest.last_commit_unix_nanos = last_timestamp;
235        write_manifest(&archive, &manifest)?;
236        sync_dir(&archive)?;
237        Ok(PitrArchiveReport {
238            archive,
239            from_epoch,
240            through_epoch,
241            records: chunk.records.len(),
242        })
243    }
244}
245
246pub fn read_pitr_manifest(archive: impl AsRef<Path>) -> Result<PitrArchiveManifest> {
247    let archive = archive.as_ref();
248    let manifest: PitrArchiveManifest =
249        serde_json::from_slice(&std::fs::read(archive.join(MANIFEST_FILE))?)
250            .map_err(|error| MongrelError::InvalidArgument(format!("PITR manifest: {error}")))?;
251    if manifest.format_version != FORMAT_VERSION {
252        return Err(MongrelError::InvalidArgument(format!(
253            "unsupported PITR archive version {}",
254            manifest.format_version
255        )));
256    }
257    Ok(manifest)
258}
259
260/// Restore an archive to a new directory at an epoch or timestamp cutoff.
261pub fn restore_pitr(
262    archive: impl AsRef<Path>,
263    destination: impl AsRef<Path>,
264    target: PitrTarget,
265    credentials: PitrCredentials<'_>,
266) -> Result<u64> {
267    let archive = archive.as_ref().canonicalize()?;
268    let manifest = read_pitr_manifest(&archive)?;
269    verify_backup(archive.join("base"))?;
270    let target_epoch = resolve_target_epoch(&manifest, target)?;
271    let (destination, parent, stage) = prepare_destination(destination.as_ref(), "pitr-restore")?;
272    std::fs::create_dir(&stage)?;
273    let outcome = (|| {
274        copy_tree(&archive.join("base"), &stage)?;
275        let meta = stage.join("_meta");
276        std::fs::create_dir_all(&meta)?;
277        write_synced(&meta.join("replica"), b"PITR restore staging\n")?;
278        crate::replication::write_replica_epoch(&stage, manifest.base_epoch)?;
279
280        let records = load_records_through(&archive, &manifest, target_epoch)?;
281        if !records.is_empty() {
282            let replica = open_with_credentials(&stage, credentials)?;
283            let applied = replica.append_replication_batch(&records)?;
284            drop(replica);
285            let recovered = open_with_credentials(&stage, credentials)?;
286            if recovered.visible_epoch().0 < applied || recovered.visible_epoch().0 < target_epoch {
287                return Err(MongrelError::Other(format!(
288                    "PITR recovery stopped at epoch {}, expected {target_epoch}",
289                    recovered.visible_epoch().0
290                )));
291            }
292            drop(recovered);
293        }
294        let _ = std::fs::remove_file(meta.join("replica"));
295        let _ = std::fs::remove_file(meta.join("repl_epoch"));
296        sync_dir(&meta)?;
297        crate::backup::sync_directories(&stage)?;
298        if destination.exists() {
299            return Err(MongrelError::Conflict(format!(
300                "PITR destination already exists: {}",
301                destination.display()
302            )));
303        }
304        std::fs::rename(&stage, &destination)?;
305        sync_dir(&parent)?;
306        Ok(target_epoch)
307    })();
308    if outcome.is_err() && stage.exists() {
309        let _ = std::fs::remove_dir_all(stage);
310    }
311    outcome
312}
313
314fn materialize_spilled_records(db: &Database, records: Vec<Record>) -> Result<Vec<Record>> {
315    let table_names: HashMap<u64, String> = db
316        .catalog_snapshot()
317        .tables
318        .into_iter()
319        .filter(|entry| matches!(entry.state, TableState::Live))
320        .map(|entry| (entry.table_id, entry.name))
321        .collect();
322    let commit_epochs: HashMap<u64, u64> = records
323        .iter()
324        .filter_map(|record| match record.op {
325            Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
326            _ => None,
327        })
328        .collect();
329    let mut output = Vec::with_capacity(records.len());
330    for record in records {
331        let Op::TxnCommit { epoch, added_runs } = &record.op else {
332            output.push(record);
333            continue;
334        };
335        for run in added_runs {
336            let name = table_names.get(&run.table_id).ok_or_else(|| {
337                MongrelError::Conflict(format!(
338                    "PITR cannot materialize spilled run {} for unavailable table {}",
339                    run.run_id, run.table_id
340                ))
341            })?;
342            let handle = db.table(name)?;
343            let table = handle.lock();
344            let mut reader = table.open_reader(run.run_id)?;
345            let mut rows = reader.all_rows()?;
346            for row in &mut rows {
347                row.committed_epoch = Epoch(*epoch);
348            }
349            output.push(Record::new(
350                record.seq,
351                record.txn_id,
352                Op::Put {
353                    table_id: run.table_id,
354                    rows: bincode::serialize(&rows)?,
355                },
356            ));
357        }
358        let mut commit = record;
359        if let Op::TxnCommit { added_runs, .. } = &mut commit.op {
360            added_runs.clear();
361        }
362        output.push(commit);
363    }
364    let complete: HashSet<u64> = output
365        .iter()
366        .filter_map(|record| match record.op {
367            Op::TxnCommit { .. } => Some(record.txn_id),
368            _ => None,
369        })
370        .collect();
371    if commit_epochs
372        .keys()
373        .any(|txn_id| !complete.contains(txn_id))
374    {
375        return Err(MongrelError::Conflict(
376            "PITR conversion lost a transaction commit".into(),
377        ));
378    }
379    Ok(output)
380}
381
382fn resolve_target_epoch(manifest: &PitrArchiveManifest, target: PitrTarget) -> Result<u64> {
383    match target {
384        PitrTarget::Latest => Ok(manifest.archived_through_epoch),
385        PitrTarget::Epoch(epoch)
386            if epoch >= manifest.base_epoch && epoch <= manifest.archived_through_epoch =>
387        {
388            Ok(epoch)
389        }
390        PitrTarget::Epoch(epoch) => Err(MongrelError::InvalidArgument(format!(
391            "PITR epoch {epoch} outside archive range {}..={}",
392            manifest.base_epoch, manifest.archived_through_epoch
393        ))),
394        PitrTarget::TimestampNanos(timestamp) => {
395            if timestamp < manifest.base_unix_nanos {
396                return Err(MongrelError::InvalidArgument(
397                    "PITR timestamp predates base backup".into(),
398                ));
399            }
400            let mut epoch = manifest.base_epoch;
401            for commit in manifest.chunks.iter().flat_map(|chunk| &chunk.commits) {
402                if commit.unix_nanos > timestamp {
403                    break;
404                }
405                epoch = commit.epoch;
406            }
407            Ok(epoch)
408        }
409    }
410}
411
412fn load_records_through(
413    archive: &Path,
414    manifest: &PitrArchiveManifest,
415    target_epoch: u64,
416) -> Result<Vec<Record>> {
417    let mut records = Vec::new();
418    for reference in &manifest.chunks {
419        if reference.from_epoch >= target_epoch {
420            break;
421        }
422        let bytes = std::fs::read(archive.join(&reference.file))?;
423        if bytes.len() as u64 != reference.bytes || sha256_bytes(&bytes) != reference.sha256 {
424            return Err(MongrelError::Other(format!(
425                "PITR chunk {} checksum mismatch",
426                reference.file
427            )));
428        }
429        let chunk: PitrChunk = bincode::deserialize(&bytes)?;
430        if chunk.format_version != FORMAT_VERSION {
431            return Err(MongrelError::InvalidArgument(format!(
432                "unsupported PITR chunk version {}",
433                chunk.format_version
434            )));
435        }
436        let selected: HashSet<u64> = chunk
437            .records
438            .iter()
439            .filter_map(|record| match record.op {
440                Op::TxnCommit { epoch, .. } if epoch <= target_epoch => Some(record.txn_id),
441                _ => None,
442            })
443            .collect();
444        records.extend(
445            chunk
446                .records
447                .into_iter()
448                .filter(|record| selected.contains(&record.txn_id)),
449        );
450        if reference.through_epoch >= target_epoch {
451            break;
452        }
453    }
454    Ok(records)
455}
456
457fn open_with_credentials(path: &Path, credentials: PitrCredentials<'_>) -> Result<Database> {
458    match credentials {
459        PitrCredentials::None => Database::open(path),
460        #[cfg(feature = "encryption")]
461        PitrCredentials::Encryption(passphrase) => Database::open_encrypted(path, passphrase),
462        #[cfg(not(feature = "encryption"))]
463        PitrCredentials::Encryption(_) => Err(MongrelError::Encryption(
464            "encryption feature is disabled".into(),
465        )),
466        PitrCredentials::User { username, password } => {
467            Database::open_with_credentials(path, username, password)
468        }
469        #[cfg(feature = "encryption")]
470        PitrCredentials::EncryptionAndUser {
471            passphrase,
472            username,
473            password,
474        } => Database::open_encrypted_with_credentials(path, passphrase, username, password),
475        #[cfg(not(feature = "encryption"))]
476        PitrCredentials::EncryptionAndUser { .. } => Err(MongrelError::Encryption(
477            "encryption feature is disabled".into(),
478        )),
479    }
480}
481
482fn prepare_destination(path: &Path, label: &str) -> Result<(PathBuf, PathBuf, PathBuf)> {
483    if path.exists() {
484        return Err(MongrelError::Conflict(format!(
485            "destination already exists: {}",
486            path.display()
487        )));
488    }
489    let name = path
490        .file_name()
491        .ok_or_else(|| MongrelError::InvalidArgument("invalid destination".into()))?;
492    let requested_parent = path
493        .parent()
494        .filter(|parent| !parent.as_os_str().is_empty())
495        .unwrap_or_else(|| Path::new("."));
496    std::fs::create_dir_all(requested_parent)?;
497    let parent = requested_parent.canonicalize()?;
498    let destination = parent.join(name);
499    let stage = parent.join(format!(
500        ".{}.{}-{}-{}",
501        name.to_string_lossy(),
502        label,
503        std::process::id(),
504        unix_nanos()
505    ));
506    if stage.exists() {
507        return Err(MongrelError::Conflict(
508            "PITR staging path already exists".into(),
509        ));
510    }
511    Ok((destination, parent, stage))
512}
513
514fn copy_tree(source: &Path, destination: &Path) -> Result<()> {
515    std::fs::create_dir_all(destination)?;
516    let mut entries = std::fs::read_dir(source)?.collect::<std::io::Result<Vec<_>>>()?;
517    entries.sort_by_key(std::fs::DirEntry::file_name);
518    for entry in entries {
519        let file_type = entry.file_type()?;
520        if file_type.is_symlink() {
521            return Err(MongrelError::InvalidArgument(format!(
522                "PITR restore refuses symlink {}",
523                entry.path().display()
524            )));
525        }
526        let target = destination.join(entry.file_name());
527        if file_type.is_dir() {
528            copy_tree(&entry.path(), &target)?;
529        } else if file_type.is_file() {
530            crate::backup::copy_file_synced(&entry.path(), &target)?;
531        }
532    }
533    Ok(())
534}
535
536fn write_manifest(root: &Path, manifest: &PitrArchiveManifest) -> Result<()> {
537    let bytes = serde_json::to_vec_pretty(manifest)
538        .map_err(|error| MongrelError::Other(format!("PITR manifest encode: {error}")))?;
539    let temporary = root.join(format!(".{MANIFEST_FILE}.tmp"));
540    write_synced(&temporary, &bytes)?;
541    std::fs::rename(temporary, root.join(MANIFEST_FILE))?;
542    sync_dir(root)
543}
544
545fn write_synced(path: &Path, bytes: &[u8]) -> Result<()> {
546    let mut file = std::fs::File::create(path)?;
547    file.write_all(bytes)?;
548    file.sync_all()?;
549    Ok(())
550}
551
552fn sync_dir(path: &Path) -> Result<()> {
553    std::fs::File::open(path)?.sync_all()?;
554    Ok(())
555}
556
557fn sha256_bytes(bytes: &[u8]) -> String {
558    Sha256::digest(bytes)
559        .iter()
560        .map(|byte| format!("{byte:02x}"))
561        .collect()
562}
563
564fn unix_nanos() -> u64 {
565    std::time::SystemTime::now()
566        .duration_since(std::time::UNIX_EPOCH)
567        .unwrap_or_default()
568        .as_nanos() as u64
569}