1use 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 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 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 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
260pub 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}