From 91dd2c1cbff704aef09e270f2d2cda20b4bcf1cc Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Thu, 6 Jun 2024 18:41:32 +1000
Subject: [PATCH] Implement DB upgrade

---
 beacon_node/beacon_chain/src/schema_change.rs |  9 ++++
 .../src/schema_change/migration_schema_v21.rs | 49 +++++++++++++++++++
 .../src/validator_pubkey_cache.rs             |  6 +--
 beacon_node/store/src/metadata.rs             |  2 +-
 4 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 beacon_node/beacon_chain/src/schema_change/migration_schema_v21.rs

diff --git a/beacon_node/beacon_chain/src/schema_change.rs b/beacon_node/beacon_chain/src/schema_change.rs
index 06d189a8c0..3fe75e348c 100644
--- a/beacon_node/beacon_chain/src/schema_change.rs
+++ b/beacon_node/beacon_chain/src/schema_change.rs
@@ -3,6 +3,7 @@ mod migration_schema_v17;
 mod migration_schema_v18;
 mod migration_schema_v19;
 mod migration_schema_v20;
+mod migration_schema_v21;
 
 use crate::beacon_chain::BeaconChainTypes;
 use crate::types::ChainSpec;
@@ -87,6 +88,14 @@ pub fn migrate_schema<T: BeaconChainTypes>(
             let ops = migration_schema_v20::downgrade_from_v20::<T>(db.clone(), log)?;
             db.store_schema_version_atomically(to, ops)
         }
+        (SchemaVersion(20), SchemaVersion(21)) => {
+            let ops = migration_schema_v21::upgrade_to_v21::<T>(db.clone(), log)?;
+            db.store_schema_version_atomically(to, ops)
+        }
+        (SchemaVersion(21), SchemaVersion(20)) => {
+            let ops = migration_schema_v21::downgrade_from_v21::<T>(db.clone(), log)?;
+            db.store_schema_version_atomically(to, ops)
+        }
         // Anything else is an error.
         (_, _) => Err(HotColdDBError::UnsupportedSchemaVersion {
             target_version: to,
diff --git a/beacon_node/beacon_chain/src/schema_change/migration_schema_v21.rs b/beacon_node/beacon_chain/src/schema_change/migration_schema_v21.rs
new file mode 100644
index 0000000000..0485852551
--- /dev/null
+++ b/beacon_node/beacon_chain/src/schema_change/migration_schema_v21.rs
@@ -0,0 +1,49 @@
+use crate::beacon_chain::BeaconChainTypes;
+use crate::validator_pubkey_cache::DatabasePubkey;
+use slog::{debug, Logger};
+use ssz::Decode;
+use std::sync::Arc;
+use store::{DBColumn, Error, HotColdDB, KeyValueStore, KeyValueStoreOp, StoreItem};
+use types::{Hash256, PublicKey};
+
+const LOG_EVERY: usize = 100_000;
+
+pub fn upgrade_to_v21<T: BeaconChainTypes>(
+    db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
+    log: Logger,
+) -> Result<Vec<KeyValueStoreOp>, Error> {
+    let mut ops = vec![];
+
+    debug!(log, "Migrating from v20 to v21");
+
+    // Iterate through all pubkeys and decompress them.
+    for (i, res) in db
+        .hot_db
+        .iter_column::<Hash256>(DBColumn::PubkeyCache)
+        .enumerate()
+    {
+        let (key, value) = res?;
+        let pubkey = PublicKey::from_ssz_bytes(&value)?;
+        let compressed = DatabasePubkey::from_pubkey(&pubkey);
+        ops.push(compressed.as_kv_store_op(key));
+
+        if i > 0 && i % LOG_EVERY == 0 {
+            debug!(
+                log,
+                "Public key decompression in progress";
+                "keys_decompressed" => i
+            );
+        }
+    }
+    debug!(log, "Public key decompression complete");
+
+    Ok(ops)
+}
+
+pub fn downgrade_from_v21<T: BeaconChainTypes>(
+    _db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
+    _log: Logger,
+) -> Result<Vec<KeyValueStoreOp>, Error> {
+    // TODO(sproul): impl downgrade
+    Ok(vec![])
+}
diff --git a/beacon_node/beacon_chain/src/validator_pubkey_cache.rs b/beacon_node/beacon_chain/src/validator_pubkey_cache.rs
index 3c9685e3d4..576fbf0fd1 100644
--- a/beacon_node/beacon_chain/src/validator_pubkey_cache.rs
+++ b/beacon_node/beacon_chain/src/validator_pubkey_cache.rs
@@ -169,7 +169,7 @@ impl<T: BeaconChainTypes> ValidatorPubkeyCache<T> {
 ///
 /// Keyed by the validator index as `Hash256::from_low_u64_be(index)`.
 #[derive(Encode, Decode)]
-struct DatabasePubkey {
+pub struct DatabasePubkey {
     pubkey: SmallVec<[u8; PUBLIC_KEY_UNCOMPRESSED_BYTES_LEN]>,
 }
 
@@ -192,13 +192,13 @@ impl DatabasePubkey {
         Hash256::from_low_u64_be(index as u64)
     }
 
-    fn from_pubkey(pubkey: &PublicKey) -> Self {
+    pub fn from_pubkey(pubkey: &PublicKey) -> Self {
         Self {
             pubkey: pubkey.serialize_uncompressed().into(),
         }
     }
 
-    fn as_pubkey(&self) -> Result<(PublicKey, PublicKeyBytes), BeaconChainError> {
+    pub fn as_pubkey(&self) -> Result<(PublicKey, PublicKeyBytes), BeaconChainError> {
         let pubkey = PublicKey::deserialize_uncompressed(&self.pubkey)
             .map_err(BeaconChainError::InvalidValidatorPubkeyBytes)?;
         let pubkey_bytes = pubkey.compress();
diff --git a/beacon_node/store/src/metadata.rs b/beacon_node/store/src/metadata.rs
index 116926ad3f..a22dc4aab4 100644
--- a/beacon_node/store/src/metadata.rs
+++ b/beacon_node/store/src/metadata.rs
@@ -4,7 +4,7 @@ use ssz::{Decode, Encode};
 use ssz_derive::{Decode, Encode};
 use types::{Checkpoint, Hash256, Slot};
 
-pub const CURRENT_SCHEMA_VERSION: SchemaVersion = SchemaVersion(20);
+pub const CURRENT_SCHEMA_VERSION: SchemaVersion = SchemaVersion(21);
 
 // All the keys that get stored under the `BeaconMeta` column.
 //