refactor(rootstore): own nodestore logic (#1495)

As mentioned in #1470, there exists the following circular dependency: 

```mermaid
stateDiagram-v2
    s1: Revision Manager
    s2: RootStore

    s1 --> s2: Wants historical revision X
    s2 --> s1: Wants latest revision
```

This PR eliminates the circular dependency above by giving `RootStore` a
reference to the file backend, which allows it to construct revisions
without having to rely on the storage of another `NodeStore`
(previously, we needed the latest `NodeStore` whose storage we would
clone).

This PR does the following:
- `NodeStore` construction with a specified root hash + address now
takes in a storage reference rather than another `NodeStore.
- Revision construction is now handled by `RootStore` instead of the
RevisionManager
- Moves revision caching from the RevisionManager to `RootStore`
- Adds tests for `RootStore` caching logic

Note: there exists an earlier version of this PR (#1480). However,
#1480 (comment)
pointed out that having to pass in a `NodeStore` to
`NodeStore::with_root()` seemed odd. After some thought, I agree with
this point and found that passing in a storage reference instead removed
the need to have to share revisions between the `RevisionManager` and
`RootStore`. Hence, I went with creating a new PR rather than clobbering
the existing one.

Author: RodrigoVillar

firewood/src/db.rs

@@ -1047,8 +1047,11 @@ mod test {
         let latest_value = latest_revision.val(key).unwrap().unwrap();
         assert_eq!(new_value, latest_value.as_ref());
 
-        let node_store =
-            NodeStore::with_root(root_hash.into_hash_type(), root_address, latest_revision);
+        let node_store = NodeStore::with_root(
+            root_hash.into_hash_type(),
+            root_address,
+            latest_revision.get_storage(),
+        );
 
         let retrieved_value = node_store.val(key).unwrap().unwrap();
         assert_eq!(value, retrieved_value.as_ref());

firewood/src/manager.rs

@@ -14,22 +14,21 @@ use parking_lot::{Mutex, RwLock};
 use std::collections::{HashMap, VecDeque};
 use std::num::NonZero;
 use std::path::PathBuf;
-use std::sync::{Arc, OnceLock, Weak};
+use std::sync::{Arc, OnceLock};
 
 use firewood_storage::logger::{trace, warn};
 use metrics::gauge;
 use rayon::{ThreadPool, ThreadPoolBuilder};
 use typed_builder::TypedBuilder;
-use weak_table::WeakValueHashMap;
 
 use crate::merkle::Merkle;
 use crate::root_store::RootStore;
 use crate::v2::api::{ArcDynDbView, HashKey, OptionalHashKeyExt};
 
 pub use firewood_storage::CacheReadStrategy;
 use firewood_storage::{
-    BranchNode, Committed, FileBacked, FileIoError, HashedNodeReader, ImmutableProposal,
-    IntoHashType, NodeStore, TrieHash,
+    BranchNode, Committed, FileBacked, FileIoError, HashedNodeReader, ImmutableProposal, NodeStore,
+    TrieHash,
 };
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, TypedBuilder)]
@@ -69,7 +68,7 @@ pub struct ConfigManager {
     pub manager: RevisionManagerConfig,
 }
 
-type CommittedRevision = Arc<NodeStore<Committed, FileBacked>>;
+pub type CommittedRevision = Arc<NodeStore<Committed, FileBacked>>;
 type ProposedRevision = Arc<NodeStore<Arc<ImmutableProposal>, FileBacked>>;
 
 #[derive(Debug)]
@@ -83,7 +82,6 @@ pub(crate) struct RevisionManager {
     proposals: Mutex<Vec<ProposedRevision>>,
     // committing_proposals: VecDeque<Arc<ProposedImmutable>>,
     by_hash: RwLock<HashMap<TrieHash, CommittedRevision>>,
-    by_rootstore: Mutex<WeakValueHashMap<TrieHash, Weak<NodeStore<Committed, FileBacked>>>>,
     threadpool: OnceLock<ThreadPool>,
     root_store: Option<RootStore>,
 }
@@ -122,20 +120,20 @@ impl RevisionManager {
         // from opening the same database simultaneously
         fb.lock()?;
 
-        let root_store: Option<RootStore> = match config.root_store_dir {
-            Some(path) => Some(RootStore::new(path).map_err(RevisionManagerError::RootStoreError)?),
-            None => None,
-        };
-
         let storage = Arc::new(fb);
         let nodestore = Arc::new(NodeStore::open(storage.clone())?);
+        let root_store = config
+            .root_store_dir
+            .map(|path| RootStore::new(path, storage.clone()))
+            .transpose()
+            .map_err(RevisionManagerError::RootStoreError)?;
+
         let manager = Self {
             max_revisions: config.manager.max_revisions,
             historical: RwLock::new(VecDeque::from([nodestore.clone()])),
             by_hash: RwLock::new(Default::default()),
             proposals: Mutex::new(Default::default()),
             // committing_proposals: Default::default(),
-            by_rootstore: Mutex::new(WeakValueHashMap::new()),
             threadpool: OnceLock::new(),
             root_store,
         };
@@ -322,47 +320,28 @@ impl RevisionManager {
     /// Retrieve a committed revision by its root hash.
     /// To retrieve a revision involves a few steps:
     /// 1. Check the in-memory revision manager.
-    /// 2. Check the in-memory `RootStore` cache.
-    /// 3. Check the persistent `RootStore`.
+    /// 2. Check `RootStore` (if it exists).
     pub fn revision(&self, root_hash: HashKey) -> Result<CommittedRevision, RevisionManagerError> {
         // 1. Check the in-memory revision manager.
         if let Some(revision) = self.by_hash.read().get(&root_hash).cloned() {
             return Ok(revision);
         }
 
-        let mut cache_guard = self.by_rootstore.lock();
-
-        // 2. Check the in-memory `RootStore` cache.
-        if let Some(nodestore) = cache_guard.get(&root_hash) {
-            return Ok(nodestore);
-        }
-
-        // 3. Check the persistent `RootStore`.
-        // If the revision exists, get its root address and construct a NodeStore for it.
-        let root_address = match &self.root_store {
-            Some(store) => store
-                .get(&root_hash)
-                .map_err(RevisionManagerError::RootStoreError)?
+        // 2. Check `RootStore` (if it exists).
+        let root_store =
+            self.root_store
+                .as_ref()
                 .ok_or(RevisionManagerError::RevisionNotFound {
                     provided: root_hash.clone(),
-                })?,
-            None => {
-                return Err(RevisionManagerError::RevisionNotFound {
-                    provided: root_hash.clone(),
-                });
-            }
-        };
-
-        let nodestore = Arc::new(NodeStore::with_root(
-            root_hash.clone().into_hash_type(),
-            root_address,
-            self.current_revision(),
-        ));
-
-        // Cache the nodestore (stored as a weak reference).
-        cache_guard.insert(root_hash, nodestore.clone());
+                })?;
+        let revision = root_store
+            .get(&root_hash)
+            .map_err(RevisionManagerError::RootStoreError)?
+            .ok_or(RevisionManagerError::RevisionNotFound {
+                provided: root_hash.clone(),
+            })?;
 
-        Ok(nodestore)
+        Ok(revision)
     }
 
     pub fn root_hash(&self) -> Result<Option<HashKey>, RevisionManagerError> {

firewood/src/root_store.rs

@@ -2,10 +2,17 @@
 // See the file LICENSE.md for licensing terms.
 
 use fjall::{Config, Keyspace, PartitionCreateOptions, PartitionHandle, PersistMode};
-use std::path::Path;
+use parking_lot::Mutex;
+use std::{
+    path::Path,
+    sync::{Arc, Weak},
+};
+use weak_table::WeakValueHashMap;
 
 use derive_where::derive_where;
-use firewood_storage::{LinearAddress, TrieHash};
+use firewood_storage::{Committed, FileBacked, IntoHashType, LinearAddress, NodeStore, TrieHash};
+
+use crate::manager::CommittedRevision;
 
 const FJALL_PARTITION_NAME: &str = "firewood";
 
@@ -14,25 +21,36 @@ const FJALL_PARTITION_NAME: &str = "firewood";
 pub struct RootStore {
     keyspace: Keyspace,
     items: PartitionHandle,
+    storage: Arc<FileBacked>,
+    /// Cache of reconstructed revisions by hash.
+    revision_cache: Mutex<WeakValueHashMap<TrieHash, Weak<NodeStore<Committed, FileBacked>>>>,
 }
 
 impl RootStore {
     /// Creates or opens an instance of `RootStore`.
     ///
     /// Args:
-    /// - path: the directory where `RootStore` will write to.
+    /// - `path`: the directory where `RootStore` will write to.
+    /// - `storage`: the underlying store to create nodestores from.
     ///
     /// # Errors
     ///
     /// Will return an error if unable to create or open an instance of `RootStore`.
     pub fn new<P: AsRef<Path>>(
         path: P,
+        storage: Arc<FileBacked>,
     ) -> Result<RootStore, Box<dyn std::error::Error + Send + Sync>> {
         let keyspace = Config::new(path).open()?;
         let items =
             keyspace.open_partition(FJALL_PARTITION_NAME, PartitionCreateOptions::default())?;
+        let revision_cache = Mutex::new(WeakValueHashMap::new());
 
-        Ok(Self { keyspace, items })
+        Ok(Self {
+            keyspace,
+            items,
+            storage,
+            revision_cache,
+        })
     }
 
     /// `add_root` persists a revision's address to `RootStore`.
@@ -57,24 +75,127 @@ impl RootStore {
         Ok(())
     }
 
-    /// `get` returns the address of a revision.
+    /// `get` retrieves a committed revision by its hash.
+    ///
+    /// To retrieve a committed revision involves a few steps:
+    /// 1. Check if the committed revision is cached.
+    /// 2. If the committed revision is not cached, query the underlying
+    ///    datastore for the revision's root address.
+    /// 3. Construct the committed revision.
     ///
     /// Args:
     /// - hash: the hash of the revision
     ///
     /// # Errors
     ///
-    ///  Will return an error if unable to query the underlying datastore.
+    ///  Will return an error if unable to query the underlying datastore or if
+    ///  the stored address is invalid.
+    ///
+    /// # Panics
+    ///
+    ///  Will panic if the latest revision does not exist.
     pub fn get(
         &self,
         hash: &TrieHash,
-    ) -> Result<Option<LinearAddress>, Box<dyn std::error::Error + Send + Sync>> {
+    ) -> Result<Option<CommittedRevision>, Box<dyn std::error::Error + Send + Sync>> {
+        // Obtain the lock to prevent multiple threads from caching the same result.
+        let mut revision_cache = self.revision_cache.lock();
+
+        // 1. Check if the committed revision is cached.
+        if let Some(v) = revision_cache.get(hash) {
+            return Ok(Some(v));
+        }
+
+        // 2. If the committed revision is not cached, query the underlying
+        //    datastore for the revision's root address.
         let Some(v) = self.items.get(**hash)? else {
             return Ok(None);
         };
 
         let array: [u8; 8] = v.as_ref().try_into()?;
+        let addr = LinearAddress::new(u64::from_be_bytes(array))
+            .ok_or("invalid address: empty address")?;
+
+        // 3. Construct the committed revision.
+        let nodestore = Arc::new(NodeStore::with_root(
+            hash.clone().into_hash_type(),
+            addr,
+            self.storage.clone(),
+        ));
+
+        // Cache for future lookups.
+        revision_cache.insert(hash.clone(), nodestore.clone());
+
+        Ok(Some(nodestore))
+    }
+}
+
+#[cfg(test)]
+#[allow(clippy::unwrap_used)]
+mod tests {
+    use super::*;
+    use firewood_storage::{CacheReadStrategy, FileBacked, NodeStore};
+    use std::num::NonZero;
+    use std::sync::Arc;
+
+    #[test]
+    fn test_cache_hit() {
+        let tmpdir = tempfile::tempdir().unwrap();
+
+        let db_path = tmpdir.as_ref().join("testdb");
+        let file_backed = Arc::new(
+            FileBacked::new(
+                db_path,
+                NonZero::new(1024).unwrap(),
+                NonZero::new(1024).unwrap(),
+                false,
+                true,
+                CacheReadStrategy::WritesOnly,
+            )
+            .unwrap(),
+        );
+
+        let root_store_dir = tmpdir.as_ref().join("root_store");
+        let root_store = RootStore::new(root_store_dir, file_backed.clone()).unwrap();
+
+        // Create a revision to cache.
+        let revision = Arc::new(NodeStore::new_empty_committed(file_backed.clone()));
+
+        let hash = TrieHash::from_bytes([1; 32]);
+        root_store
+            .revision_cache
+            .lock()
+            .insert(hash.clone(), revision.clone());
+
+        // Since the underlying datastore is empty, this should get the revision
+        // from the cache.
+        let retrieved_revision = root_store.get(&hash).unwrap().unwrap();
+
+        assert!(Arc::ptr_eq(&revision, &retrieved_revision));
+    }
+
+    #[test]
+    fn test_nonexistent_revision() {
+        let tmpdir = tempfile::tempdir().unwrap();
+
+        let db_path = tmpdir.as_ref().join("testdb");
+        let file_backed = Arc::new(
+            FileBacked::new(
+                db_path,
+                NonZero::new(1024).unwrap(),
+                NonZero::new(1024).unwrap(),
+                false,
+                true,
+                CacheReadStrategy::WritesOnly,
+            )
+            .unwrap(),
+        );
+
+        let root_store_dir = tmpdir.as_ref().join("root_store");
+        let root_store = RootStore::new(root_store_dir, file_backed.clone()).unwrap();
 
-        Ok(LinearAddress::new(u64::from_be_bytes(array)))
+        // Try to get a hash that doesn't exist in the cache nor in the underlying datastore.
+        let nonexistent_hash = TrieHash::from_bytes([1; 32]);
+        assert!(root_store.get(&nonexistent_hash).unwrap().is_none());
     }
 }

storage/src/nodestore/mod.rs

@@ -160,21 +160,15 @@ impl<S: ReadableStorage> NodeStore<Committed, S> {
     ///
     /// This constructor is used when you have an existing root node at a known
     /// address and hash, typically when reconstructing a [`NodeStore`] from
-    /// a committed state. The `latest_nodestore` provides access to the underlying
-    /// storage backend containing the persisted trie data.
+    /// a committed state.
     ///
     /// ## Panics
     ///
     /// Panics in debug builds if the hash of the node at `root_address` does
     /// not equal `root_hash`.
     #[must_use]
-    pub fn with_root(
-        root_hash: HashType,
-        root_address: LinearAddress,
-        latest_nodestore: Arc<NodeStore<Committed, S>>,
-    ) -> Self {
+    pub fn with_root(root_hash: HashType, root_address: LinearAddress, storage: Arc<S>) -> Self {
         let header = NodeStoreHeader::with_root(Some(root_address));
-        let storage = latest_nodestore.storage.clone();
 
         let nodestore = NodeStore {
             header,
@@ -201,6 +195,15 @@ impl<S: ReadableStorage> NodeStore<Committed, S> {
     }
 }
 
+impl<S: ReadableStorage> NodeStore<Committed, S> {
+    /// Get the underlying storage for a `NodeStore`.
+    #[cfg(any(test, feature = "test_utils"))]
+    #[must_use]
+    pub fn get_storage(&self) -> Arc<S> {
+        self.storage.clone()
+    }
+}
+
 /// Some nodestore kinds implement Parentable.
 ///
 /// This means that the nodestore can have children.