src/backup/backup_info.rs: improve prune algorithm

src/backup/backup_info.rs

@@ -3,7 +3,7 @@ use crate::tools;
 use failure::*;
 use regex::Regex;
 use std::os::unix::io::RawFd;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
 use chrono::{DateTime, Datelike, TimeZone, SecondsFormat, Utc, Local};
 
@@ -44,6 +44,8 @@ pub struct BackupGroup {
     backup_id: String,
 }
 
+enum PruneMark { Keep, Remove }
+
 impl BackupGroup {
 
     pub fn new<T: Into<String>, U: Into<String>>(backup_type: T, backup_id: U) -> Self {
@@ -102,20 +104,26 @@ impl BackupGroup {
     }
 
     fn mark_selections<F: Fn(DateTime<Local>, &BackupInfo) -> String> (
-        mark: &mut HashSet<PathBuf>,
+        mark: &mut HashMap<PathBuf, PruneMark>,
         list: &Vec<BackupInfo>,
         keep: usize,
         select_id: F,
     ) {
+
         let mut hash = HashSet::new();
         for info in list {
-            let local_time = info.backup_dir.backup_time().with_timezone(&Local);
-            if hash.len() >= keep as usize { break; }
             let backup_id = info.backup_dir.relative_path();
+            if let Some(_) = mark.get(&backup_id) {
+                continue;
+            }
+            let local_time = info.backup_dir.backup_time().with_timezone(&Local);
             let sel_id: String = select_id(local_time, &info);
             if !hash.contains(&sel_id) {
+                if hash.len() >= keep { break; }
                 hash.insert(sel_id);
-                mark.insert(backup_id);
+                mark.insert(backup_id, PruneMark::Keep);
+            } else {
+                mark.insert(backup_id, PruneMark::Remove);
             }
         }
     }
@@ -129,14 +137,17 @@ impl BackupGroup {
         keep_yearly: Option<u64>,
     ) -> Result<Vec<BackupInfo>, Error> {
 
-        let mut mark = HashSet::new();
+        let mut mark = HashMap::new();
         
         BackupInfo::sort_list(&mut list, false);
         
         if let Some(keep_last) = keep_last {
-            list.iter().take(keep_last as usize).for_each(|info| {
+            for _ in 0..keep_last {
-                mark.insert(info.backup_dir.relative_path());
+                if list.is_empty() { break; }
-            });
+                let info = list.remove(0);
+                let backup_id = info.backup_dir.relative_path();
+                mark.insert(backup_id, PruneMark::Keep);
+            }
         }
 
         if let Some(keep_daily) = keep_daily {
@@ -164,7 +175,14 @@ impl BackupGroup {
         }
 
         let mut remove_list: Vec<BackupInfo> = list.into_iter()
-            .filter(|info| !mark.contains(&info.backup_dir.relative_path())).collect();
+            .filter(|info| {
+                let backup_id = info.backup_dir.relative_path();
+                match mark.get(&backup_id) {
+                    Some(PruneMark::Keep) => false,
+                    _ => true,
+                }
+            })
+            .collect();
 
         BackupInfo::sort_list(&mut remove_list, true);
 

tests/prune.rs

@@ -81,20 +81,16 @@ fn test_prune_simple() -> Result<(), Error> {
     // keep-last, keep-daily mixed
     let list = orig_list.clone();
     let remove_list = get_prune_list(list, Some(2), Some(2), None, None, None);
-    let expect: Vec<PathBuf> = vec![
+    let expect: Vec<PathBuf> = vec![];
-        PathBuf::from("host/elsa/2019-12-02T11:59:15Z"),
-    ];
     assert_eq!(remove_list, expect);
 
     // keep-daily test
-
     let list = orig_list.clone();
     let remove_list = get_prune_list(list, None, Some(3), None, None, None);
     let expect: Vec<PathBuf> = vec![PathBuf::from("host/elsa/2019-12-04T11:59:15Z")];
     assert_eq!(remove_list, expect);
 
     // keep-daily test
-
     let list = orig_list.clone();
     let remove_list = get_prune_list(list, None, Some(2), None, None, None);
     let expect: Vec<PathBuf> = vec![
@@ -106,6 +102,7 @@ fn test_prune_simple() -> Result<(), Error> {
     // keep-weekly
     let list = orig_list.clone();
     let remove_list = get_prune_list(list, None, None, Some(5), None, None);
+    // all backup are within the same week, so we only keep a single file
     let expect: Vec<PathBuf> = vec![
         PathBuf::from("host/elsa/2019-12-02T11:59:15Z"),
         PathBuf::from("host/elsa/2019-12-03T11:59:15Z"),
@@ -113,9 +110,19 @@ fn test_prune_simple() -> Result<(), Error> {
     ];
     assert_eq!(remove_list, expect);
 
+    // keep-daily + keep-weekly
+    let list = orig_list.clone();
+    let remove_list = get_prune_list(list, None, Some(1), Some(5), None, None);
+    let expect: Vec<PathBuf> = vec![
+        PathBuf::from("host/elsa/2019-12-02T11:59:15Z"),
+        PathBuf::from("host/elsa/2019-12-04T11:59:15Z"),
+    ];
+    assert_eq!(remove_list, expect);
+
     // keep-monthly
     let list = orig_list.clone();
     let remove_list = get_prune_list(list, None, None, None, Some(6), None);
+    // all backup are within the same month, so we only keep a single file
     let expect: Vec<PathBuf> = vec![
         PathBuf::from("host/elsa/2019-12-02T11:59:15Z"),
         PathBuf::from("host/elsa/2019-12-03T11:59:15Z"),
@@ -126,6 +133,18 @@ fn test_prune_simple() -> Result<(), Error> {
     // keep-yearly
     let list = orig_list.clone();
     let remove_list = get_prune_list(list, None, None, None, None, Some(7));
+    // all backup are within the same year, so we only keep a single file
+    let expect: Vec<PathBuf> = vec![
+        PathBuf::from("host/elsa/2019-12-02T11:59:15Z"),
+        PathBuf::from("host/elsa/2019-12-03T11:59:15Z"),
+        PathBuf::from("host/elsa/2019-12-04T11:59:15Z"),
+    ];
+    assert_eq!(remove_list, expect);
+
+    // keep-weekly + keep-monthly + keep-yearly
+    let list = orig_list.clone();
+    let remove_list = get_prune_list(list, None, None, Some(5), Some(6), Some(7));
+    // all backup are within one week, so we only keep a single file
     let expect: Vec<PathBuf> = vec![
         PathBuf::from("host/elsa/2019-12-02T11:59:15Z"),
         PathBuf::from("host/elsa/2019-12-03T11:59:15Z"),