Remove publisher finalize and shut down property a merge pipeline wit… (#2263)

* Remove publisher finalize and shut down property a merge pipeline with no corresponding indexing pipelines.

* Fixing one unit test.

The indexing pipeline updates its statistics on its finalize.

* Using quickwit-actors with the testsuite feature when running tests.

* Using the HEARTBEAT as the delete pipeline observation period.

* Apply suggestions from code review

Co-authored-by: Paul Masurel <[email protected]>

Author: fmassot

quickwit/quickwit-actors/src/actor.rs

@@ -22,7 +22,7 @@ use std::convert::Infallible;
 use std::fmt;
 use std::ops::Deref;
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
@@ -239,7 +239,7 @@ pub struct ActorContextInner<A: Actor> {
     // This counter is useful to unsure that obsolete WakeUp
     // events do not effect ulterior `sleep`.
     sleep_count: AtomicUsize,
-    observable_state_tx: Mutex<watch::Sender<A::ObservableState>>,
+    observable_state_tx: watch::Sender<A::ObservableState>,
 }
 
 /// Internal command used to resume an actor that was paused using
@@ -285,7 +285,7 @@ impl<A: Actor> ActorContext<A> {
                 registry,
                 actor_state: AtomicState::default(),
                 sleep_count: AtomicUsize::default(),
-                observable_state_tx: Mutex::new(observable_state_tx),
+                observable_state_tx,
             }
             .into(),
         }
@@ -408,11 +408,7 @@ impl<A: Actor> ActorContext<A> {
 
     pub(crate) fn observe(&self, actor: &mut A) -> A::ObservableState {
         let obs_state = actor.observable_state();
-        let _ = self
-            .observable_state_tx
-            .lock()
-            .unwrap()
-            .send(obs_state.clone());
+        let _ = self.observable_state_tx.send(obs_state.clone());
         obs_state
     }
 

quickwit/quickwit-actors/src/lib.rs

@@ -68,7 +68,7 @@ pub use self::supervisor::{Supervisor, SupervisorState};
 /// If an actor does not advertise a progress within an interval of duration `HEARTBEAT`,
 /// its supervisor will consider it as blocked and will proceed to kill it, as well
 /// as all of the actors all the actors that share the killswitch.
-pub const HEARTBEAT: Duration = if cfg!(test) {
+pub const HEARTBEAT: Duration = if cfg!(any(test, feature = "testsuite")) {
     // Right now some unit test end when we detect that a
     // pipeline has terminated, which can require waiting
     // for a heartbeat.

quickwit/quickwit-actors/src/spawn_builder.rs

@@ -290,6 +290,7 @@ async fn actor_loop<A: Actor>(actor: A, inbox: Inbox<A>, ctx: ActorContext<A>) -
     };
 
     let final_exit_status = actor_env.finalize(after_process_exit_status).await;
+    // The last observation is collected on `ActorExecutionEnv::Drop`.
     actor_env.process_exit_status(&final_exit_status);
     final_exit_status
 }

quickwit/quickwit-cli/src/index.rs

@@ -917,7 +917,8 @@ pub async fn ingest_docs_cli(args: IngestDocsArgs) -> anyhow::Result<()> {
         quickwit_storage_uri_resolver().clone(),
     )
     .await?;
-    let (indexing_server_mailbox, _) = universe.spawn_builder().spawn(indexing_server);
+    let (indexing_server_mailbox, indexing_server_handle) =
+        universe.spawn_builder().spawn(indexing_server);
     let pipeline_id = indexing_server_mailbox
         .ask_for_res(SpawnPipeline {
             index_id: args.index_id.clone(),
@@ -942,6 +943,11 @@ pub async fn ingest_docs_cli(args: IngestDocsArgs) -> anyhow::Result<()> {
     }
     let statistics =
         start_statistics_reporting_loop(pipeline_handle, args.input_path_opt.is_none()).await?;
+    // Shutdown the indexing server.
+    universe
+        .send_exit_with_success(&indexing_server_mailbox)
+        .await?;
+    indexing_server_handle.join().await;
     if statistics.num_published_splits > 0 {
         println!(
             "Now, you can query the index with the following command:\nquickwit index search \

quickwit/quickwit-common/src/net.rs

@@ -210,7 +210,6 @@ impl Display for HostAddr {
 /// Finds a random available TCP port.
 ///
 /// This function induces a race condition, use it only in unit tests.
-#[cfg(any(test, feature = "testsuite"))]
 pub fn find_available_tcp_port() -> anyhow::Result<u16> {
     let socket: SocketAddr = ([127, 0, 0, 1], 0u16).into();
     let listener = TcpListener::bind(socket)?;

quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs

@@ -117,6 +117,35 @@ impl Actor for IndexingPipeline {
         self.handle(Supervise, ctx).await?;
         Ok(())
     }
+
+    async fn finalize(
+        &mut self,
+        _exit_status: &ActorExitStatus,
+        _ctx: &ActorContext<Self>,
+    ) -> anyhow::Result<()> {
+        // We update the observation to ensure our last "black box" observation
+        // is up to date.
+        if let Some(handles) = &self.handles {
+            let (doc_processor_counters, indexer_counters, uploader_counters, publisher_counters) = join!(
+                handles.doc_processor.observe(),
+                handles.indexer.observe(),
+                handles.uploader.observe(),
+                handles.publisher.observe(),
+            );
+            self.statistics = self
+                .previous_generations_statistics
+                .clone()
+                .add_actor_counters(
+                    &*doc_processor_counters,
+                    &*indexer_counters,
+                    &*uploader_counters,
+                    &*publisher_counters,
+                )
+                .set_generation(self.statistics.generation)
+                .set_num_spawn_attempts(self.statistics.num_spawn_attempts);
+        }
+        Ok(())
+    }
 }
 
 impl IndexingPipeline {
@@ -475,13 +504,15 @@ mod tests {
     use std::path::PathBuf;
     use std::sync::Arc;
 
-    use quickwit_actors::Universe;
-    use quickwit_config::{IndexingSettings, SourceParams};
+    use quickwit_actors::{Command, Universe};
+    use quickwit_config::{IndexingSettings, SourceParams, VoidSourceParams};
     use quickwit_doc_mapper::default_doc_mapper_for_test;
     use quickwit_metastore::{IndexMetadata, MetastoreError, MockMetastore};
     use quickwit_storage::RamStorage;
 
     use super::{IndexingPipeline, *};
+    use crate::actors::merge_pipeline::{MergePipeline, MergePipelineParams};
+    use crate::merge_policy::default_merge_policy;
     use crate::models::IndexingDirectory;
 
     #[test]
@@ -674,4 +705,82 @@ mod tests {
         assert_eq!(pipeline_statistics.num_published_splits, 1);
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_merge_pipeline_does_not_stop_on_indexing_pipeline_failure() -> anyhow::Result<()>
+    {
+        let mut metastore = MockMetastore::default();
+        metastore
+            .expect_index_metadata()
+            .withf(|index_id| index_id == "test-index")
+            .returning(|_| {
+                Ok(IndexMetadata::for_test(
+                    "test-index",
+                    "ram:///indexes/test-index",
+                ))
+            });
+        metastore
+            .expect_list_splits()
+            .returning(|_, _, _, _| Ok(Vec::new()));
+        let universe = Universe::new();
+        let node_id = "test-node";
+        let metastore = Arc::new(metastore);
+        let doc_mapper = Arc::new(default_doc_mapper_for_test());
+        let pipeline_id = IndexingPipelineId {
+            index_id: "test-index".to_string(),
+            source_id: "test-source".to_string(),
+            node_id: node_id.to_string(),
+            pipeline_ord: 0,
+        };
+        let source_config = SourceConfig {
+            source_id: "test-source".to_string(),
+            num_pipelines: 1,
+            enabled: true,
+            source_params: SourceParams::Void(VoidSourceParams),
+        };
+        let storage = Arc::new(RamStorage::default());
+        let split_store = IndexingSplitStore::create_without_local_store(storage.clone());
+        let merge_pipeline_params = MergePipelineParams {
+            pipeline_id: pipeline_id.clone(),
+            doc_mapper: doc_mapper.clone(),
+            indexing_directory: IndexingDirectory::for_test().await,
+            metastore: metastore.clone(),
+            split_store: split_store.clone(),
+            merge_policy: default_merge_policy(),
+            max_concurrent_split_uploads: 2,
+            merge_max_io_num_bytes_per_sec: None,
+        };
+        let merge_pipeline = MergePipeline::new(merge_pipeline_params);
+        let merge_planner_mailbox = merge_pipeline.merge_planner_mailbox().clone();
+        let (_merge_pipeline_mailbox, merge_pipeline_handler) =
+            universe.spawn_builder().spawn(merge_pipeline);
+        let indexing_pipeline_params = IndexingPipelineParams {
+            pipeline_id,
+            doc_mapper,
+            source_config,
+            indexing_directory: IndexingDirectory::for_test().await,
+            indexing_settings: IndexingSettings::for_test(),
+            metastore: metastore.clone(),
+            queues_dir_path: PathBuf::from("./queues"),
+            storage,
+            split_store,
+            max_concurrent_split_uploads_index: 4,
+            max_concurrent_split_uploads_merge: 5,
+            merge_planner_mailbox: merge_planner_mailbox.clone(),
+        };
+        let indexing_pipeline = IndexingPipeline::new(indexing_pipeline_params);
+        let (_indexing_pipeline_mailbox, indexing_pipeline_handler) =
+            universe.spawn_builder().spawn(indexing_pipeline);
+        assert_eq!(indexing_pipeline_handler.observe().await.generation, 1);
+        // Let's shutdown the indexer, this will trigger the the indexing pipeline failure and the
+        // restart.
+        let indexer = universe.get::<Indexer>().into_iter().next().unwrap();
+        indexer.send_message(Command::Quit).await.unwrap();
+        tokio::time::sleep(Duration::from_secs(2)).await;
+        // Check indexing pipeline has restarted.
+        assert_eq!(indexing_pipeline_handler.observe().await.generation, 2);
+        // Check that the merge pipeline is still up.
+        assert_eq!(merge_pipeline_handler.health(), Health::Healthy);
+        Ok(())
+    }
 }

quickwit/quickwit-indexing/src/actors/indexing_service.rs

@@ -90,12 +90,13 @@ pub struct IndexingServiceState {
     pub num_running_pipelines: usize,
     pub num_successful_pipelines: usize,
     pub num_failed_pipelines: usize,
+    pub num_running_merge_pipelines: usize,
 }
 
 type IndexId = String;
 type SourceId = String;
 
-#[derive(Hash, Eq, PartialEq)]
+#[derive(Clone, Hash, Eq, PartialEq)]
 struct MergePipelineId {
     index_id: String,
     source_id: String,
@@ -385,19 +386,40 @@ impl IndexingService {
                     }
                 },
             );
-        // Evict merge pipelines that are not needed or failing.
+        // Evict and kill merge pipelines that are not needed.
         let needed_merge_pipeline_ids: HashSet<MergePipelineId> = self
             .indexing_pipeline_handles
             .keys()
             .map(MergePipelineId::from)
             .collect();
+        let current_merge_pipeline_ids: HashSet<MergePipelineId> =
+            self.merge_pipeline_handles.keys().cloned().collect();
+        for merge_pipeline_id_to_shut_down in
+            current_merge_pipeline_ids.difference(&needed_merge_pipeline_ids)
+        {
+            if let Some((_, merge_pipeline_handle)) = self
+                .merge_pipeline_handles
+                .remove_entry(merge_pipeline_id_to_shut_down)
+            {
+                // We kill the merge pipeline to avoid waiting a merge operation to finish as it can
+                // be long.
+                info!(
+                    index_id=%merge_pipeline_id_to_shut_down.index_id,
+                    source_id=%merge_pipeline_id_to_shut_down.source_id,
+                    "No more indexing pipeline on this index and source, killing merge pipeline."
+                );
+                merge_pipeline_handle.handle.kill().await;
+            }
+        }
+        // Finally remove the merge pipelien with an exit status.
         self.merge_pipeline_handles
-            .retain(|merge_pipeline_id, merge_pipeline_mailbox_handle| {
+            .retain(|_, merge_pipeline_mailbox_handle| {
                 match merge_pipeline_mailbox_handle.handle.health() {
-                    Health::Healthy => needed_merge_pipeline_ids.contains(merge_pipeline_id),
+                    Health::Healthy => true,
                     Health::FailureOrUnhealthy | Health::Success => false,
                 }
             });
+        self.state.num_running_merge_pipelines = self.merge_pipeline_handles.len();
         Ok(())
     }
 
@@ -617,7 +639,7 @@ impl Handler<ShutdownPipeline> for IndexingService {
 mod tests {
     use std::time::Duration;
 
-    use quickwit_actors::{ObservationType, Universe};
+    use quickwit_actors::{ObservationType, Universe, HEARTBEAT};
     use quickwit_common::rand::append_random_suffix;
     use quickwit_common::uri::Uri;
     use quickwit_config::{SourceConfig, VecSourceParams};
@@ -853,4 +875,84 @@ mod tests {
         }
         panic!("Sleep");
     }
+
+    #[tokio::test]
+    async fn test_indexing_service_shut_down_merge_pipeline_when_no_indexing_pipeline() {
+        quickwit_common::setup_logging_for_tests();
+        let metastore_uri = Uri::from_well_formed("ram:///metastore".to_string());
+        let metastore = quickwit_metastore_uri_resolver()
+            .resolve(&metastore_uri)
+            .await
+            .unwrap();
+
+        let index_id = append_random_suffix("test-indexing-service");
+        let index_uri = format!("ram:///indexes/{index_id}");
+        let index_metadata = IndexMetadata::for_test(&index_id, &index_uri);
+
+        let source_config = SourceConfig {
+            source_id: "test-indexing-service--source".to_string(),
+            num_pipelines: 1,
+            enabled: true,
+            source_params: SourceParams::void(),
+        };
+        metastore.create_index(index_metadata).await.unwrap();
+        metastore
+            .add_source(&index_id, source_config.clone())
+            .await
+            .unwrap();
+
+        // Test `IndexingService::new`.
+        let temp_dir = tempfile::tempdir().unwrap();
+        let data_dir_path = temp_dir.path().to_path_buf();
+        let indexer_config = IndexerConfig::for_test().unwrap();
+        let storage_resolver = StorageUriResolver::for_test();
+        let universe = Universe::new();
+        let queues_dir_path = data_dir_path.join(QUEUES_DIR_NAME);
+        init_ingest_api(&universe, &queues_dir_path).await.unwrap();
+        let indexing_server = IndexingService::new(
+            "test-node".to_string(),
+            data_dir_path,
+            indexer_config,
+            metastore.clone(),
+            storage_resolver.clone(),
+        )
+        .await
+        .unwrap();
+        let (indexing_server_mailbox, indexing_server_handle) =
+            universe.spawn_builder().spawn(indexing_server);
+        indexing_server_mailbox
+            .ask_for_res(SpawnPipelines {
+                index_id: index_id.clone(),
+            })
+            .await
+            .unwrap();
+        let observation = indexing_server_handle.observe().await;
+        assert_eq!(observation.num_running_pipelines, 1);
+        assert_eq!(observation.num_failed_pipelines, 0);
+        assert_eq!(observation.num_successful_pipelines, 0);
+
+        // Test `shutdown_pipeline`
+        indexing_server_mailbox
+            .ask_for_res(ShutdownPipelines {
+                index_id: index_id.clone(),
+                source_id: None,
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            indexing_server_handle.observe().await.num_running_pipelines,
+            0
+        );
+        assert_eq!(
+            indexing_server_handle
+                .observe()
+                .await
+                .num_running_merge_pipelines,
+            0
+        );
+        universe.simulate_time_shift(HEARTBEAT).await;
+        // Check that the merge pipeline is also shut down as they are no more indexing pipeilne on
+        // the index.
+        assert!(universe.get_one::<MergePipeline>().is_none());
+    }
 }