feat: store merkle paths and mark as verified after submitting proof

Author: MarcosNicolau

aggregation_mode/db/migrations/001_init.sql

@@ -1,4 +1,4 @@
-CREATE TYPE task_status AS ENUM ('pending', 'processing', 'verified', 'submitted');
+CREATE TYPE task_status AS ENUM ('pending', 'processing', 'verified');
 
 CREATE TABLE tasks (
     task_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

aggregation_mode/db/src/types.rs

@@ -8,9 +8,8 @@ use sqlx::{
 #[sqlx(type_name = "task_status", rename_all = "lowercase")]
 pub enum TaskStatus {
     Pending,
-    Running,
-    Done,
-    Failed,
+    Processing,
+    Verified,
 }
 
 #[derive(Debug, Clone, FromRow)]

aggregation_mode/proof_aggregator/src/backend/db.rs

@@ -1,5 +1,5 @@
 use db::types::Task;
-use sqlx::{postgres::PgPoolOptions, Pool, Postgres};
+use sqlx::{postgres::PgPoolOptions, types::Uuid, Pool, Postgres};
 
 #[derive(Clone, Debug)]
 pub struct Db {
@@ -29,7 +29,18 @@ impl Db {
         limit: i64,
     ) -> Result<Vec<Task>, DbError> {
         sqlx::query_as::<_, Task>(
-            "SELECT * FROM tasks WHERE status = 'pending' AND proving_system_id = $1 LIMIT $2",
+            "WITH selected AS (
+                    SELECT task_id
+                    FROM tasks
+                    WHERE proving_system_id = $1 AND status = 'pending'
+                    LIMIT $2
+                    FOR UPDATE SKIP LOCKED
+                )
+                UPDATE tasks t
+                SET status = 'processing'
+                FROM selected s
+                WHERE t.task_id = s.task_id
+                RETURNING t.*;",
         )
         .bind(proving_system_id)
         .bind(limit)
@@ -38,11 +49,40 @@ impl Db {
         .map_err(|e| DbError::Query(e.to_string()))
     }
 
-    pub async fn mark_tasks_as_pending(&self) {}
+    pub async fn insert_tasks_merkle_path_and_mark_them_as_submitted(
+        &self,
+        updates: Vec<(Uuid, Vec<u8>)>,
+    ) -> Result<(), DbError> {
+        let mut tx = self
+            .pool
+            .begin()
+            .await
+            .map_err(|e| DbError::Query(e.to_string()))?;
+
+        for (task_id, merkle_path) in updates {
+            if let Err(e) = sqlx::query(
+                "UPDATE tasks SET merkle_path = $1, status = 'verified' WHERE task_id = $2",
+            )
+            .bind(merkle_path)
+            .bind(task_id)
+            .execute(&mut *tx)
+            .await
+            {
+                tx.rollback()
+                    .await
+                    .map_err(|e| DbError::Query(e.to_string()))?;
+                tracing::error!("Error while updating task merkle path and status {}", e);
+                return Err(DbError::Query(e.to_string()));
+            }
+        }
 
-    pub async fn mark_tasks_as_processing(&self) {}
+        tx.commit()
+            .await
+            .map_err(|e| DbError::Query(e.to_string()))?;
 
-    pub async fn mark_tasks_as_verified(&self) {}
+        Ok(())
+    }
 
-    pub async fn mark_tasks_as_submitted(&self) {}
+    // TODO: this should be used when rolling back processing proofs on unexpected errors
+    pub async fn mark_tasks_as_pending(&self) {}
 }

aggregation_mode/proof_aggregator/src/backend/fetcher.rs

@@ -5,7 +5,8 @@ use crate::{
     },
     backend::db::{Db, DbError},
 };
-use rayon::prelude::*;
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+use sqlx::types::Uuid;
 use tracing::{error, info};
 
 #[derive(Debug)]
@@ -26,48 +27,56 @@ impl ProofsFetcher {
         &self,
         engine: ZKVMEngine,
         limit: i64,
-    ) -> Result<Vec<AlignedProof>, ProofsFetcherError> {
+    ) -> Result<(Vec<AlignedProof>, Vec<Uuid>), ProofsFetcherError> {
         let tasks = self
             .db
             .get_pending_tasks_and_mark_them_as_processed(engine.proving_system_id() as i64, limit)
             .await
             .map_err(ProofsFetcherError::Query)?;
 
-        let proofs_to_aggregate: Vec<AlignedProof> = match engine {
-            ZKVMEngine::SP1 => tasks
-                .into_par_iter()
-                .filter_map(|task| {
-                    let vk = bincode::deserialize(&task.program_commitment).ok()?;
-                    let proof_with_pub_values = bincode::deserialize(&task.proof).ok()?;
-                    let sp1_proof = SP1ProofWithPubValuesAndVk::new(proof_with_pub_values, vk);
+        let (tasks_id, proofs_to_aggregate): (Vec<Uuid>, Vec<AlignedProof>) = match engine {
+            ZKVMEngine::SP1 => {
+                let pairs: Vec<(Uuid, AlignedProof)> = tasks
+                    .into_par_iter()
+                    .filter_map(|task| {
+                        let vk = bincode::deserialize(&task.program_commitment).ok()?;
+                        let proof_with_pub_values = bincode::deserialize(&task.proof).ok()?;
 
-                    match sp1_proof {
-                        Ok(proof) => Some(AlignedProof::SP1(proof.into())),
-                        Err(err) => {
-                            error!("Could not add proof, verification failed: {:?}", err);
-                            None
+                        match SP1ProofWithPubValuesAndVk::new(proof_with_pub_values, vk) {
+                            Ok(proof) => Some((task.task_id, AlignedProof::SP1(proof.into()))),
+                            Err(err) => {
+                                error!("Could not add proof, verification failed: {:?}", err);
+                                None
+                            }
                         }
-                    }
-                })
-                .collect(),
-            ZKVMEngine::RISC0 => tasks
-                .into_par_iter()
-                .filter_map(|task| {
-                    let mut image_id = [0u8; 32];
-                    image_id.copy_from_slice(&task.program_commitment);
-                    let receipt: risc0_zkvm::Receipt = bincode::deserialize(&task.proof).ok()?;
+                    })
+                    .collect();
 
-                    let risc0_proof = Risc0ProofReceiptAndImageId::new(image_id, receipt);
+                pairs.into_iter().unzip()
+            }
+            ZKVMEngine::RISC0 => {
+                let pairs: Vec<(Uuid, AlignedProof)> = tasks
+                    .into_par_iter()
+                    .filter_map(|task| {
+                        let mut image_id = [0u8; 32];
+                        image_id.copy_from_slice(&task.program_commitment);
+                        // we are inside a for_each callback so it returns for this particular iteration only
+                        let receipt = bincode::deserialize(&task.proof).ok()?;
 
-                    match risc0_proof {
-                        Ok(proof) => Some(AlignedProof::Risc0(proof.into())),
-                        Err(err) => {
-                            error!("Could not add proof, verification failed: {:?}", err);
-                            None
+                        let risc0_proof = Risc0ProofReceiptAndImageId::new(image_id, receipt);
+
+                        match risc0_proof {
+                            Ok(proof) => Some((task.task_id, AlignedProof::Risc0(proof.into()))),
+                            Err(err) => {
+                                error!("Could not add proof, verification failed: {:?}", err);
+                                None
+                            }
                         }
-                    }
-                })
-                .collect(),
+                    })
+                    .collect();
+
+                pairs.into_iter().unzip()
+            }
         };
 
         info!(
@@ -76,6 +85,6 @@ impl ProofsFetcher {
             proofs_to_aggregate.len()
         );
 
-        Ok(proofs_to_aggregate)
+        Ok((proofs_to_aggregate, tasks_id))
     }
 }

aggregation_mode/proof_aggregator/src/backend/mod.rs

@@ -6,7 +6,7 @@ mod types;
 
 use crate::{
     aggregators::{AlignedProof, ProofAggregationError, ZKVMEngine},
-    backend::db::Db,
+    backend::db::{Db, DbError},
 };
 
 use alloy::{
@@ -39,6 +39,7 @@ pub enum AggregatedProofSubmissionError {
     ZKVMAggregation(ProofAggregationError),
     BuildingMerkleRoot,
     MerkleRootMisMatch,
+    StoringMerklePaths(DbError),
 }
 
 pub struct ProofAggregator {
@@ -115,10 +116,11 @@ impl ProofAggregator {
         }
     }
 
+    // TODO: on failure, mark proofs as pending again
     async fn aggregate_and_submit_proofs_on_chain(
         &mut self,
     ) -> Result<(), AggregatedProofSubmissionError> {
-        let proofs = self
+        let (proofs, tasks_id) = self
             .fetcher
             .query(self.engine.clone(), self.config.total_proofs_limit)
             .await
@@ -167,7 +169,27 @@ impl ProofAggregator {
             receipt.transaction_hash
         );
 
-        // TODO: mark them as verified and store their merkle paths
+        info!("Storing merkle paths for each task...",);
+        let mut merkle_paths_for_tasks = vec![];
+        for task_id in tasks_id {
+            let Some(proof) = merkle_tree.get_proof_by_pos(0) else {
+                warn!("Proof not found for task id {task_id}");
+                continue;
+            };
+            let proof_bytes = proof
+                .merkle_path
+                .iter()
+                .map(|e| e.to_vec())
+                .flatten()
+                .collect::<Vec<_>>();
+
+            merkle_paths_for_tasks.push((task_id, proof_bytes))
+        }
+        self.db
+            .insert_tasks_merkle_path_and_mark_them_as_submitted(merkle_paths_for_tasks)
+            .await
+            .map_err(AggregatedProofSubmissionError::StoringMerklePaths)?;
+        info!("Merkle path inserted sucessfully",);
 
         Ok(())
     }