pd: use commit_in_place to write app safeguard (#4954)

**Describe your changes**

This PR:
- fixes a bug in the app version safeguard logic that would cause
discrepancies between JMT state versions and block height
- refactors some of the inner logic to avoid writing any state unless it
is necessary (i.e, no safeguard was found)

**Testing this:**

To observe the state discrepancy and test remediation:

1. Checkout the commit that only contains the extra tracing statement:
`git checkout a8c51f0`
2. Create a local devnet
3. Restart the devnet N times
4. Observe that the state version and block height have offset=N

To test that this fixes it:
1. Run this patch
2. Restart the node
3. Observe that the state version stays stable.


## Checklist before requesting a review

- [x] I have added guiding text to explain how a reviewer should test
these changes.

- [x] If this code contains consensus-breaking changes, I have added the
"consensus-breaking" label. Otherwise, I declare my belief that there
are not consensus-breaking changes, for the following reason:

  > Not consensus breaking.

crates/bin/pd/src/main.rs

@@ -19,8 +19,8 @@ use pd::{
         join::network_join,
     },
 };
-use penumbra_app::app_version::assert_latest_app_version;
-use penumbra_app::SUBSTORE_PREFIXES;
+use penumbra_app::app_version::check_and_update_app_version;
+use penumbra_app::{APP_VERSION, SUBSTORE_PREFIXES};
 use rand::Rng;
 use rand_core::OsRng;
 use tendermint_config::net::Address as TendermintAddress;
@@ -103,9 +103,10 @@ async fn main() -> anyhow::Result<()> {
                 .context(
                     "Unable to initialize RocksDB storage - is there another `pd` process running?",
                 )?;
-            assert_latest_app_version(storage.clone()).await?;
+            check_and_update_app_version(storage.clone()).await?;
 
             tracing::info!(
+                APP_VERSION,
                 ?abci_bind,
                 ?grpc_bind,
                 ?grpc_auto_https,

crates/core/app/src/app_version.rs

@@ -5,6 +5,6 @@ pub const APP_VERSION: u64 = 8;
 cfg_if::cfg_if! {
     if #[cfg(feature="component")] {
         mod component;
-        pub use component::{assert_latest_app_version, migrate_app_version};
+        pub use component::{check_and_update_app_version, migrate_app_version};
     }
 }

crates/core/app/src/app_version/component.rs

@@ -26,11 +26,14 @@ enum CheckContext {
     Migration,
 }
 
+/// Check that the expected version matches the found version, if it set.
+/// This will return an error if the versions do not match.
 fn check_version(ctx: CheckContext, expected: u64, found: Option<u64>) -> anyhow::Result<()> {
-    let found = match (expected, found) {
-        (x, Some(y)) if x != y => y,
-        _ => return Ok(()),
-    };
+    let found = found.unwrap_or(expected);
+    if found == expected {
+        return Ok(());
+    }
+
     match ctx {
         CheckContext::Running => {
             let expected_name = version_to_software_version(expected);
@@ -79,6 +82,7 @@ fn check_version(ctx: CheckContext, expected: u64, found: Option<u64>) -> anyhow
     }
 }
 
+/// Read the app version safeguard from nonverifiable storage.
 async fn read_app_version_safeguard<S: StateReadProto>(s: &S) -> anyhow::Result<Option<u64>> {
     let out = s
         .nonverifiable_get_proto(crate::app::state_key::app_version::safeguard().as_bytes())
@@ -87,38 +91,45 @@ async fn read_app_version_safeguard<S: StateReadProto>(s: &S) -> anyhow::Result<
     Ok(out)
 }
 
-// Neither async nor a result are needed, but only right now, so I'm putting these here
-// to reserve the right to change them later.
-async fn write_app_version_safeguard<S: StateWriteProto>(s: &mut S, x: u64) -> anyhow::Result<()> {
+/// Write the app version safeguard to nonverifiable storage.
+fn write_app_version_safeguard<S: StateWriteProto>(s: &mut S, x: u64) {
     s.nonverifiable_put_proto(
         crate::app::state_key::app_version::safeguard()
             .as_bytes()
             .to_vec(),
         x,
-    );
-    Ok(())
+    )
 }
 
-/// Assert that the app version saved in the state is the correct one.
+/// Ensure that the app version safeguard is `APP_VERSION`, or update it if it is missing.
 ///
-/// You should call this before starting a node.
+/// # Errors
+/// This method errors if the app version safeguard is different than `APP_VERSION`.
 ///
-/// This will succeed if no app version was found in local storage, or if the app version saved matches
-/// exactly.
+/// # Usage
+/// This should be called on startup. This method short-circuits if the database
+/// is uninitialized (pregenesis).
 ///
-/// This will also result in the current app version being stored, so that future
-/// calls to this function will be checked against this state.
-pub async fn assert_latest_app_version(s: Storage) -> anyhow::Result<()> {
+/// # UIP:
+/// More context is available in the UIP-6 document: https://uips.penumbra.zone/uip-6.html
+pub async fn check_and_update_app_version(s: Storage) -> anyhow::Result<()> {
     // If the storage is not initialized, avoid touching it at all,
     // to avoid complaints about it already being initialized before the first genesis.
     if s.latest_version() == u64::MAX {
         return Ok(());
     }
     let mut delta = StateDelta::new(s.latest_snapshot());
-    let found = read_app_version_safeguard(&delta).await?;
-    check_version(CheckContext::Running, APP_VERSION, found)?;
-    write_app_version_safeguard(&mut delta, APP_VERSION).await?;
-    s.commit(delta).await?;
+
+    // If the safeguard is not set, set it to the current version.
+    // Otherwise, ensure that it matches the current version.
+    match read_app_version_safeguard(&delta).await? {
+        None => {
+            tracing::debug!(?APP_VERSION, "version safeguard not found, initializing");
+            write_app_version_safeguard(&mut delta, APP_VERSION);
+            s.commit_in_place(delta).await?;
+        }
+        Some(found) => check_version(CheckContext::Running, APP_VERSION, Some(found))?,
+    }
     Ok(())
 }
 
@@ -132,7 +143,7 @@ pub async fn migrate_app_version<S: StateWriteProto>(s: &mut S, to: u64) -> anyh
     anyhow::ensure!(to > 1, "you can't migrate to the first penumbra version!");
     let found = read_app_version_safeguard(s).await?;
     check_version(CheckContext::Migration, to - 1, found)?;
-    write_app_version_safeguard(s, to).await?;
+    write_app_version_safeguard(s, to);
     Ok(())
 }
 

crates/core/app/src/server/consensus.rs

@@ -219,7 +219,8 @@ impl Consensus {
     }
 
     async fn end_block(&mut self, end_block: request::EndBlock) -> response::EndBlock {
-        tracing::info!(height = ?end_block.height, "ending block");
+        let latest_state_version = self.storage.latest_version();
+        tracing::info!(height = ?end_block.height, ?latest_state_version, "ending block");
         let events = self.app.end_block(&end_block).await;
         trace_events(&events);