From 5822764240f8b4e8cfeca4bccf7d399a0bf71897 Mon Sep 17 00:00:00 2001
From: Jiayu Ye <yejiayu.fe@gmail.com>
Date: Tue, 4 Aug 2020 21:10:42 +0800
Subject: [PATCH] perf(metrics): Add metrics of state (#397)

* perf(metrics): Add metrics of state

* fix get state of metrics

* fix lint

* fix

* fix
---
 common/apm/src/metrics/storage.rs      | 15 +++++++++++++++
 framework/src/binding/state/trie_db.rs | 22 +++++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/common/apm/src/metrics/storage.rs b/common/apm/src/metrics/storage.rs
index 1a6c4e56d..c5243854b 100644
--- a/common/apm/src/metrics/storage.rs
+++ b/common/apm/src/metrics/storage.rs
@@ -15,6 +15,7 @@ make_auto_flush_static_metric! {
     signed_tx,
     wal,
     hash_height,
+    state,
   }
 
   pub struct StoragePutCfTimeUsageVec: LocalCounter {
@@ -72,6 +73,20 @@ lazy_static! {
         auto_flush_from!(STORAGE_GET_CF_COUNTER_VEC, StorageGetCfTotalVec);
 }
 
+pub fn on_storage_get_state(duration: Duration, keys: i64) {
+    let seconds = duration_to_sec(duration);
+
+    STORAGE_GET_CF_TIME_USAGE.state.inc_by(seconds);
+    STORAGE_GET_CF_COUNTER.state.inc_by(keys);
+}
+
+pub fn on_storage_put_state(duration: Duration, size: i64) {
+    let seconds = duration_to_sec(duration);
+
+    STORAGE_PUT_CF_TIME_USAGE.state.inc_by(seconds);
+    STORAGE_PUT_CF_BYTES_COUNTER.state.inc_by(size);
+}
+
 pub fn on_storage_get_cf(sc: StorageCategory, duration: Duration, keys: i64) {
     let seconds = duration_to_sec(duration);
 
diff --git a/framework/src/binding/state/trie_db.rs b/framework/src/binding/state/trie_db.rs
index 4a91eef91..8d8340e3b 100644
--- a/framework/src/binding/state/trie_db.rs
+++ b/framework/src/binding/state/trie_db.rs
@@ -1,10 +1,12 @@
 use std::path::Path;
 use std::sync::Arc;
+use std::time::Instant;
 
 use bytes::Bytes;
 use derive_more::{Display, From};
 use rocksdb::{Options, WriteBatch, DB};
 
+use common_apm::metrics::storage::{on_storage_get_state, on_storage_put_state};
 use protocol::{ProtocolError, ProtocolErrorKind, ProtocolResult};
 
 pub struct RocksTrieDB {
@@ -32,7 +34,12 @@ impl cita_trie::DB for RocksTrieDB {
     type Error = RocksTrieDBError;
 
     fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
-        Ok(self.db.get(key).map_err(to_store_err)?.map(|v| v.to_vec()))
+        let inst = Instant::now();
+
+        let res = self.db.get(key).map_err(to_store_err)?.map(|v| v.to_vec());
+
+        on_storage_get_state(inst.elapsed(), 1i64);
+        Ok(res)
     }
 
     fn contains(&self, key: &[u8]) -> Result<bool, Self::Error> {
@@ -40,25 +47,38 @@ impl cita_trie::DB for RocksTrieDB {
     }
 
     fn insert(&self, key: Vec<u8>, value: Vec<u8>) -> Result<(), Self::Error> {
+        let inst = Instant::now();
+        let size = key.len() + value.len();
+
         self.db
             .put(Bytes::from(key), Bytes::from(value))
             .map_err(to_store_err)?;
+
+        on_storage_put_state(inst.elapsed(), size as i64);
         Ok(())
     }
 
     fn insert_batch(&self, keys: Vec<Vec<u8>>, values: Vec<Vec<u8>>) -> Result<(), Self::Error> {
+        let inst = Instant::now();
+
         if keys.len() != values.len() {
             return Err(RocksTrieDBError::BatchLengthMismatch);
         }
 
+        let mut total_size = 0;
         let mut batch = WriteBatch::default();
         for i in 0..keys.len() {
             let key = &keys[i];
             let value = &values[i];
+
+            total_size += key.len();
+            total_size += value.len();
             batch.put(key, value);
         }
 
         self.db.write(batch).map_err(to_store_err)?;
+
+        on_storage_put_state(inst.elapsed(), total_size as i64);
         Ok(())
     }