pd: implement pd upgrade

crates/bin/pd/Cargo.toml

@@ -19,22 +19,28 @@ std = ["ibc-types/std"]
 
 [dependencies]
 # Workspace dependencies
-penumbra-proto            = { path = "../../proto" }
-penumbra-storage          = { path = "../../storage" }
-penumbra-asset           = { path = "../../core/asset" }
-penumbra-keys           = { path = "../../core/keys" }
-penumbra-shielded-pool    = { path = "../../core/component/shielded-pool", features = ["parallel"] }
-penumbra-stake            = { path = "../../core/component/stake", features = ["parallel"] }
-penumbra-sct              = { path = "../../core/component/sct" }
-penumbra-dex              = { path = "../../core/component/dex", features = ["parallel"] }
-penumbra-governance       = { path = "../../core/component/governance", features = ["parallel"]}
-penumbra-ibc              = { path = "../../core/component/ibc" }
-penumbra-compact-block    = { path = "../../core/component/compact-block" }
-penumbra-chain            = { path = "../../core/component/chain" }
-penumbra-transaction      = { path = "../../core/transaction" }
-penumbra-app              = { path = "../../core/app" }
-penumbra-wallet           = { path = "../../wallet" }
-penumbra-tower-trace      = { path = "../../util/tower-trace" }
+penumbra-proto = { path = "../../proto" }
+penumbra-storage = { path = "../../storage", features = ["migration"] }
+penumbra-asset = { path = "../../core/asset" }
+penumbra-keys = { path = "../../core/keys" }
+penumbra-shielded-pool = { path = "../../core/component/shielded-pool", features = [
+    "parallel",
+] }
+penumbra-stake = { path = "../../core/component/stake", features = [
+    "parallel",
+] }
+penumbra-sct = { path = "../../core/component/sct" }
+penumbra-dex = { path = "../../core/component/dex", features = ["parallel"] }
+penumbra-governance = { path = "../../core/component/governance", features = [
+    "parallel",
+] }
+penumbra-ibc = { path = "../../core/component/ibc" }
+penumbra-compact-block = { path = "../../core/component/compact-block" }
+penumbra-chain = { path = "../../core/component/chain" }
+penumbra-transaction = { path = "../../core/transaction" }
+penumbra-app = { path = "../../core/app" }
+penumbra-wallet = { path = "../../wallet" }
+penumbra-tower-trace = { path = "../../util/tower-trace" }
 penumbra-tendermint-proxy = { path = "../../util/tendermint-proxy" }
 
 # Penumbra dependencies
@@ -51,7 +57,9 @@ tendermint = "0.32.0"
 tendermint-light-client-verifier = "0.32.0"
 ibc-types = "0.3.0"
 
-ibc-proto = { version = "0.31.0", default-features = false, features = ["server"] }
+ibc-proto = { version = "0.31.0", default-features = false, features = [
+    "server",
+] }
 prost = "0.11"
 toml = "0.5"
 # We don't need this crate at all, but its upstream published a breaking change as
@@ -126,6 +134,6 @@ dist = true
 # Update the guide docs to use the most recent version, to reduce tedium.
 # Preview the changes with `cargo release replace`.
 pre-release-replacements = [
-    { file="../../../docs/guide/src/pcli/install.md", search="(cd penumbra && git fetch && git checkout) .*", replace="$1 {{tag_name}}", exactly=1 },
-    { file="../../../docs/guide/src/pcli/update.md", search="(cd penumbra && git fetch && git checkout) .*", replace="$1 {{tag_name}}", exactly=1 },
+    { file = "../../../docs/guide/src/pcli/install.md", search = "(cd penumbra && git fetch && git checkout) .*", replace = "$1 {{tag_name}}", exactly = 1 },
+    { file = "../../../docs/guide/src/pcli/update.md", search = "(cd penumbra && git fetch && git checkout) .*", replace = "$1 {{tag_name}}", exactly = 1 },
 ]

crates/bin/pd/src/lib.rs

@@ -11,6 +11,7 @@ mod snapshot;
 pub mod auto_https;
 pub mod events;
 pub mod testnet;
+pub mod upgrade;
 
 pub use crate::metrics::register_metrics;
 pub use consensus::Consensus;

crates/bin/pd/src/main.rs

@@ -16,6 +16,7 @@ use pd::testnet::{
     generate::TestnetConfig,
     join::testnet_join,
 };
+use pd::upgrade::{self, Upgrade};
 use penumbra_proto::client::v1alpha1::{
     oblivious_query_service_server::ObliviousQueryServiceServer,
     specific_query_service_server::SpecificQueryServiceServer,
@@ -132,6 +133,13 @@ enum RootCommand {
         #[clap(long, display_order = 300)]
         prune: bool,
     },
+    /// Run a migration on the exported storage state of the full node,
+    /// and create a genesis file.
+    Upgrade {
+        /// The directory containing the exported state.
+        #[clap(long, display_order = 200)]
+        export_path: PathBuf,
+    },
 }
 
 #[derive(Debug, Subcommand)]
@@ -574,6 +582,7 @@ async fn main() -> anyhow::Result<()> {
                 ?export_path,
                 "copying from data dir to export dir",
             );
+            std::fs::create_dir_all(&export_path)?;
             fs_extra::copy_items(&from, export_path.as_path(), &copy_opts)?;
 
             tracing::info!("done copying");
@@ -591,6 +600,12 @@ async fn main() -> anyhow::Result<()> {
             // - apply checks: root hash, size, etc.
             todo!()
         }
+        RootCommand::Upgrade { export_path } => {
+            tracing::info!("upgrading state from {}", export_path.display());
+            let _ = upgrade::migrate(export_path.clone(), Upgrade::Testnet60)
+                .await
+                .unwrap();
+        }
     }
     Ok(())
 }

crates/bin/pd/src/testnet/generate.rs

@@ -200,7 +200,9 @@ impl TestnetConfig {
     }
 
     /// Build Tendermint genesis data, based on Penumbra initial application state.
-    fn make_genesis(app_state: genesis::AppState) -> anyhow::Result<Genesis<genesis::AppState>> {
+    pub(crate) fn make_genesis(
+        app_state: genesis::AppState,
+    ) -> anyhow::Result<Genesis<genesis::AppState>> {
         // Use now as genesis time
         let genesis_time = Time::from_unix_timestamp(
             SystemTime::now()

crates/bin/pd/src/upgrade.rs

@@ -0,0 +1,76 @@
+use std::path::PathBuf;
+
+use penumbra_chain::{
+    component::{AppHash, StateReadExt},
+    genesis::AppState,
+};
+use penumbra_stake::StateReadExt as _;
+use penumbra_storage::{StateDelta, StateWrite, Storage};
+
+use crate::testnet::generate::TestnetConfig;
+
+pub enum Upgrade {
+    /// No-op migration
+    Noop,
+    /// Testnet 60 migration
+    Testnet60,
+}
+
+pub async fn migrate(path_to_export: PathBuf, upgrade: Upgrade) -> anyhow::Result<()> {
+    match upgrade {
+        Upgrade::Noop => (),
+        Upgrade::Testnet60 => {
+            let mut db_path = path_to_export.clone();
+            db_path.push("rocksdb");
+            let storage = Storage::load(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();
+
+            /* --------- writing to the jmt  ------------ */
+            tracing::info!(?app_hash_pre_migration, "app hash pre upgrade");
+            let mut delta = StateDelta::new(export_state);
+            delta.put_raw("testnet_60_forked".to_string(), "done".into());
+            let root_hash = storage.commit_in_place(delta).await?;
+            let app_hash_post_migration: AppHash = root_hash.into();
+            tracing::info!(?app_hash_post_migration, "app hash post upgrade");
+
+            /* --------- collecting genesis data -------- */
+            tracing::info!("generating genesis");
+            let migrated_state = storage.latest_snapshot();
+            let root_hash = migrated_state.root_hash().await.expect("can get root hash");
+            let app_hash: AppHash = root_hash.into();
+            tracing::info!(?root_hash, "root hash post upgrade2");
+            let height = migrated_state
+                .get_block_height()
+                .await
+                .expect("can get block height");
+            let next_height = height + 1;
+            let chain_params = migrated_state
+                .get_chain_params()
+                .await
+                .expect("can get chain params");
+
+            /* ---------- genereate genesis ------------  */
+            let validators = migrated_state.validator_list().await?;
+            let mut app_state = AppState::default();
+            app_state.chain_params = chain_params;
+            app_state.validators = validators.into_iter().map(Into::into).collect();
+            let mut genesis = TestnetConfig::make_genesis(app_state.clone()).unwrap();
+            genesis.app_hash = app_hash
+                .0
+                .to_vec()
+                .try_into()
+                .expect("infaillible conversion");
+            genesis.initial_height = next_height as i64;
+            genesis.genesis_time = tendermint::time::Time::now();
+
+            let genesis_json = serde_json::to_string(&genesis).expect("can serialize genesis");
+            tracing::info!("genesis: {}", genesis_json);
+            let mut genesis_path = path_to_export.clone();
+            genesis_path.push("genesis.json");
+            std::fs::write(genesis_path, genesis_json).expect("can write genesis");
+        }
+    }
+    Ok(())
+}

crates/storage/Cargo.toml

@@ -5,6 +5,9 @@ edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+migration = []
+
 [dependencies]
 jmt = "0.6"
 tokio = { version = "1.21.1", features = ["full", "tracing"] }
@@ -19,7 +22,7 @@ hex = "0.4"
 metrics = "0.19.0"
 parking_lot = "0.12"
 pin-project = "1.0.12"
-smallvec = { version = "1.10" , features = ["union" , "const_generics"] }
+smallvec = { version = "1.10", features = ["union", "const_generics"] }
 
 # Tendermint/IBC crates 
 ics23 = "0.10.1"

crates/storage/src/storage.rs

@@ -156,10 +156,15 @@ impl Storage {
         self.0.snapshots.read().get(version)
     }
 
+    /// Commits the provided [`StateDelta`] to persistent storage as the latest
+    /// version of the chain state. If `write_to_snapshot_cache` is `false`, the
+    /// snapshot will not be written to the snapshot cache, and no subscribers
+    /// will be notified.
     async fn commit_inner(
         &self,
         cache: Cache,
         new_version: jmt::Version,
+        write_to_snapshot_cache: bool,
     ) -> Result<crate::RootHash> {
         let span = Span::current();
         let inner = self.0.clone();
@@ -237,6 +242,10 @@ impl Storage {
                         };
                     }
 
+                    if !write_to_snapshot_cache {
+                        return Ok(root_hash);
+                    }
+
                     let latest_snapshot = Snapshot::new(inner.db.clone(), new_version);
                     // Obtain a write lock to the snapshot cache, and push the latest snapshot
                     // available. The lock guard is implicitly dropped immediately.
@@ -271,7 +280,16 @@ impl Storage {
             anyhow::bail!("version mismatch in commit: expected state forked from version {} but found state forked from version {}", old_version, snapshot.version());
         }
 
-        self.commit_inner(changes, new_version).await
+        self.commit_inner(changes, new_version, true).await
+    }
+
+    #[cfg(feature = "migration")]
+    /// Commits the provided [`StateDelta`] to persistent storage without increasing the version
+    /// of the chain state.
+    pub async fn commit_in_place(&self, delta: StateDelta<Snapshot>) -> Result<crate::RootHash> {
+        let (_, changes) = delta.flatten();
+        let old_version = self.latest_version();
+        self.commit_inner(changes, old_version, false).await
     }
 
     /// Returns the internal handle to RocksDB, this is useful to test adjacent storage crates.