Refactor cluster sandbox wait_for methods (#3237)

Switches retry logic in cluster sandbox wait_for methods to use
quickwit_common::test_utils::wait_for_server_ready.

Fixes #3230
Fixes #3214

Author: imotov

quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs

@@ -27,17 +27,19 @@ use futures_util::{future, Future};
 use itertools::Itertools;
 use quickwit_actors::ActorExitStatus;
 use quickwit_common::new_coolid;
-use quickwit_common::test_utils::wait_for_server_ready;
+use quickwit_common::test_utils::{wait_for_server_ready, wait_until_predicate};
 use quickwit_common::uri::Uri as QuickwitUri;
 use quickwit_config::service::QuickwitService;
 use quickwit_config::QuickwitConfig;
 use quickwit_metastore::SplitState;
-use quickwit_rest_client::rest_client::{QuickwitClient, Transport, DEFAULT_BASE_URL};
+use quickwit_rest_client::models::IngestSource;
+use quickwit_rest_client::rest_client::{CommitType, QuickwitClient, Transport, DEFAULT_BASE_URL};
 use quickwit_serve::{serve_quickwit, ListSplitsQueryParams};
 use reqwest::Url;
 use tempfile::TempDir;
 use tokio::sync::watch::{self, Receiver, Sender};
 use tokio::task::JoinHandle;
+use tracing::debug;
 
 /// Configuration of a node made of a [`QuickwitConfig`] and a
 /// set of services.
@@ -94,6 +96,43 @@ fn transport_url(addr: SocketAddr) -> Url {
     url
 }
 
+#[macro_export]
+macro_rules! ingest_json {
+    ($($json:tt)+) => {
+        quickwit_rest_client::models::IngestSource::Bytes(json!($($json)+).to_string().into())
+    };
+}
+
+pub async fn ingest_with_retry(
+    client: &QuickwitClient,
+    index_id: &str,
+    ingest_source: IngestSource,
+    commit_type: CommitType,
+) -> anyhow::Result<()> {
+    wait_until_predicate(
+        || {
+            let commit_type_clone = commit_type.clone();
+            let ingest_source_clone = ingest_source.clone();
+            async move {
+                // Index one record.
+                if let Err(err) = client
+                    .ingest(index_id, ingest_source_clone, None, commit_type_clone)
+                    .await
+                {
+                    debug!("Failed to index into {} due to error: {}", index_id, err);
+                    false
+                } else {
+                    true
+                }
+            }
+        },
+        Duration::from_secs(10),
+        Duration::from_millis(100),
+    )
+    .await?;
+    Ok(())
+}
+
 impl ClusterSandbox {
     // Starts one node that runs all the services.
     pub async fn start_standalone_node() -> anyhow::Result<Self> {
@@ -172,19 +211,31 @@ impl ClusterSandbox {
         &self,
         expected_num_alive_nodes: usize,
     ) -> anyhow::Result<()> {
-        let mut num_attempts = 0;
-        let max_num_attempts = 3;
-        while num_attempts < max_num_attempts {
-            tokio::time::sleep(Duration::from_millis(100 * (num_attempts + 1))).await;
-            let cluster_snapshot = self.indexer_rest_client.cluster().snapshot().await?;
-            if cluster_snapshot.ready_nodes.len() == expected_num_alive_nodes {
-                return Ok(());
-            }
-            num_attempts += 1;
-        }
-        if num_attempts == max_num_attempts {
-            anyhow::bail!("Too many attempts to get expected num members.");
-        }
+        wait_until_predicate(
+            || async move {
+                match self.indexer_rest_client.cluster().snapshot().await {
+                    Ok(result) => {
+                        if result.live_nodes.len() != expected_num_alive_nodes {
+                            debug!(
+                                "wait_for_cluster_num_ready_nodes expected {} alive nodes, got {}",
+                                expected_num_alive_nodes,
+                                result.live_nodes.len()
+                            );
+                            false
+                        } else {
+                            true
+                        }
+                    }
+                    Err(err) => {
+                        debug!("wait_for_cluster_num_ready_nodes error {err}");
+                        false
+                    }
+                }
+            },
+            Duration::from_secs(10),
+            Duration::from_millis(100),
+        )
+        .await?;
         Ok(())
     }
 
@@ -193,28 +244,30 @@ impl ClusterSandbox {
         &self,
         required_pipeline_num: usize,
     ) -> anyhow::Result<()> {
-        let mut num_attempts = 0;
-        let max_num_attempts = 3;
-        while num_attempts < max_num_attempts {
-            if num_attempts > 0 {
-                tokio::time::sleep(Duration::from_millis(100 * (num_attempts))).await;
-            }
-            if self
-                .indexer_rest_client
-                .node_stats()
-                .indexing()
-                .await
-                .unwrap()
-                .num_running_pipelines
-                == required_pipeline_num
-            {
-                return Ok(());
-            }
-            num_attempts += 1;
-        }
-        if num_attempts == max_num_attempts {
-            anyhow::bail!("Too many attempts to get expected number of pipelines.");
-        }
+        wait_until_predicate(
+            || async move {
+                match self.indexer_rest_client.node_stats().indexing().await {
+                    Ok(result) => {
+                        if result.num_running_pipelines != required_pipeline_num {
+                            debug!(
+                                "wait_for_indexing_pipelines expected {} pipelines, got {}",
+                                required_pipeline_num, result.num_running_pipelines
+                            );
+                            false
+                        } else {
+                            true
+                        }
+                    }
+                    Err(err) => {
+                        debug!("wait_for_cluster_num_ready_nodes error {err}");
+                        false
+                    }
+                }
+            },
+            Duration::from_secs(10),
+            Duration::from_millis(100),
+        )
+        .await?;
         Ok(())
     }
 
@@ -225,29 +278,43 @@ impl ClusterSandbox {
         split_states: Option<Vec<SplitState>>,
         required_splits_num: usize,
     ) -> anyhow::Result<()> {
-        let mut num_attempts = 0;
-        let max_num_attempts = 3;
-        while num_attempts < max_num_attempts {
-            if num_attempts > 0 {
-                tokio::time::sleep(Duration::from_millis(100 * (num_attempts))).await;
-            }
-            if self
-                .indexer_rest_client
-                .splits(index_id)
-                .list(ListSplitsQueryParams {
+        wait_until_predicate(
+            || {
+                let splits_query_params = ListSplitsQueryParams {
                     split_states: split_states.clone(),
                     ..Default::default()
-                })
-                .await
-                .unwrap()
-                .len()
-                == required_splits_num
-            {
-                return Ok(());
-            }
-            num_attempts += 1;
-        }
-        anyhow::bail!("Too many attempts to get expected number of published splits.");
+                };
+                async move {
+                    match self
+                        .indexer_rest_client
+                        .splits(index_id)
+                        .list(splits_query_params)
+                        .await
+                    {
+                        Ok(result) => {
+                            if result.len() != required_splits_num {
+                                debug!(
+                                    "wait_for_published_splits expected {} splits, got {}",
+                                    required_splits_num,
+                                    result.len()
+                                );
+                                false
+                            } else {
+                                true
+                            }
+                        }
+                        Err(err) => {
+                            debug!("wait_for_cluster_num_ready_nodes error {err}");
+                            false
+                        }
+                    }
+                }
+            },
+            Duration::from_secs(10),
+            Duration::from_millis(100),
+        )
+        .await?;
+        Ok(())
     }
 
     pub async fn shutdown(self) -> Result<Vec<HashMap<String, ActorExitStatus>>, anyhow::Error> {

quickwit/quickwit-integration-tests/src/test_utils/mod.rs

@@ -19,4 +19,4 @@
 
 mod cluster_sandbox;
 
-pub use cluster_sandbox::{build_node_configs, ClusterSandbox};
+pub use cluster_sandbox::{build_node_configs, ingest_with_retry, ClusterSandbox};

quickwit/quickwit-integration-tests/src/tests/basic_tests.rs

@@ -28,7 +28,7 @@ use quickwit_rest_client::models::IngestSource;
 use quickwit_rest_client::rest_client::CommitType;
 use quickwit_serve::SearchRequestQueryString;
 
-use crate::test_utils::ClusterSandbox;
+use crate::test_utils::{ingest_with_retry, ClusterSandbox};
 
 fn get_ndjson_filepath(ndjson_dataset_filename: &str) -> String {
     format!(
@@ -223,16 +223,14 @@ async fn test_multi_nodes_cluster() {
     // Check that ingest request send to searcher is forwarded to indexer and thus indexed.
     let ndjson_filepath = get_ndjson_filepath("documents_to_ingest.json");
     let ingest_source = IngestSource::File(PathBuf::from_str(&ndjson_filepath).unwrap());
-    sandbox
-        .searcher_rest_client
-        .ingest(
-            "my-new-multi-node-index",
-            ingest_source,
-            None,
-            CommitType::Auto,
-        )
-        .await
-        .unwrap();
+    ingest_with_retry(
+        &sandbox.searcher_rest_client,
+        "my-new-multi-node-index",
+        ingest_source,
+        CommitType::Auto,
+    )
+    .await
+    .unwrap();
     // Wait until split is commited and search.
     tokio::time::sleep(Duration::from_secs(4)).await;
     let search_response_one_hit = sandbox

quickwit/quickwit-integration-tests/src/tests/index_tests.rs

@@ -19,12 +19,12 @@
 
 use bytes::Bytes;
 use quickwit_metastore::SplitState;
-use quickwit_rest_client::models::IngestSource;
 use quickwit_rest_client::rest_client::CommitType;
 use quickwit_serve::SearchRequestQueryString;
 use serde_json::json;
 
-use crate::test_utils::ClusterSandbox;
+use crate::ingest_json;
+use crate::test_utils::{ingest_with_retry, ClusterSandbox};
 
 #[tokio::test]
 async fn test_restarting_standalone_server() {
@@ -74,16 +74,14 @@ async fn test_restarting_standalone_server() {
         .incarnation_id;
 
     // Index one record.
-    sandbox
-        .indexer_rest_client
-        .ingest(
-            index_id,
-            IngestSource::Bytes(json!({"body": "first record"}).to_string().into()),
-            None,
-            CommitType::Force,
-        )
-        .await
-        .unwrap();
+    ingest_with_retry(
+        &sandbox.indexer_rest_client,
+        index_id,
+        ingest_json!({"body": "first record"}),
+        CommitType::Force,
+    )
+    .await
+    .unwrap();
 
     // Delete the indexq
     sandbox
@@ -113,33 +111,29 @@ async fn test_restarting_standalone_server() {
     assert_ne!(old_incarnation_id, new_incarnation_id);
 
     // Index a couple of records to create 2 additional splits.
+    ingest_with_retry(
+        &sandbox.indexer_rest_client,
+        index_id,
+        ingest_json!({"body": "second record"}),
+        CommitType::Force,
+    )
+    .await
+    .unwrap();
     sandbox
         .indexer_rest_client
         .ingest(
             index_id,
-            IngestSource::Bytes(json!({"body": "second record"}).to_string().into()),
-            None,
-            CommitType::Force,
-        )
-        .await
-        .unwrap();
-
-    sandbox
-        .indexer_rest_client
-        .ingest(
-            index_id,
-            IngestSource::Bytes(json!({"body": "third record"}).to_string().into()),
+            ingest_json!({"body": "third record"}),
             None,
             CommitType::Force,
         )
         .await
         .unwrap();
-
     sandbox
         .indexer_rest_client
         .ingest(
             index_id,
-            IngestSource::Bytes(json!({"body": "fourth record"}).to_string().into()),
+            ingest_json!({"body": "fourth record"}),
             None,
             CommitType::Force,
         )

quickwit/quickwit-rest-client/src/models.rs

@@ -71,6 +71,7 @@ impl ApiResponse {
     }
 }
 
+#[derive(Clone)]
 pub enum IngestSource {
     Bytes(Bytes),
     File(PathBuf),