feat(stats): query http gateway, wait until gateway up to collect discord stats

services/gateway/src/cache_api.rs

@@ -5,12 +5,15 @@ use axum::{
     routing::get,
     Router,
 };
-use serde_json::to_string;
+use serde_json::{json, to_string};
 use tracing::{error, info};
 use twilight_model::guild::Permissions;
 use twilight_model::id::Id;
 
-use crate::discord::cache::{dm_channel, DiscordCache, DM_PERMISSIONS};
+use crate::discord::{
+    cache::{dm_channel, DiscordCache, DM_PERMISSIONS},
+    gateway::cluster_config,
+};
 use std::sync::Arc;
 
 fn status_code(code: StatusCode, body: String) -> Response {
@@ -156,6 +159,17 @@ pub async fn run_server(cache: Arc<DiscordCache>) -> anyhow::Result<()> {
             })
         )
 
+        .route("/stats", get(|State(cache): State<Arc<DiscordCache>>| async move {
+            let cluster = cluster_config();
+            let has_been_up = cache.2.read().await.len() as u32 == if cluster.total_shards > 16 {16} else {cluster.total_shards};
+            let stats = json!({
+                "guild_count": cache.0.stats().guilds(),
+                "channel_count": cache.0.stats().channels(),
+                "up": has_been_up,
+            });
+            status_code(StatusCode::FOUND, to_string(&stats).unwrap())
+        }))
+
         .layer(axum::middleware::from_fn(crate::logger::logger))
         .with_state(cache);
 

services/gateway/src/discord/cache.rs

@@ -1,6 +1,7 @@
 use anyhow::format_err;
 use lazy_static::lazy_static;
 use std::sync::Arc;
+use tokio::sync::RwLock;
 use twilight_cache_inmemory::{
     model::CachedMember,
     permission::{MemberRoles, RootError},
@@ -110,10 +111,14 @@ pub fn new() -> DiscordCache {
             .build(),
     );
 
-    DiscordCache(cache, client)
+    DiscordCache(cache, client, RwLock::new(Vec::new()))
 }
 
-pub struct DiscordCache(pub Arc<InMemoryCache>, pub Arc<twilight_http::Client>);
+pub struct DiscordCache(
+    pub Arc<InMemoryCache>,
+    pub Arc<twilight_http::Client>,
+    pub RwLock<Vec<u32>>,
+);
 
 impl DiscordCache {
     pub async fn guild_permissions(

services/gateway/src/discord/gateway.rs

@@ -1,3 +1,4 @@
+use libpk::_config::ClusterSettings;
 use std::sync::{mpsc::Sender, Arc};
 use tracing::{info, warn};
 use twilight_gateway::{
@@ -13,6 +14,18 @@ use crate::discord::identify_queue::{self, RedisQueue};
 
 use super::{cache::DiscordCache, shard_state::ShardStateManager};
 
+pub fn cluster_config() -> ClusterSettings {
+    libpk::config
+        .discord
+        .cluster
+        .clone()
+        .unwrap_or(libpk::_config::ClusterSettings {
+            node_id: 0,
+            total_shards: 1,
+            total_nodes: 1,
+        })
+}
+
 pub fn create_shards(redis: fred::pool::RedisPool) -> anyhow::Result<Vec<Shard<RedisQueue>>> {
     let intents = Intents::GUILDS
         | Intents::DIRECT_MESSAGES
@@ -23,16 +36,7 @@ pub fn create_shards(redis: fred::pool::RedisPool) -> anyhow::Result<Vec<Shard<R
 
     let queue = identify_queue::new(redis);
 
-    let cluster_settings =
-        libpk::config
-            .discord
-            .cluster
-            .clone()
-            .unwrap_or(libpk::_config::ClusterSettings {
-                node_id: 0,
-                total_shards: 1,
-                total_nodes: 1,
-            });
+    let cluster_settings = cluster_config();
 
     let (start_shard, end_shard): (u32, u32) = if cluster_settings.total_shards < 16 {
         warn!("we have less than 16 shards, assuming single gateway process");
@@ -77,6 +81,11 @@ pub async fn runner(
                 {
                     tracing::warn!(?error, "error updating redis state")
                 }
+                if let Event::Ready(_) = event {
+                    if !cache.2.read().await.contains(&shard.id().number()) {
+                        cache.2.write().await.push(shard.id().number());
+                    }
+                }
                 cache.0.update(&event);
                 //if let Err(error) = tx.send((shard.id(), event)) {
                 //    tracing::warn!(?error, "error sending event to global handler: {error}",);