From b77390c0cdd087775d50c4c5c94f7d302c7b43cb Mon Sep 17 00:00:00 2001
From: alyssa <alyssa@mailc.net>
Date: Sat, 8 Mar 2025 13:13:14 +0000
Subject: [PATCH] feat: add metric for remaining gateway sessions

---
 Cargo.lock                          |  1 +
 crates/scheduled_tasks/Cargo.toml   |  1 +
 crates/scheduled_tasks/src/main.rs  | 25 +++++++++++++++++++++++++
 crates/scheduled_tasks/src/tasks.rs |  8 +++++++-
 4 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 63502687..ee181b20 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2929,6 +2929,7 @@ dependencies = [
  "sqlx",
  "tokio",
  "tracing",
+ "twilight-http",
 ]
 
 [[package]]
diff --git a/crates/scheduled_tasks/Cargo.toml b/crates/scheduled_tasks/Cargo.toml
index e554edd3..7420e087 100644
--- a/crates/scheduled_tasks/Cargo.toml
+++ b/crates/scheduled_tasks/Cargo.toml
@@ -16,6 +16,7 @@ serde_json = { workspace = true }
 sqlx = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
+twilight-http = { workspace = true }
 
 croner = "2.1.0"
 num-format = "0.4.4"
diff --git a/crates/scheduled_tasks/src/main.rs b/crates/scheduled_tasks/src/main.rs
index fc5e68ed..266325ed 100644
--- a/crates/scheduled_tasks/src/main.rs
+++ b/crates/scheduled_tasks/src/main.rs
@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use chrono::Utc;
 use croner::Cron;
 use fred::prelude::RedisPool;
@@ -14,15 +16,38 @@ pub struct AppCtx {
     pub messages: PgPool,
     pub stats: PgPool,
     pub redis: RedisPool,
+
+    pub discord: Arc<twilight_http::Client>,
 }
 
 libpk::main!("scheduled_tasks");
 async fn real_main() -> anyhow::Result<()> {
+    let mut client_builder = twilight_http::Client::builder().token(
+        libpk::config
+            .discord
+            .as_ref()
+            .expect("missing discord config")
+            .bot_token
+            .clone(),
+    );
+
+    if let Some(base_url) = libpk::config
+        .discord
+        .as_ref()
+        .expect("missing discord config")
+        .api_base_url
+        .clone()
+    {
+        client_builder = client_builder.proxy(base_url, true);
+    }
+
     let ctx = AppCtx {
         data: libpk::db::init_data_db().await?,
         messages: libpk::db::init_messages_db().await?,
         stats: libpk::db::init_stats_db().await?,
         redis: libpk::db::init_redis().await?,
+
+        discord: Arc::new(client_builder.build()),
     };
 
     info!("starting scheduled tasks runner");
diff --git a/crates/scheduled_tasks/src/tasks.rs b/crates/scheduled_tasks/src/tasks.rs
index 13622281..e0247768 100644
--- a/crates/scheduled_tasks/src/tasks.rs
+++ b/crates/scheduled_tasks/src/tasks.rs
@@ -25,7 +25,13 @@ pub async fn update_prometheus(ctx: AppCtx) -> anyhow::Result<()> {
 
     gauge!("pluralkit_image_cleanup_queue_length").set(count.count as f64);
 
-    // todo: remaining shard session_start_limit
+    let gateway = ctx.discord.gateway().authed().await?.model().await?;
+
+    gauge!("pluralkit_gateway_sessions_remaining")
+        .set(gateway.session_start_limit.remaining as f64);
+    gauge!("pluralkit_gateway_sessions_reset_after")
+        .set(gateway.session_start_limit.reset_after as f64);
+
     Ok(())
 }