From 41788c3f2749610ac346a34bec233456748d53ca Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 10 Nov 2020 11:11:58 +1000
Subject: [PATCH] Add metrics for the non-finalized state and queue (#1263)

---
 .../memory_state/non_finalized_state.rs       | 40 +++++++++++++++++--
 .../src/service/memory_state/queued_blocks.rs | 16 ++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/zebra-state/src/service/memory_state/non_finalized_state.rs b/zebra-state/src/service/memory_state/non_finalized_state.rs
index 0536e239..5df57399 100644
--- a/zebra-state/src/service/memory_state/non_finalized_state.rs
+++ b/zebra-state/src/service/memory_state/non_finalized_state.rs
@@ -50,11 +50,13 @@ impl NonFinalizedState {
             }
         }
 
+        self.update_metrics_for_chains();
+
         // return the finalized block
         finalized_block
     }
 
-    /// Commit block to the non-finalize state.
+    /// Commit block to the non-finalized state.
     pub fn commit_block(&mut self, block: Arc<Block>) {
         let parent_hash = block.header.previous_block_hash;
 
@@ -68,16 +70,18 @@ impl NonFinalizedState {
             })
             .expect("commit_block is only called with blocks that are ready to be commited");
 
-        parent_chain.push(block);
+        parent_chain.push(block.clone());
         self.chain_set.insert(parent_chain);
+        self.update_metrics_for_committed_block(block);
     }
 
     /// Commit block to the non-finalized state as a new chain where its parent
     /// is the finalized tip.
     pub fn commit_new_chain(&mut self, block: Arc<Block>) {
         let mut chain = Chain::default();
-        chain.push(block);
+        chain.push(block.clone());
         self.chain_set.insert(Box::new(chain));
+        self.update_metrics_for_committed_block(block);
     }
 
     /// Returns the length of the non-finalized portion of the current best chain.
@@ -182,4 +186,34 @@ impl NonFinalizedState {
             .next_back()
             .map(|box_chain| box_chain.deref())
     }
+
+    /// Update the metrics after `block` is committed
+    fn update_metrics_for_committed_block(&self, block: Arc<Block>) {
+        let height = block.coinbase_height().unwrap();
+
+        metrics::counter!("state.memory.committed.block.count", 1);
+        metrics::gauge!("state.memory.committed.block.height", height.0 as _);
+
+        if self
+            .best_chain()
+            .unwrap()
+            .blocks
+            .iter()
+            .next_back()
+            .unwrap()
+            .1
+            == &block
+        {
+            metrics::counter!("state.memory.best.committed.block.count", 1);
+            metrics::gauge!("state.memory.best.committed.block.height", height.0 as _);
+        }
+
+        self.update_metrics_for_chains();
+    }
+
+    /// Update the metrics after `self.chain_set` is modified
+    fn update_metrics_for_chains(&self) {
+        metrics::gauge!("state.memory.chain.count", self.chain_set.len() as _);
+        metrics::gauge!("state.memory.best.chain.length", self.best_chain_len() as _);
+    }
 }
diff --git a/zebra-state/src/service/memory_state/queued_blocks.rs b/zebra-state/src/service/memory_state/queued_blocks.rs
index d8e74839..dc772185 100644
--- a/zebra-state/src/service/memory_state/queued_blocks.rs
+++ b/zebra-state/src/service/memory_state/queued_blocks.rs
@@ -49,6 +49,7 @@ impl QueuedBlocks {
         assert!(inserted, "hashes must be unique");
 
         tracing::trace!(num_blocks = %self.blocks.len(), %parent_hash, ?new_height,  "Finished queueing a new block");
+        self.update_metrics();
     }
 
     /// Dequeue and return all blocks that were waiting for the arrival of
@@ -73,6 +74,7 @@ impl QueuedBlocks {
         }
 
         tracing::trace!(num_blocks = %self.blocks.len(), "Finished dequeuing blocks waiting for parent hash",);
+        self.update_metrics();
 
         queued_children
     }
@@ -99,10 +101,24 @@ impl QueuedBlocks {
                 .expect("parent is present")
                 .remove(&hash);
         }
+
+        tracing::trace!(num_blocks = %self.blocks.len(), "Finished pruning blocks at or beneath the finalized tip height",);
+        self.update_metrics();
     }
 
     /// Return the queued block if it has already been registered
     pub fn get_mut(&mut self, hash: &block::Hash) -> Option<&mut QueuedBlock> {
         self.blocks.get_mut(&hash)
     }
+
+    /// Update metrics after the queue is modified
+    fn update_metrics(&self) {
+        if let Some(max_height) = self.by_height.keys().next_back() {
+            metrics::gauge!("state.memory.queued.max.height", max_height.0 as i64);
+        } else {
+            // use -1 as a sentinel value for "None", because 0 is a valid height
+            metrics::gauge!("state.memory.queued.max.height", -1);
+        }
+        metrics::gauge!("state.memory.queued.block.count", self.blocks.len() as _);
+    }
 }