From 73c444c31ddc50ac9ca868c3c7ba87295f27853d Mon Sep 17 00:00:00 2001 From: alyssa Date: Sat, 26 Oct 2024 03:30:58 +0900 Subject: [PATCH] feat: avatar cleanup service --- ci/Dockerfile.rust | 3 +- ci/rust-docker-target.sh | 2 +- go.work.sum | 1 + lib/libpk/src/_config.rs | 5 + lib/libpk/src/db/repository/avatars.rs | 7 ++ services/avatars/Cargo.toml | 4 + services/avatars/src/cleanup.rs | 146 +++++++++++++++++++++++++ services/avatars/src/init.sql | 4 +- services/scheduled_tasks/go.mod | 11 +- services/scheduled_tasks/go.sum | 20 +++- services/scheduled_tasks/main.go | 29 ++++- services/scheduled_tasks/repo.go | 9 ++ services/scheduled_tasks/tasks.go | 30 +++++ 13 files changed, 256 insertions(+), 15 deletions(-) create mode 100644 services/avatars/src/cleanup.rs diff --git a/ci/Dockerfile.rust b/ci/Dockerfile.rust index f019273e..9a620e22 100644 --- a/ci/Dockerfile.rust +++ b/ci/Dockerfile.rust @@ -33,10 +33,11 @@ COPY services/avatars/ /build/services/avatars RUN cargo build --bin api --release --target x86_64-unknown-linux-musl RUN cargo build --bin gateway --release --target x86_64-unknown-linux-musl RUN cargo build --bin avatars --release --target x86_64-unknown-linux-musl +RUN cargo build --bin avatar_cleanup --release --target x86_64-unknown-linux-musl FROM scratch COPY --from=binary-builder /build/target/x86_64-unknown-linux-musl/release/api /api COPY --from=binary-builder /build/target/x86_64-unknown-linux-musl/release/gateway /gateway COPY --from=binary-builder /build/target/x86_64-unknown-linux-musl/release/avatars /avatars -COPY --from=binary-builder /build/target/x86_64-unknown-linux-musl/release/avatars /avatars +COPY --from=binary-builder /build/target/x86_64-unknown-linux-musl/release/avatar_cleanup /avatar_cleanup diff --git a/ci/rust-docker-target.sh b/ci/rust-docker-target.sh index 0bd8f038..19748c4c 100755 --- a/ci/rust-docker-target.sh +++ b/ci/rust-docker-target.sh @@ -39,4 +39,4 @@ EOF # add rust binaries here to build build api build gateway -build avatars +build avatars "COPY .docker-bin/avatar_cleanup /bin/avatar_cleanup" diff --git a/go.work.sum b/go.work.sum index c0f58b37..20bb40cf 100644 --- a/go.work.sum +++ b/go.work.sum @@ -1,2 +1,3 @@ github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= +golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4= google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w= diff --git a/lib/libpk/src/_config.rs b/lib/libpk/src/_config.rs index 298e603c..d112eacd 100644 --- a/lib/libpk/src/_config.rs +++ b/lib/libpk/src/_config.rs @@ -61,6 +61,11 @@ pub struct AvatarsConfig { #[serde(default)] pub migrate_worker_count: u32, + + #[serde(default)] + pub cloudflare_zone_id: Option, + #[serde(default)] + pub cloudflare_token: Option, } #[derive(Deserialize, Clone, Debug)] diff --git a/lib/libpk/src/db/repository/avatars.rs b/lib/libpk/src/db/repository/avatars.rs index cb77ee8c..6732ac67 100644 --- a/lib/libpk/src/db/repository/avatars.rs +++ b/lib/libpk/src/db/repository/avatars.rs @@ -2,6 +2,13 @@ use sqlx::{PgPool, Postgres, Transaction}; use crate::db::types::avatars::*; +pub async fn get_by_id(pool: &PgPool, id: String) -> anyhow::Result> { + Ok(sqlx::query_as("select * from images where id = $1") + .bind(id) + .fetch_optional(pool) + .await?) +} + pub async fn get_by_original_url( pool: &PgPool, original_url: &str, diff --git a/services/avatars/Cargo.toml b/services/avatars/Cargo.toml index bb664f6e..b4f9e22b 100644 --- a/services/avatars/Cargo.toml +++ b/services/avatars/Cargo.toml @@ -3,6 +3,10 @@ name = "avatars" version = "0.1.0" edition = "2021" +[[bin]] +name = "avatar_cleanup" +path = "src/cleanup.rs" + [dependencies] libpk = { path = "../../lib/libpk" } anyhow = { workspace = true } diff --git a/services/avatars/src/cleanup.rs b/services/avatars/src/cleanup.rs new file mode 100644 index 00000000..eaa29527 --- /dev/null +++ b/services/avatars/src/cleanup.rs @@ -0,0 +1,146 @@ +use anyhow::Context; +use reqwest::{ClientBuilder, StatusCode}; +use sqlx::prelude::FromRow; +use std::{sync::Arc, time::Duration}; +use tracing::{error, info}; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + libpk::init_logging("avatar_cleanup")?; + libpk::init_metrics()?; + info!("hello world"); + + let config = libpk::config + .avatars + .as_ref() + .expect("missing avatar service config"); + + let bucket = { + let region = s3::Region::Custom { + region: "s3".to_string(), + endpoint: config.s3.endpoint.to_string(), + }; + + let credentials = s3::creds::Credentials::new( + Some(&config.s3.application_id), + Some(&config.s3.application_key), + None, + None, + None, + ) + .unwrap(); + + let bucket = s3::Bucket::new(&config.s3.bucket, region, credentials)?; + + Arc::new(bucket) + }; + + let pool = libpk::db::init_data_db().await?; + + loop { + // no infinite loops + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + match cleanup_job(pool.clone(), bucket.clone()).await { + Ok(()) => {} + Err(err) => { + error!("failed to run avatar cleanup job: {}", err); + // sentry + } + } + } +} + +#[derive(FromRow)] +struct CleanupJobEntry { + id: String, +} + +async fn cleanup_job(pool: sqlx::PgPool, bucket: Arc) -> anyhow::Result<()> { + let mut tx = pool.begin().await?; + + let image_id: Option = + sqlx::query_as("select id from image_cleanup_jobs for update skip locked limit 1;") + .fetch_optional(&mut *tx) + .await?; + if image_id.is_none() { + info!("no job to run, sleeping for 1 minute"); + tokio::time::sleep(tokio::time::Duration::from_secs(60)).await; + return Ok(()); + } + let image_id = image_id.unwrap().id; + info!("got image {image_id}, cleaning up..."); + + let image_data = libpk::db::repository::avatars::get_by_id(&pool, image_id.clone()).await?; + if image_data.is_none() { + info!("image {image_id} was already deleted, skipping"); + sqlx::query("delete from image_cleanup_jobs where id = $1") + .bind(image_id) + .execute(&mut *tx) + .await?; + return Ok(()); + } + let image_data = image_data.unwrap(); + + let config = libpk::config + .avatars + .as_ref() + .expect("missing avatar service config"); + + let path = image_data + .url + .strip_prefix(config.cdn_url.as_str()) + .unwrap(); + + let s3_resp = bucket.delete_object(path).await?; + match s3_resp.status_code() { + 204 => { + info!("successfully deleted image {image_id} from s3"); + } + _ => { + anyhow::bail!("s3 returned bad error code {}", s3_resp.status_code()); + } + } + + if let Some(zone_id) = config.cloudflare_zone_id.as_ref() { + let client = ClientBuilder::new() + .connect_timeout(Duration::from_secs(3)) + .timeout(Duration::from_secs(3)) + .build() + .context("error making client")?; + + let cf_resp = client + .post(format!( + "https://api.cloudflare.com/client/v4/zones/{zone_id}/purge_cache" + )) + .header( + "Authorization", + format!("Bearer {}", config.cloudflare_token.as_ref().unwrap()), + ) + .body(format!(r#"{{"files":["{}"]}}"#, image_data.url)) + .send() + .await?; + + match cf_resp.status() { + StatusCode::OK => { + info!( + "successfully purged url {} from cloudflare cache", + image_data.url + ); + } + _ => { + let status = cf_resp.status(); + println!("{:#?}", cf_resp.text().await?); + anyhow::bail!("cloudflare returned bad error code {}", status); + } + } + } + + sqlx::query("delete from images where id = $1") + .bind(image_id.clone()) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + Ok(()) +} diff --git a/services/avatars/src/init.sql b/services/avatars/src/init.sql index 854c065b..09e7ed54 100644 --- a/services/avatars/src/init.sql +++ b/services/avatars/src/init.sql @@ -21,4 +21,6 @@ create index if not exists images_uploaded_by_account_idx on images (uploaded_by create table if not exists image_queue (itemid serial primary key, url text not null, kind text not null); alter table images add column if not exists uploaded_by_system uuid; -alter table images add column if not exists content_type text default 'image/webp'; \ No newline at end of file +alter table images add column if not exists content_type text default 'image/webp'; + +create table image_cleanup_jobs(id text references images(id) on delete cascade); diff --git a/services/scheduled_tasks/go.mod b/services/scheduled_tasks/go.mod index be1ba87e..7cea5024 100644 --- a/services/scheduled_tasks/go.mod +++ b/services/scheduled_tasks/go.mod @@ -6,13 +6,14 @@ require ( github.com/getsentry/sentry-go v0.15.0 github.com/go-redis/redis/v8 v8.11.5 github.com/jackc/pgx/v4 v4.16.1 + github.com/prometheus/client_golang v1.20.5 golang.org/x/text v0.16.0 ) require ( + github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect - github.com/google/go-cmp v0.6.0 // indirect github.com/jackc/chunkreader/v2 v2.0.1 // indirect github.com/jackc/pgconn v1.12.1 // indirect github.com/jackc/pgio v1.0.0 // indirect @@ -21,8 +22,12 @@ require ( github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b // indirect github.com/jackc/pgtype v1.11.0 // indirect github.com/jackc/puddle v1.2.1 // indirect - github.com/stretchr/testify v1.9.0 // indirect + github.com/klauspost/compress v1.17.9 // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/prometheus/client_model v0.6.1 // indirect + github.com/prometheus/common v0.55.0 // indirect + github.com/prometheus/procfs v0.15.1 // indirect golang.org/x/crypto v0.24.0 // indirect - golang.org/x/net v0.26.0 // indirect golang.org/x/sys v0.22.0 // indirect + google.golang.org/protobuf v1.34.2 // indirect ) diff --git a/services/scheduled_tasks/go.sum b/services/scheduled_tasks/go.sum index 0bf4962f..e4e20bdf 100644 --- a/services/scheduled_tasks/go.sum +++ b/services/scheduled_tasks/go.sum @@ -1,5 +1,7 @@ github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/Masterminds/semver/v3 v3.1.1/go.mod h1:VPu/7SZ7ePZ3QOrcuXROw5FAcLl4a0cBrbBpGY/8hQs= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I= @@ -24,7 +26,6 @@ github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/me github.com/gofrs/uuid v4.0.0+incompatible h1:1SD/1F5pU8p29ybwgQSwpQk+mwdRrXCYuPhW6m+TnJw= github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= -github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI= github.com/jackc/chunkreader v1.0.0/go.mod h1:RT6O25fNZIuasFJRyZ4R/Y2BbhasbmZXF9QQ7T3kePo= github.com/jackc/chunkreader/v2 v2.0.0/go.mod h1:odVSm741yZoC3dpHEUXIqA9tQRhFrgOHwnPIn9lDKlk= @@ -75,12 +76,15 @@ github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dv github.com/jackc/puddle v1.2.1 h1:gI8os0wpRXFd4FiAY2dWiqRK037tjj3t7rKFeO4X5iw= github.com/jackc/puddle v1.2.1/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= +github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA= +github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/konsorten/go-windows-terminal-sequences v1.0.2/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/pty v1.1.8/go.mod h1:O1sed60cT9XZ5uDucP5qwvh+TE3NnUj51EiZO/lmSfw= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/lib/pq v1.0.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo= github.com/lib/pq v1.1.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo= github.com/lib/pq v1.2.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo= @@ -91,6 +95,8 @@ github.com/mattn/go-colorable v0.1.6/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope github.com/mattn/go-isatty v0.0.5/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s= github.com/mattn/go-isatty v0.0.7/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s= github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE= github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= github.com/onsi/gomega v1.18.1 h1:M1GfJqGRrBrrGGsbxzV5dqM2U2ApXefZCQpkukxYRLE= @@ -99,6 +105,14 @@ github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus/client_golang v1.20.5 h1:cxppBPuYhUnsO6yo/aoRol4L7q7UFfdm+bR9r+8l63Y= +github.com/prometheus/client_golang v1.20.5/go.mod h1:PIEt8X02hGcP8JWbeHyeZ53Y/jReSnHgO035n//V5WE= +github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p8ais2e9E= +github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY= +github.com/prometheus/common v0.55.0 h1:KEi6DK7lXW/m7Ig5i47x0vRzuBsHuvJdi5ee6Y3G1dc= +github.com/prometheus/common v0.55.0/go.mod h1:2SECS4xJG1kd8XF9IcM1gMX6510RAEL65zxzNImwdc8= +github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc= +github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk= github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4= github.com/rs/xid v1.2.1/go.mod h1:+uKXf+4Djp6Md1KODXJxgGQPKngRmWyn10oCKFzNHOQ= github.com/rs/zerolog v1.13.0/go.mod h1:YbFCdg8HfsridGWAh22vktObvhZbQsZXe4/zB0OKkWU= @@ -118,7 +132,6 @@ github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81P github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= -github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/zenazn/goji v0.9.0/go.mod h1:7S9M489iMyHBNxwZnk9/EHS098H4/F6TATF2mIxtB1Q= go.uber.org/atomic v1.3.2/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE= go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE= @@ -151,7 +164,6 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL golang.org/x/net v0.0.0-20190813141303-74dc4d7220e7/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ= -golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -190,6 +202,8 @@ golang.org/x/xerrors v0.0.0-20190513163551-3ee3066db522/go.mod h1:I/5z698sn9Ka8T golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= +google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= diff --git a/services/scheduled_tasks/main.go b/services/scheduled_tasks/main.go index 240bbc6a..5e337c62 100644 --- a/services/scheduled_tasks/main.go +++ b/services/scheduled_tasks/main.go @@ -7,12 +7,24 @@ import ( "runtime/debug" "strings" "time" + "net/http" "github.com/getsentry/sentry-go" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + "github.com/prometheus/client_golang/prometheus/promhttp" ) var set_guild_count = false +var ( + cleanupQueueLength = promauto.NewGauge(prometheus.GaugeOpts{ + Name: "pluralkit_image_cleanup_queue_length", + Help: "Remaining image cleanup jobs", + }) +) + func main() { if _, ok := os.LookupEnv("SET_GUILD_COUNT"); ok { set_guild_count = true @@ -29,14 +41,19 @@ func main() { connect_dbs() log.Println("starting scheduled tasks runner") - wait_until_next_minute() + go func() { + wait_until_next_minute() - go doforever(time.Minute, withtime("stats updater", update_db_meta)) - go doforever(time.Minute*10, withtime("message stats updater", update_db_message_meta)) - go doforever(time.Minute, withtime("discord stats updater", update_discord_stats)) + go doforever(time.Minute, withtime("stats updater", update_db_meta)) + go doforever(time.Minute*10, withtime("message stats updater", update_db_message_meta)) + go doforever(time.Minute, withtime("discord stats updater", update_discord_stats)) + go doforever(time.Minute*30, withtime("queue deleted image cleanup job", queue_deleted_image_cleanup)) + }() - // block main thread - select{} + go doforever(time.Second * 10, withtime("prometheus updater", update_prom)) + log.Println("listening for prometheus on :9000") + http.Handle("/metrics", promhttp.Handler()) + http.ListenAndServe(":9000", nil) } func wait_until_next_minute() { diff --git a/services/scheduled_tasks/repo.go b/services/scheduled_tasks/repo.go index 54e11910..ece7e0dd 100644 --- a/services/scheduled_tasks/repo.go +++ b/services/scheduled_tasks/repo.go @@ -131,6 +131,15 @@ func get_message_count() int { return count } +func get_image_cleanup_queue_length() int { + var count int + row := data_db.QueryRow(context.Background(), "select count(*) as count from image_cleanup_jobs") + if err := row.Scan(&count); err != nil { + panic(err) + } + return count +} + func run_data_stats_query() map[string]interface{} { s := map[string]interface{}{} diff --git a/services/scheduled_tasks/tasks.go b/services/scheduled_tasks/tasks.go index a480e7df..a1404682 100644 --- a/services/scheduled_tasks/tasks.go +++ b/services/scheduled_tasks/tasks.go @@ -18,6 +18,11 @@ func plural(key string) string { return key + "s" } +func update_prom() { + count := get_image_cleanup_queue_length() + cleanupQueueLength.Set(float64(count)) +} + func update_db_meta() { for _, key := range table_stat_keys { q := fmt.Sprintf("update info set %s_count = (select count(*) from %s)", key, plural(key)) @@ -71,3 +76,28 @@ func update_discord_stats() { panic(err) } } + +// MUST add new image columns here +var deletedImageCleanupQuery = ` +insert into image_cleanup_jobs +select id from images where + not exists (select from image_cleanup_jobs j where j.id = images.id) + 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) + + and not exists (select from members where avatar_url = images.url) + and not exists (select from members where banner_image = images.url) + and not exists (select from members where webhook_avatar_url = images.url) + and not exists (select from member_guild where avatar_url = images.url) + + and not exists (select from groups where icon = images.url) + and not exists (select from groups where banner_image = images.url); + ` + +func queue_deleted_image_cleanup() { + _, err := data_db.Exec(context.Background(), deletedImageCleanupQuery) + if err != nil { + panic(err) + } +}