Merge branch 'admin-command-for-cache-memory-usage' into 'master'

Add show_cache_usage admin room command

See merge request famedly/conduit!167

src/database.rs

@@ -1,21 +1,15 @@
-pub mod abstraction;
-
-pub mod account_data;
-pub mod admin;
-pub mod appservice;
-pub mod globals;
-pub mod key_backups;
-pub mod media;
-pub mod proxy;
-pub mod pusher;
-pub mod rooms;
-pub mod sending;
-pub mod transaction_ids;
-pub mod uiaa;
-pub mod users;
+use std::hash::Hash;
+use std::{
+    collections::{BTreeMap, HashMap, HashSet},
+    convert::{TryFrom, TryInto},
+    fs::{self, remove_dir_all},
+    io::Write,
+    mem::size_of,
+    ops::Deref,
+    path::Path,
+    sync::{Arc, Mutex, RwLock},
+};
 
-use crate::{utils, Error, Result};
-use abstraction::DatabaseEngine;
 use directories::ProjectDirs;
 use lru_cache::LruCache;
 use rocket::{
@@ -26,21 +20,31 @@ use rocket::{
 };
 use ruma::{DeviceId, EventId, RoomId, ServerName, UserId};
 use serde::{de::IgnoredAny, Deserialize};
-use std::{
-    collections::{BTreeMap, HashMap, HashSet},
-    convert::{TryFrom, TryInto},
-    fs::{self, remove_dir_all},
-    io::Write,
-    mem::size_of,
-    ops::Deref,
-    path::Path,
-    sync::{Arc, Mutex, RwLock},
-};
 use tokio::sync::{OwnedRwLockReadGuard, RwLock as TokioRwLock, Semaphore};
 use tracing::{debug, error, warn};
 
+use abstraction::DatabaseEngine;
+
+use crate::{utils, Error, Result};
+
 use self::proxy::ProxyConfig;
 
+pub mod abstraction;
+
+pub mod account_data;
+pub mod admin;
+pub mod appservice;
+pub mod globals;
+pub mod key_backups;
+pub mod media;
+pub mod proxy;
+pub mod pusher;
+pub mod rooms;
+pub mod sending;
+pub mod transaction_ids;
+pub mod uiaa;
+pub mod users;
+
 #[derive(Clone, Debug, Deserialize)]
 pub struct Config {
     server_name: Box<ServerName>,
@@ -132,6 +136,17 @@ pub type Engine = abstraction::sqlite::Engine;
 #[cfg(feature = "heed")]
 pub type Engine = abstraction::heed::Engine;
 
+// for each key: (memory_usage in bytes, items in cache, capacity)
+pub struct CacheUsageStatistics {
+    pdu_cache: (usize, usize, usize),
+    auth_chain_cache: (usize, usize, usize),
+    shorteventid_cache: (usize, usize, usize),
+    eventidshort_cache: (usize, usize, usize),
+    statekeyshort_cache: (usize, usize, usize),
+    shortstatekey_cache: (usize, usize, usize),
+    stateinfo_cache: (usize, usize, usize),
+}
+
 pub struct Database {
     _db: Arc<Engine>,
     pub globals: globals::Globals,
@@ -917,6 +932,66 @@ impl Database {
             }
         });
     }
+
+    /// Measures memory usage in bytes and how full the caches are in percent for all caches in the Database struct.
+    pub fn get_cache_usage(&mut self) -> Result<CacheUsageStatistics> {
+        fn memory_usage_of_locked_cache<K: Eq + Hash, V>(
+            cache: &mut Mutex<LruCache<K, V>>,
+        ) -> usize {
+            let raw_cache = cache.lock().unwrap();
+            let mut cache_items_size_sum: usize = 0;
+            for cache_item in raw_cache.iter() {
+                cache_items_size_sum += std::mem::size_of_val(&cache_item);
+            }
+            cache_items_size_sum += std::mem::size_of_val(&cache);
+            cache_items_size_sum
+        }
+
+        fn items_in_locked_cache<K: Eq + Hash, V>(cache: &mut Mutex<LruCache<K, V>>) -> usize {
+            cache.lock().unwrap().len()
+        }
+        fn capacity_of_locked_cache<K: Eq + Hash, V>(cache: &mut Mutex<LruCache<K, V>>) -> usize {
+            cache.lock().unwrap().capacity()
+        }
+
+        return Ok(CacheUsageStatistics {
+            pdu_cache: (
+                memory_usage_of_locked_cache(&mut self.rooms.pdu_cache),
+                items_in_locked_cache(&mut self.rooms.pdu_cache),
+                capacity_of_locked_cache(&mut self.rooms.pdu_cache),
+            ),
+            auth_chain_cache: (
+                memory_usage_of_locked_cache(&mut self.rooms.auth_chain_cache),
+                items_in_locked_cache(&mut self.rooms.auth_chain_cache),
+                capacity_of_locked_cache(&mut self.rooms.auth_chain_cache),
+            ),
+            shorteventid_cache: (
+                memory_usage_of_locked_cache(&mut self.rooms.shorteventid_cache),
+                items_in_locked_cache(&mut self.rooms.shorteventid_cache),
+                capacity_of_locked_cache(&mut self.rooms.shorteventid_cache),
+            ),
+            eventidshort_cache: (
+                memory_usage_of_locked_cache(&mut self.rooms.eventidshort_cache),
+                items_in_locked_cache(&mut self.rooms.eventidshort_cache),
+                capacity_of_locked_cache(&mut self.rooms.eventidshort_cache),
+            ),
+            statekeyshort_cache: (
+                memory_usage_of_locked_cache(&mut self.rooms.statekeyshort_cache),
+                items_in_locked_cache(&mut self.rooms.statekeyshort_cache),
+                capacity_of_locked_cache(&mut self.rooms.statekeyshort_cache),
+            ),
+            shortstatekey_cache: (
+                memory_usage_of_locked_cache(&mut self.rooms.shortstatekey_cache),
+                items_in_locked_cache(&mut self.rooms.shortstatekey_cache),
+                capacity_of_locked_cache(&mut self.rooms.shortstatekey_cache),
+            ),
+            stateinfo_cache: (
+                memory_usage_of_locked_cache(&mut self.rooms.stateinfo_cache),
+                items_in_locked_cache(&mut self.rooms.stateinfo_cache),
+                capacity_of_locked_cache(&mut self.rooms.stateinfo_cache),
+            ),
+        });
+    }
 }
 
 pub struct DatabaseGuard(OwnedRwLockReadGuard<Database>);

src/database/admin.rs

@@ -3,19 +3,21 @@ use std::{
     sync::Arc,
 };
 
-use crate::{pdu::PduBuilder, Database};
 use rocket::futures::{channel::mpsc, stream::StreamExt};
 use ruma::{
     events::{room::message, EventType},
     UserId,
 };
-use tokio::sync::{MutexGuard, RwLock, RwLockReadGuard};
+use tokio::sync::{MutexGuard, RwLock, RwLockWriteGuard};
 use tracing::warn;
 
+use crate::{pdu::PduBuilder, Database};
+
 pub enum AdminCommand {
     RegisterAppservice(serde_yaml::Value),
     ListAppservices,
     SendMessage(message::MessageEventContent),
+    ShowCacheUsage,
 }
 
 #[derive(Clone)]
@@ -59,7 +61,7 @@ impl Admin {
             drop(guard);
 
             let send_message = |message: message::MessageEventContent,
-                                guard: RwLockReadGuard<'_, Database>,
+                                guard: RwLockWriteGuard<'_, Database>,
                                 mutex_lock: &MutexGuard<'_, ()>| {
                 guard
                     .rooms
@@ -83,7 +85,7 @@ impl Admin {
             loop {
                 tokio::select! {
                     Some(event) = receiver.next() => {
-                        let guard = db.read().await;
+                        let mut guard = db.write().await;
                         let mutex_state = Arc::clone(
                             guard.globals
                                 .roomid_mutex_state
@@ -92,6 +94,7 @@ impl Admin {
                                 .entry(conduit_room.clone())
                                 .or_default(),
                         );
+
                         let state_lock = mutex_state.lock().await;
 
                         match event {
@@ -114,6 +117,37 @@ impl Admin {
                             AdminCommand::SendMessage(message) => {
                                 send_message(message, guard, &state_lock);
                             }
+                            AdminCommand::ShowCacheUsage => {
+
+                                fn format_cache_statistics_triple(name: String, triple: (usize, usize, usize)) -> String {
+                                    let (memory_usage, item_count, capacity) = triple;
+                                    format!(
+                                        "{0} is using {1} MB ({2} bytes) of RAM at {3:.2}% utilization.",
+                                        name,
+                                        memory_usage / 100_00,
+                                        memory_usage ,
+                                        ((item_count as f32 / capacity as f32) * 100.0)
+                                    )
+                                }
+
+                                if let Ok(cache_usage_statistics) = guard.get_cache_usage() {
+
+                                    let mut statistics_lines =  Vec::with_capacity(7);
+                                    statistics_lines.push(format_cache_statistics_triple("pdu_cache".to_string(), cache_usage_statistics.pdu_cache));
+                                    statistics_lines.push(format_cache_statistics_triple("auth_chain_cache".to_string(), cache_usage_statistics.auth_chain_cache));
+                                    statistics_lines.push(format_cache_statistics_triple("shorteventid_cache".to_string(), cache_usage_statistics.shorteventid_cache));
+                                    statistics_lines.push(format_cache_statistics_triple("eventidshort_cache".to_string(), cache_usage_statistics.eventidshort_cache));
+                                    statistics_lines.push(format_cache_statistics_triple("statekeyshort_cache".to_string(), cache_usage_statistics.statekeyshort_cache));
+                                    statistics_lines.push(format_cache_statistics_triple("shortstatekey_cache".to_string(), cache_usage_statistics.shortstatekey_cache));
+                                    statistics_lines.push(format_cache_statistics_triple("stateinfo_cache".to_string(), cache_usage_statistics.stateinfo_cache));
+
+                                    send_message(message::MessageEventContent::text_plain(statistics_lines.join("\n")), guard, &state_lock);
+                                } else {
+                                    let result_text = "Could not calculate database cache size";
+                                    send_message(message::MessageEventContent::text_plain(result_text), guard, &state_lock);
+                                }
+
+                            }
                         }
 
                         drop(state_lock);

src/database/rooms.rs

@@ -1,10 +1,12 @@
-mod edus;
-
-pub use edus::RoomEdus;
-use member::MembershipState;
+use std::{
+    collections::{BTreeMap, HashMap, HashSet},
+    convert::{TryFrom, TryInto},
+    mem::size_of,
+    sync::{Arc, Mutex},
+};
 
-use crate::{pdu::PduBuilder, utils, Database, Error, PduEvent, Result};
 use lru_cache::LruCache;
+use member::MembershipState;
 use regex::Regex;
 use ring::digest;
 use rocket::http::RawStr;
@@ -22,17 +24,17 @@ use ruma::{
     state_res::{self, RoomVersion, StateMap},
     uint, EventId, RoomAliasId, RoomId, RoomVersionId, ServerName, UserId,
 };
-use std::{
-    collections::{BTreeMap, HashMap, HashSet},
-    convert::{TryFrom, TryInto},
-    mem::size_of,
-    sync::{Arc, Mutex},
-};
 use tokio::sync::MutexGuard;
 use tracing::{error, warn};
 
+pub use edus::RoomEdus;
+
+use crate::{pdu::PduBuilder, utils, Database, Error, PduEvent, Result};
+
 use super::{abstraction::Tree, admin::AdminCommand, pusher};
 
+mod edus;
+
 /// The unique identifier of each state group.
 ///
 /// This is created when a state group is added to the database by
@@ -1566,10 +1568,13 @@ impl Rooms {
                                         ));
                                     }
                                 }
+                                "show_cache_usage" => {
+                                    db.admin.send(AdminCommand::ShowCacheUsage);
+                                }
                                 _ => {
                                     db.admin.send(AdminCommand::SendMessage(
                                         message::MessageEventContent::text_plain(format!(
-                                            "Unrecognized command: {}",
+                                            "Unrecognized command: {}\nAvailable commands:\n- register_appservice\n- list_appservices\n- get_pdu\n- show_cache_usage",
                                             command
                                         )),
                                     ));