ckb_memory_tracker/
rocksdb.rs

1use ckb_db::internal::ops::{GetColumnFamilys, GetProperty, GetPropertyCF};
2
3#[allow(dead_code)]
4#[derive(Debug, Clone)]
5enum PropertyValue<T> {
6    Value(T),
7    Null,
8    Error(String),
9}
10
11impl PropertyValue<u64> {
12    pub(crate) fn as_i64(&self) -> i64 {
13        match self {
14            Self::Value(v) => *v as i64,
15            Self::Null => -1,
16            Self::Error(_) => -2,
17        }
18    }
19}
20
21impl<T> From<Result<Option<T>, String>> for PropertyValue<T> {
22    fn from(res: Result<Option<T>, String>) -> Self {
23        match res {
24            Ok(Some(v)) => Self::Value(v),
25            Ok(None) => Self::Null,
26            Err(e) => Self::Error(e),
27        }
28    }
29}
30
31/// A trait which used to track the RocksDB memory usage.
32///
33/// References: [Memory usage in RocksDB](https://github.com/facebook/rocksdb/wiki/Memory-usage-in-RocksDB)
34pub trait TrackRocksDBMemory {
35    /// Gather memory statistics through [ckb-metrics](../../ckb_metrics/index.html)
36    fn gather_memory_stats(&self) {
37        self.gather_int_values("estimate-table-readers-mem");
38        self.gather_int_values("size-all-mem-tables");
39        self.gather_int_values("cur-size-all-mem-tables");
40        self.gather_int_values("block-cache-capacity");
41        self.gather_int_values("block-cache-usage");
42        self.gather_int_values("block-cache-pinned-usage");
43    }
44
45    /// Gather integer values through [ckb-metrics](../../ckb_metrics/index.html)
46    fn gather_int_values(&self, _: &str) {}
47}
48
49pub(crate) struct DummyRocksDB;
50
51impl TrackRocksDBMemory for DummyRocksDB {}
52
53impl<RocksDB> TrackRocksDBMemory for RocksDB
54where
55    RocksDB: GetColumnFamilys + GetProperty + GetPropertyCF,
56{
57    fn gather_int_values(&self, key: &str) {
58        let mut values = Vec::new();
59        for (cf_name, cf) in self.get_cfs() {
60            let value_col: PropertyValue<u64> = self
61                .property_int_value_cf(cf, &format!("rocksdb.{key}"))
62                .map_err(|err| format!("{err}"))
63                .into();
64            if let Some(metrics) = ckb_metrics::handle() {
65                metrics
66                    .ckb_sys_mem_rocksdb
67                    .with_label_values(&[key, cf_name])
68                    .set(value_col.as_i64());
69            }
70            values.push(value_col);
71        }
72    }
73}