diff --git a/docs/configuration/node-config.md b/docs/configuration/node-config.md index eac504e7d45..32b07cd7442 100644 --- a/docs/configuration/node-config.md +++ b/docs/configuration/node-config.md @@ -247,6 +247,7 @@ This section contains the configuration options for an indexer. The split store | `enable_otlp_endpoint` | If true, enables the OpenTelemetry exporter endpoint to ingest logs and traces via the OpenTelemetry Protocol (OTLP). | `false` | | `cpu_capacity` | Advisory parameter used by the control plane. The value can expressed be in threads (e.g. `2`) or in term of millicpus (`2000m`). The control plane will attempt to schedule indexing pipelines on the different nodes proportionally to the cpu capacity advertised by the indexer. It is NOT used as a limit. All pipelines will be scheduled regardless of whether the cluster has sufficient capacity or not. The control plane does not attempt to spread the work equally when the load is well below the `cpu_capacity`. Users who need a balanced load on all of their indexer nodes can set the `cpu_capacity` to an arbitrarily low value as long as they keep it proportional to the number of threads available. | `num threads available` | | `enable_cooperative_indexing` | Enable sharing resources more efficiently when the number of indexes actively written to is significantly higher than the number of cores but might decrease the overall indexing throughput. | `false` | +| `shutdown_drain_timeout` | Time budget granted to each indexing pipeline to drain gracefully before its remaining actors are killed, on node shutdown and when the control plane tears down a pipeline. Set it above the largest `commit_timeout_secs` plus the time to upload and publish the final splits. Beware that on shutdown, draining only starts after the ingester and compactor decommissions complete, so the deployment's shutdown grace period (e.g. `terminationGracePeriodSeconds` on Kubernetes) must exceed the largest decommission timeout *plus* this value — about 600 seconds with the default timeouts. Can be overridden with the `QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT` environment variable. | `300s` | Set the `QW_INDEXING_MAX_WRITE_THROUGHPUT` environment variable to limit the aggregate indexing IO throughput on a node. It accepts human-readable byte sizes per second, such as `500mb`, and is unlimited by default. diff --git a/quickwit/quickwit-config/src/node_config/mod.rs b/quickwit/quickwit-config/src/node_config/mod.rs index 3bd488a2c0f..df24c7c336e 100644 --- a/quickwit/quickwit-config/src/node_config/mod.rs +++ b/quickwit/quickwit-config/src/node_config/mod.rs @@ -215,9 +215,25 @@ pub struct IndexerConfig { pub enable_cooperative_indexing: bool, #[serde(default = "IndexerConfig::default_cpu_capacity")] pub cpu_capacity: CpuCapacity, + #[serde(default = "IndexerConfig::default_shutdown_drain_timeout")] + pub shutdown_drain_timeout: HumanDuration, } impl IndexerConfig { + /// Returns the shutdown drain timeout, as defined in the environment variable or in + /// the configuration, in that order (the environment variable overrides the configuration). + pub fn shutdown_drain_timeout(&self) -> Duration { + quickwit_common::get_duration_from_env( + "QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT", + Duration::from(self.shutdown_drain_timeout.clone()), + ) + } + + fn default_shutdown_drain_timeout() -> HumanDuration { + HumanDuration::try_from("300s".to_string()) + .expect("`300s` should be a valid human duration") + } + fn default_enable_cooperative_indexing() -> bool { false } @@ -265,6 +281,7 @@ impl IndexerConfig { cpu_capacity: PIPELINE_FULL_CAPACITY * 4u32, max_merge_write_throughput: None, merge_concurrency: NonZeroUsize::new(3).unwrap(), + shutdown_drain_timeout: Self::default_shutdown_drain_timeout(), }; Ok(indexer_config) } @@ -281,6 +298,7 @@ impl Default for IndexerConfig { cpu_capacity: Self::default_cpu_capacity(), merge_concurrency: Self::default_merge_concurrency(), max_merge_write_throughput: None, + shutdown_drain_timeout: Self::default_shutdown_drain_timeout(), } } } diff --git a/quickwit/quickwit-config/src/node_config/serialize.rs b/quickwit/quickwit-config/src/node_config/serialize.rs index 41936f476a0..0384e5e25ce 100644 --- a/quickwit/quickwit-config/src/node_config/serialize.rs +++ b/quickwit/quickwit-config/src/node_config/serialize.rs @@ -890,6 +890,7 @@ mod tests { cpu_capacity: IndexerConfig::default_cpu_capacity(), enable_cooperative_indexing: false, max_merge_write_throughput: Some(ByteSize::mb(100)), + shutdown_drain_timeout: IndexerConfig::default_shutdown_drain_timeout(), } ); assert_eq!( diff --git a/quickwit/quickwit-indexing/src/actors/indexer.rs b/quickwit/quickwit-indexing/src/actors/indexer.rs index 742c94c25c1..b35b193fd3d 100644 --- a/quickwit/quickwit-indexing/src/actors/indexer.rs +++ b/quickwit/quickwit-indexing/src/actors/indexer.rs @@ -595,6 +595,15 @@ impl Indexer { ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { fail_point!("indexer:batch:before"); + // An empty batch can only flush an existing workbench (drain path); + // without one there is nothing to commit and creating a workbench + // just to consume it empty would needlessly hit the metastore. + if batch.docs.is_empty() + && batch.checkpoint_delta.is_empty() + && self.indexing_workbench_opt.is_none() + { + return Ok(()); + } let force_commit = batch.force_commit; self.indexer_state .index_batch( diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs index d6643fbe0e7..9eb01ae0f0d 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs @@ -47,7 +47,7 @@ use crate::merge_policy::MergePolicy; use crate::metrics::{ACTOR_NAME, BACKPRESSURE_MICROS, INDEXING_PIPELINES}; use crate::models::{IndexingStatistics, SharedPublishToken}; use crate::source::{ - AssignShards, Assignment, SourceActor, SourceRuntime, quickwit_supported_sources, + AssignShards, Assignment, Drain, SourceActor, SourceRuntime, quickwit_supported_sources, }; use crate::split_store::IndexingSplitStore; @@ -83,10 +83,76 @@ pub(crate) struct Spawn { pub(crate) retry_count: usize, } +/// Asks a pipeline to drain and shut itself down: the source stops emitting, +/// and once the in-flight batches are flushed, published, and settled +/// (acknowledged), the source exits on its own. Past `drain_timeout`, the +/// pipeline gives up, kills its actors, and still exits: delivery degrades to +/// at-least-once. A draining pipeline is never respawned. +/// +/// Only sources opting in through +/// [`crate::source::Source::should_be_drained`] actually drain: for the +/// others the pipeline kills its actors and exits immediately. +#[derive(Debug)] +pub struct DrainPipeline { + pub drain_timeout: Duration, +} + +/// Drain progress of a pipeline supervisor (see [`DrainPipeline`]). +#[derive(Default)] +pub(crate) struct DrainState { + deadline_opt: Option, +} + +pub(crate) enum DrainAction { + /// The drain is now (or was already) in progress: keep supervising. + KeepRunning, + /// Nothing left to settle: exit successfully. + Exit, + /// The source cannot drain: kill the actors, then exit. + TerminateAndExit, +} + +impl DrainState { + /// Whether a drain was initiated. A draining pipeline must never respawn. + pub(crate) fn is_draining(&self) -> bool { + self.deadline_opt.is_some() + } + + pub(crate) async fn on_drain_request( + &mut self, + running_source_opt: Option<(&Mailbox, bool)>, + drain_timeout: Duration, + ) -> DrainAction { + if self.is_draining() { + return DrainAction::KeepRunning; + } + let Some((source_mailbox, source_should_be_drained)) = running_source_opt else { + return DrainAction::Exit; + }; + if !source_should_be_drained { + // The source settles nothing on a drain: tearing the pipeline down + // right away. + return DrainAction::TerminateAndExit; + } + self.deadline_opt = Some(Instant::now() + drain_timeout); + let _ = source_mailbox.send_message(Drain).await; + DrainAction::KeepRunning + } + + /// Whether the draining pipeline ran out of time to settle: the caller + /// kills whatever is left and exits. + pub(crate) fn is_deadline_exceeded(&self) -> bool { + self.deadline_opt + .map(|deadline| Instant::now() >= deadline) + .unwrap_or(false) + } +} + /// Handles for standard Tantivy-based indexing pipeline. struct IndexingPipelineHandles { source_mailbox: Mailbox, source_handle: ActorHandle, + source_should_be_drained: bool, doc_processor: ActorHandle, indexer: ActorHandle, index_serializer: ActorHandle, @@ -115,6 +181,7 @@ pub struct IndexingPipeline { handles_opt: Option, // Killswitch used for the actors in the pipeline. This is not the supervisor killswitch. kill_switch: KillSwitch, + drain_state: DrainState, // The set of shard is something that can change dynamically without necessarily // requiring a respawn of the pipeline. @@ -168,6 +235,7 @@ impl IndexingPipeline { previous_generations_statistics: Default::default(), handles_opt: None, kill_switch: KillSwitch::default(), + drain_state: DrainState::default(), statistics: IndexingStatistics { params_fingerprint, ..Default::default() @@ -295,6 +363,15 @@ impl IndexingPipeline { Health::Healthy => {} Health::FailureOrUnhealthy => { self.terminate().await; + if self.drain_state.is_draining() { + // A draining pipeline is never respawned: whatever could + // not be settled is redelivered (see `DrainPipeline`). + warn!( + pipeline_id=?self.params.pipeline_id, + "draining indexing pipeline failed; exiting" + ); + return Err(ActorExitStatus::Success); + } let first_retry_delay = wait_duration_before_retry(0); ctx.schedule_self_msg(first_retry_delay, Spawn { retry_count: 0 }); } @@ -442,6 +519,7 @@ impl IndexingPipeline { let source = ctx .protect_future(quickwit_supported_sources().load_source(source_runtime)) .await?; + let source_should_be_drained = source.should_be_drained(); let actor_source = SourceActor::new(source, doc_processor_mailbox); let (source_mailbox, source_handle) = ctx .spawn_actor() @@ -460,6 +538,7 @@ impl IndexingPipeline { self.handles_opt = Some(IndexingPipelineHandles { source_mailbox, source_handle, + source_should_be_drained, doc_processor: doc_processor_handle, indexer: indexer_handle, index_serializer: index_serializer_handle, @@ -496,6 +575,14 @@ impl Handler for IndexingPipeline { ) -> Result<(), ActorExitStatus> { self.perform_observe(ctx); self.perform_health_check(ctx).await?; + if self.drain_state.is_deadline_exceeded() { + warn!( + pipeline_id=?self.params.pipeline_id, + "indexing pipeline could not be drained before the deadline" + ); + self.terminate().await; + return Err(ActorExitStatus::Success); + } ctx.schedule_self_msg(SUPERVISE_INTERVAL, supervise_loop_token); Ok(()) } @@ -510,6 +597,10 @@ impl Handler for IndexingPipeline { spawn: Spawn, ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { + if self.drain_state.is_draining() { + // A draining pipeline is shutting down: never respawn it. + return Ok(()); + } if self.handles_opt.is_some() { return Ok(()); } @@ -534,6 +625,34 @@ impl Handler for IndexingPipeline { } } +#[async_trait] +impl Handler for IndexingPipeline { + type Reply = (); + + async fn handle( + &mut self, + drain: DrainPipeline, + _ctx: &ActorContext, + ) -> Result<(), ActorExitStatus> { + let running_source_opt = self + .handles_opt + .as_ref() + .map(|handles| (&handles.source_mailbox, handles.source_should_be_drained)); + match self + .drain_state + .on_drain_request(running_source_opt, drain.drain_timeout) + .await + { + DrainAction::KeepRunning => Ok(()), + DrainAction::Exit => Err(ActorExitStatus::Success), + DrainAction::TerminateAndExit => { + self.terminate().await; + Err(ActorExitStatus::Success) + } + } + } +} + #[async_trait] impl Handler for IndexingPipeline { type Reply = (); diff --git a/quickwit/quickwit-indexing/src/actors/indexing_service.rs b/quickwit/quickwit-indexing/src/actors/indexing_service.rs index 97e0cf31b5c..3b04f109a78 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_service.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_service.rs @@ -16,6 +16,7 @@ use std::collections::{HashMap, HashSet}; use std::fmt::{Debug, Formatter}; use std::path::PathBuf; use std::sync::{Arc, LazyLock}; +use std::time::Duration; use anyhow::Context; use async_trait::async_trait; @@ -23,8 +24,8 @@ use bytesize::ByteSize; use futures::TryStreamExt; use itertools::Itertools; use quickwit_actors::{ - Actor, ActorContext, ActorExitStatus, ActorHandle, ActorState, Handler, Healthz, Mailbox, - Observation, + Actor, ActorContext, ActorExitStatus, ActorHandle, ActorState, DeferableReplyHandler, Handler, + Healthz, Mailbox, Observation, }; use quickwit_cluster::Cluster; use quickwit_common::pubsub::EventBroker; @@ -42,6 +43,7 @@ use quickwit_metastore::{ ListIndexesMetadataResponseExt, ListSplitsQuery, ListSplitsRequestExt, ListSplitsResponseExt, SplitMetadata, SplitState, }; +use quickwit_proto::GrpcServiceError; use quickwit_proto::indexing::{ ApplyIndexingPlanRequest, ApplyIndexingPlanResponse, IndexingError, IndexingPipelineId, IndexingTask, MergePipelineId, PipelineMetrics, @@ -60,6 +62,7 @@ use tracing::{debug, error, info, warn}; use super::merge_pipeline::{MergePipeline, MergePipelineParams}; use super::{FinishPendingMergesAndShutdownPipeline, MergePlanner, MergeSchedulerService}; +use crate::actors::indexing_pipeline::{DrainPipeline, SUPERVISE_INTERVAL}; use crate::docs_clustering::Fingerprinter; use crate::models::{DetachIndexingPipeline, DetachMergePipeline, ObservePipeline, SpawnPipeline}; use crate::source::{AssignShards, Assignment}; @@ -96,6 +99,10 @@ struct IndexingPipelineHandle { handle: ActorHandle, } +/// Reply callback of a pending `DrainAllPipelines` request, along with the +/// pipelines it still waits on. +type DrainAllWaiter = (Vec, Box); + /// The indexing service is (single) actor service running on indexer and in charge /// of executing the indexing plans received from the control plane. /// @@ -112,6 +119,16 @@ pub struct IndexingService { ingester_pool: IngesterPool, storage_resolver: StorageResolver, indexing_pipelines: HashMap, + /// Detached pipelines draining before their teardown. They exit on their + /// own and are reaped by the supervise loop. + draining_pipelines: Vec, + /// Pending `DrainAllPipelines` replies, completed by the supervise loop + /// once the tracked pipelines have exited. + drain_all_waiters: Vec, + drain_timeout: Duration, + /// Whether the whole service is draining ahead of shutdown: no new + /// pipeline may be spawned, whatever the requester. + draining: bool, latest_indexing_plan_id: IndexingPlanId, counters: IndexingServiceCounters, max_concurrent_split_uploads: usize, @@ -176,6 +193,10 @@ impl IndexingService { storage_resolver, split_cache, indexing_pipelines: Default::default(), + draining_pipelines: Vec::new(), + drain_all_waiters: Vec::new(), + drain_timeout: indexer_config.shutdown_drain_timeout(), + draining: false, latest_indexing_plan_id: String::new(), counters: Default::default(), max_concurrent_split_uploads: indexer_config.max_concurrent_split_uploads, @@ -191,7 +212,7 @@ impl IndexingService { async fn detach_indexing_pipeline( &mut self, pipeline_uid: &PipelineUid, - ) -> Result, IndexingError> { + ) -> Result { let pipeline_handle = self .indexing_pipelines .remove(pipeline_uid) @@ -200,7 +221,7 @@ impl IndexingService { IndexingError::Internal(message) })?; self.counters.num_running_pipelines -= 1; - Ok(pipeline_handle.handle) + Ok(pipeline_handle) } async fn detach_merge_pipeline( @@ -237,6 +258,14 @@ impl IndexingService { source_config: SourceConfig, pipeline_uid: PipelineUid, ) -> Result { + if self.draining { + return Err(IndexingError::new_unavailable( + "indexing service and all of its pipelines are being drained, won't spawn new \ + pipelines" + .to_string(), + )); + } + let index_metadata = self.index_metadata(ctx, &index_id).await?; let pipeline_id = IndexingPipelineId { index_uid: index_metadata.index_uid.clone(), @@ -522,6 +551,18 @@ impl IndexingService { Ok(per_merge_pipeline_immature_splits) } + /// Whether the pipeline is still alive, attached or draining. A missing + /// pipeline counts as exited. + fn is_pipeline_alive(&self, pipeline_uid: PipelineUid) -> bool { + if let Some(pipeline_handle) = self.indexing_pipelines.get(&pipeline_uid) { + return !pipeline_handle.handle.state().is_exit(); + } + self.draining_pipelines.iter().any(|pipeline_handle| { + pipeline_handle.pipeline_id.pipeline_uid == pipeline_uid + && !pipeline_handle.handle.state().is_exit() + }) + } + async fn handle_supervise(&mut self) -> Result<(), ActorExitStatus> { self.indexing_pipelines .retain(|pipeline_uid, pipeline_handle| { @@ -543,9 +584,30 @@ impl IndexingService { } } }); + + // Draining pipelines exit on their own, reap the + // detached ones and complete the `DrainAllPipelines` replies if all + // of their pipelines are gone. + self.draining_pipelines + .retain(|pipeline_handle| !pipeline_handle.handle.state().is_exit()); + if !self.drain_all_waiters.is_empty() { + let drain_all_waiters = std::mem::take(&mut self.drain_all_waiters); + for (mut pipeline_uids, reply) in drain_all_waiters { + pipeline_uids.retain(|pipeline_uid| self.is_pipeline_alive(*pipeline_uid)); + if pipeline_uids.is_empty() { + reply(()); + } else { + self.drain_all_waiters.push((pipeline_uids, reply)); + } + } + } + + // Draining pipelines may still publish splits until they exit: their + // merge pipelines must survive them. let merge_pipelines_to_retain: HashSet = self .indexing_pipelines .values() + .chain(self.draining_pipelines.iter()) .map(|pipeline_handle| pipeline_handle.pipeline_id.merge_pipeline_id()) .collect(); @@ -664,6 +726,15 @@ impl IndexingService { plan_request: ApplyIndexingPlanRequest, ctx: &ActorContext, ) -> Result<(), IndexingError> { + if self.draining { + // Reapplied plans would otherwise respawn the pipelines that just + // drained and exited, defeating the drain. + return Err(IndexingError::new_unavailable( + "indexing service and all of its pipelines are being drained, ignoring indexing \ + plan" + .to_string(), + )); + } // Plan ids are ULIDs if plan_request.indexing_plan_id < self.latest_indexing_plan_id { info!( @@ -816,12 +887,10 @@ impl IndexingService { .flat_map(|pipeline_uid| self.indexing_pipelines.get(pipeline_uid)) .any(|pipeline_handle| pipeline_handle.pipeline_id.source_id == INGEST_API_SOURCE_ID); + let mut detached_pipeline_handles: Vec = Vec::new(); for pipeline_to_shutdown in pipelines_to_shutdown { match self.detach_indexing_pipeline(pipeline_to_shutdown).await { - Ok(pipeline_handle) => { - // Killing the pipeline ensures that all the pipeline actors will stop. - pipeline_handle.kill().await; - } + Ok(pipeline_handle) => detached_pipeline_handles.push(pipeline_handle), Err(error) => { // Just log the detach error, it can only come from a missing pipeline in the // `indexing_pipeline_handles`. @@ -833,6 +902,13 @@ impl IndexingService { } } } + + // Keep a handle to the pipeline that are being drained so that in case of + // global shutdown we still wait for these ones. + for pipeline_handle in detached_pipeline_handles { + start_pipeline_drain(&pipeline_handle, self.drain_timeout).await; + self.draining_pipelines.push(pipeline_handle); + } // If at least one ingest source has been removed, the related index has possibly been // deleted. Thus we run a garbage collect to remove queues of potentially deleted // indexes. @@ -929,6 +1005,84 @@ impl IndexingService { } } +/// Drains every running indexing pipeline: it is meant +/// to run right before the universe is torn down on node shutdown, so planned +/// shutdowns publish and settle their in-flight batches instead of dropping +/// them. The reply is deferred until the drained pipelines have all exited, +/// without blocking the service. +#[derive(Debug)] +pub struct DrainAllPipelines; + +#[async_trait] +impl DeferableReplyHandler for IndexingService { + type Reply = (); + + async fn handle_message( + &mut self, + _msg: DrainAllPipelines, + reply: impl FnOnce(()) + Send + Sync + 'static, + ctx: &ActorContext, + ) -> Result<(), ActorExitStatus> { + self.draining = true; + + let mut pipeline_uids: Vec = Vec::new(); + for (pipeline_uid, pipeline_handle) in &self.indexing_pipelines { + start_pipeline_drain(pipeline_handle, self.drain_timeout).await; + pipeline_uids.push(*pipeline_uid); + } + // Some pipelines were potentially already draining when shutting down, e.g. in case + // of new indexing plans. + for pipeline_handle in &self.draining_pipelines { + if !pipeline_handle.handle.state().is_exit() { + pipeline_uids.push(pipeline_handle.pipeline_id.pipeline_uid); + } + } + if pipeline_uids.is_empty() { + reply(()); + return Ok(()); + } + info!( + num_pipelines = pipeline_uids.len(), + "draining indexing pipelines" + ); + + // Does a drain request was already sent? + let start_drain_supervise_loop = self.drain_all_waiters.is_empty(); + self.drain_all_waiters + .push((pipeline_uids, Box::new(reply))); + // The normal supervise loop only ticks every `HEARTBEAT` (30s in + // production), which would delay the reply, and thus the node + // shutdown. Tick faster while waiters are pending. + if start_drain_supervise_loop { + ctx.schedule_self_msg(SUPERVISE_INTERVAL, DrainSuperviseLoop); + } + Ok(()) + } +} + +#[derive(Debug)] +struct DrainSuperviseLoop; + +#[async_trait] +impl Handler for IndexingService { + type Reply = (); + + async fn handle( + &mut self, + _message: DrainSuperviseLoop, + ctx: &ActorContext, + ) -> Result<(), ActorExitStatus> { + if self.drain_all_waiters.is_empty() { + return Ok(()); + } + self.handle_supervise().await?; + if !self.drain_all_waiters.is_empty() { + ctx.schedule_self_msg(SUPERVISE_INTERVAL, DrainSuperviseLoop); + } + Ok(()) + } +} + #[async_trait] impl Handler for IndexingService { type Reply = Result, IndexingError>; @@ -954,8 +1108,11 @@ impl Handler for IndexingService { _ctx: &ActorContext, ) -> Result { let pipeline_uid = msg.pipeline_id.pipeline_uid; - let detach_pipeline_result = self.detach_indexing_pipeline(&pipeline_uid).await; - Ok(detach_pipeline_result) + let detached_pipeline_result = self + .detach_indexing_pipeline(&pipeline_uid) + .await + .map(|v| v.handle); + Ok(detached_pipeline_result) } } @@ -1059,13 +1216,21 @@ struct IndexingPipelineDiff { pipelines_to_spawn: Vec, } +async fn start_pipeline_drain(pipeline: &IndexingPipelineHandle, drain_timeout: Duration) { + // await for the DrainPipeline handler so when the source does not opt into + // draining, the immediate teardown completes before replacement + // pipelines can spawn and overlap its consumers. Opted-in pipelines + // reply as soon as the drain is initiated. + let _ = pipeline.mailbox.ask(DrainPipeline { drain_timeout }).await; +} + #[cfg(test)] mod tests { use std::num::NonZeroUsize; use std::path::Path; use std::time::Duration; - use quickwit_actors::{HEARTBEAT, Health, ObservationType, Supervisable, Universe}; + use quickwit_actors::{AskError, HEARTBEAT, Health, ObservationType, Supervisable, Universe}; use quickwit_cluster::{ChitchatTransport, create_cluster_for_test}; use quickwit_common::ServiceStream; use quickwit_common::rand::append_random_suffix; @@ -1225,6 +1390,85 @@ mod tests { universe.assert_quit().await; } + #[tokio::test] + async fn test_indexing_service_drain_all_pipelines() { + quickwit_common::setup_logging_for_tests(); + let transport = ChitchatTransport::default(); + let cluster = create_cluster_for_test(Vec::new(), &["indexer"], &transport, true) + .await + .unwrap(); + let metastore = metastore_for_test(); + + let index_id = append_random_suffix("test-indexing-service"); + let index_uri = format!("ram:///indexes/{index_id}"); + let index_config = IndexConfig::for_test(&index_id, &index_uri); + let create_index_request = + CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + let index_uid: IndexUid = metastore + .create_index(create_index_request) + .await + .unwrap() + .index_uid() + .clone(); + let source_config = SourceConfig { + source_id: "test-indexing-service--source".to_string(), + num_pipelines: NonZeroUsize::MIN, + enabled: true, + source_params: SourceParams::void(), + transform_config: None, + input_format: SourceInputFormat::Json, + }; + let add_source_request = + AddSourceRequest::try_from_source_config(index_uid, &source_config).unwrap(); + metastore.add_source(add_source_request).await.unwrap(); + + let universe = Universe::with_accelerated_time(); + let temp_dir = tempfile::tempdir().unwrap(); + let (indexing_service, indexing_service_handle) = + spawn_indexing_service_for_test(temp_dir.path(), &universe, metastore, cluster).await; + indexing_service + .ask_for_res(SpawnPipeline { + index_id: index_id.clone(), + source_config: source_config.clone(), + pipeline_uid: PipelineUid::for_test(0u128), + }) + .await + .unwrap(); + + // The reply is deferred until the pipeline has exited. The void source + // does not opt into draining, so the pipeline is torn down right away. + indexing_service.ask(DrainAllPipelines).await.unwrap(); + let observation = indexing_service_handle.observe().await; + assert_eq!(observation.num_running_pipelines, 0); + assert_eq!(observation.num_successful_pipelines, 1); + + // A draining service accepts neither new pipelines nor indexing plans. + let spawn_error = indexing_service + .ask_for_res(SpawnPipeline { + index_id, + source_config, + pipeline_uid: PipelineUid::for_test(1u128), + }) + .await + .unwrap_err(); + assert!(matches!( + spawn_error, + AskError::ErrorReply(IndexingError::Unavailable(_)) + )); + let apply_plan_error = indexing_service + .ask_for_res(ApplyIndexingPlanRequest { + indexing_tasks: Vec::new(), + indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5FA1".to_string(), + }) + .await + .unwrap_err(); + assert!(matches!( + apply_plan_error, + AskError::ErrorReply(IndexingError::Unavailable(_)) + )); + universe.assert_quit().await; + } + #[tokio::test] async fn test_indexing_service_supervise_pipelines() { quickwit_common::setup_logging_for_tests(); diff --git a/quickwit/quickwit-indexing/src/actors/mod.rs b/quickwit/quickwit-indexing/src/actors/mod.rs index 06490457b68..e9e3272159c 100644 --- a/quickwit/quickwit-indexing/src/actors/mod.rs +++ b/quickwit/quickwit-indexing/src/actors/mod.rs @@ -34,7 +34,9 @@ pub use doc_processor::{DocProcessor, DocProcessorCounters}; pub use index_serializer::IndexSerializer; pub use indexer::{Indexer, IndexerCounters}; pub use indexing_pipeline::{IndexingPipeline, IndexingPipelineParams}; -pub use indexing_service::{INDEXING_DIR_NAME, IndexingService, IndexingServiceCounters}; +pub use indexing_service::{ + DrainAllPipelines, INDEXING_DIR_NAME, IndexingService, IndexingServiceCounters, +}; pub use merge_executor::{MergeExecutor, combine_partition_ids, merge_split_attrs}; pub use merge_pipeline::{ FinishPendingMergesAndShutdownPipeline, MergePipeline, MergePipelineParams, diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index 372c22c72c5..f667d23acbc 100644 --- a/quickwit/quickwit-indexing/src/source/mod.rs +++ b/quickwit/quickwit-indexing/src/source/mod.rs @@ -295,6 +295,25 @@ pub trait Source: Send + 'static { Ok(()) } + /// Whether tearing down the pipeline must go through a drain ([`Drain`]) + /// rather than a plain kill. Only acknowledgment-based sources should opt + /// in: killing them loses the acknowledgments of their in-flight messages, + /// which are then redelivered and indexed again after the broker's ack + /// timeout. + fn should_be_drained(&self) -> bool { + false + } + + /// Once the source is draining (see [`Drain`]), reports whether every + /// message it delivered has been durably accounted for, i.e. published + /// and, for acknowledgment-based sources, acknowledged. + /// + /// Sources without acknowledgment state have nothing to wait for and are + /// drained as soon as they stop emitting. + fn is_drained(&self) -> bool { + true + } + /// Finalize is called once after the actor terminates. async fn finalize( &mut self, @@ -320,6 +339,7 @@ pub trait Source: Send + 'static { pub struct SourceActor { source: Box, doc_processor_mailbox: Mailbox, + draining: bool, } impl SourceActor { @@ -327,6 +347,7 @@ impl SourceActor { SourceActor { source, doc_processor_mailbox, + draining: false, } } } @@ -382,11 +403,50 @@ impl Actor for SourceActor { } } +/// Asks the source to stop emitting batches while staying alive to handle +/// `SuggestTruncate`. An empty force-commit batch is pushed downstream so the +/// indexer flushes its pending workbench instead of waiting for the commit +/// timeout. The source actor exits with success once everything it delivered +/// has been durably accounted for (see [`Source::is_drained`]) — immediately +/// for sources without acknowledgment state, otherwise on the truncate +/// notification that settles the last in-flight message. +#[derive(Debug)] +pub struct Drain; + +#[async_trait] +impl Handler for SourceActor { + type Reply = (); + + async fn handle( + &mut self, + _message: Drain, + ctx: &SourceContext, + ) -> Result<(), ActorExitStatus> { + if self.draining { + return Ok(()); + } + self.draining = true; + // sends a last empty batch with force commit to wait for the entire pipeline flush. + let flush_batch = RawDocBatch::new(Vec::new(), SourceCheckpointDelta::default(), true); + ctx.send_message(&self.doc_processor_mailbox, flush_batch) + .await?; + + if self.source.is_drained() { + return Err(ActorExitStatus::Success); + } + Ok(()) + } +} + #[async_trait] impl Handler for SourceActor { type Reply = (); async fn handle(&mut self, _message: Loop, ctx: &SourceContext) -> Result<(), ActorExitStatus> { + // If we are draining we should stop emit_batches so a source can flush gracefully. + if self.draining { + return Ok(()); + } let wait_for = self .source .emit_batches(&self.doc_processor_mailbox, ctx) @@ -520,6 +580,10 @@ impl Handler for SourceActor { // kill the source nor the indexing pipeline, but we log the error. error!(%error, "failed to process suggest truncate"); } + if self.draining && self.source.is_drained() { + // This truncate settled the last in-flight message (see `Drain`). + return Err(ActorExitStatus::Success); + } Ok(()) } } @@ -698,6 +762,47 @@ mod tests { } } + #[tokio::test] + async fn test_source_actor_drain() { + use quickwit_actors::Universe; + + let universe = Universe::with_accelerated_time(); + let source_config = SourceConfig { + source_id: "void".to_string(), + num_pipelines: NonZeroUsize::MIN, + enabled: true, + source_params: SourceParams::void(), + transform_config: None, + input_format: SourceInputFormat::Json, + }; + let source_runtime = + SourceRuntimeBuilder::new(IndexUid::new_with_random_ulid("test-index"), source_config) + .build(); + let source = quickwit_supported_sources() + .load_source(source_runtime) + .await + .unwrap(); + let (doc_processor_mailbox, doc_processor_inbox) = + universe.create_test_mailbox::(); + let source_actor = SourceActor::new(source, doc_processor_mailbox); + let (source_mailbox, source_handle) = universe.spawn_builder().spawn(source_actor); + + source_mailbox.send_message(Drain).await.unwrap(); + // A source without ack state drains immediately: the actor exits with + // success on its own. + let (exit_status, _state) = source_handle.join().await; + assert!(matches!(exit_status, ActorExitStatus::Success)); + + // Draining pushes an empty force-commit batch to flush the indexer. + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + assert_eq!(batches.len(), 1); + assert!(batches[0].docs.is_empty()); + assert!(batches[0].checkpoint_delta.is_empty()); + assert!(batches[0].force_commit); + + universe.assert_quit().await; + } + #[tokio::test] async fn test_check_source_connectivity() -> anyhow::Result<()> { { diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index 3fcbb2d5593..c8b08927674 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -84,7 +84,7 @@ use quickwit_config::{ClusterConfig, IngestApiConfig, NodeConfig, disable_ingest use quickwit_control_plane::control_plane::{ControlPlane, ControlPlaneEventSubscriber}; use quickwit_control_plane::{IndexerPool, IndexerPoolEntry}; use quickwit_index_management::{IndexService as IndexManager, IndexServiceError}; -use quickwit_indexing::actors::{IndexingService, MergeSchedulerService}; +use quickwit_indexing::actors::{DrainAllPipelines, IndexingService, MergeSchedulerService}; use quickwit_indexing::models::ShardPositionsService; use quickwit_indexing::{IndexingSplitCache, start_indexing_service}; use quickwit_ingest::{ @@ -505,6 +505,8 @@ async fn shutdown_signal_handler( shutdown_signal: BoxFutureInfaillible<()>, shutdown_token: CancellationToken, universe: Universe, + indexing_service_opt: Option>, + indexer_shutdown_drain_timeout: Duration, ingester_opt: Option, ingester_decommission_timeout: Duration, compactor_service_opt: Option>, @@ -548,6 +550,26 @@ async fn shutdown_signal_handler( if let Err(error) = compactor_result { error!("failed to decommission compactor gracefully: {:?}", error); } + + if let Some(indexing_service) = &indexing_service_opt { + match tokio::time::timeout( + indexer_shutdown_drain_timeout, + indexing_service.ask(DrainAllPipelines), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + error!("failed to drain indexing pipelines gracefully: {:?}", error); + } + Err(_elapsed) => { + error!( + timeout=?indexer_shutdown_drain_timeout, + "timed out draining the indexing pipelines" + ); + } + } + } let actor_exit_statuses = universe.quit().await; if grpc_shutdown_trigger_tx.send(()).is_err() { @@ -903,6 +925,7 @@ pub async fn serve_quickwit( let grpc_listen_addr = node_config.grpc_listen_addr; let rest_listen_addr = node_config.rest_config.listen_addr; + let indexer_shutdown_drain_timeout = node_config.indexer_config.shutdown_drain_timeout(); let ingester_decommission_timeout = node_config.ingest_api_config.decommission_timeout(); let compactor_decommission_timeout = node_config.compactor_config.decommission_timeout(); let quickwit_services: Arc = Arc::new(QuickwitServices { @@ -1015,6 +1038,8 @@ pub async fn serve_quickwit( shutdown_signal, shutdown_token.clone(), universe, + quickwit_services.indexing_service_opt.clone(), + indexer_shutdown_drain_timeout, ingester_opt, ingester_decommission_timeout, compactor_service_opt,