Add memory stats for kvdb-s

kvdb-memorydb/Cargo.toml

@@ -8,5 +8,6 @@ license = "GPL-3.0"
 edition = "2018"
 
 [dependencies]
+parity-util-mem = { path = "../parity-util-mem", version = "0.3" }
 parking_lot = "0.9.0"
 kvdb = { version = "0.2", path = "../kvdb" }

kvdb-memorydb/src/lib.rs

@@ -15,6 +15,7 @@
 // along with Parity.  If not, see <http://www.gnu.org/licenses/>.
 
 use kvdb::{DBOp, DBTransaction, DBValue, KeyValueDB};
+use parity_util_mem::MallocSizeOf;
 use parking_lot::RwLock;
 use std::{
 	collections::{BTreeMap, HashMap},
@@ -23,7 +24,7 @@ use std::{
 
 /// A key-value database fulfilling the `KeyValueDB` trait, living in memory.
 /// This is generally intended for tests and is not particularly optimized.
-#[derive(Default)]
+#[derive(Default, MallocSizeOf)]
 pub struct InMemory {
 	columns: RwLock<HashMap<u32, BTreeMap<Vec<u8>, DBValue>>>,
 }

kvdb-rocksdb/Cargo.toml

@@ -22,6 +22,7 @@ parking_lot = "0.9.0"
 regex = "1.3.1"
 rocksdb = { version = "0.13", features = ["snappy"], default-features = false }
 owning_ref = "0.4.0"
+parity-util-mem = { path = "../parity-util-mem", version = "0.3" }
 
 [dev-dependencies]
 alloc_counter = "0.0.4"

kvdb-rocksdb/src/lib.rs

@@ -18,6 +18,7 @@ mod iter;
 
 use std::{cmp, collections::HashMap, convert::identity, error, fs, io, mem, path::Path, result};
 
+use parity_util_mem::MallocSizeOf;
 use parking_lot::{Mutex, MutexGuard, RwLock};
 use rocksdb::{
 	BlockBasedOptions, ColumnFamily, ColumnFamilyDescriptor, Error, Options, ReadOptions, WriteBatch, WriteOptions, DB,
@@ -57,6 +58,7 @@ pub const DB_DEFAULT_COLUMN_MEMORY_BUDGET_MB: MiB = 128;
 /// The default memory budget in MiB.
 pub const DB_DEFAULT_MEMORY_BUDGET_MB: MiB = 512;
 
+#[derive(MallocSizeOf)]
 enum KeyState {
 	Insert(DBValue),
 	Delete,
@@ -229,19 +231,43 @@ struct DBAndColumns {
 	column_names: Vec<String>,
 }
 
+fn static_property_or_warn(db: &DB, prop: &str) -> usize {
+	match db.property_int_value(prop) {
+		Ok(Some(v)) => v as usize,
+		_ => {
+			warn!("Cannot read expected static property of RocksDb database: {}", prop);
+			0
+		}
+	}
+}
+
+impl MallocSizeOf for DBAndColumns {
+	fn size_of(&self, ops: &mut parity_util_mem::MallocSizeOfOps) -> usize {
+		self.column_names.size_of(ops)
+			+ static_property_or_warn(&self.db, "rocksdb.estimate-table-readers-mem")
+			+ static_property_or_warn(&self.db, "rocksdb.cur-size-all-mem-tables")
+			+ static_property_or_warn(&self.db, "rocksdb.block-cache-usage")
+	}
+}
+
 impl DBAndColumns {
 	fn cf(&self, i: usize) -> &ColumnFamily {
 		self.db.cf_handle(&self.column_names[i]).expect("the specified column name is correct; qed")
 	}
 }
 
 /// Key-Value database.
+#[derive(MallocSizeOf)]
 pub struct Database {
 	db: RwLock<Option<DBAndColumns>>,
+	#[ignore_malloc_size_of = "insignificant"]
 	config: DatabaseConfig,
 	path: String,
+	#[ignore_malloc_size_of = "insignificant"]
 	write_opts: WriteOptions,
+	#[ignore_malloc_size_of = "insignificant"]
 	read_opts: ReadOptions,
+	#[ignore_malloc_size_of = "insignificant"]
 	block_opts: BlockBasedOptions,
 	// Dirty values added with `write_buffered`. Cleaned on `flush`.
 	overlay: RwLock<Vec<HashMap<DBKey, KeyState>>>,
@@ -759,6 +785,36 @@ mod tests {
 		assert_eq!(&*db.get(0, key1.as_bytes()).unwrap().unwrap(), b"horse");
 	}
 
+	#[test]
+	fn mem_tables_size() {
+		let tempdir = TempDir::new("").unwrap();
+
+		let config = DatabaseConfig {
+			max_open_files: 512,
+			memory_budget: HashMap::new(),
+			compaction: CompactionProfile::default(),
+			columns: 11,
+			keep_log_file_num: 1,
+		};
+
+		let db = Database::open(&config, tempdir.path().to_str().unwrap()).unwrap();
+
+		let mut batch = db.transaction();
+		for i in 0u32..10000u32 {
+			batch.put(i / 1000 + 1, &i.to_le_bytes(), &(i * 17).to_le_bytes());
+		}
+		db.write(batch).unwrap();
+
+		db.flush().unwrap();
+
+		{
+			let db = db.db.read();
+			db.as_ref().map(|db| {
+				assert!(super::static_property_or_warn(&db.db, "rocksdb.cur-size-all-mem-tables") > 512);
+			});
+		}
+	}
+
 	#[test]
 	fn kvdb() {
 		let tempdir = TempDir::new("").unwrap();

kvdb-web/Cargo.toml

@@ -16,6 +16,7 @@ kvdb-memorydb = { version = "0.2", path = "../kvdb-memorydb" }
 futures = "0.3"
 log = "0.4.8"
 send_wrapper = "0.3.0"
+parity-util-mem = { path = "../parity-util-mem", version = "0.3" }
 
 [dependencies.web-sys]
 version = "0.3.31"

kvdb-web/src/lib.rs

@@ -45,6 +45,9 @@ pub struct Database {
 	indexed_db: SendWrapper<IdbDatabase>,
 }
 
+// TODO: implement when web-based implementation need memory stats
+parity_util_mem::malloc_size_of_is_0!(Database);
+
 impl Database {
 	/// Opens the database with the given name,
 	/// and the specified number of columns (not including the default one).

kvdb/Cargo.toml

@@ -10,3 +10,4 @@ edition = "2018"
 [dependencies]
 smallvec = "1.0.0"
 bytes = { package = "parity-bytes", version = "0.1", path = "../parity-bytes" }
+parity-util-mem = { path = "../parity-util-mem", version = "0.3" }

kvdb/src/lib.rs

@@ -107,7 +107,7 @@ impl DBTransaction {
 ///
 /// The API laid out here, along with the `Sync` bound implies interior synchronization for
 /// implementation.
-pub trait KeyValueDB: Sync + Send {
+pub trait KeyValueDB: Sync + Send + parity_util_mem::MallocSizeOf {
 	/// Helper to create a new transaction.
 	fn transaction(&self) -> DBTransaction {
 		DBTransaction::new()

parity-util-mem/src/malloc_size.rs

@@ -425,6 +425,12 @@ where
 	}
 }
 
+impl<I: MallocSizeOf> MallocSizeOf for rstd::cmp::Reverse<I> {
+	fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
+		self.0.size_of(ops)
+	}
+}
+
 #[cfg(feature = "std")]
 impl<K, V, S> MallocShallowSizeOf for std::collections::HashMap<K, V, S> {
 	fn shallow_size_of(&self, ops: &mut MallocSizeOfOps) -> usize {

parity-util-mem/tests/derive.rs

@@ -29,8 +29,8 @@ fn derive_vec() {
 	assert!(t.malloc_size_of() > 1000);
 }
 
-#[test]
 #[cfg(feature = "std")]
+#[test]
 fn derive_hashmap() {
 	#[derive(MallocSizeOf, Default)]
 	struct Trivia {
@@ -58,6 +58,5 @@ fn derive_ignore() {
 
 	t.hm.insert(1, vec![0u8; 2048]);
 	t.v = vec![0u8; 1024];
-
 	assert!(t.malloc_size_of() < 3000);
 }