feat(coprocessor): sns-worker, accept missing keys

Author: rudy-6-4

coprocessor/fhevm-engine/sns-worker/src/executor.rs

@@ -154,6 +154,18 @@ impl SwitchNSquashService {
     }
 }
 
+async fn get_keyset(pool: PgPool,
+    keys_cache: Arc<RwLock<lru::LruCache<String, KeySet>>>,
+    tenant_api_key: &String
+) -> Result<Option<KeySet>, ExecutionError> {
+    let t = telemetry::tracer("worker_loop_init");
+    let s = t.child_span("fetch_keyset");
+    let keys: Option<KeySet> = fetch_keyset(&keys_cache, &pool, tenant_api_key).await?;
+    telemetry::end_span(s);
+    t.end();
+    Ok(keys)
+}
+
 /// Executes the worker logic for the SnS task.
 pub(crate) async fn run_loop(
     conf: Config,
@@ -173,13 +185,16 @@ pub(crate) async fn run_loop(
         .listen_all(conf.db.listen_channels.iter().map(|v| v.as_str()))
         .await?;
 
-    let t = telemetry::tracer("worker_loop_init");
-    let s = t.child_span("fetch_keyset");
-    let keys: KeySet = fetch_keyset(&keys_cache, &pool, tenant_api_key).await?;
-    telemetry::end_span(s);
-    t.end();
-
-    info!(tenant_api_key = tenant_api_key, "Fetched keyset");
+    let mut keys = match get_keyset(pool.clone(), keys_cache.clone(), tenant_api_key).await? {
+        Some(keys) => {
+            info!(tenant_api_key = tenant_api_key, "Fetched keyset");
+            Some(keys)
+        }
+        None => {
+            warn!(tenant_api_key = tenant_api_key, "No keys found for the given tenant_api_key");
+            None
+        }
+    };
 
     let mut gc_ticker = interval(conf.db.cleanup_interval);
     let mut gc_timestamp = SystemTime::now();
@@ -189,6 +204,16 @@ pub(crate) async fn run_loop(
         // Continue looping until the service is cancelled or a critical error occurs
         update_last_active(last_active_at.clone()).await;
 
+        let Some(keys) = keys.as_ref() else {
+            warn!(tenant_api_key = tenant_api_key, "No keys available, retrying in 5 seconds");
+            tokio::time::sleep(Duration::from_secs(5)).await;
+            if token.is_cancelled() {
+                return Ok(());
+            }
+            keys = get_keyset(pool.clone(), keys_cache.clone(), tenant_api_key).await?;
+            continue;
+        };
+
         let maybe_remaining = fetch_and_execute_sns_tasks(&pool, &tx, &keys, &conf, &token).await?;
         if maybe_remaining {
             if token.is_cancelled() {

coprocessor/fhevm-engine/sns-worker/src/keyset.rs

@@ -16,22 +16,24 @@ pub(crate) async fn fetch_keyset(
     cache: &Arc<RwLock<lru::LruCache<String, KeySet>>>,
     pool: &PgPool,
     tenant_api_key: &String,
-) -> Result<KeySet, ExecutionError> {
+) -> Result<Option<KeySet>, ExecutionError> {
     let mut cache = cache.write().await;
     if let Some(keys) = cache.get(tenant_api_key) {
         info!(tenant_api_key, "Cache hit");
-        return Ok(keys.clone());
+        return Ok(Some(keys.clone()));
     }
 
     info!(tenant_api_key, "Cache miss");
-    let (client_key, server_key) = fetch_keys(pool, tenant_api_key).await?;
+    let Some((client_key, server_key)) = fetch_keys(pool, tenant_api_key).await? else {
+        return Ok(None);
+    };
     let key_set: KeySet = KeySet {
         client_key,
         server_key,
     };
 
     cache.push(tenant_api_key.clone(), key_set.clone());
-    Ok(key_set)
+    Ok(Some(key_set))
 }
 
 /// Retrieve both the ClientKey and ServerKey from the tenants table
@@ -44,7 +46,7 @@ pub(crate) async fn fetch_keyset(
 pub async fn fetch_keys(
     pool: &PgPool,
     tenant_api_key: &String,
-) -> anyhow::Result<(Option<tfhe::ClientKey>, tfhe::ServerKey)> {
+) -> anyhow::Result<Option<(Option<tfhe::ClientKey>, tfhe::ServerKey)>> {
     let blob = read_keys_from_large_object(
         pool,
         tenant_api_key,
@@ -53,6 +55,9 @@ pub async fn fetch_keys(
     )
     .await?;
     info!(bytes_len = blob.len(), "Retrieved sns_pk");
+    if blob.is_empty() {
+        return Ok(None);
+    }
 
     let server_key: tfhe::ServerKey = safe_deserialize_sns_key(&blob)?;
 
@@ -70,9 +75,9 @@ pub async fn fetch_keys(
         if !cks.is_empty() {
             info!(bytes_len = cks.len(), "Retrieved cks");
             let client_key: tfhe::ClientKey = safe_deserialize_sns_key(&cks)?;
-            return Ok((Some(client_key), server_key));
+            return Ok(Some((Some(client_key), server_key)));
         }
     }
 
-    Ok((None, server_key))
+    Ok(Some((None, server_key)))
 }