diff --git a/src/service/rooms/state_compressor/mod.rs b/src/service/rooms/state_compressor/mod.rs index bf90d5c4..6b520ad3 100644 --- a/src/service/rooms/state_compressor/mod.rs +++ b/src/service/rooms/state_compressor/mod.rs @@ -1,18 +1,22 @@ use std::{ - collections::HashSet, + collections::{HashMap, HashSet}, fmt::Write, mem::size_of, sync::{Arc, Mutex}, }; -use conduit::{checked, err, expected, utils, utils::math::usize_from_f64, Result}; +use conduit::{ + at, checked, err, expected, utils, + utils::{bytes, math::usize_from_f64}, + Result, +}; use database::Map; use lru_cache::LruCache; use ruma::{EventId, RoomId}; use crate::{ rooms, - rooms::short::{ShortStateHash, ShortStateKey}, + rooms::short::{ShortId, ShortStateHash, ShortStateKey}, Dep, }; @@ -53,12 +57,13 @@ pub struct HashSetCompressStateEvent { pub removed: Arc, } -pub(crate) type CompressedState = HashSet; -pub(crate) type CompressedStateEvent = [u8; 2 * size_of::()]; type StateInfoLruCache = LruCache; type ShortStateInfoVec = Vec; type ParentStatesVec = Vec; +pub(crate) type CompressedState = HashSet; +pub(crate) type CompressedStateEvent = [u8; 2 * size_of::()]; + impl crate::Service for Service { fn build(args: crate::Args<'_>) -> Result> { let config = &args.server.config; @@ -75,9 +80,28 @@ impl crate::Service for Service { })) } - fn memory_usage(&self, out: &mut dyn Write) -> Result<()> { - let stateinfo_cache = self.stateinfo_cache.lock().expect("locked").len(); - writeln!(out, "stateinfo_cache: {stateinfo_cache}")?; + fn memory_usage(&self, out: &mut dyn Write) -> Result { + let (cache_len, ents) = { + let cache = self.stateinfo_cache.lock().expect("locked"); + let ents = cache + .iter() + .map(at!(1)) + .flat_map(|vec| vec.iter()) + .fold(HashMap::new(), |mut ents, ssi| { + ents.insert(Arc::as_ptr(&ssi.added), compressed_state_size(&ssi.added)); + ents.insert(Arc::as_ptr(&ssi.removed), compressed_state_size(&ssi.removed)); + ents.insert(Arc::as_ptr(&ssi.full_state), compressed_state_size(&ssi.full_state)); + ents + }); + + (cache.len(), ents) + }; + + let ents_len = ents.len(); + let bytes = ents.values().copied().fold(0_usize, usize::saturating_add); + + let bytes = bytes::pretty(bytes); + writeln!(out, "stateinfo_cache: {cache_len} {ents_len} ({bytes})")?; Ok(()) } @@ -435,3 +459,11 @@ impl Service { .insert(&shortstatehash.to_be_bytes(), &value); } } + +#[inline] +fn compressed_state_size(compressed_state: &CompressedState) -> usize { + compressed_state + .len() + .checked_mul(size_of::()) + .expect("CompressedState size overflow") +}