diff --git a/src/service/globals/mod.rs b/src/service/globals/mod.rs index 3f98babe..3a514d5c 100644 --- a/src/service/globals/mod.rs +++ b/src/service/globals/mod.rs @@ -7,7 +7,7 @@ use std::{ time::Instant, }; -use conduwuit::{error, Config, Result}; +use conduwuit::{error, utils::bytes::pretty, Config, Result}; use data::Data; use regex::RegexSet; use ruma::{OwnedEventId, OwnedRoomAliasId, OwnedServerName, OwnedUserId, ServerName, UserId}; @@ -93,13 +93,18 @@ impl crate::Service for Service { Ok(Arc::new(s)) } - fn memory_usage(&self, out: &mut dyn Write) -> Result<()> { - let bad_event_ratelimiter = self - .bad_event_ratelimiter - .read() - .expect("locked for reading") - .len(); - writeln!(out, "bad_event_ratelimiter: {bad_event_ratelimiter}")?; + fn memory_usage(&self, out: &mut dyn Write) -> Result { + let (ber_count, ber_bytes) = self.bad_event_ratelimiter.read()?.iter().fold( + (0_usize, 0_usize), + |(mut count, mut bytes), (event_id, _)| { + bytes = bytes.saturating_add(event_id.capacity()); + bytes = bytes.saturating_add(size_of::()); + count = count.saturating_add(1); + (count, bytes) + }, + ); + + writeln!(out, "bad_event_ratelimiter: {ber_count} ({})", pretty(ber_bytes))?; Ok(()) } diff --git a/src/service/resolver/cache.rs b/src/service/resolver/cache.rs index 3e961f4c..e309a129 100644 --- a/src/service/resolver/cache.rs +++ b/src/service/resolver/cache.rs @@ -6,7 +6,10 @@ use std::{ }; use arrayvec::ArrayVec; -use conduwuit::{trace, utils::rand}; +use conduwuit::{ + trace, + utils::{math::Expected, rand}, +}; use ruma::{OwnedServerName, ServerName}; use super::fed::FedDest; @@ -113,6 +116,15 @@ impl CachedDest { pub(crate) fn default_expire() -> SystemTime { rand::timepoint_secs(60 * 60 * 18..60 * 60 * 36) } + + #[inline] + #[must_use] + pub fn size(&self) -> usize { + self.dest + .size() + .expected_add(self.host.len()) + .expected_add(size_of_val(&self.expire)) + } } impl CachedOverride { @@ -126,4 +138,8 @@ impl CachedOverride { pub(crate) fn default_expire() -> SystemTime { rand::timepoint_secs(60 * 60 * 6..60 * 60 * 12) } + + #[inline] + #[must_use] + pub fn size(&self) -> usize { size_of_val(self) } } diff --git a/src/service/resolver/fed.rs b/src/service/resolver/fed.rs index 3986db8e..76fc6894 100644 --- a/src/service/resolver/fed.rs +++ b/src/service/resolver/fed.rs @@ -5,6 +5,7 @@ use std::{ }; use arrayvec::ArrayString; +use conduwuit::utils::math::Expected; #[derive(Clone, Debug, PartialEq, Eq)] pub enum FedDest { @@ -76,6 +77,15 @@ impl FedDest { pub fn default_port() -> PortString { PortString::from(DEFAULT_PORT).expect("default port string") } + + #[inline] + #[must_use] + pub fn size(&self) -> usize { + match self { + | Self::Literal(saddr) => size_of_val(saddr), + | Self::Named(host, port) => host.len().expected_add(port.capacity()), + } + } } impl fmt::Display for FedDest { diff --git a/src/service/resolver/mod.rs b/src/service/resolver/mod.rs index e18cf0bd..6a6289b6 100644 --- a/src/service/resolver/mod.rs +++ b/src/service/resolver/mod.rs @@ -6,7 +6,7 @@ mod tests; use std::{fmt::Write, sync::Arc}; -use conduwuit::{Result, Server}; +use conduwuit::{utils, utils::math::Expected, Result, Server}; use self::{cache::Cache, dns::Resolver}; use crate::{client, Dep}; @@ -36,22 +36,25 @@ impl crate::Service for Service { })) } - fn memory_usage(&self, out: &mut dyn Write) -> Result<()> { - let resolver_overrides_cache = self - .cache - .overrides - .read() - .expect("locked for reading") - .len(); - writeln!(out, "resolver_overrides_cache: {resolver_overrides_cache}")?; + fn memory_usage(&self, out: &mut dyn Write) -> Result { + use utils::bytes::pretty; - let resolver_destinations_cache = self - .cache - .destinations - .read() - .expect("locked for reading") - .len(); - writeln!(out, "resolver_destinations_cache: {resolver_destinations_cache}")?; + let (oc_count, oc_bytes) = self.cache.overrides.read()?.iter().fold( + (0_usize, 0_usize), + |(count, bytes), (key, val)| { + (count.expected_add(1), bytes.expected_add(key.len()).expected_add(val.size())) + }, + ); + + let (dc_count, dc_bytes) = self.cache.destinations.read()?.iter().fold( + (0_usize, 0_usize), + |(count, bytes), (key, val)| { + (count.expected_add(1), bytes.expected_add(key.len()).expected_add(val.size())) + }, + ); + + writeln!(out, "resolver_overrides_cache: {oc_count} ({})", pretty(oc_bytes))?; + writeln!(out, "resolver_destinations_cache: {dc_count} ({})", pretty(dc_bytes))?; Ok(()) } diff --git a/src/service/rooms/state_accessor/mod.rs b/src/service/rooms/state_accessor/mod.rs index fd132798..d89c8835 100644 --- a/src/service/rooms/state_accessor/mod.rs +++ b/src/service/rooms/state_accessor/mod.rs @@ -10,7 +10,11 @@ use std::{ use conduwuit::{ err, error, pdu::PduBuilder, - utils::{math::usize_from_f64, ReadyExt}, + utils, + utils::{ + math::{usize_from_f64, Expected}, + ReadyExt, + }, Err, Error, PduEvent, Result, }; use futures::StreamExt; @@ -84,12 +88,35 @@ impl crate::Service for Service { })) } - fn memory_usage(&self, out: &mut dyn Write) -> Result<()> { - let server_visibility_cache = self.server_visibility_cache.lock().expect("locked").len(); - writeln!(out, "server_visibility_cache: {server_visibility_cache}")?; + fn memory_usage(&self, out: &mut dyn Write) -> Result { + use utils::bytes::pretty; - let user_visibility_cache = self.user_visibility_cache.lock().expect("locked").len(); - writeln!(out, "user_visibility_cache: {user_visibility_cache}")?; + let (svc_count, svc_bytes) = self.server_visibility_cache.lock()?.iter().fold( + (0_usize, 0_usize), + |(count, bytes), (key, _)| { + ( + count.expected_add(1), + bytes + .expected_add(key.0.capacity()) + .expected_add(size_of_val(&key.1)), + ) + }, + ); + + let (uvc_count, uvc_bytes) = self.user_visibility_cache.lock()?.iter().fold( + (0_usize, 0_usize), + |(count, bytes), (key, _)| { + ( + count.expected_add(1), + bytes + .expected_add(key.0.capacity()) + .expected_add(size_of_val(&key.1)), + ) + }, + ); + + writeln!(out, "server_visibility_cache: {svc_count} ({})", pretty(svc_bytes))?; + writeln!(out, "user_visibility_cache: {uvc_count} ({})", pretty(uvc_bytes))?; Ok(()) }