chore(coprocessor): observe query timings in tfhe-worker

- add --dcid_ttl_sec config
- add otel traces for dcid

Author: goshawk-3

coprocessor/fhevm-engine/tfhe-worker/benches/utils.rs

@@ -109,6 +109,7 @@ async fn start_coprocessor(rx: Receiver<bool>, app_port: u16, db_url: &str) {
         metric_rerand_batch_latency: MetricsConfig::default(),
         metric_fhe_batch_latency: MetricsConfig::default(),
         worker_id: None,
+        dcid_ttl_sec: 30,
     };
 
     std::thread::spawn(move || {

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

@@ -90,6 +90,11 @@ pub struct Args {
     #[arg(long, value_parser = clap::value_parser!(Uuid))]
     pub worker_id: Option<Uuid>,
 
+    /// Time-to-live in seconds for dependence chain locks
+    /// Defaults to 30 seconds if not provided
+    #[arg(long, value_parser = clap::value_parser!(u32), default_value_t = 30)]
+    pub dcid_ttl_sec: u32,
+
     /// Log level for the application
     #[arg(
         long,

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

@@ -1,9 +1,37 @@
 use chrono::{DateTime, Utc};
+use prometheus::{register_histogram, register_int_counter, Histogram, IntCounter};
 use sqlx::Postgres;
-use std::fmt;
+use std::{fmt, sync::LazyLock, time::SystemTime};
 use tracing::{debug, error, info, warn};
 use uuid::Uuid;
 
+pub(crate) static ACQUIRED_DEPENDENCE_CHAIN_ID_COUNTER: LazyLock<IntCounter> =
+    LazyLock::new(|| {
+        register_int_counter!(
+            "coprocessor_tfhe_worker_dcid_counter",
+            "Number of acquired dependence chain IDs in tfhe-worker"
+        )
+        .unwrap()
+    });
+
+pub static ACQUIRE_DEPENDENCE_CHAIN_ID_QUERY_HISTOGRAM: LazyLock<Histogram> = LazyLock::new(|| {
+    register_histogram!(
+        "coprocessor_tfhe_worker_query_acquire_dcid_seconds",
+        "Histogram of query-time spent acquiring dependence chain IDs in tfhe-worker",
+        vec![0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 1.0, 2.0, 5.0, 10.0]
+    )
+    .unwrap()
+});
+
+pub static EXTEND_DEPENDENCE_CHAIN_ID_QUERY_HISTOGRAM: LazyLock<Histogram> = LazyLock::new(|| {
+    register_histogram!(
+        "coprocessor_tfhe_worker_query_extend_dcid_seconds",
+        "Histogram of query-time spent extending dependence_chain lock in tfhe-worker",
+        vec![0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 1.0, 2.0, 5.0, 10.0]
+    )
+    .unwrap()
+});
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum LockingReason {
     UpdatedUnowned, // Normal lock acquisition
@@ -30,7 +58,7 @@ pub struct LockMngr {
     pool: sqlx::Pool<Postgres>,
     worker_id: Uuid,
     lock: Option<DatabaseChainLock>,
-    expiration_duration_secs: i64,
+    lock_ttl_sec: i64,
 }
 
 /// Dependence chain lock data
@@ -68,17 +96,13 @@ impl LockMngr {
             worker_id,
             pool,
             lock: None,
-            expiration_duration_secs: 30,
+            lock_ttl_sec: 30,
         }
     }
 
-    pub fn new_with_expiry(
-        worker_id: Uuid,
-        pool: sqlx::Pool<Postgres>,
-        expiration_duration_secs: i64,
-    ) -> Self {
+    pub fn new_with_ttl(worker_id: Uuid, pool: sqlx::Pool<Postgres>, lock_ttl_sec: u32) -> Self {
         let mut mgr = Self::new(worker_id, pool);
-        mgr.expiration_duration_secs = expiration_duration_secs;
+        mgr.lock_ttl_sec = lock_ttl_sec as i64;
         mgr
     }
 
@@ -88,6 +112,7 @@ impl LockMngr {
     pub async fn acquire_next_lock(
         &mut self,
     ) -> Result<(Option<Vec<u8>>, LockingReason), sqlx::Error> {
+        let started_at = SystemTime::now();
         let row = sqlx::query_as::<_, DatabaseChainLock>(
             r#"
             WITH candidate AS (
@@ -124,7 +149,7 @@ impl LockMngr {
         "#,
         )
         .bind(self.worker_id)
-        .bind(self.expiration_duration_secs)
+        .bind(self.lock_ttl_sec)
         .fetch_optional(&self.pool)
         .await?;
 
@@ -135,7 +160,15 @@ impl LockMngr {
         };
 
         self.lock.replace(row.clone());
-        info!(?row, "Acquired lock");
+
+        ACQUIRED_DEPENDENCE_CHAIN_ID_COUNTER.inc();
+
+        let elapsed = started_at.elapsed().map(|d| d.as_secs_f64()).unwrap_or(0.0);
+        if elapsed > 0.0 {
+            ACQUIRE_DEPENDENCE_CHAIN_ID_QUERY_HISTOGRAM.observe(elapsed);
+        }
+
+        info!(?row, query_elapsed = %elapsed, "Acquired lock");
 
         Ok((
             Some(row.dependence_chain_id),
@@ -253,6 +286,7 @@ impl LockMngr {
     pub async fn extend_current_lock(
         &mut self,
     ) -> Result<Option<(Vec<u8>, LockingReason)>, sqlx::Error> {
+        let started_at = SystemTime::now();
         let dependence_chain_id = match &self.lock {
             Some(lock) => lock.dependence_chain_id.clone(),
             None => {
@@ -272,14 +306,15 @@ impl LockMngr {
         "#,
             dependence_chain_id,
             self.worker_id,
-            self.expiration_duration_secs as f64
+            self.lock_ttl_sec as f64
         )
         .fetch_optional(&self.pool)
         .await?;
 
         let lock_expires_at = match row {
             Some(r) => r,
             None => {
+                self.lock.take();
                 error!(dcid = %hex::encode(&dependence_chain_id), "No lock extended");
                 return Ok(None);
             }
@@ -291,6 +326,11 @@ impl LockMngr {
             info!(dcid = %hex::encode(&dependence_chain_id), expires_at = ?lock.lock_expires_at, "Extended lock");
         }
 
+        let elapsed = started_at.elapsed().map(|d| d.as_secs_f64()).unwrap_or(0.0);
+        if elapsed > 0.0 {
+            EXTEND_DEPENDENCE_CHAIN_ID_QUERY_HISTOGRAM.observe(elapsed);
+        }
+
         Ok(Some((dependence_chain_id, LockingReason::ExtendedLock)))
     }
 

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

@@ -20,7 +20,7 @@ async fn test_acquire_next_lock() {
 
     for dependence_chain_id in dependence_chain_ids.iter() {
         info!(target: "deps_chain", ?dependence_chain_id, "Testing acquire_next_lock");
-        let mut mgr = LockMngr::new_with_expiry(Uuid::new_v4(), pool.clone(), 3600);
+        let mut mgr = LockMngr::new_with_ttl(Uuid::new_v4(), pool.clone(), 3600);
 
         let (acquired, locking) = mgr.acquire_next_lock().await.unwrap();
         assert_eq!(acquired, Some(dependence_chain_id.clone()));
@@ -65,7 +65,7 @@ async fn test_work_stealing() {
         info!(?dependence_chain_id, "Testing acquire_next_lock");
 
         let worker = Uuid::new_v4();
-        let mut mgr = LockMngr::new_with_expiry(worker, pool.clone(), expiration_duration_secs);
+        let mut mgr = LockMngr::new_with_ttl(worker, pool.clone(), expiration_duration_secs);
         let acquired = mgr.acquire_next_lock().await.unwrap().0;
         assert_eq!(acquired, Some(dependence_chain_id.clone()));
 

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

@@ -118,6 +118,7 @@ async fn start_coprocessor(rx: Receiver<bool>, app_port: u16, db_url: &str) {
         metric_rerand_batch_latency: MetricsConfig::default(),
         metric_fhe_batch_latency: MetricsConfig::default(),
         worker_id: None,
+        dcid_ttl_sec: 30,
     };
 
     std::thread::spawn(move || {

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

@@ -8,13 +8,14 @@ use itertools::Itertools;
 use lazy_static::lazy_static;
 use opentelemetry::trace::{Span, TraceContextExt, Tracer};
 use opentelemetry::KeyValue;
-use prometheus::{register_int_counter, IntCounter};
+use prometheus::{register_histogram, register_int_counter, Histogram, IntCounter};
 use scheduler::dfg::types::{DFGTxInput, SchedulerError};
 use scheduler::dfg::{build_component_nodes, ComponentNode, DFComponentGraph, DFGOp};
 use scheduler::dfg::{scheduler::Scheduler, types::DFGTaskInput};
 use sqlx::types::Uuid;
 use sqlx::Postgres;
 use sqlx::{postgres::PgListener, query, Acquire};
+use std::time::SystemTime;
 use std::{
     collections::{BTreeSet, HashMap},
     num::NonZeroUsize,
@@ -55,12 +56,20 @@ lazy_static! {
         "work items successfully processed and stored in the database"
     )
     .unwrap();
+    static ref WORK_ITEMS_QUERY_HISTOGRAM: Histogram = register_histogram!(
+        "coprocessor_tfhe_worker_query_work_items_seconds",
+        "Histogram of time spent querying work items in tfhe-worker",
+        vec![0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 1.0, 2.0, 5.0, 10.0]
+    )
+    .unwrap();
 }
 
 pub async fn run_tfhe_worker(
     args: crate::daemon_cli::Args,
     health_check: crate::health_check::HealthCheck,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    // Determine worker ID to use for the lifetime of this process
+    // In case of a failure in tfhe_worker_cycle, the same id must be reused to quickly unlock any held locks
     let worker_id = args.worker_id.unwrap_or(Uuid::new_v4());
     info!(target: "tfhe_worker", worker_id = %worker_id, "Starting tfhe-worker service");
     loop {
@@ -92,10 +101,11 @@ async fn tfhe_worker_cycle(
     let mut listener = PgListener::connect_with(&pool).await?;
     listener.listen("work_available").await?;
 
-    let mut deps_chain_mngr = dependence_chain::LockMngr::new(worker_id, pool.clone());
+    let mut dcid_mngr =
+        dependence_chain::LockMngr::new_with_ttl(worker_id, pool.clone(), args.dcid_ttl_sec);
 
     // Release all owned locks on startup to avoid stale locks
-    deps_chain_mngr.release_all_owned_locks().await?;
+    dcid_mngr.release_all_owned_locks().await?;
 
     #[cfg(feature = "bench")]
     populate_cache_with_tenant_keys(vec![1i32], &pool, &tenant_key_cache).await?;
@@ -131,18 +141,26 @@ async fn tfhe_worker_cycle(
             args,
             &health_check,
             &mut trx,
-            &mut deps_chain_mngr,
+            &mut dcid_mngr,
             &tracer,
             &loop_ctx,
         )
         .await?;
         if transactions.is_empty() {
-            deps_chain_mngr.release_current_lock().await?;
+            dcid_mngr.release_current_lock().await?;
 
             // Lock another dependence chain if available and
             // continue processing without waiting for notification
-            let (lock, _) = deps_chain_mngr.acquire_next_lock().await?;
-            immedially_poll_more_work = lock.is_some();
+            let mut s = tracer.start_with_context("query_dependence_chain", &loop_ctx);
+
+            let (dependence_chain_id, _) = dcid_mngr.acquire_next_lock().await?;
+            immedially_poll_more_work = dependence_chain_id.is_some();
+
+            s.set_attribute(KeyValue::new(
+                "dependence_chain_id",
+                format!("{:?}", dependence_chain_id.as_ref().map(hex::encode)),
+            ));
+            s.end();
 
             continue;
         } else {
@@ -161,7 +179,12 @@ async fn tfhe_worker_cycle(
 
         // Execute transactions segregated by tenant
         for (tenant_id, ref mut tenant_txs) in transactions.iter_mut() {
-            deps_chain_mngr.extend_current_lock().await?;
+            if dcid_mngr.extend_current_lock().await?.is_none() {
+                // best-effort attempt to extend the lock and prevent other replicas from trying to lock the same DCID.
+                // Worst-case scenario, it returns None if the lock has expired.
+                // However, the worker has already secured exclusive access to the txn computations in the Computations table.
+                warn!(target: "tfhe_worker", tenant_id = %tenant_id, "Lost dcid lock while processing transactions");
+            }
 
             let mut tx_graph = build_transaction_graph_and_execute(
                 tenant_id,
@@ -178,7 +201,7 @@ async fn tfhe_worker_cycle(
                 &mut tx_graph,
                 &mut unneeded_handles,
                 &mut trx,
-                &mut deps_chain_mngr,
+                &mut dcid_mngr,
                 &tracer,
                 &loop_ctx,
             )
@@ -324,8 +347,7 @@ async fn query_for_work<'a>(
     (Vec<(i32, Vec<ComponentNode>)>, Vec<(Handle, Handle)>),
     Box<dyn std::error::Error + Send + Sync>,
 > {
-    // 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 mut s = tracer.start_with_context("query_dependence_chain", loop_ctx);
 
     // Lock dependence chain
     let (dependence_chain_id, locking_reason) = match deps_chain_mngr.extend_current_lock().await? {
@@ -343,7 +365,11 @@ async fn query_for_work<'a>(
         "dependence_chain_id",
         format!("{:?}", dependence_chain_id.as_ref().map(hex::encode)),
     ));
+    s.end();
 
+    // 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();
     let the_work = query!(
         "
 WITH selected_computations AS (
@@ -398,6 +424,8 @@ FOR UPDATE SKIP LOCKED            ",
         error!(target: "tfhe_worker", { error = %err }, "error while querying work items");
         err
     })?;
+
+    WORK_ITEMS_QUERY_HISTOGRAM.observe(started_at.elapsed().unwrap_or_default().as_secs_f64());
     s.set_attribute(KeyValue::new("count", the_work.len() as i64));
     s.end();
     health_check.update_db_access();