From 5edf51eb449711be03b47f1920ce6d42310d6715 Mon Sep 17 00:00:00 2001
From: LightQuantum <self@lightquantum.me>
Date: Sat, 11 Feb 2023 00:30:57 +0800
Subject: [PATCH] feat(gc): delete partial index

---
 rsync-gc/src/index.rs |  4 ++++
 rsync-gc/src/main.rs  | 17 +++++++++++++++--
 rsync-gc/src/opts.rs  |  3 +++
 rsync-gc/src/plan.rs  | 12 ++++++++++--
 4 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/rsync-gc/src/index.rs b/rsync-gc/src/index.rs
index d3e6e3f..7b7ac89 100644
--- a/rsync-gc/src/index.rs
+++ b/rsync-gc/src/index.rs
@@ -63,6 +63,7 @@ pub async fn commit_gc(
     conn: &mut (impl aio::ConnectionLike + Send),
     namespace: &str,
     index: &[u64],
+    delete_partial: bool,
 ) -> Result<()> {
     let mut pipe = redis::pipe();
     pipe.atomic();
@@ -71,6 +72,9 @@ pub async fn commit_gc(
     for i in index {
         pipe.del(format!("{namespace}:stale:{i}"));
     }
+    if delete_partial {
+        pipe.del(format!("{namespace}:partial"));
+    }
     pipe.del(format!("{namespace}:partial-stale"));
 
     let result: Vec<bool> = pipe.query_async(conn).await?;
diff --git a/rsync-gc/src/main.rs b/rsync-gc/src/main.rs
index cf00e41..7d076ac 100644
--- a/rsync-gc/src/main.rs
+++ b/rsync-gc/src/main.rs
@@ -59,11 +59,24 @@ async fn main() -> Result<()> {
     info!("deleted {} objects", deleted);
 
     info!("deleting stale objects...");
-    let hashes = hashes_to_remove(&mut redis_conn, &redis_opts.namespace, &delete, &keep).await?;
+    let hashes = hashes_to_remove(
+        &mut redis_conn,
+        &redis_opts.namespace,
+        &delete,
+        &keep,
+        opts.delete_partial,
+    )
+    .await?;
     bulk_delete_objs(&s3, &s3_opts, &hashes).await?;
     info!("deleted {} objects", hashes.len());
 
-    commit_gc(&mut redis_conn, &redis_opts.namespace, &delete).await?;
+    commit_gc(
+        &mut redis_conn,
+        &redis_opts.namespace,
+        &delete,
+        opts.delete_partial,
+    )
+    .await?;
 
     Ok(())
 }
diff --git a/rsync-gc/src/opts.rs b/rsync-gc/src/opts.rs
index 890ffa4..52ff609 100644
--- a/rsync-gc/src/opts.rs
+++ b/rsync-gc/src/opts.rs
@@ -12,6 +12,9 @@ pub struct Opts {
     /// How many revisions to keep.
     #[clap(short, long)]
     pub keep: usize,
+    /// Whether to delete partial files.
+    #[clap(long, default_value = "false")]
+    pub delete_partial: bool,
     /// S3 endpoint url.
     /// For specifying authentication, use environment variables:
     /// AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY
diff --git a/rsync-gc/src/plan.rs b/rsync-gc/src/plan.rs
index 221f328..8ea43c2 100644
--- a/rsync-gc/src/plan.rs
+++ b/rsync-gc/src/plan.rs
@@ -2,6 +2,7 @@ use std::collections::HashSet;
 use std::iter;
 
 use eyre::Result;
+use itertools::Either;
 use redis::{aio, AsyncCommands};
 
 use rsync_core::metadata::{MetaExtra, Metadata};
@@ -11,12 +12,19 @@ pub async fn hashes_to_remove(
     namespace: &str,
     stale_indices: &[u64],
     alive_indices: &[u64],
+    delete_partial: bool,
 ) -> Result<HashSet<[u8; 20]>> {
-    // We also need to remove partial-stale index.
     let stale_indices = stale_indices
         .iter()
         .map(|index| format!("{namespace}:stale:{index}"))
-        .chain(iter::once(format!("{namespace}:partial-stale")));
+        // We also need to remove partial-stale index.
+        .chain(iter::once(format!("{namespace}:partial-stale")))
+        // May need to delete partial.
+        .chain(if delete_partial {
+            Either::Left(iter::once(format!("{namespace}:partial")))
+        } else {
+            Either::Right(iter::empty())
+        });
     let alive_indices = alive_indices
         .iter()
         .map(|index| format!("{namespace}:index:{index}"));