task-manager: add status check

taikoxyz · May 15, 2024 · 9305c18 · 9305c18
1 parent 1d9a9ab
commit 9305c18
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 7 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -91,6 +91,7 @@ serde = { version = "1.0", default-features = false, features = [
 ] }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
 serde_with = { version = "3.0" }
+num_enum = "0.7.2"
 
 # tracing
 tracing = "0.1"

diff --git a/task_manager/Cargo.toml b/task_manager/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021" # { workspace = true }
 [dependencies]
 raiko-primitives = { workspace = true }
 rusqlite = { workspace = true }
+num_enum = { workspace = true }
 
 [dev-dependencies]
 rand = "0.9.0-alpha.1" # This is an alpha version, that has rng.gen_iter::<T>()

diff --git a/task_manager/src/lib.rs b/task_manager/src/lib.rs
@@ -163,6 +163,8 @@ use raiko_primitives::{BlockNumber, ChainId, B256};
 use rusqlite::{named_params, Statement};
 use rusqlite::{Connection, OpenFlags};
 
+use num_enum::{IntoPrimitive, FromPrimitive};
+
 // Types
 // ----------------------------------------------------------------
 
@@ -208,7 +210,8 @@ pub enum TaskProofsys {
 
 #[allow(non_camel_case_types)]
 #[rustfmt::skip]
-#[derive(Debug, Copy, Clone)]
+#[derive(PartialEq, Debug, Copy, Clone, IntoPrimitive, FromPrimitive)]
+#[repr(i32)]
 pub enum TaskStatus {
     Success                   =     0,
     Registered                =  1000,
@@ -222,6 +225,8 @@ pub enum TaskStatus {
     CancellationInProgress    = -3210,
     InvalidOrUnsupportedBlock = -4000,
     UnspecifiedFailureReason  = -9999,
+    #[num_enum(default)]
+    SqlDbCorruption           = -99999,
 }
 
 // Implementation
@@ -611,8 +616,8 @@ impl TaskDb {
         let get_task_proving_status = conn.prepare(
             "
             SELECT
-                thirdparty_desc,
-                id_status,
+                t3p.thirdparty_desc,
+                ts.id_status,
                 MAX(timestamp)
             FROM
                 task_status ts
@@ -625,7 +630,7 @@ impl TaskDb {
                 AND t.blockhash = :blockhash
                 AND t.id_proofsys = :id_proofsys
             GROUP BY
-                ts.id_thirdparty
+                t3p.id_thirdparty
             ORDER BY
                 ts.timestamp DESC;
             ")?;
@@ -710,12 +715,27 @@ impl<'db> TaskManager<'db> {
             ":blockhash": blockhash.as_slice(),
             ":id_proofsys": proof_system as u8,
             ":fulfiller": fulfiller,
-            ":id_status": status as isize,
+            ":id_status": status as i32,
             ":proof": proof
         })?;
         Ok(())
     }
 
+    pub fn get_task_proving_status(
+        &mut self,
+        chain_id: ChainId,
+        blockhash: &B256,
+        proof_system: TaskProofsys,
+    ) -> Result<TaskStatus, TaskManagerError> {
+        let proving_status = self.get_task_proving_status.query_row(named_params! {
+            ":chain_id": chain_id as u64,
+            ":blockhash": blockhash.as_slice(),
+            ":id_proofsys": proof_system as u8,
+        }, |r| r.get::<_, i32>(0).map(TaskStatus::from))?;
+
+        Ok(proving_status)
+    }
+
     pub fn get_task_proof(
         &mut self,
         chain_id: ChainId,

diff --git a/task_manager/tests/main.rs b/task_manager/tests/main.rs
@@ -164,6 +164,11 @@ mod tests {
                 &payload,
             ).unwrap();
 
+            assert_eq!(
+                TaskStatus::Registered,
+                tama.get_task_proving_status(chain_id, &blockhash, proofsys).unwrap()
+            );
+
             tasks.push((
                 chain_id,
                 blockhash,
@@ -180,6 +185,11 @@ mod tests {
             None
         ).unwrap();
 
+        assert_eq!(
+            TaskStatus::Cancelled_NeverStarted,
+            tama.get_task_proving_status(tasks[0].0, &tasks[0].1, tasks[0].2).unwrap()
+        );
+
         // -----------------------
 
         tama.update_task_progress(
@@ -191,6 +201,11 @@ mod tests {
             None
         ).unwrap();
 
+        assert_eq!(
+            TaskStatus::WorkInProgress,
+            tama.get_task_proving_status(tasks[1].0, &tasks[1].1, tasks[1].2).unwrap()
+        );
+
         tama.update_task_progress(
             tasks[1].0,
             &tasks[1].1,
@@ -200,6 +215,11 @@ mod tests {
             None
         ).unwrap();
 
+        assert_eq!(
+            TaskStatus::CancellationInProgress,
+            tama.get_task_proving_status(tasks[1].0, &tasks[1].1, tasks[1].2).unwrap()
+        );
+
         tama.update_task_progress(
             tasks[1].0,
             &tasks[1].1,
@@ -209,6 +229,11 @@ mod tests {
             None
         ).unwrap();
 
+        assert_eq!(
+            TaskStatus::Cancelled,
+            tama.get_task_proving_status(tasks[1].0, &tasks[1].1, tasks[1].2).unwrap()
+        );
+
         // -----------------------
 
         tama.update_task_progress(
@@ -220,6 +245,11 @@ mod tests {
             None
         ).unwrap();
 
+        assert_eq!(
+            TaskStatus::WorkInProgress,
+            tama.get_task_proving_status(tasks[2].0, &tasks[2].1, tasks[2].2).unwrap()
+        );
+
         let proof: Vec<_> = (&mut rng).gen_iter::<u8>().take(128).collect();
         tama.update_task_progress(
             tasks[2].0,
@@ -230,11 +260,16 @@ mod tests {
             Some(&proof)
         ).unwrap();
 
-        // -----------------------
+        assert_eq!(
+            TaskStatus::Success,
+            tama.get_task_proving_status(tasks[2].0, &tasks[2].1, tasks[2].2).unwrap()
+        );
+
         assert_eq!(
             proof,
             tama.get_task_proof(tasks[2].0, &tasks[2].1, tasks[2].2).unwrap()
-        )
+        );
+
     }
 
 }