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

Added a gauge to track the number of unpersisted nodes at any time.
These can get used later by the revision manager to decide whether or
not to persist a proposal (among other possible strategies).

It's also helpful to know how many nodes are still in memory that could
be persisted.

When increasing this counter, we do it all at once, but we decrement one
by one (or by the number of entries from the kernel's completion queue).
This is because the writes are slower than the adds, and there is some
integer to floating point math that happens each time the counter is
adjusted.

Author: rkuris

storage/src/nodestore/hash.rs

@@ -68,15 +68,16 @@ where
     }
 
     /// Hashes `node`, which is at the given `path_prefix`, and its children recursively.
-    /// Returns the hashed node and its hash.
+    /// Returns the hashed node, its hash, and the count of nodes processed.
     pub(super) fn hash_helper(
         #[cfg(feature = "ethhash")] &self,
         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")]
@@ -157,10 +158,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:?}");
@@ -188,6 +192,6 @@ 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))
     }
 }

storage/src/nodestore/mod.rs

@@ -49,9 +49,11 @@ use crate::logger::trace;
 use crate::node::branch::ReadSerializable as _;
 use arc_swap::ArcSwap;
 use arc_swap::access::DynAccess;
+use metrics::gauge;
 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};
@@ -125,6 +127,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,
         };
@@ -154,6 +157,7 @@ impl<S: ReadableStorage> NodeStore<Committed, S> {
                 deleted: Box::default(),
                 root_hash: None,
                 root: None,
+                unwritten_nodes: AtomicUsize::new(0),
             },
         })
     }
@@ -354,11 +358,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)]
@@ -390,6 +410,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 {
@@ -407,6 +429,17 @@ 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)]
+            gauge!("firewood.nodes.unwritten").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.
@@ -465,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,
         }
     }
 }
@@ -499,6 +541,7 @@ impl NodeStore<Arc<ImmutableProposal>, FileBacked> {
                 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(),
         }
@@ -524,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,
         };
@@ -536,20 +580,28 @@ 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)]
+        gauge!("firewood.nodes.unwritten").increment(unwritten_count as f64);
+
         Ok(nodestore)
     }
 }

storage/src/nodestore/persist.rs

@@ -31,7 +31,7 @@ use std::iter::FusedIterator;
 
 use crate::linear::FileIoError;
 use coarsetime::Instant;
-use metrics::counter;
+use metrics::{counter, gauge};
 
 #[cfg(feature = "io-uring")]
 use crate::logger::trace;
@@ -229,6 +229,9 @@ impl NodeStore<Committed, FileBacked> {
                 .write(persisted_address.get(), serialized.as_slice())?;
             node.persist_at(persisted_address);
 
+            // Decrement gauge immediately after node is written to storage
+            gauge!("firewood.nodes.unwritten").decrement(1.0);
+
             // Move the arc to a vector of persisted nodes for caching
             // we save them so we don't have to lock the cache while we write them
             // If we ever persist out of band, we might have a race condition, so
@@ -272,6 +275,11 @@ impl NodeStore<Committed, FileBacked> {
         // Finally persist the header
         self.flush_header()?;
 
+        // Reset unwritten nodes counter to zero since all nodes are now persisted
+        self.kind
+            .unwritten_nodes
+            .store(0, std::sync::atomic::Ordering::Relaxed);
+
         Ok(())
     }
 
@@ -292,11 +300,13 @@ impl NodeStore<Committed, FileBacked> {
         }
 
         /// Helper function to handle completion queue entries and check for errors
+        /// Returns the number of completed operations
         fn handle_completion_queue(
             storage: &FileBacked,
             completion_queue: io_uring::cqueue::CompletionQueue<'_>,
             saved_pinned_buffers: &mut [PinnedBufferEntry],
-        ) -> Result<(), FileIoError> {
+        ) -> Result<usize, FileIoError> {
+            let mut completed_count = 0;
             for entry in completion_queue {
                 let item = entry.user_data() as usize;
                 let pbe = saved_pinned_buffers
@@ -322,8 +332,9 @@ impl NodeStore<Committed, FileBacked> {
                     ));
                 }
                 pbe.offset = None;
+                completed_count += 1;
             }
-            Ok(())
+            Ok(completed_count)
         }
 
         const RINGSIZE: usize = FileBacked::RINGSIZE as usize;
@@ -404,11 +415,16 @@ impl NodeStore<Committed, FileBacked> {
                 })?;
                 let completion_queue = ring.completion();
                 trace!("competion queue length: {}", completion_queue.len());
-                handle_completion_queue(
+                let completed_writes = handle_completion_queue(
                     &self.storage,
                     completion_queue,
                     &mut saved_pinned_buffers,
                 )?;
+
+                // Decrement gauge for writes that have actually completed
+                if completed_writes > 0 {
+                    gauge!("firewood.nodes.unwritten").decrement(completed_writes as f64);
+                }
             }
 
             // Mark node as persisted and collect for cache
@@ -424,7 +440,13 @@ impl NodeStore<Committed, FileBacked> {
                 .file_io_error(e, 0, Some("io-uring final submit_and_wait".to_string()))
         })?;
 
-        handle_completion_queue(&self.storage, ring.completion(), &mut saved_pinned_buffers)?;
+        let final_completed_writes =
+            handle_completion_queue(&self.storage, ring.completion(), &mut saved_pinned_buffers)?;
+
+        // Decrement gauge for final batch of writes that completed
+        if final_completed_writes > 0 {
+            gauge!("firewood.nodes.unwritten").decrement(final_completed_writes as f64);
+        }
 
         debug_assert!(
             !saved_pinned_buffers.iter().any(|pbe| pbe.offset.is_some()),