chore(coprocessor): solve dcid unit-tests issue

Author: goshawk-3

coprocessor/fhevm-engine/tfhe-worker/src/dependence_chain.rs

@@ -127,7 +127,7 @@ impl LockMngr {
         &mut self,
     ) -> Result<(Option<Vec<u8>>, LockingReason), sqlx::Error> {
         if self.disable_locking {
-            warn!("Locking is disabled");
+            debug!("Locking is disabled");
             return Ok((None, LockingReason::Missing));
         }
 

coprocessor/fhevm-engine/tfhe-worker/src/tests/dependence_chain.rs

@@ -1,7 +1,7 @@
 use crate::dependence_chain::{LockMngr, LockingReason};
 use serial_test::serial;
 use sqlx::postgres::PgPoolOptions;
-use test_harness::instance::ImportMode;
+use test_harness::instance::{DBInstance, ImportMode};
 use tokio::time::{sleep, Duration};
 use tracing::info;
 use uuid::Uuid;
@@ -11,7 +11,13 @@ const NUM_SAMPLE_CHAINS: usize = 10;
 #[tokio::test]
 #[serial(db)]
 async fn test_acquire_next_lock() {
-    let pool = setup().await;
+    let instance = setup().await;
+    let pool = sqlx::postgres::PgPoolOptions::new()
+        .max_connections(2)
+        .connect(instance.db_url.as_str())
+        .await
+        .expect("Failed to connect to the database");
+
     let dependence_chain_ids = insert_dependence_chains(&pool, NUM_SAMPLE_CHAINS)
         .await
         .expect("inserted chains");
@@ -52,7 +58,12 @@ async fn test_acquire_next_lock() {
 #[tokio::test]
 #[serial(db)]
 async fn test_work_stealing() {
-    let pool = setup().await;
+    let instance = setup().await;
+    let pool = sqlx::postgres::PgPoolOptions::new()
+        .max_connections(2)
+        .connect(instance.db_url.as_str())
+        .await
+        .expect("Failed to connect to the database");
 
     let dependence_chain_ids = insert_dependence_chains(&pool, NUM_SAMPLE_CHAINS)
         .await
@@ -178,7 +189,9 @@ async fn insert_dependence_chains(
     let mut out = Vec::with_capacity(num_chains);
 
     for i in 0..num_chains {
+        info!("Inserting dependence chain {}", i);
         let dependence_chain_id = i.to_le_bytes().to_vec();
+
         sqlx::query!(
             r#"
             INSERT INTO dependence_chain (dependence_chain_id, status, last_updated_at)
@@ -199,7 +212,7 @@ async fn insert_dependence_chains(
     Ok(out)
 }
 
-async fn setup() -> sqlx::PgPool {
+async fn setup() -> DBInstance {
     let _ = tracing_subscriber::fmt().json().with_level(true).try_init();
     let test_instance = test_harness::instance::setup_test_db(ImportMode::None)
         .await
@@ -216,5 +229,5 @@ async fn setup() -> sqlx::PgPool {
         .await
         .unwrap();
 
-    pool
+    test_instance
 }

coprocessor/fhevm-engine/tfhe-worker/src/tfhe_worker.rs

@@ -379,6 +379,8 @@ async fn query_for_work<'a>(
     ));
     s.end();
 
+    let transaction_batch_size = args.work_items_batch_size;
+
     // This query locks our work items so other worker doesn't select them.
     let mut s = tracer.start_with_context("query_work_items", loop_ctx);
     let started_at = SystemTime::now();
@@ -424,7 +426,7 @@ WHERE c.transaction_id IN (
   )
 FOR UPDATE SKIP LOCKED            ",
         dependence_chain_id,
-        args.work_items_batch_size as i32,
+        transaction_batch_size as i32,
     )
     .fetch_all(trx.as_mut())
     .await
@@ -439,14 +441,14 @@ FOR UPDATE SKIP LOCKED            ",
     health_check.update_db_access();
     if the_work.is_empty() {
         if let Some(dependence_chain_id) = &dependence_chain_id {
-            warn!(target: "tfhe_worker", dcid = %hex::encode(dependence_chain_id), locking = ?locking_reason, "No work items found to process");
+            warn!(target: "tfhe_worker", dcid = %hex::encode(dependence_chain_id), locking = ?locking_reason, "No computations found to process");
         }
         health_check.update_activity();
         return Ok((vec![], vec![]));
     }
     WORK_ITEMS_FOUND_COUNTER.inc_by(the_work.len() as u64);
     info!(target: "tfhe_worker", { count = the_work.len(), dcid = ?dependence_chain_id.as_ref().map(hex::encode),
-         locking = ?locking_reason }, "Processing work items");
+         locking = ?locking_reason }, "Processing computations");
     // Make sure we process each tenant independently to avoid
     // setting different keys from different tenants in the worker
     // threads