feat(host): create initial cancel handling (#316)

* [WIP](host): connect cancel API * feat(host): create initial cancel handling * refactor(host,task_manager): remove duplicate struct and impls * refactor(lib): remove reference where not needed * refactor(host,task_manager): replace struct with type * chore(host): use get instead of post for reporting * fix(task_manager): clear hash map for in memory db * refactor(host,task_manager): clean up function params * refactor(host): minimize fetching and type complexity * refactor(host): skip task if getting task data is unsuccessful * feat(host): pass error from proving to task actor * feat(host): make getting task manager more ergonomic * refactor(core,lib): clean up preflight code * feat(host): simplify proof request config reading * fix(lib): fix typo * refactor(host): extract cache utils to separate module * fix(core): handle errors for blob checking * fix(task_manager): handle proof retrieval better * refactor(host): make proof handler more ergonomic * fix(tasks): pass closure instead of variable * fix(tasks): handle inconsistency in status triplet * chore(core,lib): clean up dependency declaration * feat(core,host,lib,risc0,sp1): unify proof output * fix(clippy): fix lint * refactor(host): add error trace * refactor(tasks): use type instead of struct * refactor(tasks): move struct definition * feat(host): use single channel for tasks * feat(all): pass a store to provers to cancel a remote task * fix(all): fix some async errors * fix(core,host,lib): handle async boundary * fix(provers): fix types on trait implementors * fix(core): fix lints * fix(provers): fix lints * fix(core): fix test mock * fix(core): use provers instead of drivers * fix(lib,tasks): account for multi prover environment * fix(tasks): fix clippy lints * feat(risc0): add cancellation to bonsai (#320) * feat(risc0): add cancellation to bonsai * fix(risc0): remove id after cancellation * fix(risc0): add prover code to proof key * feat(sp1): add cancellation support (#322) * feat(sp1): add cancellation support * fix(sp1): add prover code to proof key * refactor(risc0,sp1): rename arguments * fix(clippy): remove unused imports * fix(clippy): rename unused args and unwrap result
taikoxyz · Jul 22, 2024 · f6d02b3 · f6d02b3
1 parent a017242
commit f6d02b3
Show file tree

Hide file tree

Showing 33 changed files with 1,417 additions and 987 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -39,7 +39,10 @@ raiko-core = { path = "./core" }
 raiko-tasks = { path = "./tasks" }
 
 # reth
-reth-primitives = { git = "https://github.com/taikoxyz/taiko-reth.git", branch = "v1.0.0-rc.2-taiko", default-features = false, features = ["alloy-compat", "taiko"] }
+reth-primitives = { git = "https://github.com/taikoxyz/taiko-reth.git", branch = "v1.0.0-rc.2-taiko", default-features = false, features = [
+    "alloy-compat",
+    "taiko",
+] }
 reth-evm-ethereum = { git = "https://github.com/taikoxyz/taiko-reth.git", branch = "v1.0.0-rc.2-taiko", default-features = false }
 reth-evm = { git = "https://github.com/taikoxyz/taiko-reth.git", branch = "v1.0.0-rc.2-taiko", default-features = false }
 reth-rpc-types = { git = "https://github.com/taikoxyz/taiko-reth.git", branch = "v1.0.0-rc.2-taiko", default-features = false }
@@ -77,16 +80,14 @@ alloy-network = { version = "0.1", default-features = false, features = [
     "k256",
 ] }
 alloy-contract = { version = "0.1", default-features = false }
-alloy-eips = { version = "0.1", default-features = false, features = [
-    "serde",
-] }
+alloy-eips = { version = "0.1", default-features = false, features = ["serde"] }
 alloy-provider = { version = "0.1", default-features = false, features = [
     "reqwest",
 ] }
-alloy-transport-http = { version = "0.1",default-features = false, features = [
+alloy-transport-http = { version = "0.1", default-features = false, features = [
     "reqwest",
 ] }
-alloy-signer = { version = "0.1", default-features = false}
+alloy-signer = { version = "0.1", default-features = false }
 alloy-signer-local = { version = "0.1", default-features = false }
 
 # ethers (TODO: remove)
@@ -135,6 +136,7 @@ utoipa = { version = "4.2.0", features = ["axum_extras"] }
 structopt = "0.3.24"
 prometheus = { version = "0.13.3", features = ["process"] }
 tokio = { version = "^1.23", features = ["full"] }
+tokio-util = { version = "0.7.11" }
 reqwest = { version = "0.11.22", features = ["json"] }
 url = "2.5.0"
 async-trait = "0.1.80"

diff --git a/core/Cargo.toml b/core/Cargo.toml
@@ -14,11 +14,11 @@ sgx-prover = { path = "../provers/sgx/prover", optional = true }
 raiko-lib = { workspace = true }
 
 # reth
-reth-primitives.workspace = true
-reth-evm-ethereum.workspace = true
-reth-evm.workspace = true
-reth-revm.workspace = true
-reth-provider.workspace = true
+reth-primitives = { workspace = true }
+reth-evm-ethereum = { workspace = true }
+reth-evm = { workspace = true }
+reth-revm = { workspace = true }
+reth-provider = { workspace = true }
 
 # alloy
 alloy-rlp = { workspace = true }

diff --git a/core/src/interfaces.rs b/core/src/interfaces.rs
@@ -4,7 +4,7 @@ use clap::{Args, ValueEnum};
 use raiko_lib::{
     input::{BlobProofType, GuestInput, GuestOutput},
     primitives::eip4844::{calc_kzg_proof, commitment_to_version_hash, kzg_proof_to_bytes},
-    prover::{Proof, Prover, ProverError},
+    prover::{IdStore, IdWrite, Proof, ProofKey, Prover, ProverError},
 };
 use reth_primitives::hex;
 use serde::{Deserialize, Serialize};
@@ -160,30 +160,31 @@ impl ProofType {
         input: GuestInput,
         output: &GuestOutput,
         config: &Value,
+        store: Option<&mut dyn IdWrite>,
     ) -> RaikoResult<Proof> {
         let mut proof = match self {
-            ProofType::Native => NativeProver::run(input.clone(), output, config)
+            ProofType::Native => NativeProver::run(input.clone(), output, config, store)
                 .await
-                .map_err(<ProverError as Into<RaikoError>>::into),
+                .map_err(|e| e.into()),
             ProofType::Sp1 => {
                 #[cfg(feature = "sp1")]
-                return sp1_driver::Sp1Prover::run(input.clone(), output, config)
+                return sp1_driver::Sp1Prover::run(input.clone(), output, config, store)
                     .await
                     .map_err(|e| e.into());
                 #[cfg(not(feature = "sp1"))]
                 Err(RaikoError::FeatureNotSupportedError(*self))
             }
             ProofType::Risc0 => {
                 #[cfg(feature = "risc0")]
-                return risc0_driver::Risc0Prover::run(input.clone(), output, config)
+                return risc0_driver::Risc0Prover::run(input.clone(), output, config, store)
                     .await
                     .map_err(|e| e.into());
                 #[cfg(not(feature = "risc0"))]
                 Err(RaikoError::FeatureNotSupportedError(*self))
             }
             ProofType::Sgx => {
                 #[cfg(feature = "sgx")]
-                return sgx_prover::SgxProver::run(input.clone(), output, config)
+                return sgx_prover::SgxProver::run(input.clone(), output, config, store)
                     .await
                     .map_err(|e| e.into());
                 #[cfg(not(feature = "sgx"))]
@@ -195,18 +196,55 @@ impl ProofType {
         if let Some(blob_commitment) = input.taiko.blob_commitment.clone() {
             let kzg_proof = calc_kzg_proof(
                 &input.taiko.tx_data,
-                &commitment_to_version_hash(&blob_commitment.try_into().unwrap()),
+                &commitment_to_version_hash(&blob_commitment.try_into().map_err(|_| {
+                    RaikoError::Conversion(
+                        "Could not convert blob commitment to version hash".to_owned(),
+                    )
+                })?),
             )
-            .unwrap();
-            let kzg_proof_hex = hex::encode(kzg_proof_to_bytes(&kzg_proof));
-            proof
-                .as_object_mut()
-                .unwrap()
-                .insert("kzg_proof".to_string(), Value::String(kzg_proof_hex));
+            .map_err(|e| anyhow::anyhow!(e))?;
+            proof.kzg_proof = Some(hex::encode(kzg_proof_to_bytes(&kzg_proof)));
         }
 
         Ok(proof)
     }
+
+    pub async fn cancel_proof(
+        &self,
+        proof_key: ProofKey,
+        read: Box<&mut dyn IdStore>,
+    ) -> RaikoResult<()> {
+        match self {
+            ProofType::Native => NativeProver::cancel(proof_key, read)
+                .await
+                .map_err(|e| e.into()),
+            ProofType::Sp1 => {
+                #[cfg(feature = "sp1")]
+                return sp1_driver::Sp1Prover::cancel(proof_key, read)
+                    .await
+                    .map_err(|e| e.into());
+                #[cfg(not(feature = "sp1"))]
+                Err(RaikoError::FeatureNotSupportedError(*self))
+            }
+            ProofType::Risc0 => {
+                #[cfg(feature = "risc0")]
+                return risc0_driver::Risc0Prover::cancel(proof_key, read)
+                    .await
+                    .map_err(|e| e.into());
+                #[cfg(not(feature = "risc0"))]
+                Err(RaikoError::FeatureNotSupportedError(*self))
+            }
+            ProofType::Sgx => {
+                #[cfg(feature = "sgx")]
+                return sgx_prover::SgxProver::cancel(proof_key, read)
+                    .await
+                    .map_err(|e| e.into());
+                #[cfg(not(feature = "sgx"))]
+                Err(RaikoError::FeatureNotSupportedError(*self))
+            }
+        }?;
+        Ok(())
+    }
 }
 
 #[serde_as]

diff --git a/core/src/lib.rs b/core/src/lib.rs
@@ -5,11 +5,20 @@ use crate::{
 };
 use alloy_primitives::Address;
 use alloy_rpc_types::EIP1186AccountProofResponse;
-use raiko_lib::builder::{create_mem_db, RethBlockBuilder};
-use raiko_lib::consts::{ChainSpec, VerifierType};
-use raiko_lib::input::{GuestInput, GuestOutput, TaikoProverData};
 use raiko_lib::protocol_instance::ProtocolInstance;
 use raiko_lib::prover::Proof;
+use raiko_lib::{
+    builder::{create_mem_db, RethBlockBuilder},
+    prover::ProofKey,
+};
+use raiko_lib::{
+    consts::{ChainSpec, VerifierType},
+    prover::IdStore,
+};
+use raiko_lib::{
+    input::{GuestInput, GuestOutput, TaikoProverData},
+    prover::IdWrite,
+};
 use reth_primitives::Header;
 use serde_json::Value;
 use std::{collections::HashMap, hint::black_box};
@@ -92,13 +101,26 @@ impl Raiko {
         }
     }
 
-    pub async fn prove(&self, input: GuestInput, output: &GuestOutput) -> RaikoResult<Proof> {
-        let data = serde_json::to_value(&self.request)?;
+    pub async fn prove(
+        &self,
+        input: GuestInput,
+        output: &GuestOutput,
+        store: Option<&mut dyn IdWrite>,
+    ) -> RaikoResult<Proof> {
+        let config = serde_json::to_value(&self.request)?;
         self.request
             .proof_type
-            .run_prover(input, output, &data)
+            .run_prover(input, output, &config, store)
             .await
     }
+
+    pub async fn cancel(
+        &self,
+        proof_key: ProofKey,
+        read: Box<&mut dyn IdStore>,
+    ) -> RaikoResult<()> {
+        self.request.proof_type.cancel_proof(proof_key, read).await
+    }
 }
 
 fn check_header(exp: &Header, header: &Header) -> Result<(), RaikoError> {
@@ -268,7 +290,7 @@ mod tests {
             .expect("input generation failed");
         let output = raiko.get_output(&input).expect("output generation failed");
         let _proof = raiko
-            .prove(input, &output)
+            .prove(input, &output, None)
             .await
             .expect("proof generation failed");
     }