Expose a method to return the memory usage of all dynamic atoms.

Signed-off-by: Josh Matthews <[email protected]>

Author: jdm

Cargo.toml

@@ -21,11 +21,13 @@ name = "string_cache"
 [features]
 serde_support = ["serde"]
 default = ["serde_support"]
+malloc_size_of = ["dep:malloc_size_of", "malloc_size_of_derive"]
 
 [dependencies]
 precomputed-hash = "0.1"
 serde = { version = "1", optional = true }
 malloc_size_of = { version = "0.1", default-features = false, optional = true }
+malloc_size_of_derive = { version = "0.1", optional = true }
 phf_shared = "0.11"
 new_debug_unreachable = "1.0.2"
 parking_lot = "0.12"

src/dynamic_set.rs

@@ -7,6 +7,8 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use malloc_size_of_derive::MallocSizeOf;
+use malloc_size_of::{MallocSizeOf as _, MallocSizeOfOps, MallocShallowSizeOf};
 use parking_lot::Mutex;
 use std::borrow::Cow;
 use std::mem;
@@ -22,6 +24,20 @@ pub(crate) struct Set {
     buckets: Box<[Mutex<Option<Box<Entry>>>]>,
 }
 
+impl Set {
+    #[cfg(feature = "malloc_size_of")]
+    pub(crate) fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
+        self.buckets.shallow_size_of(ops) +
+            self.buckets
+                .iter()
+                .map(|bucket| {
+                    bucket.lock().as_ref().map_or(0, |entry| entry.size_of(ops))
+                })
+                .sum::<usize>()
+    }
+}
+
+#[cfg_attr(feature = "malloc_size_of", derive(MallocSizeOf))]
 pub(crate) struct Entry {
     pub(crate) string: Box<str>,
     pub(crate) hash: u32,

src/lib.rs

@@ -127,6 +127,12 @@ pub use static_sets::{EmptyStaticAtomSet, PhfStrSet, StaticAtomSet};
 /// Use this if you don’t care about static atoms.
 pub type DefaultAtom = Atom<EmptyStaticAtomSet>;
 
+/// Measure the memory usage of the set of dynamic atoms.
+#[cfg(feature = "malloc_size_of")]
+pub fn size_of_dynamic_atoms(ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
+    dynamic_set::dynamic_set().size_of(ops)
+}
+
 // Some minor tests of internal layout here.
 // See ../integration-tests for much more.