add bundle-uri-cleaner to cleanup old bundles

Summary: We need something to cleanup old bundles. This is the thing that does it. There will be one task in tsp_global for it. Unsharded.

Reviewed By: RajivTS

Differential Revision: D75516918

fbshipit-source-id: a715cc0a51cf526b94881fa4b0fffa8b96d30d4c

Author: mzr

eden/mononoke/facebook/git_bundle_cleaner/BUCK

@@ -0,0 +1,43 @@
+load("@fbcode_macros//build_defs:rust_binary.bzl", "rust_binary")
+load("@fbcode_macros//build_defs/lib:rust_oss.bzl", "rust_oss")
+load("//fbpkg:fbpkg.bzl", "fbpkg")
+load("//tupperware/cm/antlir2/bzl:buck2_container.bzl", "tw")
+
+oncall("scm_server_infra")
+
+rust_binary(
+    name = "git_bundle_cleaner",
+    srcs = glob(["src/**/*.rs"]),
+    named_deps = {
+        "slog": "//common/rust/shed/tracing_slog_compat:tracing_slog_compat",
+    } | ({} if rust_oss.is_oss_build() else {
+    }),
+    test_deps = [],
+    deps = [
+        "fbsource//third-party/rust:anyhow",
+        "fbsource//third-party/rust:clap",
+        "fbsource//third-party/rust:tokio",
+        "//common/rust/shed/facet:facet",
+        "//common/rust/shed/fbinit:fbinit",
+        "//eden/mononoke/cmdlib/mononoke_app:mononoke_app",
+        "//eden/mononoke/git/bundle_uri:bundle_uri",
+        "//eden/mononoke/metaconfig:metaconfig_types",
+        "//eden/mononoke/repo_attributes/repo_identity:repo_identity",
+        "//manifold/clients/rust:manifold_client",
+    ],
+)
+
+tw.container(
+    name = "tw.mononoke.git_bundle_cleaner",
+    features = [
+        tw.build_fbpkg(
+            fbpkg.builder(
+                name = "mononoke.git_bundle_cleaner",
+                path_actions = {
+                    "git_bundle_cleaner": ":git_bundle_cleaner",
+                },
+            ),
+        ),
+    ],
+    oncall = "scm_server_infra",
+)

eden/mononoke/git/bundle_uri/src/lib.rs

@@ -86,6 +86,12 @@ pub trait GitBundleUri: Send + Sync {
     /// There might be None.
     async fn get_latest_bundle_list(&self) -> Result<Option<BundleList>>;
 
+    /// Get all available bundle lists for a repo.
+    async fn get_bundle_lists(&self) -> Result<Vec<BundleList>>;
+
+    /// Remove all bundles in a given bundle list from the metadata db.
+    async fn remove_bundle_list_from_metadata_db(&self, bundle_list_num: u64) -> Result<()>;
+
     async fn get_url_for_bundle_handle(
         &self,
         ctx: &CoreContext,
@@ -196,6 +202,16 @@ impl<U: Clone + Send + GitBundleUrlGenerator + Sync> GitBundleUri for BundleUri<
         self.bundle_metadata_storage.get_latest_bundle_list().await
     }
 
+    async fn get_bundle_lists(&self) -> Result<Vec<BundleList>> {
+        self.bundle_metadata_storage.get_bundle_lists().await
+    }
+
+    async fn remove_bundle_list_from_metadata_db(&self, bundle_list_num: u64) -> Result<()> {
+        self.bundle_metadata_storage
+            .remove_bundle_list(bundle_list_num)
+            .await
+    }
+
     async fn add_new_bundles(&self, bundles: &[Bundle]) -> Result<u64> {
         self.bundle_metadata_storage.add_new_bundles(bundles).await
     }

eden/mononoke/git/bundle_uri/src/sql.rs

@@ -33,6 +33,16 @@ mononoke_queries! {
     "
     }
 
+    read GetBundleListsForRepo(repo_id: RepositoryId) -> (
+        String, u64, u64, String, u64
+    ) {
+    "SELECT bundle_handle, in_bundle_list_order, bundle_list, bundle_fingerprint, generation_start_timestamp
+    FROM git_bundles
+    WHERE repo_id = {repo_id}
+    ORDER BY bundle_list DESC, in_bundle_list_order ASC;
+    "
+    }
+
     read GetLatestBundleListNumForRepo(repo_id: RepositoryId) -> (
         u64
     ) {
@@ -65,6 +75,11 @@ mononoke_queries! {
         none,
     "INSERT INTO git_bundles (repo_id, bundle_handle, bundle_list, in_bundle_list_order, bundle_fingerprint, generation_start_timestamp) VALUES {values}"
     }
+
+    write RemoveBundleList(repo_id: RepositoryId, bundle_list: u64) {
+        none,
+    "DELETE FROM git_bundles WHERE repo_id = {repo_id} and bundle_list = {bundle_list}"
+    }
 }
 
 #[derive(Clone)]
@@ -173,6 +188,81 @@ impl SqlGitBundleMetadataStorage {
 
         Ok(Some(bundle_list))
     }
+
+    pub async fn remove_bundle_list(&self, bundle_list_num: u64) -> Result<()> {
+        RemoveBundleList::query(
+            &self.connections.write_connection,
+            &self.repo_id,
+            &bundle_list_num,
+        )
+        .await?;
+        Ok(())
+    }
+
+    pub async fn get_bundle_lists(&self) -> Result<Vec<BundleList>> {
+        let rows =
+            GetBundleListsForRepo::query(&self.connections.read_connection, &self.repo_id).await?;
+
+        // +----------------------+-------------+
+        // | in_bundle_list_order | bundle_list |
+        // +----------------------+-------------+
+        // |                    1 |           3 |
+        // |                    1 |           2 |
+        // |                    2 |           2 |
+        // |                    3 |           2 |
+        // |                    1 |           1 |
+        // +----------------------+-------------+
+        let mut bundle_lists = vec![];
+        let mut rows_iter = rows.into_iter().peekable();
+        while let Some((
+            handle,
+            in_bundle_list_order,
+            first_seen_bundle_list,
+            fingerprint,
+            generation_start_timestamp,
+        )) = rows_iter.next()
+        {
+            let current_bundle_list_num = first_seen_bundle_list;
+            // First bundle for a new bundle-list.
+            let mut bundles = vec![Bundle {
+                in_bundle_list_order,
+                handle,
+                fingerprint,
+                generation_start_timestamp,
+            }];
+
+            // Rest of the bundles for the new bundle-list.
+            while let Some((
+                handle,
+                in_bundle_list_order,
+                bundle_list,
+                fingerprint,
+                generation_start_timestamp,
+            )) = rows_iter.peek()
+            {
+                if current_bundle_list_num == *bundle_list {
+                    // This bundle belongs to the current bundle-list.
+                    bundles.push(Bundle {
+                        in_bundle_list_order: *in_bundle_list_order,
+                        handle: handle.clone(),
+                        fingerprint: fingerprint.clone(),
+                        generation_start_timestamp: generation_start_timestamp.clone(),
+                    });
+                    rows_iter.next();
+                } else {
+                    // This bundle is the first elem of the next bundle-list. Do not consume it.
+                    // Break the loop to finish processing current bundle.
+                    break;
+                }
+            }
+            bundle_lists.push(BundleList {
+                bundle_list_num: current_bundle_list_num,
+                bundles,
+            });
+        }
+
+        Ok(bundle_lists)
+    }
 }
 
 #[cfg(test)]
@@ -277,4 +367,49 @@ mod test {
 
         Ok(())
     }
+
+    #[mononoke::fbinit_test]
+    async fn test_get_bundle_lists(_: FacebookInit) -> Result<()> {
+        let repo_id = RepositoryId::new(2137);
+        let storage = SqlGitBundleMetadataStorageBuilder::with_sqlite_in_memory()?.build(repo_id);
+
+        storage.add_new_bundles(&TEST_BUNDLE_LIST_2[..]).await?;
+        storage.add_new_bundles(&TEST_BUNDLE_LIST_3[..]).await?;
+
+        let bundle_lists = storage.get_bundle_lists().await?;
+        assert_eq!(bundle_lists.len(), 2);
+        for bundle_list in bundle_lists.iter() {
+            for (p, n) in bundle_list.bundles.iter().tuple_windows() {
+                assert!(p.in_bundle_list_order < n.in_bundle_list_order)
+            }
+        }
+        for (p, n) in bundle_lists.iter().tuple_windows() {
+            assert!(p.bundle_list_num > n.bundle_list_num)
+        }
+
+        Ok(())
+    }
+
+    #[mononoke::fbinit_test]
+    async fn test_remove_bundle_list(_: FacebookInit) -> Result<()> {
+        let repo_id = RepositoryId::new(2137);
+        let storage = SqlGitBundleMetadataStorageBuilder::with_sqlite_in_memory()?.build(repo_id);
+
+        let bundle_list_num_2 = storage.add_new_bundles(&TEST_BUNDLE_LIST_2[..]).await?;
+        let bundle_list_num_3 = storage.add_new_bundles(&TEST_BUNDLE_LIST_3[..]).await?;
+
+        let bundle_lists = storage.get_bundle_lists().await?;
+        assert_eq!(bundle_lists.len(), 2);
+        assert_eq!(bundle_lists[0].bundle_list_num, bundle_list_num_3);
+        assert_eq!(bundle_lists[1].bundle_list_num, bundle_list_num_2);
+        storage.remove_bundle_list(bundle_list_num_2).await?;
+        let bundle_lists = storage.get_bundle_lists().await?;
+        assert_eq!(bundle_lists.len(), 1);
+        assert_eq!(bundle_lists[0].bundle_list_num, bundle_list_num_3);
+        storage.remove_bundle_list(bundle_list_num_3).await?;
+        let bundle_lists = storage.get_bundle_lists().await?;
+        assert_eq!(bundle_lists.len(), 0);
+
+        Ok(())
+    }
 }