wip: Move digest fetching into package version filtering

Author: sondrelg

src/client/builder.rs

@@ -13,7 +13,6 @@ use tracing::debug;
 use crate::cli::models::{Account, Token};
 use crate::client::client::PackagesClient;
 use crate::client::urls::Urls;
-use base64::prelude::*;
 pub type RateLimitedService = Arc<Mutex<ConcurrencyLimit<RateLimit<Client>>>>;
 
 #[derive(Debug)]

src/client/client.rs

@@ -6,7 +6,6 @@ use chrono::{DateTime, Utc};
 use color_eyre::eyre::{eyre, Result};
 use reqwest::header::HeaderMap;
 use reqwest::{Client, Method, Request, StatusCode};
-use serde_json::Value;
 use tokio::time::sleep;
 use tower::{Service, ServiceExt};
 use tracing::{debug, error, info, Span};
@@ -481,7 +480,11 @@ impl PackagesClient {
         ))
     }
 
-    pub async fn fetch_image_manifest(&self, tag: &str) -> Result<Vec<String>> {
+    pub async fn fetch_image_manifest(
+        &self,
+        package_name: String,
+        tag: String,
+    ) -> Result<(String, String, Vec<String>)> {
         debug!(tag = tag, "Retrieving image manifest");
 
         let url = format!("https://ghcr.io/v2/snok%2Fcontainer-retention-policy/manifests/{tag}");
@@ -494,16 +497,21 @@ impl PackagesClient {
             Ok(t) => t,
             Err(e) => {
                 println!("{}", raw_json);
-                return Ok(vec![]);
+                return Err(eyre!(
+                    "Failed to fetch image manifest for \x1b[34m{package_name}\x1b[0m:\x1b[32m{tag}\x1b[0m: {e}"
+                ));
             }
         };
 
-        Ok(resp
-            .manifests
-            .unwrap_or(vec![])
-            .iter()
-            .map(|manifest| manifest.digest.to_string())
-            .collect())
+        Ok((
+            package_name,
+            tag,
+            resp.manifests
+                .unwrap_or(vec![])
+                .iter()
+                .map(|manifest| manifest.digest.to_string())
+                .collect(),
+        ))
     }
 }
 

src/core/select_package_versions.rs

@@ -8,7 +8,7 @@ use chrono::Utc;
 use color_eyre::Result;
 use humantime::Duration as HumantimeDuration;
 use indicatif::ProgressStyle;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::task::JoinSet;
@@ -287,16 +287,104 @@ pub async fn select_package_versions(
         );
     }
 
-    let mut package_version_map = HashMap::new();
+    let mut all_package_versions = vec![];
+    let mut fetch_digest_set = JoinSet::new();
 
     debug!("Fetching package versions");
+
     while let Some(r) = set.join_next().await {
         // Get all the package versions for a package
         let (package_name, mut package_versions) = r??;
 
+        // Queue fetching of digests for each tag
+        for package_version in &package_versions.tagged {
+            for tag in &package_version.metadata.container.tags {
+                fetch_digest_set.spawn(client.fetch_image_manifest(package_name.clone(), tag.clone()));
+            }
+        }
+
+        all_package_versions.push((package_name, package_versions));
+    }
+
+    debug!("Fetching package versions");
+    let mut digests = HashSet::new();
+    let mut digest_tag = HashMap::new();
+
+    while let Some(r) = fetch_digest_set.join_next().await {
+        // Get all the digests for the package
+        let (package_name, tag, package_digests) = r??;
+
+        if package_digests.is_empty() {
+            debug!(
+                package_name = package_name,
+                "Found {} associated digests for \x1b[34m{package_name}\x1b[0m:\x1b[32m{tag}\x1b[0m",
+                package_digests.len()
+            );
+        } else {
+            info!(
+                package_name = package_name,
+                "Found {} associated digests for \x1b[34m{package_name}\x1b[0m:\x1b[32m{tag}\x1b[0m",
+                package_digests.len()
+            );
+        }
+
+        digests.extend(package_digests.clone());
+        for digest in package_digests.into_iter() {
+            digest_tag.insert(digest, format!("\x1b[34m{package_name}\x1b[0m:\x1b[32m{tag}\x1b[0m"));
+        }
+    }
+
+    let mut package_version_map = HashMap::new();
+
+    for (package_name, mut package_versions) in all_package_versions {
+        package_versions.untagged = package_versions
+            .untagged
+            .into_iter()
+            .filter_map(|package_version| {
+                if digests.contains(&package_version.name) {
+                    let x: String = package_version.name.clone();
+                    let association: &String = digest_tag.get(&x as &str).unwrap();
+                    debug!(
+                        "Skipping deletion of {} because it's associated with {association}",
+                        package_version.name
+                    );
+                    None
+                } else {
+                    Some(package_version)
+                }
+            })
+            .collect();
+        let count_before = package_versions.tagged.len();
+        package_versions.tagged = package_versions
+            .tagged
+            .into_iter()
+            .filter(|package_version| {
+                if digests.contains(&package_version.name) {
+                    let association = digest_tag.get(&*(package_version.name)).unwrap();
+                    debug!(
+                        "Skipping deletion of {} because it's associated with {association}",
+                        package_version.name
+                    );
+                    false
+                } else {
+                    true
+                }
+            })
+            .collect();
+
+        let adjusted_keep_n_most_recent =
+            if keep_n_most_recent as i64 - (count_before as i64 - package_versions.tagged.len() as i64) < 0 {
+                0
+            } else {
+                keep_n_most_recent as i64 - (count_before as i64 - package_versions.tagged.len() as i64)
+            };
+
         // Keep n package versions per package, if specified
-        package_versions.tagged =
-            handle_keep_n_most_recent(package_versions.tagged, keep_n_most_recent, timestamp_to_use);
+        package_versions.tagged = handle_keep_n_most_recent(
+            package_versions.tagged,
+            adjusted_keep_n_most_recent as u32,
+            timestamp_to_use,
+        );
 
         info!(
             package_name = package_name,

src/main.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use color_eyre::eyre::Result;
 use tokio::sync::RwLock;
-use tracing::{debug, error, info, info_span, trace, Instrument};
+use tracing::{debug, error, info_span, trace, Instrument};
 use tracing_indicatif::IndicatifLayer;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
@@ -143,27 +143,16 @@ async fn main() -> Result<()> {
         *counts.remaining_requests.read().await
     );
 
-    for (package, package_versions) in package_version_map.iter() {
-        info!("Print package {package}");
-        for package_version in &package_versions.tagged {
-            for tag in &package_version.metadata.container.tags {
-                info!("Print tag {tag}");
-                let digests = client.fetch_image_manifest(tag).await.unwrap();
+    let (deleted_packages, failed_packages) =
+        delete_package_versions(package_version_map, client, counts.clone(), input.dry_run)
+            .instrument(info_span!("deleting package versions"))
+            .await;
 
-                println!("digests: {digests:?}");
-            }
-        }
-    }
-    // let (deleted_packages, failed_packages) =
-    //     delete_package_versions(package_version_map, client, counts.clone(), input.dry_run)
-    //         .instrument(info_span!("deleting package versions"))
-    //         .await;
-    //
-    // let mut github_output = env::var("GITHUB_OUTPUT").unwrap_or_default();
-    //
-    // github_output.push_str(&format!("deleted={}", deleted_packages.join(",")));
-    // github_output.push_str(&format!("failed={}", failed_packages.join(",")));
-    // env::set_var("GITHUB_OUTPUT", github_output);
+    let mut github_output = env::var("GITHUB_OUTPUT").unwrap_or_default();
+
+    github_output.push_str(&format!("deleted={}", deleted_packages.join(",")));
+    github_output.push_str(&format!("failed={}", failed_packages.join(",")));
+    env::set_var("GITHUB_OUTPUT", github_output);
 
     Ok(())
 }