Implement MallocSizeOf for Rkv internals

.github/workflows/ci.yml

@@ -46,7 +46,7 @@ jobs:
       - run: cargo build --verbose
 
       - name: Test with all features
-        run: cargo test --all --verbose
+        run: cargo test --all --all-features --verbose
 
       - name: Test with no default features
         run: cargo test --lib --no-default-features --verbose

Cargo.toml

@@ -24,6 +24,7 @@ db-dup-sort = []
 db-int-key = []
 default = ["db-dup-sort", "db-int-key"]
 no-canonicalize-path = []
+malloc-size-of = ["dep:malloc_size_of", "dep:malloc_size_of_derive"]
 
 [dependencies]
 arrayref = "0.3"
@@ -38,6 +39,8 @@ serde_derive = "1.0"
 thiserror = "2.0"
 url = "2.0"
 uuid = "1.0"
+malloc_size_of_derive = { version = "0.1.3", optional = true }
+malloc_size_of = { version = "0.2.2", package = "wr_malloc_size_of", default-features = false, optional = true }
 
 [dev-dependencies]
 byteorder = "1"

src/backend/common.rs

@@ -45,6 +45,7 @@ pub enum WriteFlags {
 
 /// Strategy to use when corrupted data is detected while opening a database.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "malloc-size-of", derive(malloc_size_of_derive::MallocSizeOf))]
 pub enum RecoveryStrategy {
     /// Bubble up the error on detecting a corrupted data file. The default.
     Error,

src/backend/impl_safe/database.rs

@@ -17,9 +17,18 @@ use crate::backend::traits::BackendDatabase;
 #[derive(Debug, Eq, PartialEq, Copy, Clone, Hash)]
 pub struct DatabaseImpl(pub(crate) Id<Database>);
 
+#[cfg(feature = "malloc-size-of")]
+impl malloc_size_of::MallocSizeOf for DatabaseImpl {
+    fn size_of(&self, _ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
+        // Id<T> does not allocate
+        0
+    }
+}
+
 impl BackendDatabase for DatabaseImpl {}
 
 #[derive(Debug, Serialize, Deserialize)]
+#[cfg_attr(feature = "malloc-size-of", derive(malloc_size_of_derive::MallocSizeOf))]
 pub struct Database {
     snapshot: Snapshot,
 }

src/backend/impl_safe/environment.rs

@@ -16,6 +16,8 @@ use std::{
     path::{Path, PathBuf},
     sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
 };
+#[cfg(feature = "malloc-size-of")]
+use std::mem;
 
 use id_arena::Arena;
 use log::warn;
@@ -34,6 +36,7 @@ type DatabaseArena = Arena<Database>;
 type DatabaseNameMap = HashMap<Option<String>, DatabaseImpl>;
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
+#[cfg_attr(feature = "malloc-size-of", derive(malloc_size_of_derive::MallocSizeOf))]
 pub struct EnvironmentBuilderImpl {
     flags: EnvironmentFlagsImpl,
     max_readers: Option<usize>,
@@ -119,6 +122,22 @@ pub(crate) struct EnvironmentDbs {
     pub(crate) name_map: DatabaseNameMap,
 }
 
+#[cfg(feature = "malloc-size-of")]
+impl malloc_size_of::MallocSizeOf for EnvironmentDbs {
+    fn size_of(&self, ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
+        let mut n = 0;
+
+        // Approximate allocations for the inner vec itself
+        let len = self.arena.len();
+        n += len * mem::size_of::<Database>();
+
+        for (_id, s) in self.arena.iter() {
+             n += s.size_of(ops)
+        }
+        n + self.name_map.size_of(ops)
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct EnvironmentDbsRefMut<'a> {
     pub(crate) arena: &'a mut DatabaseArena,
@@ -143,6 +162,20 @@ pub struct EnvironmentImpl {
     rw_txns: Arc<()>,
 }
 
+#[cfg(feature = "malloc-size-of")]
+impl malloc_size_of::MallocSizeOf for EnvironmentImpl {
+    fn size_of(&self, ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
+        let mut n = 0;
+        n += self.path.size_of(ops);
+        n += self.max_dbs.size_of(ops);
+        if let Ok(dbs) = self.dbs.read() {
+            n +=  (*dbs).size_of(ops);
+        }
+
+        n
+    }
+}
+
 impl EnvironmentImpl {
     fn serialize(&self) -> Result<Vec<u8>, ErrorImpl> {
         let dbs = self.dbs.read().map_err(|_| ErrorImpl::EnvPoisonError)?;

src/backend/impl_safe/flags.rs

@@ -23,6 +23,13 @@ bitflags! {
     }
 }
 
+#[cfg(feature = "malloc-size-of")]
+impl malloc_size_of::MallocSizeOf for EnvironmentFlagsImpl {
+    fn size_of(&self, _ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
+        0
+    }
+}
+
 impl BackendFlags for EnvironmentFlagsImpl {
     fn empty() -> EnvironmentFlagsImpl {
         EnvironmentFlagsImpl::empty()
@@ -64,6 +71,13 @@ bitflags! {
     }
 }
 
+#[cfg(feature = "malloc-size-of")]
+impl malloc_size_of::MallocSizeOf for DatabaseFlagsImpl {
+    fn size_of(&self, _ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
+        0
+    }
+}
+
 impl BackendFlags for DatabaseFlagsImpl {
     fn empty() -> DatabaseFlagsImpl {
         DatabaseFlagsImpl::empty()

src/backend/impl_safe/snapshot.rs

@@ -29,6 +29,21 @@ pub struct Snapshot {
     map: Arc<BTreeMap<Key, BTreeSet<Value>>>,
 }
 
+#[cfg(feature = "malloc-size-of")]
+impl malloc_size_of::MallocSizeOf for Snapshot {
+    fn size_of(&self, ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
+        let mut n = 0;
+        n += self.flags.size_of(ops);
+
+        // All other clones of the `Snapshot` map are transient in a transaction
+        // and thus not reachable from the `Database` object,
+        // so we're safe counting them here.
+        n += (*self.map).size_of(ops);
+
+        n
+    }
+}
+
 impl Snapshot {
     pub(crate) fn new(flags: Option<DatabaseFlagsImpl>) -> Snapshot {
         Snapshot {

src/env.rs

@@ -41,6 +41,7 @@ pub static DEFAULT_MAX_DBS: c_uint = 5;
 
 /// Wrapper around an `Environment` (e.g. a `SafeMode` environment).
 #[derive(Debug)]
+#[cfg_attr(feature = "malloc-size-of", derive(malloc_size_of_derive::MallocSizeOf))]
 pub struct Rkv<E> {
     _path: PathBuf,
     env: E,

src/store/single.rs

@@ -21,6 +21,7 @@ use crate::{
 type EmptyResult = Result<(), StoreError>;
 
 #[derive(Debug, Eq, PartialEq, Copy, Clone)]
+#[cfg_attr(feature = "malloc-size-of", derive(malloc_size_of_derive::MallocSizeOf))]
 pub struct SingleStore<D> {
     db: D,
 }