Add database cache info to cache_info command

Cargo.toml

@@ -30,7 +30,7 @@ strum_macros = "0.25"
 arrayvec = "0.7.4"
 bitflags = "2.4.1"
 paste = "1.0.14"
-typesize = "0.1.2"
+typesize = { version = "0.1.2", features = ["arrayvec"] }
 futures-util = { version = "0.3.29", default-features = false }
 
 [dependencies.symphonia]
@@ -70,3 +70,6 @@ branch = "current"
 [dependencies.songbird]
 version = "0.4"
 features = ["builtin-queue"]
+
+[patch.crates-io]
+typesize = { git = "https://github.com/GnomedDev/typesize" }

src/commands/owner.rs

@@ -17,6 +17,7 @@
 use std::{
     borrow::Cow,
     collections::{HashMap, HashSet},
+    hash::Hash,
     sync::atomic::Ordering::SeqCst,
 };
 
@@ -27,6 +28,8 @@ use rand::{rngs::ThreadRng, Rng};
 use typesize::TypeSize;
 
 use crate::{
+    database,
+    database_models::Compact,
     funcs::dm_generic,
     opt_ext::OptionTryUnwrap,
     structs::{Command, CommandResult, Context, PrefixContext, TTSModeChoice},
@@ -314,7 +317,7 @@ fn choose_random<'a, T>(rng: &mut ThreadRng, container: &'a [T]) -> &'a T {
     &container[index]
 }
 
-fn choose_random_map<'a, K: std::hash::Hash + Eq, V>(
+fn choose_random_map<'a, K: Hash + Eq, V>(
     rng: &mut ThreadRng,
     map: &'a HashMap<K, V>,
 ) -> Option<&'a V> {
@@ -328,6 +331,22 @@ fn random_guild<'a>(
     cache.guild(choose_random(rng, &cache.guilds()))
 }
 
+fn get_db_info<CacheKey, RowT>(
+    name: &'static str,
+    handler: &database::Handler<CacheKey, RowT>,
+) -> typesize::Field
+where
+    CacheKey: Eq + Hash + TypeSize,
+    RowT::Compacted: TypeSize,
+    RowT: Compact,
+{
+    typesize::Field {
+        name,
+        size: handler.get_size(),
+        collection_items: handler.get_collection_item_count(),
+    }
+}
+
 #[poise::command(prefix_command, owners_only)]
 pub async fn cache_info(ctx: Context<'_>, kind: Option<String>) -> CommandResult {
     struct Field {
@@ -339,27 +358,47 @@ pub async fn cache_info(ctx: Context<'_>, kind: Option<String>) -> CommandResult
 
     let serenity_cache = ctx.cache();
     let cache_stats = {
+        let data = ctx.data();
         let mut rng = rand::thread_rng();
         match kind.as_deref() {
-            Some("guild") => random_guild(&mut rng, serenity_cache)
-                .try_unwrap()?
-                .get_size_details(),
+            Some("db") => Some(vec![
+                get_db_info("guild db", &data.guilds_db),
+                get_db_info("userinfo db", &data.userinfo_db),
+                get_db_info("nickname db", &data.nickname_db),
+                get_db_info("user voice db", &data.user_voice_db),
+                get_db_info("guild voice db", &data.guild_voice_db),
+            ]),
+            Some("guild") => Some(
+                random_guild(&mut rng, serenity_cache)
+                    .try_unwrap()?
+                    .get_size_details(),
+            ),
             Some("channel") => {
                 let guild = random_guild(&mut rng, serenity_cache).try_unwrap()?;
-                choose_random_map(&mut rng, &guild.channels)
-                    .try_unwrap()?
-                    .get_size_details()
+                Some(
+                    choose_random_map(&mut rng, &guild.channels)
+                        .try_unwrap()?
+                        .get_size_details(),
+                )
             }
             Some("role") => {
                 let guild = random_guild(&mut rng, serenity_cache).try_unwrap()?;
-                choose_random_map(&mut rng, &guild.roles)
-                    .try_unwrap()?
-                    .get_size_details()
+                Some(
+                    choose_random_map(&mut rng, &guild.roles)
+                        .try_unwrap()?
+                        .get_size_details(),
+                )
             }
-            _ => serenity_cache.get_size_details(),
+            Some(_) => None,
+            None => Some(serenity_cache.get_size_details()),
         }
     };
 
+    let Some(cache_stats) = cache_stats else {
+        ctx.say("Unknown cache!").await?;
+        return Ok(());
+    };
+
     let mut fields = Vec::new();
     for field in cache_stats {
         let name = format!("`{}`", field.name);

src/database.rs

@@ -14,9 +14,10 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-use std::sync::Arc;
+use std::{hash::Hash, sync::Arc};
 
 use dashmap::DashMap;
+use typesize::TypeSize;
 
 pub use crate::database_models::*;
 use crate::structs::{Result, TTSMode};
@@ -25,7 +26,7 @@ type PgArguments<'a> = <sqlx::Postgres as sqlx::database::HasArguments<'a>>::Arg
 type QueryAs<'a, R> = sqlx::query::QueryAs<'a, sqlx::Postgres, R, PgArguments<'a>>;
 type Query<'a> = sqlx::query::Query<'a, sqlx::Postgres, PgArguments<'a>>;
 
-pub trait CacheKeyTrait: std::cmp::Eq + std::hash::Hash {
+pub trait CacheKeyTrait: std::cmp::Eq + Hash {
     fn bind_query(self, query: Query<'_>) -> Query<'_>;
     fn bind_query_as<R>(self, query: QueryAs<'_, R>) -> QueryAs<'_, R>;
 }
@@ -57,12 +58,11 @@ impl CacheKeyTrait for (i64, TTSMode) {
     }
 }
 
-pub struct Handler<
-    CacheKey: CacheKeyTrait,
-    RowT: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Compact,
-> {
+type OwnedArc<T> = typesize::ptr::SizableArc<T, typesize::ptr::Owned>;
+
+pub struct Handler<CacheKey, RowT: Compact> {
     pool: sqlx::PgPool,
-    cache: DashMap<CacheKey, Arc<RowT::Compacted>>,
+    cache: DashMap<CacheKey, OwnedArc<RowT::Compacted>>,
 
     default_row: Arc<RowT::Compacted>,
     single_insert: &'static str,
@@ -74,7 +74,7 @@ pub struct Handler<
 impl<CacheKey, RowT> Handler<CacheKey, RowT>
 where
     CacheKey: CacheKeyTrait + Sync + Send + Copy + Default,
-    RowT: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Compact + Sync + Send + Unpin,
+    RowT: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Compact + Send + Unpin,
 {
     pub async fn new(
         pool: sqlx::PgPool,
@@ -115,7 +115,7 @@ where
             .await?
             .unwrap_or_else(|| self.default_row.clone());
 
-        self.cache.insert(identifier, row.clone());
+        self.cache.insert(identifier, row.clone().into());
         Ok(row)
     }
 
@@ -166,6 +166,23 @@ where
     }
 }
 
+impl<CacheKey: Eq + Hash + TypeSize, RowT: Compact> TypeSize for Handler<CacheKey, RowT>
+where
+    RowT::Compacted: TypeSize,
+{
+    fn extra_size(&self) -> usize {
+        self.cache.extra_size()
+    }
+
+    fn get_collection_item_count(&self) -> Option<usize> {
+        self.cache.get_collection_item_count()
+    }
+
+    fn get_size_details(&self) -> Vec<typesize::Field> {
+        self.cache.get_size_details()
+    }
+}
+
 #[macro_export]
 macro_rules! create_db_handler {
     ($pool:expr, $table_name:literal, $id_name:literal) => {{

src/database_models.rs

@@ -1,4 +1,5 @@
 use arrayvec::ArrayString;
+use typesize::derive::TypeSize;
 
 use poise::serenity_prelude::{ChannelId, GuildId, RoleId, UserId};
 
@@ -30,9 +31,11 @@ macro_rules! named_bitflags {
     (pub struct $name:ident: $flag_size:ident {
         $(const $flag_name:ident = $flag_value:expr;)*
     }) => {
+        #[derive(TypeSize)]
+        pub struct $name($flag_size);
+
         bitflags::bitflags! {
-            #[derive(Debug)]
-            pub struct $name: $flag_size {
+            impl $name: $flag_size {
                 $(const $flag_name = $flag_value;)*
             }
         }
@@ -46,6 +49,15 @@ macro_rules! named_bitflags {
                 )*
             }
         }
+
+        impl std::fmt::Debug for $name {
+            fn fmt(&self, mut f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                write!(f, concat!(stringify!($name), "("))?;
+                bitflags::parser::to_writer(self, &mut f)?;
+                write!(f, ")")?;
+                Ok(())
+            }
+        }
     };
 }
 
@@ -85,7 +97,7 @@ named_bitflags! {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, TypeSize)]
 pub struct GuildRow {
     pub flags: GuildRowFlags,
     pub channel: Option<ChannelId>,
@@ -144,7 +156,7 @@ named_bitflags! {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, TypeSize)]
 pub struct UserRow {
     pub flags: UserRowFlags,
     pub voice_mode: Option<TTSMode>,
@@ -173,7 +185,7 @@ pub struct GuildVoiceRowRaw {
     pub voice: String,
 }
 
-#[derive(Debug)]
+#[derive(Debug, TypeSize)]
 
 pub struct GuildVoiceRow {
     pub guild_id: Option<GuildId>,
@@ -200,7 +212,7 @@ pub struct UserVoiceRowRaw {
     pub speaking_rate: Option<f32>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, TypeSize)]
 pub struct UserVoiceRow {
     pub user_id: Option<UserId>,
     pub mode: TTSMode,
@@ -222,7 +234,7 @@ impl Compact for UserVoiceRowRaw {
     }
 }
 
-#[derive(Debug, sqlx::FromRow)]
+#[derive(Debug, TypeSize, sqlx::FromRow)]
 pub struct NicknameRow {
     pub name: Option<String>,
 }