storage: implement substore support

crates/bin/pd/src/main.rs

@@ -284,7 +284,7 @@ async fn main() -> anyhow::Result<()> {
             let mut rocks_path = home.clone();
             rocks_path.push("rocksdb");
 
-            let storage = Storage::load(rocks_path)
+            let storage = Storage::init(rocks_path)
                 .await
                 .context("Unable to initialize RocksDB storage")?;
 
@@ -654,7 +654,7 @@ async fn main() -> anyhow::Result<()> {
 
             tracing::info!("pruning JMT tree");
             export_path.push("rocksdb");
-            let export = Storage::load(export_path).await?;
+            let export = Storage::init(export_path).await?;
             let _ = StateDelta::new(export.latest_snapshot());
             // TODO:
             // - add utilities in `penumbra_storage` to prune a tree

crates/bin/pd/src/mempool.rs

@@ -27,7 +27,7 @@ use crate::{metrics, App};
 pub struct Mempool {
     queue: mpsc::Receiver<Message<Request, Response, tower::BoxError>>,
     app: App,
-    snapshot_rx: watch::Receiver<Snapshot>,
+    rx_snapshot: watch::Receiver<Snapshot>,
 }
 
 impl Mempool {
@@ -41,7 +41,7 @@ impl Mempool {
         Ok(Self {
             queue,
             app,
-            snapshot_rx,
+            rx_snapshot: snapshot_rx,
         })
     }
 
@@ -95,9 +95,9 @@ impl Mempool {
                 biased;
                 // Check whether the height has changed, which requires us to throw away our
                 // ephemeral mempool state, and create a new one based on the new state.
-                change = self.snapshot_rx.changed() => {
+                change = self.rx_snapshot.changed() => {
                     if let Ok(()) = change {
-                        let snapshot = self.snapshot_rx.borrow().clone();
+                        let snapshot = self.rx_snapshot.borrow().clone();
                         tracing::debug!(height = ?snapshot.version(), "resetting ephemeral mempool state");
                         self.app = App::new(snapshot).await?;
                     } else {

crates/bin/pd/src/upgrade.rs

@@ -23,7 +23,7 @@ pub async fn migrate(path_to_export: PathBuf, upgrade: Upgrade) -> anyhow::Resul
         Upgrade::Testnet60 => {
             let mut db_path = path_to_export.clone();
             db_path.push("rocksdb");
-            let storage = Storage::load(db_path).await?;
+            let storage = Storage::init(db_path).await?;
             let export_state = storage.latest_snapshot();
             let root_hash = export_state.root_hash().await.expect("can get root hash");
             let app_hash_pre_migration: AppHash = root_hash.into();

crates/narsil/narsil/src/bin/narsild.rs

@@ -132,7 +132,7 @@ async fn main() -> anyhow::Result<()> {
             let mut rocks_path = home.clone();
             rocks_path.push("rocksdb");
 
-            let storage = Storage::load(rocks_path)
+            let storage = Storage::init(rocks_path)
                 .await
                 .context("Unable to initialize RocksDB storage")?;
 

crates/narsil/narsil/src/ledger/mempool.rs

@@ -27,7 +27,7 @@ use crate::{ledger::app::App, metrics};
 pub struct Mempool {
     queue: mpsc::Receiver<Message<Request, Response, tower::BoxError>>,
     app: App,
-    snapshot_rx: watch::Receiver<Snapshot>,
+    rx_snapshot: watch::Receiver<Snapshot>,
 }
 
 impl Mempool {
@@ -41,7 +41,7 @@ impl Mempool {
         Self {
             queue,
             app,
-            snapshot_rx,
+            rx_snapshot: snapshot_rx,
         }
     }
 
@@ -95,9 +95,9 @@ impl Mempool {
                 biased;
                 // Check whether the height has changed, which requires us to throw away our
                 // ephemeral mempool state, and create a new one based on the new state.
-                change = self.snapshot_rx.changed() => {
+                change = self.rx_snapshot.changed() => {
                     if let Ok(()) = change {
-                        let snapshot = self.snapshot_rx.borrow().clone();
+                        let snapshot = self.rx_snapshot.borrow().clone();
                         tracing::debug!(height = ?snapshot.version(), "resetting ephemeral mempool state");
                         self.app = App::new(snapshot);
                     } else {

crates/storage/src/cache.rs

@@ -1,8 +1,11 @@
-use std::{any::Any, collections::BTreeMap};
+use std::{any::Any, collections::BTreeMap, sync::Arc};
 
 use tendermint::abci;
 
-use crate::StateWrite;
+use crate::{
+    store::{multistore::MultistoreConfig, substore::SubstoreConfig},
+    StateWrite,
+};
 
 /// A cache of changes to the state of the blockchain.
 ///
@@ -72,4 +75,32 @@ impl Cache {
     pub fn take_events(&mut self) -> Vec<abci::Event> {
         std::mem::take(&mut self.events)
     }
+
+    /// Consumes a `Cache` and returns a map of `SubstoreConfig` to `Cache` that
+    /// corresponds to changes belonging to each substore. The keys in each `Cache`
+    /// are truncated to remove the substore prefix.
+    pub fn shard_by_prefix(
+        self,
+        prefixes: &MultistoreConfig,
+    ) -> BTreeMap<Arc<SubstoreConfig>, Self> {
+        let mut changes_by_substore = BTreeMap::new();
+        for (key, some_value) in self.unwritten_changes.into_iter() {
+            let (truncated_key, substore_config) = prefixes.route_key_str(&key);
+            changes_by_substore
+                .entry(substore_config)
+                .or_insert_with(Cache::default)
+                .unwritten_changes
+                .insert(truncated_key.to_string(), some_value);
+        }
+
+        for (key, some_value) in self.nonverifiable_changes {
+            let (truncated_key, substore_config) = prefixes.route_key_bytes(&key);
+            changes_by_substore
+                .entry(substore_config)
+                .or_insert_with(Cache::default)
+                .nonverifiable_changes
+                .insert(truncated_key.to_vec(), some_value);
+        }
+        changes_by_substore
+    }
 }

crates/storage/src/lib.rs

@@ -47,6 +47,7 @@ mod read;
 mod snapshot;
 mod snapshot_cache;
 mod storage;
+mod store;
 #[cfg(test)]
 mod tests;
 mod utils;