fix(coprocessor): eliminate db contention between the two workers in sns-executor

Each worker performs write operations on non-shared database tables.
The executor do updates to pbs_computations and ciphertexts,
while the uploader only updates the ciphertext_digest, thus
optimizing overall performance by reducing write conflicts.

In addition, the executor does a garbage collecting for any uploaded ct128.

Author: goshawk-3

coprocessor/fhevm-engine/sns-executor/.sqlx/query-455eb6d9fe4e6ffd2fe7f45c2fd668b1e9884a2eff149dd2638089417bad1e95.json

@@ -92,7 +92,11 @@ pub(crate) async fn process_s3_uploads(
                 let mut trx = pool.begin().await?;
 
                 let item = match job {
-                    UploadJob::Normal(item) => item,
+                    UploadJob::Normal(item) => 
+                    {  
+                        item.enqueue_upload_task(&mut trx).await?;
+                        item
+                    },
                     UploadJob::DatabaseLock(item) => {
                         if let Err(err) = sqlx::query!(
                             "SELECT * FROM ciphertext_digest

coprocessor/fhevm-engine/sns-executor/.sqlx/query-5006562c943244b0e6dd278a60e76924ff0e43a08971950550f8c97a6651dfe3.json

@@ -34,6 +34,7 @@ fn construct_config() -> Config {
             batch_limit: args.work_items_batch_size,
             polling_interval: args.pg_polling_interval,
             max_connections: args.pg_pool_connections,
+            cleanup_interval: args.cleanup_interval,
         },
         s3: S3Config {
             bucket_ct128: args.bucket_name_ct128,

coprocessor/fhevm-engine/sns-executor/.sqlx/query-a006d9ccf85f04ed5c9d65eb759b6e83376343cfc56ad730b15d5fa476d5db37.json

@@ -71,6 +71,9 @@ pub struct Args {
 
     #[arg(long, default_value = "120s", value_parser = parse_duration)]
     pub s3_regular_recheck_duration: Duration,
+    
+    #[arg(long, default_value = "120s", value_parser = parse_duration)]
+    pub cleanup_interval: Duration,
 }
 
 pub fn parse_args() -> Args {

coprocessor/fhevm-engine/sns-executor/.sqlx/query-a8b64a29e5c7c917a75c5e758ad494cf75ab8291ba6d47366157827b3550cc0f.json

@@ -15,6 +15,7 @@ use std::time::SystemTime;
 use tfhe::set_server_key;
 use tokio::select;
 use tokio::sync::mpsc::Sender;
+use tokio::time::sleep;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info};
 
@@ -64,7 +65,7 @@ pub(crate) async fn run_loop(
             match acquire_connection(&pool, token.clone()).await {
                 ConnStatus::Established(conn) => conn,
                 ConnStatus::Failed => {
-                    tokio::time::sleep(RETRY_DB_CONN_INTERVAL).await;
+                    sleep(RETRY_DB_CONN_INTERVAL).await;
                     continue; // Retry to reacquire a connection
                 }
                 ConnStatus::Cancelled => break,
@@ -94,11 +95,15 @@ pub(crate) async fn run_loop(
 
             select! {
                 _ = token.cancelled() => return Ok(()),
-                notification = listener.try_recv() => {
-                    info!(target: "worker", "Received notification {:?}", notification);
+                n = listener.try_recv() => {
+                    info!(target: "worker", "Received notification {:?}", n);
                 },
-                _ = tokio::time::sleep(Duration::from_secs(conf.polling_interval.into())) => {
+                _ =  sleep(Duration::from_secs(conf.polling_interval.into())) => {
                     debug!(target: "worker", "Polling timeout, rechecking for tasks");
+                },
+                // Garbage collecting
+                _ = sleep(conf.cleanup_interval) => {
+                    garbage_collect(&mut conn).await?;
                 }
             }
         }
@@ -107,6 +112,31 @@ pub(crate) async fn run_loop(
     Ok(())
 }
 
+// Clean up the database by removing old ciphertexts128 already uploaded to S3.
+async fn garbage_collect(conn: &mut PoolConnection<Postgres>) -> Result<(), ExecutionError> {
+    let _s = telemetry::tracer("cleanup_ct128");
+    let rows = sqlx::query!(
+        "UPDATE ciphertexts c
+        SET ciphertext128 = NULL
+        WHERE ciphertext128 is NOT NULL -- needed for getting an accurate rows_affected
+        AND EXISTS (
+            SELECT 1
+            FROM ciphertext_digest d
+            WHERE d.tenant_id = c.tenant_id
+            AND d.handle = c.handle
+            AND d.ciphertext128 IS NOT NULL
+        );"
+    )
+    .execute(conn.as_mut())
+    .await?.rows_affected();
+ 
+    if rows > 0 {
+        info!(target: "worker", "Cleaning up old ciphertexts128, rows_affected: {}", rows);
+    }
+
+    Ok(())
+}
+
 /// Fetch and process SnS tasks from the database.
 async fn fetch_and_execute_sns_tasks(
     conn: &mut PoolConnection<Postgres>,
@@ -128,17 +158,19 @@ async fn fetch_and_execute_sns_tasks(
         let t = telemetry::tracer("batch_execution");
         t.set_attribute("count", tasks.len().to_string());
 
-        process_tasks(trx, &mut tasks, keys, tx).await?;
+        process_tasks(&mut tasks, keys, tx)?;
         update_computations_status(trx, &tasks).await?;
 
         let s = t.child_span("batch_store_ciphertext128");
         update_ciphertext128(trx, &tasks).await?;
         notify_ciphertext128_ready(trx, &conf.notify_channel).await?;
+
+        // Try to enqueue the tasks for upload in the DB
+        // This is a best-effort attempt, as the upload worker might not be available
+        enqueue_upload_tasks(trx, &tasks).await?;
         telemetry::end_span(s);
 
         db_txn.commit().await?;
-
-        t.end();
     } else {
         db_txn.rollback().await?;
     }
@@ -241,9 +273,18 @@ async fn get_remaining_tasks(
     Ok(records_count.unwrap_or(0))
 }
 
-/// Processes the tasks by decompressing and transforming ciphertexts.
-async fn process_tasks(
+async fn enqueue_upload_tasks(
     db_txn: &mut Transaction<'_, Postgres>,
+    tasks: &[HandleItem],
+) -> Result<(), ExecutionError> {
+    for task in tasks.iter() {
+        task.enqueue_upload_task(db_txn).await?;
+    }
+    Ok(())
+}
+
+/// Processes the tasks by decompressing and transforming ciphertexts.
+fn process_tasks(
     tasks: &mut [HandleItem],
     keys: &KeySet,
     tx: &Sender<UploadJob>,
@@ -292,9 +333,6 @@ async fn process_tasks(
             }
         };
 
-        // Enqueue the task for upload in DB
-        task.enqueue_upload_task(db_txn).await?;
-
         // Start uploading the ciphertexts as soon as the ct128 is computed
         //
         // The service must continue running the squashed noise algorithm,
@@ -336,28 +374,23 @@ async fn update_ciphertext128(
             let ciphertext128 = &task.ct128_uncompressed;
             let s = task.otel.child_span("ct128_db_insert");
 
-            // Insert the ciphertext128 into the database only if
-            // the uploader has not already put it in AWS S3
+            // Insert the ciphertext128 into the database for reliability
+            // Later on, we clean up all uploaded ct128
             let res = sqlx::query!(
-                "
-                UPDATE ciphertexts
-                SET ciphertext128 = $1
-                WHERE handle = $2
-                AND EXISTS (
-                    SELECT 1
-                    FROM ciphertext_digest
-                    WHERE handle = $2
-                    AND ciphertext128 IS NULL
-                );",
-                ciphertext128.as_ref(),
-                task.handle
-            )
-            .execute(db_txn.as_mut())
-            .await;
+                    "
+                    UPDATE ciphertexts
+                    SET ciphertext128 = $1
+                    WHERE handle = $2;",
+                    ciphertext128.as_ref(),
+                    task.handle
+                )
+                .execute(db_txn.as_mut())
+                .await;
 
             match res {
-                Ok(_) => {
-                    debug!(target: "worker", handle = ?task.handle, "Inserted ct128 in DB");
+                Ok(val) => {
+                    info!(target: "worker", handle = compact_hex(&task.handle),
+                        query_res = format!("{:?}", val),  "Inserted ct128 in DB");
                     telemetry::end_span(s);
                 }
                 Err(err) => {

coprocessor/fhevm-engine/sns-executor/.sqlx/query-fd37e7dc679caa21ba22d7724a27409c232cceb719aa41a71e7faf44d2cb8ef9.json

@@ -32,6 +32,7 @@ pub struct DBConfig {
     pub notify_channel: String,
     pub batch_limit: u32,
     pub polling_interval: u32,
+    pub cleanup_interval: Duration,
     pub max_connections: u32,
 }
 
@@ -128,19 +129,6 @@ impl HandleItem {
         .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),
@@ -229,7 +217,7 @@ pub enum UploadJob {
 }
 
 impl UploadJob {
-    pub fn handle(&self) -> &Vec<u8> {
+    pub fn handle(&self) -> &[u8] {
         match self {
             UploadJob::Normal(item) => &item.handle,
             UploadJob::DatabaseLock(item) => &item.handle,