fix(coprocessor): allow TFHE worker to lock earliest dependence chain on no progress

Author: antoniupop

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

@@ -217,6 +217,69 @@ impl LockMngr {
         ))
     }
 
+    /// Acquire the earliest dependence-chain entry for processing
+    /// sorted by last_updated_at (FIFO).
+    /// Returns the dependence_chain_id if a lock was acquired
+    pub async fn acquire_early_lock(
+        &mut self,
+    ) -> Result<(Option<Vec<u8>>, LockingReason), sqlx::Error> {
+        if self.disable_locking {
+            debug!("Locking is disabled");
+            return Ok((None, LockingReason::Missing));
+        }
+
+        let started_at = SystemTime::now();
+        let row = sqlx::query_as::<_, DatabaseChainLock>(
+            r#"
+            WITH candidate AS (
+                SELECT dependence_chain_id, 'updated_unowned' AS match_reason, dependency_count
+                FROM dependence_chain
+                WHERE
+                    status = 'updated'      -- Marked as updated by host-listener
+                    AND
+                    worker_id IS NULL       -- Ensure no other workers own it
+                ORDER BY last_updated_at ASC        -- FIFO
+                FOR UPDATE SKIP LOCKED              -- Ensure no other worker is currently trying to lock it
+                LIMIT 1
+            )
+            UPDATE dependence_chain AS dc
+            SET
+                worker_id = $1,
+                status = 'processing',
+                lock_acquired_at = NOW(),
+                lock_expires_at = NOW() + make_interval(secs => $2)
+            FROM candidate
+            WHERE dc.dependence_chain_id = candidate.dependence_chain_id
+            RETURNING dc.*, candidate.match_reason, candidate.dependency_count;
+        "#,
+        )
+        .bind(self.worker_id)
+        .bind(self.lock_ttl_sec)
+        .fetch_optional(&self.pool)
+        .await?;
+
+        let row = if let Some(row) = row {
+            row
+        } else {
+            return Ok((None, LockingReason::Missing));
+        };
+
+        self.lock.replace((row.clone(), SystemTime::now()));
+        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 on earliest DCID");
+
+        Ok((
+            Some(row.dependence_chain_id),
+            LockingReason::from(row.match_reason.as_str()),
+        ))
+    }
+
     /// Release all locks held by this worker
     ///
     /// If host-listener has marked the dependence chain as 'updated' in the meantime,

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

@@ -152,6 +152,7 @@ async fn tfhe_worker_cycle(
             &health_check,
             &mut trx,
             &mut dcid_mngr,
+            &mut no_progress_cycles,
             &tracer,
             &loop_ctx,
         )
@@ -328,6 +329,7 @@ async fn query_for_work<'a>(
     health_check: &crate::health_check::HealthCheck,
     trx: &mut sqlx::Transaction<'a, Postgres>,
     deps_chain_mngr: &mut dependence_chain::LockMngr,
+    no_progress_cycles: &mut u32,
     tracer: &opentelemetry::global::BoxedTracer,
     loop_ctx: &opentelemetry::Context,
 ) -> Result<
@@ -340,7 +342,14 @@ async fn query_for_work<'a>(
         match deps_chain_mngr.extend_or_release_current_lock(true).await? {
             // If there is a current lock, we extend it and use its dependence_chain_id
             Some((id, reason)) => (Some(id), reason),
-            None => deps_chain_mngr.acquire_next_lock().await?,
+            None => {
+                if no_progress_cycles < 10 * args.dcid_max_no_progress_cycles {
+                    deps_chain_mngr.acquire_next_lock().await?
+                } else {
+                    *no_progress_cycles = 0;
+                    deps_chain_mngr.acquire_early_lock().await?
+                }
+            }
         };
     if deps_chain_mngr.enabled() && dependence_chain_id.is_none() {
         // No dependence chain to lock, so no work to do
@@ -482,7 +491,10 @@ WHERE c.transaction_id IN (
                     inputs,
                     is_allowed: w.is_allowed,
                 });
-                if w.schedule_order < earliest_schedule_order {
+                if w.schedule_order < earliest_schedule_order && w.is_allowed {
+                    // Only account for allowed to avoid case of reorg
+                    // where trivial encrypts will be in collision in
+                    // the same transaction and old ones are re-used
                     earliest_schedule_order = w.schedule_order;
                 }
             }