chore(coprocessor): prevent row-lock contention between execution and upload tasks

Author: goshawk-3

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

@@ -1,4 +1,4 @@
-use crate::{Config, ExecutionError, HandleItem, S3Config};
+use crate::{Config, ExecutionError, HandleItem, S3Config, UploadJob};
 use aws_config::retry::RetryConfig;
 use aws_config::timeout::TimeoutConfig;
 use aws_config::BehaviorVersion;
@@ -32,8 +32,8 @@ pub const EVENT_CIPHERTEXTS_UPLOADED: &str = "event_ciphertexts_uploaded";
 /// Process the S3 uploads
 pub(crate) async fn process_s3_uploads(
     conf: &Config,
-    mut tasks: mpsc::Receiver<HandleItem>,
-    tasks_tx: mpsc::Sender<HandleItem>,
+    mut jobs: mpsc::Receiver<UploadJob>,
+    jobs_tx: mpsc::Sender<UploadJob>,
     token: CancellationToken,
 ) -> Result<(), ExecutionError> {
     // Client construction is expensive due to connection thread pool initialization, and should
@@ -56,7 +56,7 @@ pub(crate) async fn process_s3_uploads(
         let token = token.clone();
         let pool = pool.clone();
         async move {
-            do_resubmits_loop(client, pool, &conf, tasks_tx, token, is_ready)
+            do_resubmits_loop(client, pool, &conf, jobs_tx, token, is_ready)
                 .await
                 .unwrap_or_else(|err| {
                     error!("Failed to spawn do_resubmits_loop: {}", err);
@@ -71,41 +71,50 @@ pub(crate) async fn process_s3_uploads(
 
     loop {
         select! {
-            task = tasks.recv() => {
-                let task = match task {
-                    Some(task) => task,
+            job = jobs.recv() => {
+                let job = match job {
+                    Some(job) => job,
                     None => return Ok(()),
                 };
-                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::Acquire) {
                     // 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");
-
-                    // Ensure that the uploading task is enqueued in the DB
-                    trx_lock.commit().await?;
                     continue;
                 }
 
+                let mut trx = pool.begin().await?;
+
+                let item = match job {
+                    UploadJob::Normal(item) => item,
+                    UploadJob::DatabaseLock(item) => {
+                        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",
+                            item.tenant_id,
+                            item.handle,
+                        )
+                        .fetch_one(trx.as_mut())
+                        .await
+                        {
+                            warn!(
+                                "Failed to lock pending uploads {}, handle: {}",
+                                err,
+                                compact_hex(&item.handle)
+                            );
+                            trx.rollback().await?;
+                            continue;
+                        }
+                        item
+                    },
+                 };
+
+
+                debug!("Received task, handle: {}", hex::encode(&item.handle));
+
                 // Cleanup completed tasks
                 upload_jobs.retain(|h| !h.is_finished());
 
@@ -130,8 +139,8 @@ 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_lock, task, &client, &conf).await {
+                    let s = item.otel.child_span("upload_s3");
+                        if let Err(err) = upload_ciphertexts(trx, item, &client, &conf).await {
                             if let ExecutionError::S3TransientError(_) = err {
                                 ready_flag.store(false, Ordering::Release);
                                 info!("S3 setup is not ready, due to transient error: {}", err);
@@ -339,7 +348,7 @@ fn compute_digest(ct: &[u8]) -> Vec<u8> {
 async fn fetch_pending_uploads(
     db_pool: &Pool<Postgres>,
     limit: i64,
-) -> Result<Vec<HandleItem>, ExecutionError> {
+) -> Result<Vec<UploadJob>, ExecutionError> {
     let rows = sqlx::query!(
         "SELECT tenant_id, handle, ciphertext, ciphertext128
         FROM ciphertext_digest 
@@ -351,7 +360,7 @@ async fn fetch_pending_uploads(
     .fetch_all(db_pool)
     .await?;
 
-    let mut tasks = Vec::new();
+    let mut jobs = Vec::new();
 
     for row in rows {
         let mut ct64_compressed = Arc::new(Vec::new());
@@ -404,17 +413,20 @@ async fn fetch_pending_uploads(
         }
 
         if !ct64_compressed.is_empty() || !ct128_uncompressed.is_empty() {
-            tasks.push(HandleItem {
+            let item = HandleItem {
                 tenant_id: row.tenant_id,
                 handle: handle.clone(),
                 ct64_compressed,
                 ct128_uncompressed,
                 otel: telemetry::tracer_with_handle("recovery_task", handle),
-            });
+            };
+
+            // Instruct the uploader to acquire DB lock when processing the item
+            jobs.push(UploadJob::DatabaseLock(item));
         }
     }
 
-    Ok(tasks)
+    Ok(jobs)
 }
 
 /// Resubmit for uploading ciphertexts.
@@ -424,7 +436,7 @@ async fn do_resubmits_loop(
     client: Arc<aws_sdk_s3::Client>,
     pool: Arc<Pool<Postgres>>,
     conf: &Config,
-    tasks: mpsc::Sender<HandleItem>,
+    tasks: mpsc::Sender<UploadJob>,
     token: CancellationToken,
     is_ready: Arc<AtomicBool>,
 ) -> Result<(), ExecutionError> {
@@ -472,7 +484,7 @@ async fn do_resubmits_loop(
 async fn try_resubmit(
     pool: &PgPool,
     is_ready: Arc<AtomicBool>,
-    tasks: mpsc::Sender<HandleItem>,
+    tasks: mpsc::Sender<UploadJob>,
     token: CancellationToken,
 ) -> Result<(), ExecutionError> {
     if !is_ready.load(Ordering::SeqCst) {
@@ -482,18 +494,18 @@ async fn try_resubmit(
 
     match fetch_pending_uploads(pool, 10).await {
         // TODO: const, token
-        Ok(recovery_tasks) => {
+        Ok(jobs) => {
             info!(
                 target: "worker",
                 action = "retry_s3_uploads",
                 "Fetched {} pending uploads from the database",
-                recovery_tasks.len()
+                jobs.len()
             );
             // Resubmit for uploading ciphertexts
-            for task in recovery_tasks {
+            for task in jobs {
                 select! {
                     _ = tasks.send(task.clone()) => {
-                        debug!("Task sent to upload worker, handle: {}", hex::encode(&task.handle));
+                        // TODO: debug!("Task sent to upload worker, handle: {}", hex::encode(&task.handle));
                     },
                     _ = token.cancelled() => {
                         return Ok(())

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

@@ -1,7 +1,8 @@
 use fhevm_engine_common::telemetry;
 use sns_executor::{
-    compute_128bit_ct, process_s3_uploads, Config, DBConfig, HandleItem, S3Config, S3RetryPolicy,
+    compute_128bit_ct, process_s3_uploads, Config, DBConfig, S3Config, S3RetryPolicy, UploadJob,
 };
+
 use tokio::{signal::unix, spawn, sync::mpsc};
 use tokio_util::sync::CancellationToken;
 use tracing::{error, Level};
@@ -67,7 +68,7 @@ async fn main() {
     // to avoid blocking the worker
     // and to allow for some burst of uploads
     let (uploads_tx, uploads_rx) =
-        mpsc::channel::<HandleItem>(10 * config.s3.max_concurrent_uploads as usize);
+        mpsc::channel::<UploadJob>(10 * config.s3.max_concurrent_uploads as usize);
 
     if let Err(err) = telemetry::setup_otlp(&config.service_name) {
         panic!("Error while initializing tracing: {:?}", err);

coprocessor/fhevm-engine/sns-executor/src/bin/sns_worker.rs

@@ -2,6 +2,7 @@ use crate::keyset::fetch_keyset;
 use crate::squash_noise::SquashNoiseCiphertext;
 use crate::HandleItem;
 use crate::KeySet;
+use crate::UploadJob;
 use crate::{Config, DBConfig, ExecutionError};
 use fhevm_engine_common::telemetry;
 use fhevm_engine_common::utils::compact_hex;
@@ -30,7 +31,7 @@ enum ConnStatus {
 /// Executes the worker logic for the SnS task.
 pub(crate) async fn run_loop(
     conf: &Config,
-    tx: &Sender<HandleItem>,
+    tx: &Sender<UploadJob>,
     token: CancellationToken,
 ) -> Result<(), ExecutionError> {
     let tenant_api_key = &conf.tenant_api_key;
@@ -109,7 +110,7 @@ pub(crate) async fn run_loop(
 /// Fetch and process SnS tasks from the database.
 async fn fetch_and_execute_sns_tasks(
     conn: &mut PoolConnection<Postgres>,
-    tx: &Sender<HandleItem>,
+    tx: &Sender<UploadJob>,
     keys: &KeySet,
     conf: &DBConfig,
 ) -> Result<(), ExecutionError> {
@@ -124,11 +125,20 @@ async fn fetch_and_execute_sns_tasks(
     let trx = &mut db_txn;
 
     if let Some(mut tasks) = query_sns_tasks(trx, conf.batch_limit).await? {
+        let t = telemetry::tracer("batch_execution");
+        t.set_attribute("count", tasks.len().to_string());
+
         process_tasks(trx, &mut tasks, keys, tx).await?;
         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?;
+        telemetry::end_span(s);
+
         db_txn.commit().await?;
+
+        t.end();
     } else {
         db_txn.rollback().await?;
     }
@@ -236,7 +246,7 @@ async fn process_tasks(
     db_txn: &mut Transaction<'_, Postgres>,
     tasks: &mut [HandleItem],
     keys: &KeySet,
-    tx: &Sender<HandleItem>,
+    tx: &Sender<UploadJob>,
 ) -> Result<(), ExecutionError> {
     for task in tasks.iter_mut() {
         let ct64_compressed = task.ct64_compressed.as_ref();
@@ -289,7 +299,7 @@ async fn process_tasks(
         //
         // The service must continue running the squashed noise algorithm,
         // regardless of the availability of the upload worker.
-        if let Err(err) = tx.try_send(task.clone()) {
+        if let Err(err) = tx.try_send(UploadJob::Normal(task.clone())) {
             // This could happen if either we are experiencing a burst of tasks
             // or the upload worker cannot recover the connection to AWS S3
             //

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

@@ -217,10 +217,21 @@ pub enum ExecutionError {
     S3TransientError(String),
 }
 
+#[derive(Clone)]
+pub enum UploadJob {
+    /// Represents a standard upload that is dispatched immediately
+    /// after a successful squash_noise computation
+    Normal(HandleItem),
+
+    /// Represents a job that requires acquiring a database lock
+    /// before initiating the upload process.
+    DatabaseLock(HandleItem),
+}
+
 /// Runs the SnS worker loop
 pub async fn compute_128bit_ct(
     conf: &Config,
-    tx: Sender<HandleItem>,
+    tx: Sender<UploadJob>,
     token: CancellationToken,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     info!(target: "sns", "Worker started with {}", conf);
@@ -234,8 +245,8 @@ pub async fn compute_128bit_ct(
 /// Runs the uploader loop
 pub async fn process_s3_uploads(
     conf: &Config,
-    rx: mpsc::Receiver<HandleItem>,
-    tx: Sender<HandleItem>,
+    rx: mpsc::Receiver<UploadJob>,
+    tx: Sender<UploadJob>,
     token: CancellationToken,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     info!(target: "sns", "Uploader started with {:?}", conf.s3);