From ddf1e21f25646e1b03fe601ad14833b482aa71f1 Mon Sep 17 00:00:00 2001
From: alyssa <alyssa@mailc.net>
Date: Fri, 9 May 2025 12:24:49 +0000
Subject: [PATCH] feat(avatars): check if images are used before deleting

---
 crates/avatars/src/cleanup.rs       |  4 ++-
 crates/scheduled_tasks/src/tasks.rs | 54 ++++++++++++++++++++++++-----
 2 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/crates/avatars/src/cleanup.rs b/crates/avatars/src/cleanup.rs
index 0fb815e3..83327cf9 100644
--- a/crates/avatars/src/cleanup.rs
+++ b/crates/avatars/src/cleanup.rs
@@ -55,9 +55,10 @@ async fn cleanup_job(pool: sqlx::PgPool, bucket: Arc<s3::Bucket>) -> anyhow::Res
     let mut tx = pool.begin().await?;
 
     let image_id: Option<CleanupJobEntry> = sqlx::query_as(
+        // no timestamp checking here
+        // images are only added to the table after 24h
         r#"
                 select id from image_cleanup_jobs
-                where ts < now() - interval '1 day'
                 for update skip locked limit 1;"#,
     )
     .fetch_optional(&mut *tx)
@@ -72,6 +73,7 @@ async fn cleanup_job(pool: sqlx::PgPool, bucket: Arc<s3::Bucket>) -> anyhow::Res
 
     let image_data = libpk::db::repository::avatars::get_by_id(&pool, image_id.clone()).await?;
     if image_data.is_none() {
+        // unsure how this can happen? there is a FK reference
         info!("image {image_id} was already deleted, skipping");
         sqlx::query("delete from image_cleanup_jobs where id = $1")
             .bind(image_id)
diff --git a/crates/scheduled_tasks/src/tasks.rs b/crates/scheduled_tasks/src/tasks.rs
index c7308ed3..64246fc9 100644
--- a/crates/scheduled_tasks/src/tasks.rs
+++ b/crates/scheduled_tasks/src/tasks.rs
@@ -18,11 +18,18 @@ pub async fn update_prometheus(ctx: AppCtx) -> anyhow::Result<()> {
     struct Count {
         count: i64,
     }
+
+    let pending_count: Count = sqlx::query_as("select count(*) from image_cleanup_pending_jobs")
+        .fetch_one(&ctx.data)
+        .await?;
+
     let count: Count = sqlx::query_as("select count(*) from image_cleanup_jobs")
         .fetch_one(&ctx.data)
         .await?;
 
-    gauge!("pluralkit_image_cleanup_queue_length").set(count.count as f64);
+    gauge!("pluralkit_image_cleanup_queue_length", "pending" => "true")
+        .set(pending_count.count as f64);
+    gauge!("pluralkit_image_cleanup_queue_length", "pending" => "false").set(count.count as f64);
 
     let gateway = ctx.discord.gateway().authed().await?.model().await?;
 
@@ -133,14 +140,10 @@ pub async fn update_discord_stats(ctx: AppCtx) -> anyhow::Result<()> {
 }
 
 pub async fn queue_deleted_image_cleanup(ctx: AppCtx) -> anyhow::Result<()> {
-    // todo: we want to delete immediately when system is deleted, but after a
-    // delay if member is deleted
-    ctx.data
-        .execute(
-            r#"
-insert into image_cleanup_jobs
-select id, now() from images where
-        not exists (select from image_cleanup_jobs j where j.id = images.id)
+    // if an image is present on no member, add it to the pending deletion queue
+    // if it is still present on no member after 24h, actually delete it
+
+    let usage_query = r#"
     and not exists (select from systems where avatar_url = images.url)
     and not exists (select from systems where banner_image = images.url)
     and not exists (select from system_guild where avatar_url = images.url)
@@ -152,9 +155,42 @@ select id, now() from images where
 
     and not exists (select from groups where icon = images.url)
     and not exists (select from groups where banner_image = images.url);
+    "#;
+
+    ctx.data
+        .execute(
+            format!(
+                r#"
+            insert into image_cleanup_pending_jobs
+            select id, now() from images where
+            not exists (select from image_cleanup_pending_jobs j where j.id = images.id)
+            and not exists (select from image_cleanup_jobs j where j.id = images.id)
+            {}
         "#,
+                usage_query
+            )
+            .as_str(),
         )
         .await?;
+
+    ctx.data
+        .execute(
+            format!(
+                r#"
+            insert into image_cleanup_jobs
+            select image_cleanup_pending_jobs.id from image_cleanup_pending_jobs
+            left join images on images.id = image_cleanup_pending_jobs.id
+            where
+            ts < now() - '24 hours'::interval
+            and not exists (select from image_cleanup_jobs j where j.id = images.id)
+            {}
+        "#,
+                usage_query
+            )
+            .as_str(),
+        )
+        .await?;
+
     Ok(())
 }