fix(coprocessor): fix handling of cyclic dependence errors (port from #1421) (#1439)

* fix(coprocessor): fix handling of cyclic dependence errors (#1421)

* fix(coprocessor): fix handling of dependence cycles

* fix(coprocessor): log dependence cycle

* fix(coprocessor): fix uncomputable flag and schedule order on missing input

* chore(host-listener): update max message size websocket (#1410)

---------

Co-authored-by: Oba <obatirou@gmail.com>

* fix(coprocessor): host-listener, message size leftover

---------

Co-authored-by: Oba <obatirou@gmail.com>
Co-authored-by: rudy <rudy.sicard@zama.ai>

Author: 3 people

coprocessor/fhevm-engine/host-listener/src/cmd/mod.rs

@@ -3,6 +3,7 @@ use alloy::primitives::Address;
 use alloy::providers::{Provider, ProviderBuilder, WsConnect};
 use alloy::pubsub::SubscriptionStream;
 use alloy::rpc::types::{Block, BlockNumberOrTag, Filter, Header, Log};
+use alloy::transports::ws::WebSocketConfig;
 use anyhow::{anyhow, Result};
 use clap::Parser;
 use futures_util::stream::StreamExt;
@@ -166,6 +167,10 @@ mod eth_rpc_err {
     }
 }
 
+fn websocket_config() -> WebSocketConfig {
+    WebSocketConfig::default().max_message_size(Some(256 * 1024 * 1024)) // 256MB
+}
+
 impl InfiniteLogIter {
     fn new(args: &Args) -> Self {
         let mut contract_addresses = vec![];
@@ -202,7 +207,8 @@ impl InfiniteLogIter {
     }
 
     async fn get_chain_id(&self) -> anyhow::Result<ChainId> {
-        let ws = WsConnect::new(&self.url);
+        let config = websocket_config();
+        let ws = WsConnect::new(&self.url).with_config(config);
         let provider = ProviderBuilder::new().connect_ws(ws).await?;
         Ok(provider.get_chain_id().await?)
     }
@@ -245,7 +251,10 @@ impl InfiniteLogIter {
             filter = filter.address(self.contract_addresses.clone())
         }
         // we use a specific provider to not disturb the real-time one (no buffer shared)
-        let ws = WsConnect::new(&self.url).with_max_retries(0); // disabled, alloy skips events
+        let config = websocket_config();
+        let ws = WsConnect::new(&self.url)
+            .with_config(config)
+            .with_max_retries(0); // disabled, alloy skips events
         let provider = match ProviderBuilder::new().connect_ws(ws).await {
             Ok(provider) => provider,
             Err(_) => anyhow::bail!("Cannot get a provider"),
@@ -622,9 +631,11 @@ impl InfiniteLogIter {
 
     async fn new_log_stream(&mut self, not_initialized: bool) {
         let mut retry = 20;
+        let config = websocket_config();
         loop {
-            let ws = WsConnect::new(&self.url).with_max_retries(0); // disabled, alloy skips events
-
+            let ws = WsConnect::new(&self.url)
+                .with_config(config)
+                .with_max_retries(0); // disabled, alloy skips events
             match ProviderBuilder::new().connect_ws(ws).await {
                 Ok(provider) => {
                     let catch_up_from =

coprocessor/fhevm-engine/scheduler/src/dfg.rs

@@ -2,7 +2,7 @@ pub mod scheduler;
 pub mod types;
 
 use std::{collections::HashMap, sync::atomic::AtomicUsize};
-use tracing::error;
+use tracing::{error, warn};
 
 use crate::dfg::types::*;
 use anyhow::Result;
@@ -305,7 +305,12 @@ impl DFComponentGraph {
         for (consumer, tx) in self.graph.node_references() {
             for i in tx.inputs.keys() {
                 if let Some(producer) = self.allowed_map.get(i) {
-                    dependence_pairs.push((producer, consumer));
+                    if *producer == consumer {
+                        warn!(target: "scheduler", { },
+			       "Self-dependence on node");
+                    } else {
+                        dependence_pairs.push((*producer, consumer));
+                    }
                 } else {
                     self.needed_map
                         .entry(i.clone())
@@ -314,14 +319,72 @@ impl DFComponentGraph {
                 }
             }
         }
+
+        // We build a replica of the graph and map it to the
+        // underlying DiGraph so we can identify cycles.
+        let mut digraph = self.graph.map(|idx, _| idx, |_, _| ()).graph().clone();
         // Add transaction dependence edges
-        for (producer, consumer) in dependence_pairs {
-            // Error only occurs in case of cyclic dependence which
-            // shoud not be possible between transactions. In that
-            // case, the whole cycle should be put in an error state.
-            self.graph
-                .add_edge(*producer, consumer, ())
-                .map_err(|_| SchedulerError::CyclicDependence)?;
+        for (producer, consumer) in dependence_pairs.iter() {
+            digraph.add_edge(*producer, *consumer, ());
+        }
+        let mut tarjan = daggy::petgraph::algo::TarjanScc::new();
+        let mut sccs = Vec::new();
+        tarjan.run(&digraph, |scc| {
+            if scc.len() > 1 {
+                // All non-singleton SCCs in a directed graph are
+                // dependence cycles
+                sccs.push(scc.to_vec());
+            }
+        });
+        if !sccs.is_empty() {
+            for scc in sccs {
+                error!(target: "scheduler", { cycle_size = ?scc.len() },
+		       "Dependence cycle detected");
+                for idx in scc {
+                    let idx = digraph
+                        .node_weight(idx)
+                        .ok_or(SchedulerError::DataflowGraphError)?;
+                    let tx = self
+                        .graph
+                        .node_weight_mut(*idx)
+                        .ok_or(SchedulerError::DataflowGraphError)?;
+                    // Mark the node as uncomputable so we don't go
+                    // and mark as completed operations that are in
+                    // error.
+                    tx.is_uncomputable = true;
+                    error!(target: "scheduler", { transaction_id = ?hex::encode(tx.transaction_id.clone()) },
+		       "Transaction is part of a dependence cycle");
+                    for (_, op) in tx.graph.graph.node_references() {
+                        self.results.push(DFGTxResult {
+                            transaction_id: tx.transaction_id.clone(),
+                            handle: op.result_handle.to_vec(),
+                            compressed_ct: Err(SchedulerError::CyclicDependence.into()),
+                        });
+                    }
+                }
+            }
+            return Err(SchedulerError::CyclicDependence.into());
+        } else {
+            // If no dependence cycles were found, then we can
+            // complete the graph and proceed to execution
+            for (producer, consumer) in dependence_pairs.iter() {
+                // The error case here should not happen as we've
+                // already covered it by testing for SCCs in the graph
+                // first
+                if self.graph.add_edge(*producer, *consumer, ()).is_err() {
+                    let prod = self
+                        .graph
+                        .node_weight(*producer)
+                        .ok_or(SchedulerError::DataflowGraphError)?;
+                    let cons = self
+                        .graph
+                        .node_weight(*consumer)
+                        .ok_or(SchedulerError::DataflowGraphError)?;
+                    error!(target: "scheduler", { producer_id = ?hex::encode(prod.transaction_id.clone()), consumer_id = ?hex::encode(cons.transaction_id.clone()) },
+		       "Dependence cycle when adding dependence - initial cycle detection failed");
+                    return Err(SchedulerError::CyclicDependence.into());
+                }
+            }
         }
         Ok(())
     }
@@ -402,7 +465,17 @@ impl DFComponentGraph {
             .graph
             .node_weight_mut(tx_node_index)
             .ok_or(SchedulerError::DataflowGraphError)?;
+        if tx_node.is_uncomputable {
+            return Ok(());
+        }
         tx_node.is_uncomputable = true;
+        for (_idx, op) in tx_node.graph.graph.node_references() {
+            self.results.push(DFGTxResult {
+                transaction_id: tx_node.transaction_id.clone(),
+                handle: op.result_handle.to_vec(),
+                compressed_ct: Err(SchedulerError::MissingInputs.into()),
+            });
+        }
         for edge in edges.edges_directed(tx_node_index, Direction::Outgoing) {
             let dependent_tx_index = edge.target();
             self.set_uncomputable(dependent_tx_index, edges)?;

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

@@ -448,7 +448,13 @@ async fn build_transaction_graph_and_execute<'a>(
     loop_ctx: &opentelemetry::Context,
 ) -> Result<DFComponentGraph, Box<dyn std::error::Error + Send + Sync>> {
     let mut tx_graph = DFComponentGraph::default();
-    tx_graph.build(tenant_txs)?;
+    if let Err(e) = tx_graph.build(tenant_txs) {
+        // If we had an error while building the graph, we don't
+        // execute anything and return to allow any set results
+        // (essentially errors) to be set in DB.
+        warn!(target: "tfhe_worker", { error = %e }, "error while building transaction graph");
+        return Ok(tx_graph);
+    }
     let cts_to_query = tx_graph.needed_map.keys().cloned().collect::<Vec<_>>();
     let ciphertext_map =
         query_ciphertexts(&cts_to_query, *tenant_id, trx, tracer, loop_ctx).await?;