1use crate::disk;
54use crate::lease::new_instance_id;
55use crate::manifest::{crc32_hex, latest_key, object_key, parse_segment_key, Manifest};
56use crate::segment::complete_frames;
57use crate::store::{build_store_with, BackupConfig};
58use crate::{redact, BackupError};
59use mushroomdb::restore::cut_at_commit;
60use object_store::path::Path as ObjectPath;
61use object_store::ObjectStore;
62use std::path::{Path, PathBuf};
63use std::sync::Arc;
64use time::format_description::well_known::Rfc3339;
65use time::OffsetDateTime;
66use tokio::io::AsyncWriteExt as _;
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub enum RestoreTarget {
71 CommitSeq(u64),
75 Timestamp(OffsetDateTime),
78}
79
80impl RestoreTarget {
81 pub fn parse(s: &str) -> Result<Self, String> {
84 let s = s.trim();
85 if !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit()) {
86 return s
87 .parse()
88 .map(RestoreTarget::CommitSeq)
89 .map_err(|e| format!("--at {s:?}: {e}"));
90 }
91 OffsetDateTime::parse(s, &Rfc3339)
92 .map(RestoreTarget::Timestamp)
93 .map_err(|_| {
94 format!(
95 "--at {s:?} is neither a commit number nor an RFC 3339 timestamp \
96 (for example 2026-10-01T12:00:00Z)"
97 )
98 })
99 }
100}
101
102impl std::fmt::Display for RestoreTarget {
103 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104 match self {
105 RestoreTarget::CommitSeq(n) => write!(f, "commit {n}"),
106 RestoreTarget::Timestamp(t) => write!(f, "{}", rfc3339(*t)),
107 }
108 }
109}
110
111#[derive(Debug, Clone, PartialEq, Eq)]
113pub struct RestoreS3Outcome {
114 pub into: PathBuf,
116 pub manifest: String,
118 pub commits: u64,
120 pub source_lineage: String,
122 pub lineage: String,
124 pub files: Vec<String>,
126 pub segments: usize,
128 pub bytes: u64,
130 pub tail_note: Option<String>,
134 pub install_note: Option<String>,
138 pub swept: Vec<PathBuf>,
141}
142
143#[derive(Debug)]
147pub enum RestoreS3Error {
148 Config(String),
150 Store(object_store::Error),
152 Io(std::io::Error),
154 NoBackup { from: String },
156 Manifest(String),
158 Verify(String),
162 IntoNotEmpty { into: PathBuf },
164 DataDirNotEmpty { dir: PathBuf, entries: Vec<String> },
167 BeforeOldestArchive { target: u64, oldest: u64 },
169 AfterLatest {
172 target: u64,
173 commits: u64,
174 not_restorable: Option<String>,
177 },
178 CommitTimesAbsent,
181 CommitTimesEnd {
185 recorded_through: u64,
187 recorded_through_ms: i64,
189 commits: u64,
191 },
192 Unreachable(String),
194}
195
196impl std::fmt::Display for RestoreS3Error {
197 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
198 match self {
199 RestoreS3Error::Config(m) => write!(f, "{}", redact(m)),
202 RestoreS3Error::Store(e) => write!(f, "object store: {}", redact(&e.to_string())),
203 RestoreS3Error::Io(e) => write!(f, "{e}"),
204 RestoreS3Error::NoBackup { from } => {
205 write!(f, "{from} holds no backup: there is no LATEST under it")
206 }
207 RestoreS3Error::Manifest(m) => write!(f, "manifest: {m}"),
208 RestoreS3Error::Verify(m) => write!(f, "the restore point does not verify: {m}"),
209 RestoreS3Error::IntoNotEmpty { into } => write!(
210 f,
211 "{} is not empty; restore into a new or empty directory",
212 into.display()
213 ),
214 RestoreS3Error::DataDirNotEmpty { dir, entries } => write!(
215 f,
216 "{} holds no store but is not empty ({}); a restore on boot installs only \
217 into an empty data directory. Move those away and start again",
218 dir.display(),
219 entries.join(", ")
220 ),
221 RestoreS3Error::BeforeOldestArchive { target, oldest } => write!(
222 f,
223 "commit {target} is before the oldest commit the backup holds ({oldest}); \
224 restore commit {oldest} or later"
225 ),
226 RestoreS3Error::AfterLatest {
227 target,
228 commits,
229 not_restorable,
230 } => {
231 match commits.checked_sub(1) {
232 Some(newest) => write!(
233 f,
234 "commit {target} is past the newest commit in the backup ({newest})"
235 )?,
236 None => write!(f, "commit {target}: the backup holds no commits")?,
237 }
238 match not_restorable {
239 Some(why) => write!(f, "; {why}"),
240 None => Ok(()),
241 }
242 }
243 RestoreS3Error::CommitTimesAbsent => write!(
244 f,
245 "the backup records no commit times (no commit_times.bin), so a timestamp \
246 cannot be resolved; restore by commit number instead"
247 ),
248 RestoreS3Error::CommitTimesEnd {
249 recorded_through,
250 recorded_through_ms,
251 commits,
252 } => write!(
253 f,
254 "the backup records commit times only up to commit {recorded_through}, at {at}, \
255 and commits {} to {} have none, so an instant at or after {at} cannot be \
256 placed among them. Restore an instant before it, or by commit index \
257 (--at <commit index>, for example --at {recorded_through}), or without --at \
258 for the newest",
259 recorded_through + 1,
260 commits.saturating_sub(1),
261 at = ms_rfc3339(*recorded_through_ms),
262 ),
263 RestoreS3Error::Unreachable(m) => write!(f, "{m}"),
264 }
265 }
266}
267
268impl std::error::Error for RestoreS3Error {}
269
270impl From<object_store::Error> for RestoreS3Error {
271 fn from(e: object_store::Error) -> Self {
272 RestoreS3Error::Store(e)
273 }
274}
275
276impl From<std::io::Error> for RestoreS3Error {
277 fn from(e: std::io::Error) -> Self {
278 RestoreS3Error::Io(e)
279 }
280}
281
282impl From<BackupError> for RestoreS3Error {
283 fn from(e: BackupError) -> Self {
284 match e {
285 BackupError::Store(e) => RestoreS3Error::Store(e),
286 BackupError::Io(e) => RestoreS3Error::Io(e),
287 BackupError::Manifest(m) => RestoreS3Error::Manifest(m),
288 other => RestoreS3Error::Config(other.to_string()),
289 }
290 }
291}
292
293pub async fn restore_from_s3(
300 from: &str,
301 at: Option<RestoreTarget>,
302 into: &Path,
303) -> Result<RestoreS3Outcome, RestoreS3Error> {
304 restore_from_s3_with(from, at, into, |k| std::env::var(k).ok()).await
305}
306
307pub async fn restore_from_s3_with(
309 from: &str,
310 at: Option<RestoreTarget>,
311 into: &Path,
312 get: impl Fn(&str) -> Option<String>,
313) -> Result<RestoreS3Outcome, RestoreS3Error> {
314 let config = BackupConfig::new(from.trim());
315 let location = config
316 .location()
317 .map_err(|e| RestoreS3Error::Config(config_message(e, from)))?;
318 check_into(into)?;
319 let store = build_store_with(&config, get)
320 .map_err(|e| RestoreS3Error::Config(config_message(e, from)))?;
321 restore_from_store(store, &location.prefix, at, into).await
322}
323
324fn config_message(e: BackupError, from: &str) -> String {
325 match e {
326 BackupError::Config(m) => m.replace("MUSHROOMDB_BACKUP_S3_URL", "--from"),
327 other => format!("{from}: {other}"),
328 }
329}
330
331pub async fn restore_from_store(
334 store: Arc<dyn ObjectStore>,
335 prefix: &str,
336 at: Option<RestoreTarget>,
337 into: &Path,
338) -> Result<RestoreS3Outcome, RestoreS3Error> {
339 restore_with(store, prefix, at, into, Mode::Beside).await
340}
341
342pub async fn restore_in_place(
348 store: Arc<dyn ObjectStore>,
349 prefix: &str,
350 dir: &Path,
351) -> Result<RestoreS3Outcome, RestoreS3Error> {
352 restore_with(store, prefix, None, dir, Mode::InPlace).await
353}
354
355#[derive(Debug, Clone, Copy, PartialEq, Eq)]
357enum Mode {
358 Beside,
360 InPlace,
363}
364
365impl Mode {
366 fn staging_root(self, into: &Path) -> PathBuf {
368 match self {
369 Mode::Beside => parent_of(into),
370 Mode::InPlace => into.to_path_buf(),
371 }
372 }
373}
374
375async fn restore_with(
376 store: Arc<dyn ObjectStore>,
377 prefix: &str,
378 at: Option<RestoreTarget>,
379 into: &Path,
380 mode: Mode,
381) -> Result<RestoreS3Outcome, RestoreS3Error> {
382 if mode == Mode::Beside {
383 check_into(into)?;
384 }
385 let bucket = Bucket { store, prefix };
386 let Some(latest_target) = bucket.get_text(&latest_key(prefix)).await? else {
387 return Err(RestoreS3Error::NoBackup {
388 from: if prefix.is_empty() {
389 "the bucket".into()
390 } else {
391 prefix.to_string()
392 },
393 });
394 };
395 let (root, name) = (mode.staging_root(into), name_of(into));
399 let swept = blocking(move || Ok(sweep_stale_staging(&root, &name))).await?;
400 if mode == Mode::InPlace {
401 let dir = into.to_path_buf();
402 blocking(move || check_in_place(&dir)).await?;
403 }
404 let latest = bucket.manifest(&latest_target).await?;
405 let needs_tail = !matches!(at, Some(RestoreTarget::CommitSeq(n)) if n < latest.commit_seq);
408 let (tail, tail_note) = if needs_tail {
409 bucket.tail_segments(&latest).await?
410 } else {
411 (Vec::new(), None)
412 };
413 let available = latest.commit_seq + tail.iter().map(|s| s.frames).sum::<u64>();
414 let latest_commit_seq = latest.commit_seq;
415
416 let target = match at {
417 None => None,
418 Some(RestoreTarget::CommitSeq(n)) => Some(n),
419 Some(RestoreTarget::Timestamp(t)) => {
420 Some(bucket.resolve_time(&latest, t, available).await?)
421 }
422 };
423 let (key, manifest, tail, expect) = match target {
424 None => (latest_target, latest, tail, available),
425 Some(n) if n >= available => {
426 return Err(RestoreS3Error::AfterLatest {
427 target: n,
428 commits: available,
429 not_restorable: tail_note,
430 })
431 }
432 Some(n) if n >= latest.commit_seq => (latest_target, latest, tail, n + 1),
433 Some(n) => {
434 let (key, m) = bucket.manifest_holding(n, &latest_target, &latest).await?;
435 if n < m.oldest_commit_seq {
436 return Err(RestoreS3Error::BeforeOldestArchive {
437 target: n,
438 oldest: m.oldest_commit_seq,
439 });
440 }
441 (key, m, Vec::new(), n + 1)
444 }
445 };
446 let tail_note = if manifest.commit_seq < latest_commit_seq {
447 None
448 } else {
449 tail_note
450 };
451 if let Some(note) = &tail_note {
452 tracing::warn!(note = %note, "the segments after the newest manifest stop early");
453 }
454
455 let staging = {
456 let into = into.to_path_buf();
457 blocking(move || Staging::new(&into, mode)).await?
458 };
459 let mut out = RestoreS3Outcome {
460 into: into.to_path_buf(),
461 manifest: key,
462 commits: expect,
463 source_lineage: manifest.lineage.clone(),
464 lineage: new_instance_id(),
465 files: Vec::new(),
466 segments: 0,
467 bytes: 0,
468 tail_note,
469 install_note: None,
470 swept,
471 };
472 bucket
473 .download(&manifest, &tail, staging.path(), &mut out)
474 .await?;
475
476 let dir = staging.path().to_path_buf();
477 let lineage = out.lineage.clone();
478 let cut = target;
479 let manifest_key = out.manifest.clone();
480 let parent = (mode == Mode::InPlace).then(|| disk::ParentRecord {
481 lineage: out.source_lineage.clone(),
482 manifest: out.manifest.clone(),
483 commit_seq: expect,
484 });
485 blocking(move || -> Result<(), RestoreS3Error> {
486 if let Some(n) = cut {
487 cut_at_commit(&dir, n).map_err(|e| match e {
488 mushroomdb::GraphError::CommitOutOfRange { floor, .. } if n < floor => {
489 RestoreS3Error::BeforeOldestArchive {
490 target: n,
491 oldest: floor,
492 }
493 }
494 mushroomdb::GraphError::CommitOutOfRange { .. } => {
495 RestoreS3Error::Unreachable(format!(
496 "commit {n} cannot be rebuilt from the restore point {manifest_key}: \
497 the history it needs before that restore point's snapshot is not \
498 whole in the backup (a truncating snapshot or archive retention \
499 removed some of it)"
500 ))
501 }
502 other => RestoreS3Error::Io(std::io::Error::other(other.to_string())),
503 })?;
504 }
505 mushroomdb::restore::trim_commit_times(&dir, expect)
509 .map_err(|e| RestoreS3Error::Io(std::io::Error::other(e.to_string())))?;
510 disk::create_lineage(&dir, &lineage)?;
511 if let Some(parent) = &parent {
512 disk::create_parent(&dir, parent)?;
513 }
514 let opened = mushroomdb::GraphDb::open_with_options(
515 &dir,
516 mushroomdb::OpenOptions {
517 auto_migrate: false,
518 repair_wal: false,
519 read_only: true,
520 },
521 )
522 .map_err(|e| RestoreS3Error::Verify(format!("the restored copy does not open: {e}")))?;
523 let commits = opened
524 .wal_total_commits()
525 .map_err(|e| RestoreS3Error::Verify(e.to_string()))?;
526 if commits != expect {
527 return Err(RestoreS3Error::Verify(format!(
528 "the restored copy holds {commits} commits; {manifest_key} says {expect}"
529 )));
530 }
531 Ok(())
532 })
533 .await?;
534 if !out.files.iter().any(|f| f == disk::LINEAGE_FILE) {
535 out.files.push(disk::LINEAGE_FILE.to_string());
536 }
537 if mode == Mode::InPlace {
538 out.files.push(disk::PARENT_FILE.to_string());
539 }
540 let dir = staging.path().to_path_buf();
541 let target_dir = into.to_path_buf();
542 out.install_note = match mode {
543 Mode::Beside => blocking(move || install(&dir, &target_dir)).await?,
544 Mode::InPlace => blocking(move || install_in_place(&dir, &target_dir)).await?,
545 };
546 staging.installed();
547 Ok(out)
548}
549
550const DOWNLOAD_CHUNK: usize = 8 * 1024 * 1024;
555
556struct Bucket<'a> {
557 store: Arc<dyn ObjectStore>,
558 prefix: &'a str,
559}
560
561struct TailSegment {
563 key: String,
564 bytes: Vec<u8>,
565 frames: u64,
566}
567
568impl Bucket<'_> {
569 async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, RestoreS3Error> {
570 match self.store.get(&ObjectPath::from(key)).await {
571 Ok(r) => Ok(Some(r.bytes().await?.to_vec())),
572 Err(object_store::Error::NotFound { .. }) => Ok(None),
573 Err(e) => Err(e.into()),
574 }
575 }
576
577 async fn get_text(&self, key: &str) -> Result<Option<String>, RestoreS3Error> {
578 match self.get(key).await? {
579 None => Ok(None),
580 Some(b) => String::from_utf8(b)
581 .map(|s| Some(s.trim().to_string()))
582 .map_err(|_| RestoreS3Error::Manifest(format!("{key} is not UTF-8"))),
583 }
584 }
585
586 async fn manifest(&self, key: &str) -> Result<Manifest, RestoreS3Error> {
587 let bytes = self.get(key).await?.ok_or_else(|| {
588 RestoreS3Error::Manifest(format!("{key} is named but does not exist"))
589 })?;
590 Manifest::from_json(&bytes).map_err(|e| RestoreS3Error::Manifest(format!("{key}: {e}")))
591 }
592
593 async fn manifest_holding(
598 &self,
599 n: u64,
600 latest_key: &str,
601 latest: &Manifest,
602 ) -> Result<(String, Manifest), RestoreS3Error> {
603 let dir = ObjectPath::from(object_key(self.prefix, "manifest"));
604 let mut keyed: Vec<(u64, String)> = self
605 .store
606 .list_with_delimiter(Some(&dir))
607 .await?
608 .objects
609 .into_iter()
610 .filter_map(|o| {
611 let key = o.location.to_string();
612 let seq = key
613 .rsplit('/')
614 .next()?
615 .strip_suffix(".json")?
616 .parse::<u64>()
617 .ok()?;
618 (seq > n).then_some((seq, key))
619 })
620 .collect();
621 keyed.sort();
622 for (_, key) in keyed {
623 match self.manifest(&key).await {
624 Ok(m) if m.lineage == latest.lineage && m.commit_seq > n => return Ok((key, m)),
625 Ok(_) | Err(RestoreS3Error::Manifest(_)) => continue,
626 Err(e) => return Err(e),
627 }
628 }
629 Ok((latest_key.to_string(), latest.clone()))
630 }
631
632 async fn tail_segments(
635 &self,
636 latest: &Manifest,
637 ) -> Result<(Vec<TailSegment>, Option<String>), RestoreS3Error> {
638 let mut end = match latest.segments.last() {
639 Some(s) => {
640 parse_segment_key(self.prefix, &s.key)
641 .ok_or_else(|| {
642 RestoreS3Error::Verify(format!("{} is not a segment key", s.key))
643 })?
644 .2
645 }
646 None => 0,
647 };
648 let dir = ObjectPath::from(object_key(
649 self.prefix,
650 &format!("wal/{}", latest.wal_generation),
651 ));
652 let mut later: Vec<(u64, u64, String)> = self
653 .store
654 .list_with_delimiter(Some(&dir))
655 .await?
656 .objects
657 .into_iter()
658 .filter_map(|o| {
659 let key = o.location.to_string();
660 let (g, start, stop) = parse_segment_key(self.prefix, &key)?;
661 (g == latest.wal_generation && start >= end).then_some((start, stop, key))
662 })
663 .collect();
664 later.sort();
665 let mut tail = Vec::new();
666 for (i, (start, stop, key)) in later.iter().enumerate() {
667 let why = if *start != end {
668 Some(format!(
669 "the segment that should start at byte {end} is missing \
670 ({key} starts at {start})"
671 ))
672 } else {
673 match self.get(key).await? {
674 None => Some(format!("{key} was listed but is gone")),
675 Some(bytes) => {
676 let (whole, frames) = complete_frames(&bytes);
677 if bytes.len() as u64 != stop - start || whole != bytes.len() {
678 Some(format!(
679 "{key} is damaged: not {} bytes of whole, valid frames",
680 stop - start
681 ))
682 } else {
683 end = *stop;
684 tail.push(TailSegment {
685 key: key.clone(),
686 bytes,
687 frames,
688 });
689 None
690 }
691 }
692 }
693 };
694 if let Some(why) = why {
695 let unrestored = &later[i..];
696 let bytes = unrestored.last().map_or(end, |l| l.1) - end;
697 let (at_least, count) = if *start != end {
699 ("at least ", unrestored.len() + 1)
700 } else {
701 ("", unrestored.len())
702 };
703 return Ok((
704 tail,
705 Some(format!(
706 "{at_least}{count} segment(s) uploaded after the newest manifest, \
707 {bytes} bytes of WAL from byte {end} on, were not restored, so the \
708 commits they hold are not in the restored store: {why}"
709 )),
710 ));
711 }
712 }
713 Ok((tail, None))
714 }
715
716 async fn resolve_time(
718 &self,
719 latest: &Manifest,
720 t: OffsetDateTime,
721 available: u64,
722 ) -> Result<u64, RestoreS3Error> {
723 let entry = latest
724 .files
725 .iter()
726 .find(|f| f.name == "commit_times.bin")
727 .ok_or(RestoreS3Error::CommitTimesAbsent)?;
728 let bytes = self
729 .get(&entry.key)
730 .await?
731 .ok_or_else(|| RestoreS3Error::Verify(format!("{} does not exist", entry.key)))?;
732 if bytes.len() as u64 != entry.size || crc32_hex(&bytes) != entry.crc32 {
733 return Err(RestoreS3Error::Verify(format!(
734 "{} does not match its manifest entry",
735 entry.key
736 )));
737 }
738 let ms = (t.unix_timestamp_nanos().div_euclid(1_000_000)) as i64;
739 let at =
740 mushroomdb::restore::commit_at_instant(&bytes, latest.oldest_commit_seq, ms, available)
741 .map_err(|e| match e {
742 mushroomdb::GraphError::NoRecordedTime => RestoreS3Error::CommitTimesAbsent,
743 mushroomdb::GraphError::TimeBeforeFloor {
744 floor_ms,
745 floor_commit,
746 } => RestoreS3Error::Unreachable(format!(
747 "{} is before the oldest recorded commit time ({}, commit {floor_commit})",
748 rfc3339(t),
749 ms_rfc3339(floor_ms)
750 )),
751 mushroomdb::GraphError::CommitOutOfRange { commit, floor, .. } => {
752 RestoreS3Error::BeforeOldestArchive {
753 target: commit,
754 oldest: floor,
755 }
756 }
757 other => RestoreS3Error::Verify(format!("commit_times.bin: {other}")),
758 })?;
759 if at.commit == at.newest_recorded && at.newest_recorded + 1 < available {
765 return Err(RestoreS3Error::CommitTimesEnd {
766 recorded_through: at.newest_recorded,
767 recorded_through_ms: at.newest_recorded_ms,
768 commits: available,
769 });
770 }
771 Ok(at.commit)
772 }
773
774 async fn download(
776 &self,
777 m: &Manifest,
778 tail: &[TailSegment],
779 dir: &Path,
780 out: &mut RestoreS3Outcome,
781 ) -> Result<(), RestoreS3Error> {
782 for f in &m.files {
783 if !is_store_file_name(&f.name) {
784 return Err(RestoreS3Error::Verify(format!(
785 "{} restores to {:?}, which is not a store file",
786 f.key, f.name
787 )));
788 }
789 if out.files.contains(&f.name) {
790 return Err(RestoreS3Error::Verify(format!(
791 "the manifest names {} twice",
792 f.name
793 )));
794 }
795 let crc = self
796 .download_file(&f.key, f.size, &dir.join(&f.name))
797 .await?;
798 if crc != f.crc32 {
799 return Err(RestoreS3Error::Verify(format!(
800 "{} has CRC-32 {crc}; the manifest says {}",
801 f.key, f.crc32
802 )));
803 }
804 if f.name == disk::LINEAGE_FILE {
805 let named = tokio::fs::read_to_string(dir.join(&f.name)).await?;
806 if named.trim() != m.lineage {
807 return Err(RestoreS3Error::Verify(format!(
808 "{} names lineage {}, its manifest lineage {}",
809 f.key,
810 named.trim(),
811 m.lineage
812 )));
813 }
814 }
815 out.bytes += f.size;
816 out.files.push(f.name.clone());
817 }
818
819 let mut wal = tokio::fs::File::create(dir.join("wal.bin")).await?;
820 let mut offset = 0u64;
821 for s in &m.segments {
822 let (g, start, end) = parse_segment_key(self.prefix, &s.key)
823 .ok_or_else(|| RestoreS3Error::Verify(format!("{} is not a segment key", s.key)))?;
824 if g != m.wal_generation {
825 return Err(RestoreS3Error::Verify(format!(
826 "{} is of generation {g}, the manifest's is {}",
827 s.key, m.wal_generation
828 )));
829 }
830 if start != offset {
831 return Err(RestoreS3Error::Verify(format!(
832 "{} starts at {start}; the segments before it end at {offset}",
833 s.key
834 )));
835 }
836 let bytes = self.get(&s.key).await?.ok_or_else(|| {
837 RestoreS3Error::Verify(format!("{} is named but does not exist", s.key))
838 })?;
839 let (whole, frames) = complete_frames(&bytes);
840 if bytes.len() as u64 != end - start
841 || crc32_hex(&bytes) != s.crc32
842 || whole != bytes.len()
843 || frames != s.frames
844 {
845 return Err(RestoreS3Error::Verify(format!(
846 "{} does not match its manifest entry (size, CRC-32 or frames)",
847 s.key
848 )));
849 }
850 wal.write_all(&bytes).await?;
851 offset = end;
852 out.bytes += bytes.len() as u64;
853 out.segments += 1;
854 }
855 for s in tail {
856 wal.write_all(&s.bytes).await?;
857 out.bytes += s.bytes.len() as u64;
858 out.segments += 1;
859 tracing::debug!(key = %s.key, "segment after the newest manifest");
860 }
861 wal.sync_all().await?;
862 out.files.push("wal.bin".into());
863 Ok(())
864 }
865
866 async fn download_file(
869 &self,
870 key: &str,
871 size: u64,
872 path: &Path,
873 ) -> Result<String, RestoreS3Error> {
874 let location = ObjectPath::from(key);
875 let meta = match self.store.head(&location).await {
876 Ok(m) => m,
877 Err(object_store::Error::NotFound { .. }) => {
878 return Err(RestoreS3Error::Verify(format!(
879 "{key} is named but does not exist"
880 )))
881 }
882 Err(e) => return Err(e.into()),
883 };
884 if meta.size as u64 != size {
885 return Err(RestoreS3Error::Verify(format!(
886 "{key} is {} bytes; the manifest says {size}",
887 meta.size
888 )));
889 }
890 let mut file = tokio::fs::File::create(path).await?;
891 let mut crc = crc32fast::Hasher::new();
892 let mut at = 0usize;
893 let size = size as usize;
894 while at < size {
895 let stop = (at + DOWNLOAD_CHUNK).min(size);
896 let chunk = self.store.get_range(&location, at..stop).await?;
897 if chunk.len() != stop - at {
898 return Err(RestoreS3Error::Verify(format!(
899 "{key} changed while it was read"
900 )));
901 }
902 crc.update(&chunk);
903 file.write_all(&chunk).await?;
904 at = stop;
905 }
906 file.sync_all().await?;
907 Ok(crate::manifest::format_crc32(crc.finalize()))
908 }
909}
910
911fn is_store_file_name(name: &str) -> bool {
914 is_plain_name(name)
915 && !matches!(
916 name,
917 "wal.bin" | "snapshot.intent" | "LOCK" | disk::PARENT_FILE | INSTALL_MARKER
918 )
919}
920
921fn is_plain_name(name: &str) -> bool {
923 !name.is_empty() && name != "." && name != ".." && !name.contains(['/', '\\'])
924}
925
926fn check_into(into: &Path) -> Result<(), RestoreS3Error> {
930 match std::fs::read_dir(into) {
931 Ok(mut entries) => {
932 if entries.next().is_some() {
933 Err(RestoreS3Error::IntoNotEmpty {
934 into: into.to_path_buf(),
935 })
936 } else {
937 Ok(())
938 }
939 }
940 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
941 Err(e) => Err(RestoreS3Error::Io(std::io::Error::new(
942 e.kind(),
943 format!("{}: {e}", into.display()),
944 ))),
945 }
946}
947
948fn check_in_place(dir: &Path) -> Result<(), RestoreS3Error> {
954 let entries = match std::fs::read_dir(dir) {
955 Ok(e) => e,
956 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
957 Err(e) => {
958 return Err(RestoreS3Error::Io(std::io::Error::new(
959 e.kind(),
960 format!("{}: {e}", dir.display()),
961 )))
962 }
963 };
964 let name = name_of(dir);
965 let mut others = Vec::new();
966 for e in entries {
967 let e = e?;
968 let entry = e.file_name().to_string_lossy().into_owned();
969 let empty_file =
970 |n: &str| entry == n && e.metadata().is_ok_and(|m| m.is_file() && m.len() == 0);
971 let allowed = entry == "lost+found"
972 || entry == "LOCK"
973 || empty_file("wal.bin")
974 || staging_stem(&entry, &name).is_some();
975 if !allowed {
976 others.push(entry);
977 }
978 }
979 if others.is_empty() {
980 return Ok(());
981 }
982 others.sort();
983 Err(RestoreS3Error::DataDirNotEmpty {
984 dir: dir.to_path_buf(),
985 entries: others,
986 })
987}
988
989fn parent_of(into: &Path) -> PathBuf {
991 match into.parent() {
992 Some(p) if !p.as_os_str().is_empty() => p.to_path_buf(),
993 _ => PathBuf::from("."),
994 }
995}
996
997fn name_of(into: &Path) -> String {
999 into.file_name()
1000 .map(|n| n.to_string_lossy().into_owned())
1001 .unwrap_or_else(|| "restored".into())
1002}
1003
1004fn staging_stem(entry: &str, name: &str) -> Option<String> {
1009 let rest = entry.strip_prefix(&format!(".{name}.restore-"))?;
1010 let stem = rest.strip_suffix(".lock").unwrap_or(rest);
1011 let (pid, tag) = stem.split_once('-')?;
1012 let ok = !pid.is_empty()
1013 && pid.bytes().all(|b| b.is_ascii_digit())
1014 && tag.len() == 8
1015 && tag.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f'));
1016 ok.then(|| format!(".{name}.restore-{stem}"))
1017}
1018
1019pub const STALE_STAGING_AFTER: std::time::Duration = std::time::Duration::from_secs(10);
1024
1025fn older_than(path: &Path, age: std::time::Duration) -> bool {
1027 std::fs::metadata(path)
1028 .and_then(|m| m.modified())
1029 .ok()
1030 .and_then(|t| t.elapsed().ok())
1031 .is_some_and(|elapsed| elapsed >= age)
1032}
1033
1034fn sweep_stale_staging(root: &Path, name: &str) -> Vec<PathBuf> {
1043 let Ok(entries) = std::fs::read_dir(root) else {
1044 return Vec::new();
1045 };
1046 let mut stems = std::collections::BTreeSet::new();
1047 for e in entries.flatten() {
1048 if let Some(stem) = staging_stem(&e.file_name().to_string_lossy(), name) {
1049 stems.insert(stem);
1050 }
1051 }
1052 let mut removed = Vec::new();
1053 for stem in stems {
1054 let dir = root.join(&stem);
1055 let lock = root.join(format!("{stem}.lock"));
1056 let stale = match std::fs::OpenOptions::new().write(true).open(&lock) {
1057 Ok(f) => match f.try_lock() {
1058 Ok(()) => older_than(&lock, STALE_STAGING_AFTER),
1059 Err(std::fs::TryLockError::WouldBlock) => false,
1060 Err(std::fs::TryLockError::Error(e)) => {
1061 tracing::warn!(
1062 path = %dir.display(),
1063 error = %e,
1064 "cannot tell whether a restore staging directory is in use \
1065 (no advisory locks here); left in place"
1066 );
1067 false
1068 }
1069 },
1070 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
1071 older_than(&dir, STALE_STAGING_AFTER)
1072 }
1073 Err(_) => false,
1074 };
1075 if !stale {
1076 continue;
1077 }
1078 if dir.is_dir() && std::fs::remove_dir_all(&dir).is_ok() {
1079 tracing::info!(path = %dir.display(), "removed a stale restore staging directory");
1080 removed.push(dir);
1081 }
1082 if std::fs::remove_file(&lock).is_ok() {
1083 removed.push(lock);
1084 }
1085 }
1086 removed
1087}
1088
1089pub(crate) fn sweep_in_place(dir: &Path) -> Vec<PathBuf> {
1092 sweep_stale_staging(dir, &name_of(dir))
1093}
1094
1095pub(crate) fn staging_leftovers(dir: &Path) -> Vec<PathBuf> {
1098 let name = name_of(dir);
1099 let Ok(entries) = std::fs::read_dir(dir) else {
1100 return Vec::new();
1101 };
1102 let mut out: Vec<PathBuf> = entries
1103 .flatten()
1104 .filter(|e| staging_stem(&e.file_name().to_string_lossy(), &name).is_some())
1105 .map(|e| e.path())
1106 .collect();
1107 out.sort();
1108 out
1109}
1110
1111struct Staging {
1118 path: PathBuf,
1119 lock_path: PathBuf,
1120 _lock: Option<std::fs::File>,
1121 created: Vec<PathBuf>,
1123}
1124
1125impl Staging {
1126 fn new(into: &Path, mode: Mode) -> Result<Self, RestoreS3Error> {
1127 let root = mode.staging_root(into);
1128 let stem = format!(
1129 ".{}.restore-{}-{}",
1130 name_of(into),
1131 std::process::id(),
1132 &new_instance_id()[..8]
1133 );
1134 let mut staging = Self {
1136 path: root.join(&stem),
1137 lock_path: root.join(format!("{stem}.lock")),
1138 _lock: None,
1139 created: root
1140 .ancestors()
1141 .take_while(|a| !a.as_os_str().is_empty() && !a.exists())
1142 .map(Path::to_path_buf)
1143 .collect(),
1144 };
1145 std::fs::create_dir_all(&root)?;
1146 let lock = std::fs::File::create(&staging.lock_path)?;
1149 let locked = if take_fault(into, Fault::Lock) {
1150 Err(std::io::Error::other("injected lock failure"))
1151 } else {
1152 lock.lock()
1153 };
1154 if let Err(e) = locked {
1155 tracing::warn!(
1160 path = %staging.lock_path.display(),
1161 error = %e,
1162 "cannot lock the restore's lock file; restoring without it"
1163 );
1164 }
1165 staging._lock = Some(lock);
1166 std::fs::create_dir(&staging.path)?;
1167 Ok(staging)
1168 }
1169
1170 fn path(&self) -> &Path {
1171 &self.path
1172 }
1173
1174 fn installed(mut self) {
1177 self.created.clear();
1178 }
1179}
1180
1181impl Drop for Staging {
1182 fn drop(&mut self) {
1183 let _ = std::fs::remove_dir_all(&self.path);
1184 let _ = std::fs::remove_file(&self.lock_path);
1185 for dir in &self.created {
1186 let _ = std::fs::remove_dir(dir);
1188 }
1189 }
1190}
1191
1192fn install(staging: &Path, into: &Path) -> Result<Option<String>, RestoreS3Error> {
1203 remove_if_present(&staging.join("LOCK"))?;
1205 #[cfg(unix)]
1206 std::fs::File::open(staging)?.sync_all()?;
1207 check_into(into)?;
1209 if take_fault(into, Fault::BeforeRename) {
1210 return Err(injected(into));
1211 }
1212 #[cfg(not(unix))]
1213 let replaced_empty = if into.exists() {
1214 std::fs::remove_dir(into)?;
1215 true
1216 } else {
1217 false
1218 };
1219 if let Err(e) = std::fs::rename(staging, into) {
1220 #[allow(unused_mut)]
1221 let mut recreated = String::new();
1222 #[cfg(not(unix))]
1223 if replaced_empty {
1224 recreated = match std::fs::create_dir(into) {
1225 Ok(()) => "; the empty directory it replaced was created again".into(),
1226 Err(_) => "; the empty directory it was to replace was removed and could \
1227 not be created again"
1228 .into(),
1229 };
1230 }
1231 return Err(RestoreS3Error::Io(std::io::Error::new(
1232 e.kind(),
1233 format!(
1234 "moving the restored store into {} (it must be a new or empty directory, \
1235 not a mount point; restore into a directory inside one): {e}{recreated}",
1236 into.display()
1237 ),
1238 )));
1239 }
1240 let parent = parent_of(into);
1243 Ok(sync_dir(into, &parent).err().map(|e| {
1244 let note = format!(
1245 "the restored store is in place in {}, but syncing its parent directory {} \
1246 failed ({e}): its directory entry may not survive a power loss",
1247 into.display(),
1248 parent.display()
1249 );
1250 tracing::warn!(note = %note, "restore installed");
1251 note
1252 }))
1253}
1254
1255pub use mushroomdb::restore::INSTALL_MARKER;
1260
1261fn install_in_place(staging: &Path, dir: &Path) -> Result<Option<String>, RestoreS3Error> {
1275 remove_if_present(&staging.join("LOCK"))?;
1276 #[cfg(unix)]
1277 std::fs::File::open(staging)?.sync_all()?;
1278 check_in_place(dir)?;
1279 let mut names = Vec::new();
1280 for e in std::fs::read_dir(staging)? {
1281 names.push(e?.file_name().to_string_lossy().into_owned());
1282 }
1283 names.sort_by_key(|n| (matches!(n.as_str(), "snapshot.bin" | "wal.bin"), n.clone()));
1285 if take_fault(dir, Fault::BeforeRename) {
1286 return Err(injected(dir));
1287 }
1288 let mut lines = Vec::new();
1289 for name in &names {
1290 let (size, mtime) = disk::ident_of(&std::fs::metadata(staging.join(name))?);
1291 lines.push(format!("{name}\t{size}\t{mtime}"));
1292 }
1293 let marker = dir.join(INSTALL_MARKER);
1294 {
1295 use std::io::Write as _;
1296 let mut f = std::fs::File::create(&marker)?;
1297 f.write_all(lines.join("\n").as_bytes())?;
1298 f.sync_all()?;
1299 }
1300 let mut moved: Vec<&str> = Vec::new();
1301 let result = (|| -> Result<(), RestoreS3Error> {
1302 sync_dir(dir, dir)?;
1303 for name in &names {
1304 if take_fault(dir, Fault::AfterMoving(moved.len())) {
1305 return Err(injected(dir));
1306 }
1307 std::fs::rename(staging.join(name), dir.join(name)).map_err(|e| {
1308 std::io::Error::new(
1309 e.kind(),
1310 format!("moving {name} into {}: {e}", dir.display()),
1311 )
1312 })?;
1313 moved.push(name);
1314 }
1315 sync_dir(dir, dir)?;
1316 Ok(())
1317 })();
1318 if let Err(e) = result {
1319 for name in &moved {
1320 let _ = std::fs::remove_file(dir.join(name));
1321 }
1322 if sync_dir(dir, dir).is_ok() {
1325 let _ = std::fs::remove_file(&marker);
1326 let _ = sync_dir(dir, dir);
1327 }
1328 return Err(e);
1329 }
1330 if !take_fault(dir, Fault::KeepMarker) {
1331 std::fs::remove_file(&marker)?;
1332 }
1333 if let Err(e) = sync_marker_removal(dir) {
1334 return Err(RestoreS3Error::Io(std::io::Error::new(
1335 e.kind(),
1336 format!(
1337 "the restored store is in {}, but syncing the directory after removing the \
1338 install marker failed ({e}); not served, since a power loss could bring the \
1339 marker back. Start again: the next boot serves the store, or restores it \
1340 again if the marker came back",
1341 dir.display()
1342 ),
1343 )));
1344 }
1345 Ok(None)
1346}
1347
1348fn sync_marker_removal(dir: &Path) -> std::io::Result<()> {
1349 if take_fault(dir, Fault::MarkerRemovalSync) {
1350 return Err(std::io::Error::other("injected fsync failure"));
1351 }
1352 sync_dir(dir, dir)
1353}
1354
1355pub fn install_interrupted(dir: &Path) -> bool {
1359 dir.join(INSTALL_MARKER).exists()
1360}
1361
1362pub fn undo_interrupted_install(dir: &Path) -> std::io::Result<Vec<String>> {
1375 let marker = dir.join(INSTALL_MARKER);
1376 let listed = match std::fs::read_to_string(&marker) {
1377 Ok(s) => s,
1378 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
1379 Err(e) => return Err(e),
1380 };
1381 let mut entries: Vec<(&str, Option<disk::Ident>)> = Vec::new();
1382 for line in listed.lines() {
1383 let mut fields = line.split('\t').map(str::trim);
1384 let Some(name) = fields.next().filter(|n| is_plain_name(n)) else {
1385 continue;
1386 };
1387 let size = fields.next().and_then(|v| v.parse::<u64>().ok());
1388 let mtime = fields.next().and_then(|v| v.parse::<u64>().ok());
1389 entries.push((name, size.zip(mtime)));
1390 }
1391 let changed = |why: String| {
1392 std::io::Error::other(format!(
1393 "{} holds the install marker of a restore on boot ({}), but {why} since that \
1394 install: the store may have served writes, so nothing was removed. If the store \
1395 is the one you want, delete {}; to restore again, empty the directory",
1396 dir.display(),
1397 marker.display(),
1398 marker.display()
1399 ))
1400 };
1401 for (name, ident) in &entries {
1402 let Ok(meta) = std::fs::symlink_metadata(dir.join(name)) else {
1403 continue;
1404 };
1405 if let Some(staged) = ident {
1406 if meta.is_file() && meta.len() > 0 && disk::ident_of(&meta) != *staged {
1409 return Err(changed(format!("{name} has changed")));
1410 }
1411 }
1412 }
1413 for written in ["snapshot.bin", "wal.bin"] {
1414 if !entries.iter().any(|(n, _)| *n == written) && dir.join(written).exists() {
1415 return Err(changed(format!("{written} has appeared")));
1416 }
1417 }
1418 let mut removed = Vec::new();
1419 for (name, _) in &entries {
1420 let path = dir.join(name);
1421 if std::fs::symlink_metadata(&path).is_ok_and(|m| m.is_file()) {
1422 std::fs::remove_file(&path)?;
1423 removed.push(name.to_string());
1424 }
1425 }
1426 sync_dir(dir, dir)?;
1427 std::fs::remove_file(&marker)?;
1428 sync_dir(dir, dir)?;
1429 Ok(removed)
1430}
1431
1432fn remove_if_present(path: &Path) -> std::io::Result<()> {
1433 match std::fs::remove_file(path) {
1434 Ok(()) => Ok(()),
1435 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
1436 Err(e) => Err(e),
1437 }
1438}
1439
1440fn sync_dir(target: &Path, dir: &Path) -> std::io::Result<()> {
1443 if take_fault(target, Fault::DirSync) {
1444 return Err(std::io::Error::other("injected fsync failure"));
1445 }
1446 #[cfg(unix)]
1447 std::fs::File::open(dir)?.sync_all()?;
1448 #[cfg(not(unix))]
1449 let _ = dir;
1450 Ok(())
1451}
1452
1453fn injected(into: &Path) -> RestoreS3Error {
1454 RestoreS3Error::Io(std::io::Error::other(format!(
1455 "installing into {}: injected failure",
1456 into.display()
1457 )))
1458}
1459
1460#[doc(hidden)]
1462#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1463pub enum Fault {
1464 BeforeRename,
1466 AfterMoving(usize),
1468 DirSync,
1470 Lock,
1472 MarkerRemovalSync,
1474 KeepMarker,
1477}
1478
1479static FAULTS: std::sync::Mutex<Vec<(PathBuf, Fault)>> = std::sync::Mutex::new(Vec::new());
1481
1482#[doc(hidden)]
1486pub fn fail_next_install_into(into: &Path) {
1487 inject_fault(into, Fault::BeforeRename);
1488}
1489
1490#[doc(hidden)]
1492pub fn inject_fault(into: &Path, fault: Fault) {
1493 FAULTS
1494 .lock()
1495 .unwrap_or_else(|p| p.into_inner())
1496 .push((into.to_path_buf(), fault));
1497}
1498
1499fn take_fault(into: &Path, fault: Fault) -> bool {
1500 let mut faults = FAULTS.lock().unwrap_or_else(|p| p.into_inner());
1501 match faults.iter().position(|(p, f)| p == into && *f == fault) {
1502 Some(i) => {
1503 faults.remove(i);
1504 true
1505 }
1506 None => false,
1507 }
1508}
1509
1510async fn blocking<T: Send + 'static>(
1511 f: impl FnOnce() -> Result<T, RestoreS3Error> + Send + 'static,
1512) -> Result<T, RestoreS3Error> {
1513 tokio::task::spawn_blocking(f)
1514 .await
1515 .map_err(|e| RestoreS3Error::Io(std::io::Error::other(e.to_string())))?
1516}
1517
1518fn ms_rfc3339(ms: i64) -> String {
1519 OffsetDateTime::from_unix_timestamp_nanos(ms as i128 * 1_000_000)
1520 .map(rfc3339)
1521 .unwrap_or_else(|_| format!("{ms} ms"))
1522}
1523
1524fn rfc3339(t: OffsetDateTime) -> String {
1525 t.format(&Rfc3339).unwrap_or_else(|_| t.to_string())
1526}