persistent_kv/snapshot_set/
admin.rs

1use std::fs;
2
3use super::{FileSnapshotSet, SnapshotType};
4
5/// Admin operations for a snapshot set, e.g. pruning snapshots.
6#[allow(dead_code)]
7pub trait SnapshotSetAdmin: Send {
8    /// Prunes backup snapshots, keeping only the latest `max_backups_keep` full snapshots.
9    /// This is useful to limit the number of backup snapshots that are kept around.
10    /// __Warning__: This will delete files from the file system.
11    fn prune_backup_snapshots(&mut self, max_backups_keep: usize) -> Result<(), std::io::Error>;
12
13    /// Prunes snapshots that are not completed, e.g. due to a power failure or process panic.
14    /// This is useful to clean up incomplete snapshots that should not be used.
15    /// __Warning__: This will delete files from the file system.
16    fn prune_not_completed_snapshots(&mut self) -> Result<(), std::io::Error>;
17}
18
19impl SnapshotSetAdmin for FileSnapshotSet {
20    fn prune_backup_snapshots(&mut self, max_backups_keep: usize) -> Result<(), std::io::Error> {
21        let mut full_backup_snapshots = self
22            .snapshots
23            .iter()
24            .filter(|snapshot| snapshot.snapshot_type == SnapshotType::FullCompleted)
25            .cloned()
26            .collect::<Vec<_>>();
27
28        full_backup_snapshots.sort_by_key(|snapshot| snapshot.ordinal);
29        // Skip last item, this is the latest snapshot which we never consider a backup.
30        let full_backup_snapshots = &full_backup_snapshots[..full_backup_snapshots.len() - 1];
31        if max_backups_keep >= full_backup_snapshots.len() {
32            return Ok(());
33        }
34
35        let num_snapshots_to_delete = full_backup_snapshots.len() - max_backups_keep;
36        for snapshot in full_backup_snapshots.iter().take(num_snapshots_to_delete) {
37            println!("Pruning backup snapshot: {:?}", snapshot.shard_paths);
38            for path in snapshot.shard_paths.iter() {
39                fs::remove_file(path)?
40            }
41            self.snapshots.retain(|s| s.ordinal != snapshot.ordinal);
42        }
43        Ok(())
44    }
45
46    fn prune_not_completed_snapshots(&mut self) -> Result<(), std::io::Error> {
47        let not_completed_snapshots = self
48            .snapshots
49            .iter()
50            .filter(|snapshot| snapshot.snapshot_type == SnapshotType::Pending)
51            .cloned()
52            .collect::<Vec<_>>();
53
54        for snapshot in not_completed_snapshots.iter() {
55            println!(
56                "Pruning not completed snapshots: {:?}",
57                snapshot.shard_paths
58            );
59            for path in snapshot.shard_paths.iter() {
60                fs::remove_file(path)?
61            }
62            self.snapshots.retain(|s| s.ordinal != snapshot.ordinal);
63        }
64        Ok(())
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use std::{
71        fs::File,
72        path::{Path, PathBuf},
73    };
74    use tempfile::TempDir;
75
76    use super::*;
77    use crate::snapshot_set::SnapshotOrdinal;
78
79    fn create_temp_dir() -> TempDir {
80        TempDir::new().unwrap()
81    }
82
83    fn create_snapshot_file(folder: &Path, name: &str) -> PathBuf {
84        let path = folder.join(name);
85        File::create(&path).unwrap();
86        path
87    }
88
89    #[test]
90    fn prunes_backup_snapshots() {
91        let tmp_dir = create_temp_dir();
92        create_snapshot_file(tmp_dir.path(), "snapshot_1_0-of-1_full.bin"); // Backup
93        create_snapshot_file(tmp_dir.path(), "snapshot_2_0-of-1_full.bin"); // Backup
94        create_snapshot_file(tmp_dir.path(), "snapshot_3_0-of-1_diff.bin");
95        create_snapshot_file(tmp_dir.path(), "snapshot_4_0-of-1_full.bin"); // Backup
96        create_snapshot_file(tmp_dir.path(), "snapshot_5_0-of-1_full.bin"); // Latest
97
98        let mut snapshot_set = FileSnapshotSet::new(tmp_dir.path()).unwrap();
99        assert_eq!(snapshot_set.snapshots.len(), 5);
100
101        snapshot_set.prune_backup_snapshots(3).unwrap();
102
103        assert_eq!(snapshot_set.snapshots.len(), 5);
104
105        snapshot_set.prune_backup_snapshots(1).unwrap();
106
107        assert_eq!(snapshot_set.snapshots.len(), 3);
108        assert_eq!(snapshot_set.snapshots[0].ordinal, SnapshotOrdinal(3));
109        assert_eq!(snapshot_set.snapshots[1].ordinal, SnapshotOrdinal(4));
110        assert_eq!(snapshot_set.snapshots[2].ordinal, SnapshotOrdinal(5));
111
112        // Construct a new SnapShotSet to verify that the files were actually deleted.
113        drop(snapshot_set);
114        let mut snapshot_set = FileSnapshotSet::new(tmp_dir.path()).unwrap();
115        snapshot_set.prune_backup_snapshots(0).unwrap();
116
117        assert_eq!(snapshot_set.snapshots.len(), 2);
118        assert_eq!(snapshot_set.snapshots[0].ordinal, SnapshotOrdinal(3));
119        assert_eq!(snapshot_set.snapshots[1].ordinal, SnapshotOrdinal(5));
120    }
121
122    #[test]
123    fn prunes_not_completed_snapshots() {
124        let tmp_dir = create_temp_dir();
125        create_snapshot_file(tmp_dir.path(), "snapshot_3_0-of-1_pending.bin");
126        create_snapshot_file(tmp_dir.path(), "snapshot_1_0-of-1_pending.bin"); // Not completed
127        create_snapshot_file(tmp_dir.path(), "snapshot_2_0-of-1_full.bin"); // Not completed
128
129        let mut snapshot_set = FileSnapshotSet::new(tmp_dir.path()).unwrap();
130        assert_eq!(snapshot_set.snapshots.len(), 3);
131
132        snapshot_set.prune_not_completed_snapshots().unwrap();
133
134        assert_eq!(snapshot_set.snapshots.len(), 1);
135        assert_eq!(snapshot_set.snapshots[0].ordinal, SnapshotOrdinal(2));
136
137        // Construct a new SnapShotSet to verify that the files were actually deleted.
138        drop(snapshot_set);
139        let snapshot_set = FileSnapshotSet::new(tmp_dir.path()).unwrap();
140
141        assert_eq!(snapshot_set.snapshots.len(), 1);
142        assert_eq!(snapshot_set.snapshots[0].ordinal, SnapshotOrdinal(2));
143    }
144}
persistent_kv/snapshot_set/admin.rs

persistent_kv/snapshot_set/
admin.rs
