feat(deferred-persist): Part 1: unpersisted gauge

benchmark/Grafana-dashboard.json

@@ -845,6 +845,113 @@
       ],
       "title": "Insert Merkle Ops by Type",
       "type": "timeseries"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "${DS_PROMETHEUS}"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "custom": {
+            "axisBorderShow": false,
+            "axisCenteredZero": false,
+            "axisColorMode": "text",
+            "axisLabel": "",
+            "axisPlacement": "auto",
+            "barAlignment": 0,
+            "barWidthFactor": 0.6,
+            "drawStyle": "line",
+            "fillOpacity": 0,
+            "gradientMode": "none",
+            "hideFrom": {
+              "legend": false,
+              "tooltip": false,
+              "viz": false
+            },
+            "insertNulls": false,
+            "lineInterpolation": "linear",
+            "lineWidth": 1,
+            "pointSize": 5,
+            "scaleDistribution": {
+              "type": "linear"
+            },
+            "showPoints": "auto",
+            "spanNulls": false,
+            "stacking": {
+              "group": "A",
+              "mode": "none"
+            },
+            "thresholdsStyle": {
+              "mode": "off"
+            }
+          },
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "yellow",
+                "value": 1000
+              },
+              {
+                "color": "red",
+                "value": 5000
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 8,
+        "w": 12,
+        "x": 12,
+        "y": 31
+      },
+      "id": 13,
+      "options": {
+        "legend": {
+          "calcs": [
+            "min",
+            "mean",
+            "max",
+            "last"
+          ],
+          "displayMode": "table",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "hideZeros": false,
+          "mode": "single",
+          "sort": "none"
+        }
+      },
+      "pluginVersion": "11.5.1",
+      "targets": [
+        {
+          "editorMode": "code",
+          "expr": "firewood_nodes_unwritten",
+          "legendFormat": "Unwritten Nodes",
+          "range": true,
+          "refId": "A",
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          }
+        }
+      ],
+      "title": "Unwritten Nodes (In Memory)",
+      "type": "timeseries"
     }
   ],
   "refresh": "10s",
@@ -863,4 +970,4 @@
   "uid": "adxfhfmwx5ypsc",
   "version": 12,
   "weekStart": ""
-}
+}

storage/src/macros.rs

@@ -1,7 +1,8 @@
 // Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE.md for licensing terms.
+
 #[macro_export]
-/// Macro to register and use a metric with description and labels.
+/// Macro to register and use a counter metric with description and labels.
 /// This macro is a wrapper around the `metrics` crate's `counter!` and `describe_counter!`
 /// macros. It ensures that the description is registered just once.
 ///
@@ -32,3 +33,36 @@ macro_rules! firewood_counter {
         }
     };
 }
+
+#[macro_export]
+/// Macro to register and use a gauge metric with description and labels.
+/// This macro is a wrapper around the `metrics` crate's `gauge!` and `describe_gauge!`
+/// macros. It ensures that the description is registered just once.
+///
+/// Usage:
+///   `firewood_gauge!("metric_name", "description")`
+///   `firewood_gauge!("metric_name", "description", "label" => "value")`
+///
+/// Call `.increment(val)` or `.decrement(val)` on the result as appropriate.
+macro_rules! firewood_gauge {
+    // With labels
+    ($name:expr, $desc:expr, $($labels:tt)+) => {
+        {
+            static ONCE: std::sync::Once = std::sync::Once::new();
+            ONCE.call_once(|| {
+                metrics::describe_counter!($name, $desc);
+            });
+            metrics::gauge!($name, $($labels)+)
+        }
+    };
+    // No labels
+    ($name:expr, $desc:expr) => {
+        {
+            static ONCE: std::sync::Once = std::sync::Once::new();
+            ONCE.call_once(|| {
+                metrics::describe_counter!($name, $desc);
+            });
+            metrics::gauge!($name)
+        }
+    };
+}

storage/src/nodestore/hash.rs

@@ -78,9 +78,10 @@ where
         mut node: Node,
         path_prefix: &mut Path,
         #[cfg(feature = "ethhash")] fake_root_extra_nibble: Option<u8>,
-    ) -> Result<(MaybePersistedNode, HashType), FileIoError> {
+    ) -> Result<(MaybePersistedNode, HashType, usize), FileIoError> {
         // If this is a branch, find all unhashed children and recursively hash them.
         trace!("hashing {node:?} at {path_prefix:?}");
+        let mut nodes_processed = 1usize; // Count this node
         if let Node::Branch(ref mut b) = node {
             // special case code for ethereum hashes at the account level
             #[cfg(feature = "ethhash")]
@@ -160,10 +161,13 @@ where
                 path_prefix.0.push(nibble as u8);
 
                 #[cfg(feature = "ethhash")]
-                let (child_node, child_hash) =
+                let (child_node, child_hash, child_count) =
                     self.hash_helper(child_node, path_prefix, make_fake_root)?;
                 #[cfg(not(feature = "ethhash"))]
-                let (child_node, child_hash) = Self::hash_helper(child_node, path_prefix)?;
+                let (child_node, child_hash, child_count) =
+                    Self::hash_helper(child_node, path_prefix)?;
+
+                nodes_processed = nodes_processed.saturating_add(child_count);
 
                 *child = Some(Child::MaybePersisted(child_node, child_hash));
                 trace!("child now {child:?}");
@@ -191,7 +195,7 @@ where
         #[cfg(not(feature = "ethhash"))]
         let hash = hash_node(&node, path_prefix);
 
-        Ok((SharedNode::new(node).into(), hash))
+        Ok((SharedNode::new(node).into(), hash, nodes_processed))
     }
 
     #[cfg(feature = "ethhash")]

storage/src/nodestore/mod.rs

@@ -43,13 +43,15 @@ pub(crate) mod hash;
 pub(crate) mod header;
 pub(crate) mod persist;
 
+use crate::firewood_gauge;
 use crate::logger::trace;
 use crate::node::branch::ReadSerializable as _;
 use arc_swap::ArcSwap;
 use arc_swap::access::DynAccess;
 use smallvec::SmallVec;
 use std::fmt::Debug;
 use std::io::{Error, ErrorKind};
+use std::sync::atomic::AtomicUsize;
 
 // Re-export types from alloc module
 pub use alloc::{AreaIndex, LinearAddress, NodeAllocator};
@@ -121,6 +123,7 @@ impl<S: ReadableStorage> NodeStore<Committed, S> {
                 deleted: Box::default(),
                 root_hash: None,
                 root: header.root_address().map(Into::into),
+                unwritten_nodes: AtomicUsize::new(0),
             },
             storage,
         };
@@ -150,6 +153,7 @@ impl<S: ReadableStorage> NodeStore<Committed, S> {
                 deleted: Box::default(),
                 root_hash: None,
                 root: None,
+                unwritten_nodes: AtomicUsize::new(0),
             },
         })
     }
@@ -350,11 +354,27 @@ pub trait RootReader {
 }
 
 /// A committed revision of a merkle trie.
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 pub struct Committed {
     deleted: Box<[MaybePersistedNode]>,
     root_hash: Option<TrieHash>,
     root: Option<MaybePersistedNode>,
+    /// TODO: No readers of this variable yet - will be used for tracking unwritten nodes in committed revisions
+    unwritten_nodes: AtomicUsize,
+}
+
+impl Clone for Committed {
+    fn clone(&self) -> Self {
+        Self {
+            deleted: self.deleted.clone(),
+            root_hash: self.root_hash.clone(),
+            root: self.root.clone(),
+            unwritten_nodes: AtomicUsize::new(
+                self.unwritten_nodes
+                    .load(std::sync::atomic::Ordering::Relaxed),
+            ),
+        }
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -386,6 +406,8 @@ pub struct ImmutableProposal {
     root_hash: Option<TrieHash>,
     /// The root node, either in memory or on disk
     root: Option<MaybePersistedNode>,
+    /// The number of unwritten nodes in this proposal
+    unwritten_nodes: usize,
 }
 
 impl ImmutableProposal {
@@ -403,6 +425,21 @@ impl ImmutableProposal {
     }
 }
 
+impl Drop for ImmutableProposal {
+    fn drop(&mut self) {
+        // When an immutable proposal is dropped without being committed,
+        // decrement the gauge to reflect that these nodes will never be written
+        if self.unwritten_nodes > 0 {
+            #[allow(clippy::cast_precision_loss)]
+            firewood_gauge!(
+                "firewood.nodes.unwritten",
+                "current number of unwritten nodes"
+            )
+            .decrement(self.unwritten_nodes as f64);
+        }
+    }
+}
+
 /// Contains the state of a revision of a merkle trie.
 ///
 /// The first generic parameter is the type of the revision, which supports reading nodes from parent proposals.
@@ -461,14 +498,23 @@ impl<T: Into<NodeStoreParent>, S: ReadableStorage> From<NodeStore<T, S>>
 /// Commit a proposal to a new revision of the trie
 impl<S: WritableStorage> From<NodeStore<ImmutableProposal, S>> for NodeStore<Committed, S> {
     fn from(val: NodeStore<ImmutableProposal, S>) -> Self {
+        let NodeStore {
+            header,
+            kind,
+            storage,
+        } = val;
+        // Use ManuallyDrop to prevent the Drop impl from running since we're committing
+        let kind = std::mem::ManuallyDrop::new(kind);
+
         NodeStore {
-            header: val.header,
+            header,
             kind: Committed {
-                deleted: val.kind.deleted,
-                root_hash: val.kind.root_hash,
-                root: val.kind.root,
+                deleted: kind.deleted.clone(),
+                root_hash: kind.root_hash.clone(),
+                root: kind.root.clone(),
+                unwritten_nodes: AtomicUsize::new(kind.unwritten_nodes),
             },
-            storage: val.storage,
+            storage,
         }
     }
 }
@@ -495,6 +541,7 @@ impl<S: WritableStorage> NodeStore<Arc<ImmutableProposal>, S> {
                 deleted: self.kind.deleted.clone(),
                 root_hash: self.kind.root_hash.clone(),
                 root: self.kind.root.clone(),
+                unwritten_nodes: AtomicUsize::new(self.kind.unwritten_nodes),
             },
             storage: self.storage.clone(),
         }
@@ -520,6 +567,7 @@ impl<S: ReadableStorage> TryFrom<NodeStore<MutableProposal, S>>
                 parent: Arc::new(ArcSwap::new(Arc::new(kind.parent))),
                 root_hash: None,
                 root: None,
+                unwritten_nodes: 0,
             }),
             storage,
         };
@@ -532,20 +580,32 @@ impl<S: ReadableStorage> TryFrom<NodeStore<MutableProposal, S>>
 
         // Hashes the trie and returns the address of the new root.
         #[cfg(feature = "ethhash")]
-        let (root, root_hash) = nodestore.hash_helper(root, &mut Path::new(), None)?;
+        let (root, root_hash, unwritten_count) =
+            nodestore.hash_helper(root, &mut Path::new(), None)?;
         #[cfg(not(feature = "ethhash"))]
-        let (root, root_hash) =
+        let (root, root_hash, unwritten_count) =
             NodeStore::<MutableProposal, S>::hash_helper(root, &mut Path::new())?;
 
         let immutable_proposal =
             Arc::into_inner(nodestore.kind).expect("no other references to the proposal");
+        // Use ManuallyDrop to prevent Drop from running since we're replacing the proposal
+        let immutable_proposal = std::mem::ManuallyDrop::new(immutable_proposal);
         nodestore.kind = Arc::new(ImmutableProposal {
-            deleted: immutable_proposal.deleted,
-            parent: immutable_proposal.parent,
+            deleted: immutable_proposal.deleted.clone(),
+            parent: immutable_proposal.parent.clone(),
             root_hash: Some(root_hash.into_triehash()),
             root: Some(root),
+            unwritten_nodes: unwritten_count,
         });
 
+        // Track unwritten nodes in metrics
+        #[allow(clippy::cast_precision_loss)]
+        firewood_gauge!(
+            "firewood.nodes.unwritten",
+            "current number of unwritten nodes"
+        )
+        .increment(unwritten_count as f64);
+
         Ok(nodestore)
     }
 }