Integration tests timing issues

Cargo.toml

@@ -130,6 +130,9 @@ opt-level = 1
 [profile.test]
 opt-level = 1
 
+[patch.crates-io]
+iroh-blobs = { git = "https://github.com/pefontana/iroh-blobs.git", rev = "3600dbc" }
+
 [workspace.metadata.release]
 tag-name = "v{{version}}"
 sign-tag = true

architectures/decentralized/testing/src/docker_setup.rs

@@ -67,6 +67,7 @@ pub async fn e2e_testing_setup(
         init_num_clients,
         None,
         None,
+        None,
     )
     .await
 }
@@ -80,6 +81,7 @@ pub async fn e2e_testing_setup_with_min(
     min_clients: usize,
     owner_keypair_path: Option<&Path>,
     waiting_for_members_extra_time: Option<u32>,
+    round_witness_time: Option<u32>,
 ) -> DockerTestCleanup {
     remove_old_client_containers(docker_client).await;
 
@@ -88,6 +90,7 @@ pub async fn e2e_testing_setup_with_min(
         min_clients,
         owner_keypair_path,
         waiting_for_members_extra_time,
+        round_witness_time,
     )
     .unwrap();
 
@@ -284,7 +287,7 @@ pub async fn spawn_new_client_with_monitoring(
 
 // Updated spawn function
 pub fn spawn_psyche_network(init_num_clients: usize) -> Result<(), DockerWatcherError> {
-    spawn_psyche_network_with_min(init_num_clients, init_num_clients, None, None)
+    spawn_psyche_network_with_min(init_num_clients, init_num_clients, None, None, None)
 }
 
 /// Spawn the psyche network with explicit min_clients and optional owner keypair.
@@ -293,6 +296,7 @@ pub fn spawn_psyche_network_with_min(
     min_clients: usize,
     owner_keypair_path: Option<&Path>,
     waiting_for_members_extra_time: Option<u32>,
+    round_witness_time: Option<u32>,
 ) -> Result<(), DockerWatcherError> {
     #[cfg(not(feature = "python"))]
     let mut builder = ConfigBuilder::new()
@@ -309,6 +313,10 @@ pub fn spawn_psyche_network_with_min(
         builder = builder.with_waiting_for_members_extra_time(time);
     }
 
+    if let Some(time) = round_witness_time {
+        builder = builder.with_round_witness_time(time);
+    }
+
     let config_file_path = builder.build();
 
     println!("[+] Config file written to: {}", config_file_path.display());

architectures/decentralized/testing/src/utils.rs

@@ -160,6 +160,7 @@ pub struct ConfigBuilder {
     batch_size: u32,
     architecture: String,
     waiting_for_members_extra_time: Option<u32>,
+    round_witness_time: Option<u32>,
 }
 
 impl Default for ConfigBuilder {
@@ -189,6 +190,7 @@ impl ConfigBuilder {
             batch_size: 4,
             architecture: String::from("HfLlama"),
             waiting_for_members_extra_time: None,
+            round_witness_time: None,
         }
     }
 
@@ -218,6 +220,11 @@ impl ConfigBuilder {
         self
     }
 
+    pub fn with_round_witness_time(mut self, time: u32) -> Self {
+        self.round_witness_time = Some(time);
+        self
+    }
+
     pub fn build(mut self) -> PathBuf {
         // Use min_clients if set, otherwise default to num_clients
         let min_clients = self.min_clients.unwrap_or(self.num_clients);
@@ -240,6 +247,10 @@ impl ConfigBuilder {
             self.set_value("config.waiting_for_members_extra_time", time);
         }
 
+        if let Some(time) = self.round_witness_time {
+            self.set_value("config.round_witness_time", time);
+        }
+
         let config_content = toml::to_string(&self.base_config).unwrap();
         let config_file_path = PathBuf::from("../../../config/solana-test/test-config.toml");
         fs::write(&config_file_path, config_content).unwrap();

architectures/decentralized/testing/tests/integration_tests.rs

@@ -270,7 +270,26 @@ async fn test_rejoining_client_delay() {
 
     let solana_client = Arc::new(SolanaTestClient::new("test".to_string(), None).await);
 
-    tokio::time::sleep(Duration::from_secs(30)).await;
+    // Monitor client-1 to detect when training starts
+    let _monitor_client_1 = watcher
+        .monitor_container(
+            &format!("{CLIENT_CONTAINER_PREFIX}-1"),
+            vec![IntegrationTestLogMarker::StateChange],
+        )
+        .unwrap();
+
+    // Wait for the first training round to start in epoch 0.
+    // We spawn client-2 here so it joins in the next epoch when checkpoint is p2p.
+    println!("Waiting for epoch 0 training to start");
+    while let Some(response) = watcher.log_rx.recv().await {
+        if let Response::StateChange(_, _, _old_state, new_state, epoch, step) = response {
+            println!("epoch: {epoch} step: {step} - state: {new_state}");
+            if new_state == RunState::RoundTrain.to_string() {
+                println!("Training started, spawning client-2");
+                break;
+            }
+        }
+    }
 
     // Spawn client
     spawn_new_client(docker.clone(), None).await.unwrap();
@@ -333,8 +352,8 @@ async fn disconnect_client() {
     let docker = Arc::new(Docker::connect_with_socket_defaults().unwrap());
     let mut watcher = DockerWatcher::new(docker.clone());
 
-    // Initialize a Solana run with 3 clients
-    let _cleanup = e2e_testing_setup(docker.clone(), 3).await;
+    // Initialize a Solana run with 3 clients.
+    let _cleanup = e2e_testing_setup_with_min(docker.clone(), 3, 3, None, None, Some(15)).await;
 
     let _monitor_client_1 = watcher
         .monitor_container(
@@ -426,10 +445,7 @@ async fn disconnect_client() {
                     killed_client = true;
                 }
 
-                if killed_client
-                    && seen_health_checks.len() >= 2
-                    && new_state == RunState::Cooldown.to_string()
-                {
+                if killed_client && new_state == RunState::Cooldown.to_string() {
                     let epoch_clients = solana_client.get_current_epoch_clients().await;
                     assert!(
                         epoch_clients.len() <= 2,
@@ -464,11 +480,13 @@ async fn disconnect_client() {
         }
     }
 
-    // assert that two healthchecks were sent, by the alive clients
-    assert_eq!(
-        seen_health_checks.len(),
-        2,
-        "Two healthchecks should have been sent"
+    // Each alive client should detect the killed one and send a healthcheck.
+    // Ideally both alive clients send one, but due to P2P gossip timing through the
+    // relay, a training result may not reach the other peer before bloom filters are
+    // built, causing only one client to detect the dead one.
+    assert!(
+        !seen_health_checks.is_empty(),
+        "Expected at least 1 healthcheck, got 0",
     );
 
     // check how many batches where lost due to the client shutdown
@@ -492,7 +510,8 @@ async fn drop_a_client_waitingformembers_then_reconnect() {
 
     // Use extra WFM time so we have a window to kill a client during WaitingForMembers
     let _cleanup =
-        e2e_testing_setup_with_min(docker.clone(), n_clients, n_clients, None, Some(30)).await;
+        e2e_testing_setup_with_min(docker.clone(), n_clients, n_clients, None, Some(30), None)
+            .await;
 
     let solana_client = SolanaTestClient::new(run_id, None).await;
     // Monitor clients
@@ -842,7 +861,8 @@ async fn test_pause_and_resume_run() {
     // Setup with min_clients=1 but init_num_clients=0 (we spawn manually)
     // Pass owner keypair to setup script
     let _cleanup =
-        e2e_testing_setup_with_min(docker.clone(), 0, 1, Some(owner_path.as_path()), None).await;
+        e2e_testing_setup_with_min(docker.clone(), 0, 1, Some(owner_path.as_path()), None, None)
+            .await;
 
     // Create SolanaTestClient with owner keypair for set_paused
     let solana_client = SolanaTestClient::new(run_id.clone(), Some(owner_keypair.clone())).await;

config/client/.env.rpc-fallback-test

@@ -2,7 +2,6 @@
 RPC=http://nginx-proxy-1:8901
 WS_RPC=ws://nginx-proxy-1:8901/ws/
 RUN_ID=test
-CONFIG_PATH=/usr/src/psyche/config/solana-test/config.toml
 RPC_2=http://nginx-proxy-2:8902
 WS_RPC_2=ws://nginx-proxy-2:8902/ws/
 RUST_LOG=debug,psyche=trace

shared/network/src/lib.rs

@@ -11,6 +11,7 @@ use iroh_blobs::{
     BlobsProtocol,
     api::downloader::Downloader,
     store::mem::{MemStore, Options as MemStoreOptions},
+    util::connection_pool::Options as PoolOptions,
 };
 use iroh_gossip::{
     api::{GossipReceiver, GossipSender},
@@ -340,8 +341,8 @@ where
 
         let endpoint = {
             let transport_config = QuicTransportConfig::builder()
-                .max_idle_timeout(Some(Duration::from_secs(10).try_into()?))
-                .keep_alive_interval(Duration::from_secs(1))
+                .max_idle_timeout(Some(Duration::from_secs(120).try_into()?))
+                .keep_alive_interval(Duration::from_secs(5))
                 .set_max_remote_nat_traversal_addresses(50)
                 .build();
 
@@ -445,7 +446,13 @@ where
                 add_protected: None,
             }),
         });
-        let downloader = Downloader::new(&store, &endpoint);
+        let pool_options = PoolOptions {
+            idle_timeout: Duration::from_secs(60),
+            connect_timeout: Duration::from_secs(5),
+            max_connections: 1024,
+            on_connected: None,
+        };
+        let downloader = Downloader::with_pool_options(&store, &endpoint, pool_options);
         trace!("blobs store created!");
 
         trace!("creating gossip...");