feat(host,provers,tasks): handle local job cancellation

taikoxyz · Aug 15, 2024 · 65e9fcd · 65e9fcd
1 parent 5ea0c33
commit 65e9fcd
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 6 deletions.
diff --git a/host/src/proof.rs b/host/src/proof.rs
@@ -65,7 +65,7 @@ impl ProofActor {
             )
             .await
             .or_else(|e| {
-                if e.to_string().contains("no id found") {
+                if e.to_string().contains("No data for query") {
                     warn!("Task already cancelled or not yet started!");
                     Ok(())
                 } else {

diff --git a/provers/risc0/driver/src/lib.rs b/provers/risc0/driver/src/lib.rs
@@ -112,7 +112,16 @@ impl Prover for Risc0Prover {
     }
 
     async fn cancel(key: ProofKey, id_store: Box<&mut dyn IdStore>) -> ProverResult<()> {
-        let uuid = id_store.read_id(key).await?;
+        let uuid = match id_store.read_id(key).await {
+            Ok(uuid) => uuid,
+            Err(e) => {
+                if e.to_string().contains("No data for query") {
+                    return Ok(());
+                } else {
+                    return Err(ProverError::GuestError(e.to_string()));
+                }
+            }
+        };
         cancel_proof(uuid)
             .await
             .map_err(|e| ProverError::GuestError(e.to_string()))?;

diff --git a/provers/sp1/driver/src/lib.rs b/provers/sp1/driver/src/lib.rs
@@ -158,7 +158,16 @@ impl Prover for Sp1Prover {
     }
 
     async fn cancel(key: ProofKey, id_store: Box<&mut dyn IdStore>) -> ProverResult<()> {
-        let proof_id = id_store.read_id(key).await?;
+        let proof_id = match id_store.read_id(key).await {
+            Ok(proof_id) => proof_id,
+            Err(e) => {
+                if e.to_string().contains("No data for query") {
+                    return Ok(());
+                } else {
+                    return Err(ProverError::GuestError(e.to_string()));
+                }
+            }
+        };
         let private_key = env::var("SP1_PRIVATE_KEY").map_err(|_| {
             ProverError::GuestError("SP1_PRIVATE_KEY must be set for remote proving".to_owned())
         })?;

diff --git a/tasks/src/adv_sqlite.rs b/tasks/src/adv_sqlite.rs
@@ -824,14 +824,22 @@ impl TaskDb {
               1;
             "#,
         )?;
-        let query = statement.query_row(
+        let query = match statement.query_row(
             named_params! {
                 ":chain_id": chain_id,
                 ":blockhash": blockhash.to_vec(),
                 ":proofsys_id": proof_key,
             },
             |row| row.get::<_, String>(0),
-        )?;
+        ) {
+            Ok(q) => q,
+            Err(e) => {
+                return match e {
+                    rusqlite::Error::QueryReturnedNoRows => TaskManagerError::NoData,
+                    e => e.into(),
+                }
+            }
+        };
 
         Ok(query)
     }

diff --git a/tasks/src/lib.rs b/tasks/src/lib.rs
@@ -27,6 +27,8 @@ pub enum TaskManagerError {
     IOError(IOErrorKind),
     #[error("SQL Error {0}")]
     SqlError(String),
+    #[error("No data for query")]
+    NoData,
     #[error("Anyhow error: {0}")]
     Anyhow(String),
 }

diff --git a/tasks/src/mem_db.rs b/tasks/src/mem_db.rs
@@ -143,7 +143,7 @@ impl InMemoryTaskDb {
         self.store
             .get(&key)
             .cloned()
-            .ok_or_else(|| TaskManagerError::SqlError("no id found".to_owned()))
+            .ok_or_else(|| TaskManagerError::NoData)
     }
 }