feat: add task cache (#49)

.gitignore

@@ -22,6 +22,7 @@ target/
 
 config.json
 keys/
+db/
 
 # Helm chart related
 *.lock

Cargo.toml

@@ -46,3 +46,4 @@ tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 axum = "0.6.0"
 dotenv = "0.15"
+rocksdb = "0.20"

src/config.rs

@@ -8,6 +8,7 @@ use std::fs::File;
 pub struct Config {
     pub prover_name_prefix: String,
     pub keys_dir: String,
+    pub db_path: String,
     pub coordinator: CoordinatorConfig,
     pub l2geth: Option<L2GethConfig>,
     pub prover: ProverConfig,
@@ -81,11 +82,13 @@ impl Config {
     }
 
     fn get_env_var(key: &str) -> anyhow::Result<Option<String>> {
-        Ok(std::env::var_os(key).map(|val| {
-            val.to_str()
-                .ok_or_else(|| anyhow::anyhow!("{key} env var is not valid UTF-8"))
-                .map(String::from)
-        }).transpose()?)
+        Ok(std::env::var_os(key)
+            .map(|val| {
+                val.to_str()
+                    .ok_or_else(|| anyhow::anyhow!("{key} env var is not valid UTF-8"))
+                    .map(String::from)
+            })
+            .transpose()?)
     }
 
     fn override_with_env(&mut self) -> anyhow::Result<()> {
@@ -118,6 +121,9 @@ impl Config {
                 cloud.api_key = val;
             }
         }
+        if let Some(val) = Self::get_env_var("DB_PATH")? {
+            self.db_path = val;
+        }
 
         Ok(())
     }

src/coordinator_handler/coordinator_client.rs

@@ -10,7 +10,7 @@ pub struct CoordinatorClient {
     vks: Vec<String>,
     circuit_version: String,
     pub prover_name: String,
-    key_signer: KeySigner,
+    pub key_signer: KeySigner,
     api: Api,
     token: Mutex<Option<String>>,
 }

src/db.rs

@@ -1 +1,65 @@
+use crate::coordinator_handler::GetTaskResponseData;
+use rocksdb::DB;
 
+pub struct Db {
+    db: DB,
+}
+
+impl Db {
+    pub fn new(path: &str) -> anyhow::Result<Self> {
+        let db = DB::open_default(path)?;
+        Ok(Self { db })
+    }
+
+    pub fn get_coordinator_task_by_public_key(
+        &self,
+        public_key: String,
+    ) -> Option<GetTaskResponseData> {
+        self.db
+            .get(fmt_coordinator_task_key(public_key))
+            .ok()?
+            .as_ref()
+            .map(|v| serde_json::from_slice(v).ok())
+            .flatten()
+    }
+
+    pub fn get_proving_task_id_by_public_key(&self, public_key: String) -> Option<String> {
+        self.db
+            .get(fmt_proving_task_id_key(public_key))
+            .ok()?
+            .map(|v| String::from_utf8(v).ok())
+            .flatten()
+    }
+
+    pub fn set_coordinator_task_by_public_key(
+        &self,
+        public_key: String,
+        coordinator_task: &GetTaskResponseData,
+    ) {
+        let _ = serde_json::to_vec(coordinator_task)
+            .map(|bytes| self.db.put(fmt_coordinator_task_key(public_key), bytes));
+    }
+
+    pub fn set_proving_task_id_by_public_key(&self, public_key: String, proving_task_id: String) {
+        let _ = self.db.put(
+            fmt_proving_task_id_key(public_key),
+            proving_task_id.as_bytes(),
+        );
+    }
+
+    pub fn delete_coordinator_task_by_public_key(&self, public_key: String) {
+        let _ = self.db.delete(fmt_coordinator_task_key(public_key));
+    }
+
+    pub fn delete_proving_task_id_by_public_key(&self, public_key: String) {
+        let _ = self.db.delete(fmt_proving_task_id_key(public_key));
+    }
+}
+
+fn fmt_coordinator_task_key(public_key: String) -> String {
+    format!("last_coordinator_task_{}", public_key)
+}
+
+fn fmt_proving_task_id_key(public_key: String) -> String {
+    format!("last_proving_task_id_{}", public_key)
+}