component: better test infrastructure

component/src/app/mod.rs

@@ -4,7 +4,7 @@ use anyhow::Result;
 use penumbra_chain::params::FmdParameters;
 use penumbra_chain::{genesis, AppHash, StateWriteExt as _};
 use penumbra_proto::{Protobuf, StateWriteProto};
-use penumbra_storage::{State, Storage};
+use penumbra_storage::{ArcStateExt, State, Storage};
 use penumbra_transaction::Transaction;
 use tendermint::abci::{self, types::ValidatorUpdate};
 use tracing::instrument;
@@ -71,9 +71,10 @@ impl App {
         &mut self,
         begin_block: &abci::request::BeginBlock,
     ) -> Vec<abci::Event> {
-        let state =
-            Arc::get_mut(&mut self.state).expect("state Arc should not be referenced elsewhere");
-        let mut state_tx = state.begin_transaction();
+        let mut state_tx = self
+            .state
+            .try_begin_transaction()
+            .expect("state Arc should not be referenced elsewhere");
 
         // store the block height
         state_tx.put_block_height(begin_block.header.height.into());
@@ -104,12 +105,12 @@ impl App {
         tx.check_stateless(tx.clone())?;
         tx.check_stateful(self.state.clone(), tx.clone()).await?;
 
-        // We need to get a mutable reference to the State here, so we use
-        // `Arc::get_mut`. At this point, the stateful checks should have completed,
+        // At this point, the stateful checks should have completed,
         // leaving us with exclusive access to the Arc<State>.
-        let state =
-            Arc::get_mut(&mut self.state).expect("state Arc should not be referenced elsewhere");
-        let mut state_tx = state.begin_transaction();
+        let mut state_tx = self
+            .state
+            .try_begin_transaction()
+            .expect("state Arc should not be referenced elsewhere");
         tx.execute(&mut state_tx).await?;
 
         // At this point, we've completed execution successfully with no errors,
@@ -120,9 +121,10 @@ impl App {
 
     #[instrument(skip(self, end_block))]
     pub async fn end_block(&mut self, end_block: &abci::request::EndBlock) -> Vec<abci::Event> {
-        let state =
-            Arc::get_mut(&mut self.state).expect("state Arc should not be referenced elsewhere");
-        let mut state_tx = state.begin_transaction();
+        let mut state_tx = self
+            .state
+            .try_begin_transaction()
+            .expect("state Arc should not be referenced elsewhere");
 
         Staking::end_block(&mut state_tx, end_block).await;
         IBCComponent::end_block(&mut state_tx, end_block).await;
@@ -139,8 +141,6 @@ impl App {
     ///
     /// This method also resets `self` as if it were constructed
     /// as an empty state over top of the newly written storage.
-    ///
-    /// TODO: why does this return Result?
     #[instrument(skip(self, storage))]
     pub async fn commit(&mut self, storage: Storage) -> AppHash {
         // We need to extract the State we've built up to commit it.  Fill in a dummy state.

component/src/ibc/component/client.rs

@@ -485,52 +485,38 @@ mod tests {
     use std::sync::Arc;
 
     use crate::action_handler::ActionHandler;
+    use crate::TempStorageExt;
 
     use super::*;
     use ibc_proto::ibc::core::client::v1::MsgCreateClient as RawMsgCreateClient;
     use ibc_proto::ibc::core::client::v1::MsgUpdateClient as RawMsgUpdateClient;
     use penumbra_chain::StateWriteExt as _;
     use penumbra_proto::core::ibc::v1alpha1::{ibc_action::Action as IbcActionInner, IbcAction};
     use penumbra_proto::Message;
-    use penumbra_storage::Storage;
-    use penumbra_tct as tct;
-    use penumbra_transaction::{Action, Transaction, TransactionBody};
-    use tempfile::tempdir;
+    use penumbra_storage::{ArcStateExt, TempStorage};
+    use penumbra_transaction::Transaction;
     use tendermint::Time;
 
     // test that we can create and update a light client.
     #[tokio::test]
-    async fn test_create_and_update_light_client() {
+    async fn test_create_and_update_light_client() -> anyhow::Result<()> {
         // create a storage backend for testing
-        let dir = tempdir().unwrap();
-        let file_path = dir.path().join("ibc-testing.db");
+        let storage = TempStorage::new().await?.apply_default_genesis().await?;
 
-        let storage = Storage::load(file_path.clone()).await.unwrap();
         let mut state = Arc::new(storage.latest_state());
-        let state_mut =
-            Arc::get_mut(&mut state).expect("state Arc should not be referenced elsewhere");
-        let mut state_tx = state_mut.begin_transaction();
 
-        // init chain should result in client counter = 0
-        let genesis_state = genesis::AppState::default();
-        let timestamp = Time::parse_from_rfc3339("2022-02-11T17:30:50.425417198Z").unwrap();
+        // Light client verification is time-dependent.  In practice, the latest
+        // (consensus) time will be delivered in each BeginBlock and written
+        // into the state.  Here, set the block timestamp manually so it's
+        // available to the unit test.
+        let timestamp = Time::parse_from_rfc3339("2022-02-11T17:30:50.425417198Z")?;
+        let mut state_tx = state.try_begin_transaction().unwrap();
         state_tx.put_block_timestamp(timestamp);
-        state_tx.put_block_height(0);
-        Ics2Client::init_chain(&mut state_tx, &genesis_state).await;
-
         state_tx.apply();
-        storage
-            .commit(Arc::try_unwrap(state).unwrap())
-            .await
-            .unwrap();
-
-        let mut state = Arc::new(storage.latest_state());
 
         // base64 encoded MsgCreateClient that was used to create the currently in-use Stargaze
         // light client on the cosmos hub:
         // https://cosmos.bigdipper.live/transactions/13C1ECC54F088473E2925AD497DDCC092101ADE420BC64BADE67D34A75769CE9
-        //
-        //
         let msg_create_client_stargaze_raw =
             base64::decode(include_str!("../../ibc/test/create_client.msg").replace('\n', ""))
                 .unwrap();
@@ -549,87 +535,37 @@ mod tests {
         let create_client_action = IbcAction {
             action: Some(IbcActionInner::CreateClient(msg_create_stargaze_client)),
         };
-        let create_client_tx = Arc::new(Transaction {
-            transaction_body: TransactionBody {
-                actions: vec![Action::IBCAction(create_client_action)],
-                expiry_height: 0,
-                chain_id: "".to_string(),
-                fee: Default::default(),
-                fmd_clues: vec![],
-                memo: None,
-            },
-            anchor: tct::Tree::new().root(),
-            binding_sig: [0u8; 64].into(),
-        });
-
         let update_client_action = IbcAction {
             action: Some(IbcActionInner::UpdateClient(msg_update_stargaze_client)),
         };
-        let update_client_tx = Arc::new(Transaction {
-            transaction_body: TransactionBody {
-                actions: vec![Action::IBCAction(update_client_action)],
-                expiry_height: 0,
-                chain_id: "".to_string(),
-                fee: Default::default(),
-                fmd_clues: vec![],
-                memo: None,
-            },
-            binding_sig: [0u8; 64].into(),
-            anchor: tct::Tree::new().root(),
-        });
-
-        if let Action::IBCAction(inner_action) = create_client_tx.actions().collect::<Vec<_>>()[0] {
-            inner_action
-                .check_stateless(create_client_tx.clone())
-                .unwrap();
-            inner_action
-                .check_stateful(state.clone(), create_client_tx.clone())
-                .await
-                .unwrap();
-            // execute (save client)
-            let state_mut =
-                Arc::get_mut(&mut state).expect("state Arc should not be referenced elsewhere");
-            let mut state_tx = state_mut.begin_transaction();
-            inner_action.execute(&mut state_tx).await.unwrap();
-            state_tx.apply();
-            storage
-                .commit(Arc::try_unwrap(state).unwrap())
-                .await
-                .unwrap();
-        } else {
-            panic!("expected ibc action");
-        }
 
-        let mut state = Arc::new(storage.latest_state());
-        assert_eq!(state.clone().client_counter().await.unwrap().0, 1);
+        // The ActionHandler trait provides the transaction the action was part
+        // of as context available during verification.  This is used, for instance,
+        // to allow spend and output proofs to access the (transaction-wide) anchor.
+        // Since the context is not used by the IBC action handlers, we can pass a dummy transaction.
+        let dummy_context = Arc::new(Transaction::default());
+
+        create_client_action.check_stateless(dummy_context.clone())?;
+        create_client_action
+            .check_stateful(state.clone(), dummy_context.clone())
+            .await?;
+        let mut state_tx = state.try_begin_transaction().unwrap();
+        create_client_action.execute(&mut state_tx).await?;
+        state_tx.apply();
 
-        // now try update client
-        if let Action::IBCAction(inner_action) = update_client_tx.actions().collect::<Vec<_>>()[0] {
-            inner_action
-                .check_stateless(update_client_tx.clone())
-                .unwrap();
-            // verify the ClientUpdate proof
-            inner_action
-                .check_stateful(state.clone(), update_client_tx.clone())
-                .await
-                .unwrap();
-            // save the next tm state
-            let state_mut =
-                Arc::get_mut(&mut state).expect("state Arc should not be referenced elsewhere");
-            let mut state_tx = state_mut.begin_transaction();
-            inner_action.execute(&mut state_tx).await.unwrap();
-            state_tx.apply();
-            storage
-                .commit(Arc::try_unwrap(state).unwrap())
-                .await
-                .unwrap();
-        } else {
-            panic!("expected ibc action");
-        }
+        // Check that state reflects +1 client apps registered.
+        assert_eq!(state.client_counter().await.unwrap().0, 1);
 
-        let mut state = Arc::new(storage.latest_state());
+        // Now we update the client and confirm that the update landed in state.
+        update_client_action.check_stateless(dummy_context.clone())?;
+        update_client_action
+            .check_stateful(state.clone(), dummy_context.clone())
+            .await?;
+        let mut state_tx = state.try_begin_transaction().unwrap();
+        update_client_action.execute(&mut state_tx).await?;
+        state_tx.apply();
 
-        // try one more client update
+        // We've had one client update, yes. What about second client update?
         // https://cosmos.bigdipper.live/transactions/ED217D360F51E622859F7B783FEF98BDE3544AA32BBD13C6C77D8D0D57A19FFD
         let msg_update_second =
             base64::decode(include_str!("../../ibc/test/update_client_2.msg").replace('\n', ""))
@@ -640,42 +576,15 @@ mod tests {
         let second_update_client_action = IbcAction {
             action: Some(IbcActionInner::UpdateClient(second_update)),
         };
-        let second_update_client_tx = Arc::new(Transaction {
-            transaction_body: TransactionBody {
-                actions: vec![Action::IBCAction(second_update_client_action)],
-                expiry_height: 0,
-                chain_id: "".to_string(),
-                fee: Default::default(),
-                fmd_clues: vec![],
-                memo: None,
-            },
-            anchor: tct::Tree::new().root(),
-            binding_sig: [0u8; 64].into(),
-        });
-
-        if let Action::IBCAction(inner_action) =
-            second_update_client_tx.actions().collect::<Vec<_>>()[0]
-        {
-            inner_action
-                .check_stateless(second_update_client_tx.clone())
-                .unwrap();
-            // verify the ClientUpdate proof
-            inner_action
-                .check_stateful(state.clone(), second_update_client_tx.clone())
-                .await
-                .unwrap();
-            // save the next tm state
-            let state_mut =
-                Arc::get_mut(&mut state).expect("state Arc should not be referenced elsewhere");
-            let mut state_tx = state_mut.begin_transaction();
-            inner_action.execute(&mut state_tx).await.unwrap();
-            state_tx.apply();
-            storage
-                .commit(Arc::try_unwrap(state).unwrap())
-                .await
-                .unwrap();
-        } else {
-            panic!("expected ibc action");
-        }
+
+        second_update_client_action.check_stateless(dummy_context.clone())?;
+        second_update_client_action
+            .check_stateful(state.clone(), dummy_context.clone())
+            .await?;
+        let mut state_tx = state.try_begin_transaction().unwrap();
+        second_update_client_action.execute(&mut state_tx).await?;
+        state_tx.apply();
+
+        Ok(())
     }
 }

component/src/lib.rs

@@ -4,6 +4,10 @@ use penumbra_storage::StateTransaction;
 use tendermint::abci;
 
 mod action_handler;
+
+mod temp_storage_ext;
+pub use temp_storage_ext::TempStorageExt;
+
 pub mod app;
 pub mod dex;
 pub mod governance;

component/src/temp_storage_ext.rs

@@ -0,0 +1,34 @@
+use std::ops::Deref;
+
+use async_trait::async_trait;
+use penumbra_chain::genesis;
+use penumbra_storage::TempStorage;
+
+use crate::app::App;
+
+#[async_trait]
+pub trait TempStorageExt: Sized {
+    async fn apply_genesis(self, genesis: genesis::AppState) -> anyhow::Result<Self>;
+    async fn apply_default_genesis(self) -> anyhow::Result<Self>;
+}
+
+#[async_trait]
+impl TempStorageExt for TempStorage {
+    async fn apply_genesis(self, genesis: genesis::AppState) -> anyhow::Result<Self> {
+        // Check that we haven't already applied a genesis state:
+        if self.latest_version() != u64::MAX {
+            return Err(anyhow::anyhow!("database already initialized"));
+        }
+
+        // Apply the genesis state to the storage
+        let mut app = App::new(self.latest_state());
+        app.init_chain(&genesis).await;
+        app.commit(self.deref().clone()).await;
+
+        Ok(self)
+    }
+
+    async fn apply_default_genesis(self) -> anyhow::Result<Self> {
+        self.apply_genesis(Default::default()).await
+    }
+}

storage/src/lib.rs

@@ -45,5 +45,5 @@ mod storage;
 pub use crate::metrics::register_metrics;
 pub use jmt::{ics23_spec, RootHash};
 pub use snapshot::Snapshot;
-pub use state::{State, StateRead, StateTransaction, StateWrite};
-pub use storage::{StateNotification, Storage};
+pub use state::{ArcStateExt, State, StateRead, StateTransaction, StateWrite};
+pub use storage::{StateNotification, Storage, TempStorage};

storage/src/state.rs

@@ -1,4 +1,4 @@
-use std::{any::Any, collections::BTreeMap, pin::Pin};
+use std::{any::Any, collections::BTreeMap, pin::Pin, sync::Arc};
 
 use anyhow::Result;
 use async_trait::async_trait;
@@ -179,3 +179,15 @@ impl StateRead for State {
         nonconsensus_prefix_raw_with_cache(&self.snapshot, &self.nonconsensus_changes, prefix)
     }
 }
+
+/// Extension trait providing `try_begin_transaction()` on `Arc<State>`.
+pub trait ArcStateExt: Sized {
+    /// Attempts to begin a transaction on this `Arc<State>`, returning `None` if the `Arc` is shared.
+    fn try_begin_transaction<'a>(&'a mut self) -> Option<StateTransaction<'a>>;
+}
+
+impl ArcStateExt for Arc<State> {
+    fn try_begin_transaction<'a>(&'a mut self) -> Option<StateTransaction<'a>> {
+        Arc::get_mut(self).map(|state| state.begin_transaction())
+    }
+}

storage/src/storage.rs

@@ -14,6 +14,9 @@ use crate::snapshot::Snapshot;
 use crate::snapshot_cache::SnapshotCache;
 use crate::State;
 
+mod temp;
+pub use temp::TempStorage;
+
 /// A handle for a storage instance, backed by RocksDB.
 ///
 /// The handle is cheaply clonable; all clones share the same backing data store.