add migration

EspressoSystems · Oct 15, 2024 · a9a7605 · a9a7605
1 parent 409a06d
commit a9a7605
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 3 deletions.
diff --git a/crates/example-types/src/storage_types.rs b/crates/example-types/src/storage_types.rs
@@ -32,7 +32,6 @@ type VidShares<TYPES> = HashMap<
     <TYPES as NodeType>::Time,
     HashMap<<TYPES as NodeType>::SignatureKey, Proposal<TYPES, VidDisperseShare<TYPES>>>,
 >;
-
 #[derive(Clone, Debug)]
 pub struct TestStorageState<TYPES: NodeType> {
     vids: VidShares<TYPES>,
@@ -223,4 +222,22 @@ impl<TYPES: NodeType> Storage<TYPES> for TestStorage<TYPES> {
 
         Ok(())
     }
+
+    async fn migrate_consensus(
+        &self,
+        _convert_leaf: fn(Leaf<TYPES>) -> Leaf2<TYPES>,
+        convert_proposal: fn(
+            Proposal<TYPES, QuorumProposal<TYPES>>,
+        ) -> Proposal<TYPES, QuorumProposal2<TYPES>>,
+    ) -> Result<()> {
+        let mut storage_writer = self.inner.write().await;
+
+        for (view, proposal) in storage_writer.proposals.clone().iter() {
+            storage_writer
+                .proposals2
+                .insert(*view, convert_proposal(proposal.clone()));
+        }
+
+        Ok(())
+    }
 }
diff --git a/crates/hotshot/src/lib.rs b/crates/hotshot/src/lib.rs
@@ -46,9 +46,9 @@ pub use hotshot_types::error::HotShotError;
 use hotshot_types::{
     consensus::{Consensus, ConsensusMetricsValue, OuterConsensus, View, ViewInner},
     constants::{EVENT_CHANNEL_SIZE, EXTERNAL_EVENT_CHANNEL_SIZE},
-    data::{Leaf, QuorumProposal},
+    data::{Leaf, Leaf2, QuorumProposal, QuorumProposal2},
     event::{EventType, LeafInfo},
-    message::{DataMessage, Message, MessageKind, Proposal},
+    message::{convert_proposal, DataMessage, Message, MessageKind, Proposal},
     simple_certificate::{QuorumCertificate, UpgradeCertificate},
     traits::{
         consensus_api::ConsensusApi,
@@ -57,6 +57,7 @@ use hotshot_types::{
         node_implementation::{ConsensusTime, NodeType},
         signature_key::SignatureKey,
         states::ValidatedState,
+        storage::Storage,
         EncodeBytes,
     },
     HotShotConfig,
@@ -188,6 +189,10 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> SystemContext<T
     /// well.
     ///
     /// Use this instead of `init` if you want to start the tasks manually
+    ///
+    /// # Panics
+    ///
+    /// Panics if storage migration fails.
     #[allow(clippy::too_many_arguments)]
     pub async fn new(
         public_key: TYPES::SignatureKey,
@@ -201,6 +206,20 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> SystemContext<T
         storage: I::Storage,
         marketplace_config: MarketplaceConfig<TYPES, I>,
     ) -> Arc<Self> {
+        #[allow(clippy::panic)]
+        match storage
+            .migrate_consensus(
+                Into::<Leaf2<TYPES>>::into,
+                convert_proposal::<TYPES, QuorumProposal<TYPES>, QuorumProposal2<TYPES>>,
+            )
+            .await
+        {
+            Ok(()) => {}
+            Err(e) => {
+                panic!("Failed to migrate consensus storage: {e}");
+            }
+        }
+
         let interal_chan = broadcast(EVENT_CHANNEL_SIZE);
         let external_chan = broadcast(EXTERNAL_EVENT_CHANNEL_SIZE);
 

diff --git a/crates/types/src/data.rs b/crates/types/src/data.rs
@@ -403,6 +403,21 @@ impl<TYPES: NodeType> From<QuorumProposal<TYPES>> for QuorumProposal2<TYPES> {
     }
 }
 
+impl<TYPES: NodeType> From<Leaf<TYPES>> for Leaf2<TYPES> {
+    fn from(leaf: Leaf<TYPES>) -> Self {
+        let bytes: [u8; 32] = leaf.parent_commitment.into();
+
+        Self {
+            view_number: leaf.view_number,
+            justify_qc: LeafCertificate::Quorum(leaf.justify_qc),
+            parent_commitment: Commitment::from_raw(bytes),
+            block_header: leaf.block_header,
+            upgrade_certificate: leaf.upgrade_certificate,
+            block_payload: leaf.block_payload,
+        }
+    }
+}
+
 impl<TYPES: NodeType> HasViewNumber<TYPES> for DaProposal<TYPES> {
     fn view_number(&self) -> TYPES::Time {
         self.view_number

diff --git a/crates/types/src/message.rs b/crates/types/src/message.rs
@@ -349,6 +349,22 @@ pub struct Proposal<TYPES: NodeType, PROPOSAL: HasViewNumber<TYPES> + Deserializ
     pub _pd: PhantomData<TYPES>,
 }
 
+/// Convert a `Proposal` by converting the underlying proposal type
+pub fn convert_proposal<TYPES, PROPOSAL, PROPOSAL2>(
+    proposal: Proposal<TYPES, PROPOSAL>,
+) -> Proposal<TYPES, PROPOSAL2>
+where
+    TYPES: NodeType,
+    PROPOSAL: HasViewNumber<TYPES> + DeserializeOwned,
+    PROPOSAL2: HasViewNumber<TYPES> + DeserializeOwned + From<PROPOSAL>,
+{
+    Proposal {
+        data: proposal.data.into(),
+        signature: proposal.signature,
+        _pd: proposal._pd,
+    }
+}
+
 impl<TYPES> Proposal<TYPES, QuorumProposal<TYPES>>
 where
     TYPES: NodeType,

diff --git a/crates/types/src/traits/storage.rs b/crates/types/src/traits/storage.rs
@@ -63,4 +63,12 @@ pub trait Storage<TYPES: NodeType>: Send + Sync + Clone {
         &self,
         decided_upgrade_certificate: Option<UpgradeCertificate<TYPES>>,
     ) -> Result<()>;
+    /// Migrate leaves from `Leaf` to `Leaf2`, and proposals from `QuorumProposal` to `QuorumProposal2`
+    async fn migrate_consensus(
+        &self,
+        convert_leaf: fn(Leaf<TYPES>) -> Leaf2<TYPES>,
+        convert_proposal: fn(
+            Proposal<TYPES, QuorumProposal<TYPES>>,
+        ) -> Proposal<TYPES, QuorumProposal2<TYPES>>,
+    ) -> Result<()>;
 }