fix(coprocessor): persist upload tasks and mark existing ciphertexts as uploaded

goshawk-3 · goshawk-3 · commit 384f7888733a · 2025-05-22T18:17:14.000+03:00
- Always persist upload tasks to the ciphertext_digest table
- If ciphertexts are already uploaded to S3, mark them as uploaded in the database to keep state consistent
diff --git a/coprocessor/fhevm-engine/sns-executor/src/aws_upload.rs b/coprocessor/fhevm-engine/sns-executor/src/aws_upload.rs
@@ -74,15 +74,31 @@ pub(crate) async fn process_s3_uploads(
                     Some(task) => task,
                     None => return Ok(()),
                 };
-
-                let trx = insert_and_lock(task.tenant_id, task.handle.clone(), &pool).await?;
+                debug!("Received task, handle: {}", hex::encode(&task.handle));
+
+                let mut trx_lock = pool.begin().await?;
+                task.enqueue_upload_task(&mut trx_lock).await?;
+
+                if let Err(err) = sqlx::query!(
+                    "SELECT * FROM ciphertext_digest
+                    WHERE handle = $2 AND tenant_id = $1 AND
+                    (ciphertext128 IS NULL OR ciphertext IS NULL)
+                    FOR UPDATE SKIP LOCKED",
+                    task.tenant_id,
+                    task.handle,
+                )
+                .fetch_one(trx_lock.as_mut())
+                .await {
+                    error!("Failed to lock pending uploads {}, handle: {}", err, compact_hex(&task.handle));
+                    trx_lock.rollback().await?;
+                    continue;
+                }
 
                 if !is_ready.load(Ordering::SeqCst) {
                     // If the S3 setup is not ready, we need to wait for its ready status
                     // before we can continue spawning uploading job
                     info!("Upload task skipped, S3 connection still not ready");
-                    // Queue the uploading job in the database
-                    trx.commit().await?;
+                    trx_lock.commit().await?;
                     continue;
                 }
 
@@ -111,7 +127,7 @@ pub(crate) async fn process_s3_uploads(
                 // Spawn a new task to upload the ciphertexts
                 let h = tokio::spawn(async move {
                     let s = task.otel.child_span("upload_s3");
-                        if let Err(err) = upload_ciphertexts(trx, task, &client, &conf).await {
+                        if let Err(err) = upload_ciphertexts(trx_lock, task, &client, &conf).await {
                             if let ExecutionError::S3TransientError(_) = err {
                                 ready_flag.store(false, Ordering::SeqCst );
                                 info!("S3 setup is not ready, due to transient error: {}", err);
@@ -146,40 +162,6 @@ pub(crate) async fn process_s3_uploads(
     }
 }
 
-async fn insert_and_lock(
-    tenant_id: i32,
-    handle: Vec<u8>,
-    pool: &PgPool,
-) -> Result<Transaction<'static, Postgres>, ExecutionError> {
-    let mut trx = pool.begin().await?;
-
-    sqlx::query!(
-        "INSERT INTO ciphertext_digest (tenant_id, handle)
-        VALUES ($1, $2) ON CONFLICT DO NOTHING",
-        tenant_id,
-        &handle,
-    )
-    .execute(trx.as_mut())
-    .await?;
-
-    trx.commit().await?;
-
-    let mut trx = pool.begin().await?;
-
-    sqlx::query!(
-        "SELECT * FROM ciphertext_digest
-        WHERE handle = $2 AND tenant_id = $1 AND
-        (ciphertext128 IS NULL OR ciphertext IS NULL)
-        FOR UPDATE SKIP LOCKED",
-        tenant_id,
-        handle,
-    )
-    .fetch_all(trx.as_mut())
-    .await?;
-
-    Ok(trx)
-}
-
 enum UploadResult {
     CtType128((Vec<u8>, BoxedSpan)),
     CtType64((Vec<u8>, BoxedSpan)),
@@ -230,6 +212,14 @@ async fn upload_ciphertexts(
                     .send(),
                 UploadResult::CtType128((ct128_digest.clone(), span)),
             ));
+        } else {
+            info!(
+                "ct128 already exists in S3, handle: {}, digest: {}",
+                handle_as_hex,
+                hex::encode(&ct128_digest)
+            );
+
+            task.update_ct128_uploaded(&mut trx, ct128_digest).await?;
         }
     }
 
@@ -261,8 +251,14 @@ async fn upload_ciphertexts(
                     .send(),
                 UploadResult::CtType64((ct64_digest.clone(), span)),
             ));
+        } else {
+            info!(
+                "ct64 already exists in S3, handle: {}, digest: {}",
+                handle_as_hex,
+                hex::encode(&ct64_digest)
+            );
 
-            // TODO: Update DB
+            task.update_ct64_uploaded(&mut trx, ct64_digest).await?;
         }
     }
 
@@ -287,35 +283,7 @@ async fn upload_ciphertexts(
                     telemetry::end_span_with_err(span, err.to_string());
                     transient_error = Some(ExecutionError::S3TransientError(err.to_string()));
                 } else {
-                    sqlx::query!(
-                        "UPDATE ciphertext_digest
-                        SET ciphertext128 = $1
-                        WHERE handle = $2",
-                        digest,
-                        task.handle
-                    )
-                    .execute(trx.as_mut())
-                    .await?;
-
-                    // Reset ciphertext128 as the ct128 has been successfully uploaded to S3
-                    // NB: For reclaiming the disk-space in DB, we rely on auto vacuuming in
-                    // Postgres
-
-                    sqlx::query!(
-                        "UPDATE ciphertexts
-                        SET ciphertext128 = NULL
-                        WHERE handle = $1",
-                        task.handle
-                    )
-                    .execute(trx.as_mut())
-                    .await?;
-
-                    info!(
-                        "Uploaded ct128, handle: {}, digest: {}",
-                        handle_as_hex,
-                        compact_hex(&digest)
-                    );
-
+                    task.update_ct128_uploaded(&mut trx, digest).await?;
                     telemetry::end_span_with_timestamp(span, finish_time);
                 }
             }
@@ -329,28 +297,13 @@ async fn upload_ciphertexts(
                     telemetry::end_span_with_err(span, err.to_string());
                     transient_error = Some(ExecutionError::S3TransientError(err.to_string()));
                 } else {
-                    sqlx::query!(
-                        "UPDATE ciphertext_digest
-                            SET ciphertext = $1
-                            WHERE handle = $2",
-                        digest,
-                        task.handle
-                    )
-                    .execute(trx.as_mut())
-                    .await?;
-                    info!(
-                        "Uploaded ct64, handle: {}, digest: {}",
-                        handle_as_hex,
-                        compact_hex(&digest)
-                    );
-
+                    task.update_ct64_uploaded(&mut trx, digest).await?;
                     telemetry::end_span_with_timestamp(span, finish_time);
                 }
             }
         }
     }
 
-    // TODO: Move this notify in DB query
     sqlx::query("SELECT pg_notify($1, '')")
         .bind(EVENT_CIPHERTEXTS_UPLOADED)
         .execute(trx.as_mut())
diff --git a/coprocessor/fhevm-engine/sns-executor/src/executor.rs b/coprocessor/fhevm-engine/sns-executor/src/executor.rs
@@ -120,11 +120,13 @@ async fn fetch_and_execute_sns_tasks(
         }
     };
 
-    if let Some(mut tasks) = query_sns_tasks(&mut db_txn, conf.batch_limit).await? {
-        process_tasks(&mut tasks, keys, tx)?;
-        update_computations_status(&mut db_txn, &tasks).await?;
-        update_ciphertext128(&mut db_txn, &tasks).await?;
-        notify_ciphertext128_ready(&mut db_txn, &conf.notify_channel).await?;
+    let trx = &mut db_txn;
+
+    if let Some(mut tasks) = query_sns_tasks(trx, conf.batch_limit).await? {
+        process_tasks(trx, &mut tasks, keys, tx).await?;
+        update_computations_status(trx, &tasks).await?;
+        update_ciphertext128(trx, &tasks).await?;
+        notify_ciphertext128_ready(trx, &conf.notify_channel).await?;
         db_txn.commit().await?;
     } else {
         db_txn.rollback().await?;
@@ -229,13 +231,12 @@ async fn get_remaining_tasks(
 }
 
 /// Processes the tasks by decompressing and transforming ciphertexts.
-fn process_tasks(
+async fn process_tasks(
+    db_txn: &mut Transaction<'_, Postgres>,
     tasks: &mut [HandleItem],
     keys: &KeySet,
     tx: &Sender<HandleItem>,
 ) -> Result<(), ExecutionError> {
-    set_server_key(keys.server_key.clone());
-
     for task in tasks.iter_mut() {
         let ct64_compressed = task.ct64_compressed.as_ref().ok_or_else(|| {
             ExecutionError::MissingCiphertext128(format!(
@@ -245,6 +246,7 @@ fn process_tasks(
         })?;
 
         let s = task.otel.child_span("decompress_ct64");
+        set_server_key(keys.server_key.clone());
         let ct = decompress_ct(&task.handle, ct64_compressed)?;
         telemetry::end_span(s);
 
@@ -277,7 +279,10 @@ fn process_tasks(
             }
         };
 
-        // Start uploading the ciphertexts sooner than later
+        // Enqueue the task for upload in DB
+        task.enqueue_upload_task(db_txn).await?;
+
+        // Start uploading the ciphertexts
         //
         // The service must continue running the squashed noise algorithm,
         // regardless of the availability of the upload worker.
@@ -294,8 +299,6 @@ fn process_tasks(
 
             error!({target = "worker", action = "review"},  "Failed to send task to upload worker: {err}");
             telemetry::end_span_with_err(task.otel.child_span("send_task"), err.to_string());
-
-            // TODO: Insert ciphertext
         }
     }
 
diff --git a/coprocessor/fhevm-engine/sns-executor/src/lib.rs b/coprocessor/fhevm-engine/sns-executor/src/lib.rs
@@ -8,8 +8,9 @@ mod tests;
 
 use std::time::Duration;
 
-use fhevm_engine_common::{telemetry::OtelTracer, types::FhevmError};
+use fhevm_engine_common::{telemetry::OtelTracer, types::FhevmError, utils::compact_hex};
 use serde::{Deserialize, Serialize};
+use sqlx::{Postgres, Transaction};
 use thiserror::Error;
 use tokio::sync::mpsc::{self, Sender};
 use tokio_util::sync::CancellationToken;
@@ -91,6 +92,89 @@ pub struct HandleItem {
     pub otel: OtelTracer,
 }
 
+impl HandleItem {
+    /// Enqueues the upload task into the database
+    ///
+    /// If inserted into the `ciphertext_digest` table means that the both (ct64 and ct128)
+    /// ciphertexts are ready to be uploaded to S3.
+    pub(crate) async fn enqueue_upload_task(
+        &self,
+        db_txn: &mut Transaction<'_, Postgres>,
+    ) -> Result<(), ExecutionError> {
+        sqlx::query!(
+            "INSERT INTO ciphertext_digest (tenant_id, handle)
+            VALUES ($1, $2) ON CONFLICT DO NOTHING",
+            self.tenant_id,
+            &self.handle,
+        )
+        .execute(db_txn.as_mut())
+        .await?;
+
+        Ok(())
+    }
+
+    pub(crate) async fn update_ct128_uploaded(
+        &self,
+        trx: &mut Transaction<'_, Postgres>,
+        digest: Vec<u8>,
+    ) -> Result<(), ExecutionError> {
+        sqlx::query!(
+            "UPDATE ciphertext_digest
+            SET ciphertext128 = $1
+            WHERE handle = $2",
+            digest,
+            self.handle
+        )
+        .execute(trx.as_mut())
+        .await?;
+
+        // Reset ciphertext128 as the ct128 has been successfully uploaded to S3
+        // NB: For reclaiming the disk-space in DB, we rely on auto vacuuming in
+        // Postgres
+
+        sqlx::query!(
+            "UPDATE ciphertexts
+             SET ciphertext128 = NULL
+             WHERE handle = $1",
+            self.handle
+        )
+        .execute(trx.as_mut())
+        .await?;
+
+        info!(
+            "Mark ct128 as uploaded, handle: {}, digest: {}",
+            compact_hex(&self.handle),
+            compact_hex(&digest)
+        );
+
+        Ok(())
+    }
+
+    pub(crate) async fn update_ct64_uploaded(
+        &self,
+        trx: &mut Transaction<'_, Postgres>,
+        digest: Vec<u8>,
+    ) -> Result<(), ExecutionError> {
+        sqlx::query!(
+            "UPDATE ciphertext_digest
+             SET ciphertext = $1
+             WHERE handle = $2",
+            digest,
+            self.handle
+        )
+        .execute(trx.as_mut())
+        .await?;
+
+        info!(
+            "Mark ct64 as uploaded, handle: {}, digest: {}",
+            compact_hex(&self.handle),
+            compact_hex(&digest)
+        );
+
+        Ok(())
+    }
+}
+
 #[derive(Error, Debug)]
 pub enum ExecutionError {
     #[error("Conversion error: {0}")]
@@ -126,8 +210,8 @@ pub enum ExecutionError {
     #[error("Deserialization error: {0}")]
     DeserializationError(String),
 
-    #[error("Bucket S3 upload: {0}")]
-    BucketNotExist(String),
+    #[error("Bucket not found {0}")]
+    BucketNotFound(String),
 
     #[error("S3 Transient error: {0}")]
     S3TransientError(String),
