Jonas/service lifecycle (#293)

* introduce server that manages joint tasks

* add comment

* fix container name

* update comment

* comment

* clippy

* clippy

* error handling

* re-initialize on task failur

* remove wrapper handle

* reinitialize in 'normal' mode

* join service tasks

* Enhance: Finalize Hyperlane Snapshots (#292)

Author: jonas089

crates/ev-prover/README.md

@@ -114,3 +114,19 @@ To lint the Protobuf definitions:
 cd crates/ev-prover/proto
 buf lint
 ```
+
+## Hyperlane Message Finality
+
+```mermaid
+    sequenceDiagram;
+        Reth->>MessageProver: Transactions;
+        SnapshotStore->>MessageProver: Trusted Snapshot finalized==false;
+        MessageProver->>SnapshotStore: New Snapshot after Inserts finalized==false;
+        MessageProver->>SnapshotStore: Trusted Snapshot finalized==true;
+```
+
+Finality serves two distinct purposes. Firstly, it can be used to expose the status of message proofs to the user. However the primary purpose is such that errors can be detected when the finalization status is not updated for a snapshot. Currently this will not break the system, as new messages will be proven from the next snapshot onwards, but messages could get lost until a proper retry mechanism is implemented. 
+
+So long as the finality status of all snapshots flips to `true` there is nothing to worry about. All snapshots, except for the most recent one, should be finalized at all times. Gaps in finalization indicate that message proof submission was unsuccessful, or that the DB corrupted post submission. 
+
+Note that the message prover always takes the last known snapshot, which is expected to be unfinalized, but no such check is enforced and generates a proof using all messages that occurred from the height of the trusted snapshot + 1 all the way to the trusted EV height in ZKISM, aka `committed_height`.

crates/ev-prover/src/prover/programs/combined.rs

@@ -289,7 +289,7 @@ impl EvCombinedProver {
         }
 
         let namespace = self.app.namespace;
-        for height in scan_start..latest_head {
+        for height in scan_start..=latest_head {
             if !self.is_empty_block(height, namespace).await? {
                 // Ensure batch size stays within allowed range
                 let blocks_elapsed = height.saturating_sub(trusted_celestia_height);
@@ -323,10 +323,13 @@ impl EvCombinedProver {
         let msg = MsgUpdateZkExecutionIsm::new(id, proof.bytes(), public_values, signer);
 
         info!("Updating ZKISM on Celestia...");
-        let res = self.app.ism_client.send_tx(msg).await?;
-        assert!(res.success);
+        let response = self.app.ism_client.send_tx(msg).await?;
+        if !response.success {
+            error!("Failed to submit state transition proof to ZKISM: {:?}", response);
+            return Err(anyhow::anyhow!("Failed to submit state transition proof to ZKISM"));
+        }
 
-        info!("[Done] Proof tx submitted to ism with hash: {}", res.tx_hash);
+        info!("[Done] Proof tx submitted to ism with hash: {}", response.tx_hash);
 
         Ok(())
     }

crates/ev-prover/src/prover/programs/message.rs

@@ -221,15 +221,16 @@ impl HyperlaneMessageProver {
         &self,
         evm_provider: &DefaultProvider,
         indexer: &mut HyperlaneIndexer,
-        height: u64,
+        committed_height: u64,
         proof: EIP1186AccountProofResponse,
         state_root: FixedBytes<32>,
         ism_client: &CelestiaIsmClient,
     ) -> Result<()> {
         // generate a new proof for all messages that occurred since the last trusted height, inserting into the last snapshot
         // then save new snapshot
-        let mut snapshot = self.snapshot_store.get_snapshot(self.snapshot_store.current_index()?)?;
-        if snapshot.height == height {
+        let trusted_snapshot_index = self.snapshot_store.current_index()?;
+        let mut snapshot = self.snapshot_store.get_snapshot(trusted_snapshot_index)?;
+        if snapshot.height == committed_height {
             debug!("No new ev blocks so no new messages to prove");
             return Ok(());
         }
@@ -239,15 +240,15 @@ impl HyperlaneMessageProver {
             .address(indexer.contract_address)
             .event(&Dispatch::id())
             .from_block(start_height)
-            .to_block(height);
+            .to_block(committed_height);
 
         // run the indexer to get all messages that occurred since the last trusted height
         indexer
             .index(self.message_store.clone(), Arc::new(evm_provider.clone()))
             .await?;
 
         let mut messages: Vec<StoredHyperlaneMessage> = Vec::new();
-        for block in start_height..=height {
+        for block in start_height..=committed_height {
             messages.extend(self.message_store.get_by_block(block)?);
         }
 
@@ -281,17 +282,29 @@ impl HyperlaneMessageProver {
 
         let message_proof_msg = MsgSubmitMessages::new(
             self.ctx.ism_id.clone(),
-            height,
+            committed_height,
             message_proof.0.bytes(),
             message_proof.0.public_values.as_slice().to_vec(),
             ism_client.signer_address().to_string(),
         );
 
+        // store the unfinalized snapshot
+        snapshot.height = committed_height;
+        let snapshot_index = self.snapshot_store.current_index()? + 1;
+        self.snapshot_store.insert_snapshot(snapshot_index, snapshot)?;
         info!("Submitting Hyperlane tree proof to ZKISM...");
         let response = ism_client.send_tx(message_proof_msg).await?;
 
-        assert!(response.success);
+        if !response.success {
+            error!("Failed to submit Hyperlane tree proof to ZKISM: {:?}", response);
+            return Err(anyhow::anyhow!("Failed to submit Hyperlane tree proof to ZKISM"));
+        }
         info!("[Done] ZKISM was updated successfully");
+        self.proof_store
+            .store_membership_proof(committed_height, &message_proof.0, &message_proof.1)
+            .await?;
+
+        self.snapshot_store.finalize_snapshot(trusted_snapshot_index)?;
 
         info!("Relaying verified Hyperlane messages to Celestia...");
         // submit all now verified messages to hyperlane
@@ -305,18 +318,13 @@ impl HyperlaneMessageProver {
                 message_hex,
             );
             let response = ism_client.send_tx(msg).await?;
-            assert!(response.success);
+            if !response.success {
+                error!("Failed to relay Hyperlane message to Celestia: {:?}", response);
+                return Err(anyhow::anyhow!("Failed to relay Hyperlane message to Celestia"));
+            }
         }
         info!("[Done] Tia was bridged back to Celestia");
 
-        self.proof_store
-            .store_membership_proof(height, &message_proof.0, &message_proof.1)
-            .await?;
-
-        snapshot.height = height;
-        self.snapshot_store
-            .insert_snapshot(self.snapshot_store.current_index()? + 1, snapshot)?;
-
         Ok(())
     }
 }

crates/storage/src/hyperlane/snapshot.rs

@@ -11,12 +11,23 @@ use std::sync::{Arc, RwLock};
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct HyperlaneSnapshot {
+    // the trusted EV height in the ZKISM
     pub height: u64,
+    // the Hyperlane Message Tree e.g. Snapshot
     pub tree: MerkleTree,
+    // whether this Snapshot has been finalized
+    pub finalized: bool,
 }
 impl HyperlaneSnapshot {
     pub fn new(height: u64, tree: MerkleTree) -> HyperlaneSnapshot {
-        HyperlaneSnapshot { height, tree }
+        HyperlaneSnapshot {
+            height,
+            tree,
+            finalized: false,
+        }
+    }
+    pub fn finalize(&mut self) {
+        self.finalized = true;
     }
 }
 
@@ -53,6 +64,7 @@ impl HyperlaneSnapshotStore {
         Ok(vec![ColumnFamilyDescriptor::new("snapshots", Options::default())])
     }
 
+    /// Insert a Hyperlane Snapshot into the database
     pub fn insert_snapshot(&self, index: u64, snapshot: HyperlaneSnapshot) -> Result<()> {
         // Serialize outside the lock to minimize lock duration
         let serialized = bincode::serialize(&snapshot).context("Failed to serialize snapshot")?;
@@ -70,6 +82,7 @@ impl HyperlaneSnapshotStore {
         Ok(())
     }
 
+    /// Get a Hyperlane Snapshot by index
     pub fn get_snapshot(&self, index: u64) -> Result<HyperlaneSnapshot> {
         let read_lock = self.db.read().map_err(|e| anyhow::anyhow!("lock error: {e}"))?;
         let cf = read_lock.cf_handle("snapshots").context("Missing CF")?;
@@ -88,6 +101,49 @@ impl HyperlaneSnapshotStore {
         Ok(snapshot)
     }
 
+    /// Get the latest pending snapshot, we expect only the most recent snapshot to be unfinalized
+    pub fn get_pending_snapshot(&self) -> Result<Option<(u64, HyperlaneSnapshot)>> {
+        let read_lock = self
+            .db
+            .read()
+            .map_err(|e| anyhow::anyhow!("Failed to acquire read lock: {e}"))?;
+        let cf = read_lock.cf_handle("snapshots").context("Missing CF")?;
+        let mut iter = read_lock.iterator_cf(cf, IteratorMode::End);
+        while let Some(Ok((k, v))) = iter.next() {
+            if k.len() != 8 {
+                continue;
+            }
+            let mut buf = [0u8; 8];
+            buf.copy_from_slice(&k);
+            let index = u64::from_be_bytes(buf);
+            let mut snapshot: HyperlaneSnapshot = bincode::deserialize(&v).context("Failed to deserialize snapshot")?;
+            for h in snapshot.tree.branch.iter_mut() {
+                if h.is_empty() {
+                    *h = ZERO_BYTES.to_string();
+                }
+            }
+            if !snapshot.finalized {
+                return Ok(Some((index, snapshot)));
+            }
+        }
+        Ok(None)
+    }
+
+    /// Finalize a Hyperlane Snapshot after successful proof submission
+    pub fn finalize_snapshot(&self, index: u64) -> Result<()> {
+        let mut snapshot = self
+            .get_snapshot(index)
+            .with_context(|| format!("Snapshot at index {index} not found"))?;
+        if snapshot.finalized {
+            return Err(anyhow::anyhow!(
+                "Tried to finalize a finalized snapshot at index {index}"
+            ));
+        }
+        snapshot.finalized = true;
+        self.insert_snapshot(index, snapshot)
+    }
+
+    /// Get the next insert index for the Hyperlane Snapshot store
     pub fn current_index(&self) -> Result<u64> {
         let read_lock = self
             .db
@@ -104,6 +160,7 @@ impl HyperlaneSnapshotStore {
         }
     }
 
+    /// Reset the database by dropping the snapshots column family and creating a new one
     pub fn reset_db(&self) -> Result<()> {
         let mut write_lock = self
             .db
@@ -115,3 +172,46 @@ impl HyperlaneSnapshotStore {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_insert_snapshot() {
+        let store = HyperlaneSnapshotStore::new(tempfile::tempdir().unwrap(), None).unwrap();
+        let snapshot = HyperlaneSnapshot::new(0, MerkleTree::default());
+        store.insert_snapshot(0, snapshot).unwrap();
+    }
+    #[test]
+    fn test_get_snapshot() {
+        let store = HyperlaneSnapshotStore::new(tempfile::tempdir().unwrap(), None).unwrap();
+        let snapshot = HyperlaneSnapshot::new(0, MerkleTree::default());
+        store.insert_snapshot(0, snapshot.clone()).unwrap();
+        let retrieved_snapshot = store.get_snapshot(0).unwrap();
+        assert_eq!(retrieved_snapshot, snapshot);
+    }
+    #[test]
+    fn test_get_pending_snapshot() {
+        let store = HyperlaneSnapshotStore::new(tempfile::tempdir().unwrap(), None).unwrap();
+        let first_snapshot = HyperlaneSnapshot::new(0, MerkleTree::default());
+        let second_snapshot = HyperlaneSnapshot::new(1, MerkleTree::default());
+        let third_snapshot = HyperlaneSnapshot::new(2, MerkleTree::default());
+        store.insert_snapshot(0, first_snapshot.clone()).unwrap();
+        store.insert_snapshot(1, second_snapshot.clone()).unwrap();
+        store.insert_snapshot(2, third_snapshot.clone()).unwrap();
+        store.finalize_snapshot(0).unwrap();
+        store.finalize_snapshot(1).unwrap();
+        let retrieved_snapshot = store.get_pending_snapshot().unwrap();
+        assert_eq!(retrieved_snapshot, Some((2, third_snapshot)));
+    }
+    #[test]
+    fn test_finalize_snapshot() {
+        let store = HyperlaneSnapshotStore::new(tempfile::tempdir().unwrap(), None).unwrap();
+        let snapshot = HyperlaneSnapshot::new(0, MerkleTree::default());
+        store.insert_snapshot(0, snapshot.clone()).unwrap();
+        store.finalize_snapshot(0).unwrap();
+        let retrieved_snapshot = store.get_snapshot(0).unwrap();
+        assert!(retrieved_snapshot.finalized);
+    }
+}