From fd6b5d49ec672b8ab5800f2dba2d9fd143aade93 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 17:48:05 +0200 Subject: [PATCH 01/41] feat: indexing pipeline drain Drain indexing pipelines on planned teardowns instead of killing them: the source stops emitting, flushes the in-flight batches with a forced commit, and the pipeline exits on its own once everything is published and settled. Sources opt in via source_needs_drain; others keep the plain kill semantics. Co-Authored-By: Claude Fable 5 --- .../quickwit-indexing/src/actors/indexer.rs | 9 ++ .../src/actors/indexing_pipeline.rs | 59 +++++++++- .../src/actors/indexing_service.rs | 101 ++++++++++++++-- quickwit/quickwit-indexing/src/actors/mod.rs | 3 +- .../parquet_pipeline/indexing_service_impl.rs | 3 + .../src/actors/parquet_pipeline/pipeline.rs | 61 +++++++++- .../src/actors/pipeline_shared.rs | 35 +++++- quickwit/quickwit-indexing/src/source/mod.rs | 108 ++++++++++++++++++ quickwit/quickwit-serve/src/lib.rs | 13 ++- 9 files changed, 377 insertions(+), 15 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/indexer.rs b/quickwit/quickwit-indexing/src/actors/indexer.rs index 5de2eda20f0..8dfdabeafba 100644 --- a/quickwit/quickwit-indexing/src/actors/indexer.rs +++ b/quickwit/quickwit-indexing/src/actors/indexer.rs @@ -591,6 +591,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 e134a23ed43..6b4203c54f2 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs @@ -39,7 +39,8 @@ use tracing::{debug, error, info, instrument, warn}; use super::{DocProcessor, IndexSerializer, Indexer, MergePlanner, Packager}; use crate::SplitsUpdateMailbox; use crate::actors::pipeline_shared::{ - SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, SuperviseLoop, wait_duration_before_retry, + DRAIN_GRACE_MARGIN, DrainPipeline, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, + SuperviseLoop, wait_duration_before_retry, }; use crate::actors::sequencer::Sequencer; use crate::actors::uploader::UploaderType; @@ -49,7 +50,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; @@ -85,6 +86,10 @@ pub struct IndexingPipeline { handles_opt: Option, // Killswitch used for the actors in the pipeline. This is not the supervisor killswitch. kill_switch: KillSwitch, + // Set when the pipeline is draining (see `DrainPipeline`): the pipeline is + // shutting down and must never respawn; past the deadline it gives up and + // kills whatever is left. + drain_deadline_opt: Option, // The set of shard is something that can change dynamically without necessarily // requiring a respawn of the pipeline. @@ -137,6 +142,7 @@ impl IndexingPipeline { previous_generations_statistics: Default::default(), handles_opt: None, kill_switch: KillSwitch::default(), + drain_deadline_opt: None, statistics: IndexingStatistics { params_fingerprint, ..Default::default() @@ -264,6 +270,15 @@ impl IndexingPipeline { Health::Healthy => {} Health::FailureOrUnhealthy => { self.terminate().await; + if self.drain_deadline_opt.is_some() { + // 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 }); } @@ -459,6 +474,16 @@ impl Handler for IndexingPipeline { ) -> Result<(), ActorExitStatus> { self.perform_observe(ctx); self.perform_health_check(ctx).await?; + if let Some(drain_deadline) = self.drain_deadline_opt + && Instant::now() >= drain_deadline + { + 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(()) } @@ -473,6 +498,10 @@ impl Handler for IndexingPipeline { spawn: Spawn, ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { + if self.drain_deadline_opt.is_some() { + // A draining pipeline is shutting down: never respawn it. + return Ok(()); + } if self.handles_opt.is_some() { return Ok(()); } @@ -497,6 +526,32 @@ impl Handler for IndexingPipeline { } } +#[async_trait] +impl Handler for IndexingPipeline { + type Reply = (); + + async fn handle( + &mut self, + _drain: DrainPipeline, + _ctx: &ActorContext, + ) -> Result<(), ActorExitStatus> { + if self.drain_deadline_opt.is_some() { + return Ok(()); + } + let Some(handles) = &self.handles_opt else { + // Nothing is running, so nothing is in flight. + return Err(ActorExitStatus::Success); + }; + self.drain_deadline_opt = Some( + Instant::now() + self.params.indexing_settings.commit_timeout() + DRAIN_GRACE_MARGIN, + ); + // If the source is already dead, the health check terminates the + // draining pipeline instead of respawning it. + let _ = handles.source_mailbox.send_message(Drain).await; + Ok(()) + } +} + #[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 a533d2789a0..9be098a9928 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_service.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_service.rs @@ -22,8 +22,8 @@ use async_trait::async_trait; 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; @@ -62,7 +62,7 @@ use super::pipeline_shared::{ActorPipeline, PipelineHandle}; use super::{FinishPendingMergesAndShutdownPipeline, MergePlanner, MergeSchedulerService}; use crate::docs_clustering::Fingerprinter; use crate::models::{DetachIndexingPipeline, DetachMergePipeline, ObservePipeline, SpawnPipeline}; -use crate::source::{AssignShards, Assignment}; +use crate::source::{AssignShards, Assignment, source_needs_drain}; use crate::split_store::IndexingSplitCache; use crate::{IndexingPipeline, IndexingPipelineParams, IndexingSplitStore, IndexingStatistics}; @@ -92,6 +92,10 @@ struct ParquetMergePipelineHandle { pub type BoxedPipelineHandle = Box; +/// 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. /// @@ -108,6 +112,12 @@ pub struct IndexingService { pub(crate) ingester_pool: IngesterPool, pub(crate) storage_resolver: StorageResolver, indexing_pipelines: HashMap, + /// Detached pipelines draining before their teardown. They exit on their + /// own (see `DrainPipeline`) 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, latest_indexing_plan_id: IndexingPlanId, counters: IndexingServiceCounters, pub(crate) max_concurrent_split_uploads: usize, @@ -178,6 +188,8 @@ impl IndexingService { storage_resolver, split_cache, indexing_pipelines: Default::default(), + draining_pipelines: Vec::new(), + drain_all_waiters: Vec::new(), latest_indexing_plan_id: String::new(), counters: Default::default(), max_concurrent_split_uploads: indexer_config.max_concurrent_split_uploads, @@ -409,6 +421,7 @@ impl IndexingService { ); } + let needs_drain = source_needs_drain(&source_config.source_params); let pipeline_params = IndexingPipelineParams { pipeline_id: indexing_pipeline_id.clone(), metastore: self.metastore.clone(), @@ -437,6 +450,7 @@ impl IndexingService { pipeline_id: indexing_pipeline_id, mailbox, handle, + needs_drain, })) } @@ -568,6 +582,28 @@ impl IndexingService { } } }); + // Draining pipelines exit on their own (see `DrainPipeline`): reap the + // detached ones and complete the `DrainAllPipelines` replies whose + // pipelines are all gone. + self.draining_pipelines + .retain(|pipeline_handle| pipeline_handle.state().is_running()); + 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.indexing_pipelines + .get(pipeline_uid) + .map(|pipeline_handle| pipeline_handle.state().is_running()) + .unwrap_or(false) + }); + if pipeline_uids.is_empty() { + reply(()); + } else { + self.drain_all_waiters.push((pipeline_uids, reply)); + } + } + } + let merge_pipelines_to_retain: HashSet = self .indexing_pipelines .values() @@ -965,12 +1001,10 @@ impl IndexingService { pipeline_handle.indexing_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`. @@ -982,6 +1016,20 @@ impl IndexingService { } } } + // Acknowledgment-based sources are drained before their teardown so + // their in-flight messages are published and acknowledged instead of + // redelivered (see `DrainPipeline`). A drain can take up to the commit + // timeout, so it is not awaited here: the draining pipeline exits on + // its own and the supervise loop reaps its handle. + for pipeline_handle in detached_pipeline_handles { + if pipeline_handle.needs_drain() { + pipeline_handle.start_drain().await; + self.draining_pipelines.push(pipeline_handle); + } else { + // Killing the pipeline ensures that all the pipeline actors will stop. + pipeline_handle.kill().await; + } + } // 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. @@ -1078,6 +1126,45 @@ impl IndexingService { } } +/// Drains every running indexing pipeline whose source needs it (see +/// `DrainPipeline`): 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> { + let mut pipeline_uids: Vec = Vec::new(); + for (pipeline_uid, pipeline_handle) in &self.indexing_pipelines { + if pipeline_handle.needs_drain() { + pipeline_handle.start_drain().await; + pipeline_uids.push(*pipeline_uid); + } + } + if pipeline_uids.is_empty() { + reply(()); + return Ok(()); + } + info!( + num_pipelines = pipeline_uids.len(), + "draining indexing pipelines" + ); + self.drain_all_waiters + .push((pipeline_uids, Box::new(reply))); + Ok(()) + } +} + #[async_trait] impl Handler for IndexingService { type Reply = Result, IndexingError>; diff --git a/quickwit/quickwit-indexing/src/actors/mod.rs b/quickwit/quickwit-indexing/src/actors/mod.rs index fdab5df4bcc..dc0a41724a0 100644 --- a/quickwit/quickwit-indexing/src/actors/mod.rs +++ b/quickwit/quickwit-indexing/src/actors/mod.rs @@ -39,7 +39,8 @@ pub use index_serializer::IndexSerializer; pub use indexer::{Indexer, IndexerCounters}; pub use indexing_pipeline::{IndexingPipeline, IndexingPipelineParams}; pub use indexing_service::{ - BoxedPipelineHandle, INDEXING_DIR_NAME, IndexingService, IndexingServiceCounters, + BoxedPipelineHandle, DrainAllPipelines, INDEXING_DIR_NAME, IndexingService, + IndexingServiceCounters, }; pub use log_publisher_impl::MERGE_PUBLISHER_NAME; pub(crate) use log_publisher_impl::PUBLISHER_NAME; diff --git a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/indexing_service_impl.rs b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/indexing_service_impl.rs index 857bf399838..2a658ad10d2 100644 --- a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/indexing_service_impl.rs +++ b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/indexing_service_impl.rs @@ -26,6 +26,7 @@ use quickwit_proto::indexing::{IndexingError, IndexingPipelineId}; use crate::actors::pipeline_shared::ActorPipeline; use crate::actors::{MetricsPipeline, MetricsPipelineParams}; +use crate::source::source_needs_drain; use crate::{BoxedPipelineHandle, IndexingService}; impl IndexingService { @@ -86,6 +87,7 @@ impl IndexingService { ctx, )?; + let needs_drain = source_needs_drain(&source_config.source_params); let pipeline_params = MetricsPipelineParams { pipeline_id: indexing_pipeline_id.clone(), metastore: self.metastore.clone(), @@ -111,6 +113,7 @@ impl IndexingService { pipeline_id: indexing_pipeline_id, mailbox, handle, + needs_drain, })) } diff --git a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs index 4056fbeaaa5..eaed82b6119 100644 --- a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs @@ -42,18 +42,19 @@ use quickwit_proto::indexing::IndexingPipelineId; use quickwit_proto::metastore::{MetastoreError, MetastoreServiceClient}; use quickwit_proto::types::ShardId; use quickwit_storage::{Storage, StorageResolver}; -use tracing::{debug, error, info, instrument}; +use tracing::{debug, error, info, instrument, warn}; use super::{ParquetDocProcessor, ParquetIndexer, ParquetPackager, ParquetUploader}; use crate::actors::pipeline_shared::{ - SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, SuperviseLoop, wait_duration_before_retry, + DRAIN_GRACE_MARGIN, DrainPipeline, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, + SuperviseLoop, wait_duration_before_retry, }; use crate::actors::sequencer::Sequencer; use crate::actors::{Publisher, UploaderType}; use crate::metrics::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, }; struct MetricsPipelineHandles { @@ -113,6 +114,10 @@ pub struct MetricsPipeline { statistics: IndexingStatistics, handles_opt: Option, kill_switch: KillSwitch, + // Set when the pipeline is draining (see `DrainPipeline`): the pipeline is + // shutting down and must never respawn; past the deadline it gives up and + // kills whatever is left. + drain_deadline_opt: Option, shard_ids: BTreeSet, // Id of the last indexing plan assigned to this pipeline. Kept here, like `shard_ids`, so it // can be re-sent to the source on respawn; the source adopts it as its publish token. @@ -162,6 +167,7 @@ impl MetricsPipeline { previous_generations_statistics: Default::default(), handles_opt: None, kill_switch: KillSwitch::default(), + drain_deadline_opt: None, statistics: IndexingStatistics { params_fingerprint, ..Default::default() @@ -287,6 +293,15 @@ impl MetricsPipeline { Health::Healthy => {} Health::FailureOrUnhealthy => { self.terminate().await; + if self.drain_deadline_opt.is_some() { + // A draining pipeline is never respawned: whatever could + // not be settled is redelivered (see `DrainPipeline`). + warn!( + pipeline_id=?self.params.pipeline_id, + "draining metrics 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 }); } @@ -499,6 +514,16 @@ impl Handler for MetricsPipeline { ) -> Result<(), ActorExitStatus> { self.perform_observe(ctx); self.perform_health_check(ctx).await?; + if let Some(drain_deadline) = self.drain_deadline_opt + && Instant::now() >= drain_deadline + { + warn!( + pipeline_id=?self.params.pipeline_id, + "metrics 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(()) } @@ -513,6 +538,10 @@ impl Handler for MetricsPipeline { spawn: Spawn, ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { + if self.drain_deadline_opt.is_some() { + // A draining pipeline is shutting down: never respawn it. + return Ok(()); + } if self.handles_opt.is_some() { return Ok(()); } @@ -537,6 +566,32 @@ impl Handler for MetricsPipeline { } } +#[async_trait] +impl Handler for MetricsPipeline { + type Reply = (); + + async fn handle( + &mut self, + _drain: DrainPipeline, + _ctx: &ActorContext, + ) -> Result<(), ActorExitStatus> { + if self.drain_deadline_opt.is_some() { + return Ok(()); + } + let Some(handles) = &self.handles_opt else { + // Nothing is running, so nothing is in flight. + return Err(ActorExitStatus::Success); + }; + self.drain_deadline_opt = Some( + Instant::now() + self.params.indexing_settings.commit_timeout() + DRAIN_GRACE_MARGIN, + ); + // If the source is already dead, the health check terminates the + // draining pipeline instead of respawning it. + let _ = handles.source_mailbox.send_message(Drain).await; + Ok(()) + } +} + #[async_trait] impl Handler for MetricsPipeline { type Reply = (); diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index 33bd9ebbeff..61a40a7df76 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -50,6 +50,21 @@ 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 for acknowledgment-based sources), the source — and with it +/// the whole pipeline — exits with success on its own. Past the drain deadline +/// (commit timeout + [`DRAIN_GRACE_MARGIN`]), the pipeline gives up, kills its +/// actors, and still exits: delivery degrades to at-least-once. A draining +/// pipeline is never respawned. Fire-and-forget: watch the pipeline actor's +/// state to know when it is done. +#[derive(Debug)] +pub struct DrainPipeline; + +/// Extra time granted to a draining pipeline on top of the commit timeout, +/// covering the upload and publication of the final splits. +pub(crate) const DRAIN_GRACE_MARGIN: Duration = Duration::from_secs(30); + // --------------------------------------------------------------------------- // Pipeline trait — type-erased handle for any indexing pipeline actor // --------------------------------------------------------------------------- @@ -75,6 +90,11 @@ pub trait PipelineHandle: Send + Sync { fn last_observation(&self) -> IndexingStatistics; fn check_health(&self, check_for_progress: bool) -> Health; async fn send_assign_shards(&self, message: AssignShards) -> Result<(), SendError>; + /// See [`DrainPipeline`]. Fire-and-forget: the pipeline exits on its own. + async fn start_drain(&self); + /// Whether a teardown must drain before killing + /// (see [`crate::source::source_needs_drain`]). + fn needs_drain(&self) -> bool; async fn observe(&self) -> Observation; async fn join(self: Box) -> (ActorExitStatus, IndexingStatistics); async fn quit(self: Box) -> (ActorExitStatus, IndexingStatistics); @@ -87,11 +107,14 @@ pub(crate) struct ActorPipeline> pub pipeline_id: IndexingPipelineId, pub mailbox: Mailbox, pub handle: ActorHandle, + pub needs_drain: bool, } #[async_trait] impl PipelineHandle for ActorPipeline -where A: Actor + DeferableReplyHandler +where A: Actor + + DeferableReplyHandler + + DeferableReplyHandler { fn indexing_pipeline_id(&self) -> &IndexingPipelineId { &self.pipeline_id @@ -118,6 +141,16 @@ where A: Actor + DeferableReplyHandler bool { + self.needs_drain + } + async fn observe(&self) -> Observation { self.handle.observe().await } diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index d6b76d32d48..e23f44dc732 100644 --- a/quickwit/quickwit-indexing/src/source/mod.rs +++ b/quickwit/quickwit-indexing/src/source/mod.rs @@ -296,6 +296,16 @@ pub trait Source: Send + 'static { Ok(()) } + /// 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, @@ -321,6 +331,7 @@ pub trait Source: Send + 'static { pub struct SourceActor { source: Box, source_sink: SourceSink, + draining: bool, } impl SourceActor { @@ -328,6 +339,7 @@ impl SourceActor { SourceActor { source, source_sink: source_sink.into(), + draining: false, } } } @@ -381,11 +393,62 @@ 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); + self.source_sink + .send_raw_doc_batch(flush_batch, ctx) + .await?; + if self.source.is_drained() { + return Err(ActorExitStatus::Success); + } + Ok(()) + } +} + +/// Whether tearing down a pipeline for this source must go through a drain +/// ([`Drain`]) rather than a plain kill. Only acknowledgment-based +/// sources benefit: killing them loses the acknowledgments of their in-flight +/// messages, which are then redelivered and indexed again after the broker's +/// ack timeout. Checkpoint-based sources resume from their last checkpoint, so +/// a drain buys them nothing. +pub(crate) fn source_needs_drain(_source_params: &SourceParams) -> bool { + // No built-in source opts in yet: acknowledgment-based sources add their + // variant here as they land. + false +} + #[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.source_sink, ctx).await?; if wait_for.is_zero() { ctx.send_self_message(Loop).await?; @@ -516,6 +579,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(()) } } @@ -694,6 +761,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 1bd34f3f4ac..8f5973d3423 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::{IndexerNodeInfo, IndexerPool}; 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::{ @@ -508,6 +508,7 @@ fn start_shard_positions_service( async fn shutdown_signal_handler( shutdown_signal: BoxFutureInfaillible<()>, universe: Universe, + indexing_service_opt: Option>, ingester_opt: Option, ingester_decommission_timeout: Duration, compactor_supervisor_opt: Option>, @@ -537,6 +538,15 @@ async fn shutdown_signal_handler( if let Err(error) = compactor_result { error!("failed to decommission compactor gracefully: {:?}", error); } + // Drain the indexing pipelines only after the ingester decommission: it + // waits for the persisted shard data to be indexed, which requires the + // pipelines to still be consuming. Draining publishes and settles the + // in-flight batches before the universe teardown drops them. + if let Some(indexing_service) = &indexing_service_opt + && let Err(error) = indexing_service.ask(DrainAllPipelines).await + { + error!("failed to drain indexing pipelines gracefully: {:?}", error); + } let actor_exit_statuses = universe.quit().await; if grpc_shutdown_trigger_tx.send(()).is_err() { @@ -1033,6 +1043,7 @@ pub async fn serve_quickwit( let shutdown_handle = tokio::spawn(shutdown_signal_handler( shutdown_signal, universe, + quickwit_services.indexing_service_opt.clone(), ingester_opt, ingester_decommission_timeout, compactor_supervisor_opt, From adffac87ce972555f77ca6b9a6881c488cf2f2f4 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 17:48:24 +0200 Subject: [PATCH 02/41] feat: nats source NATS JetStream source binding to a pre-provisioned durable consumer: messages are acknowledged (server-confirmed) once the split containing them is published, and the pipeline drain settles in-flight messages on planned teardowns. Ships config, docs, tutorial, and broker tests. Co-Authored-By: Claude Fable 5 --- .env.example | 1 + .github/workflows/coverage.yml | 3 + LICENSE-3rdparty.csv | 5 + .../tutorials/stackoverflow/nats-source.yaml | 11 + .../stackoverflow/send_messages_to_nats.py | 19 + docker-compose.yml | 16 + docs/configuration/source-config.md | 45 + docs/ingest-data/nats.md | 247 ++++ quickwit/Cargo.lock | 100 ++ quickwit/Cargo.toml | 5 + quickwit/quickwit-cli/Cargo.toml | 3 + quickwit/quickwit-common/src/metrics.rs | 3 + quickwit/quickwit-config/src/lib.rs | 10 +- .../quickwit-config/src/source_config/mod.rs | 196 ++++ .../src/source_config/serialize.rs | 4 +- .../src/indexing_scheduler/mod.rs | 1 + quickwit/quickwit-indexing/Cargo.toml | 6 + quickwit/quickwit-indexing/src/source/mod.rs | 35 +- .../src/source/nats_source.rs | 1017 +++++++++++++++++ .../quickwit-metastore/src/metastore/mod.rs | 1 + 20 files changed, 1717 insertions(+), 11 deletions(-) create mode 100644 config/tutorials/stackoverflow/nats-source.yaml create mode 100644 config/tutorials/stackoverflow/send_messages_to_nats.py create mode 100644 docs/ingest-data/nats.md create mode 100644 quickwit/quickwit-indexing/src/source/nats_source.rs diff --git a/.env.example b/.env.example index 9556b302b34..2cf7a40877a 100644 --- a/.env.example +++ b/.env.example @@ -19,6 +19,7 @@ MAP_HOST_POSTGRES=0.0.0.0 MAP_HOST_PULSAR=0.0.0.0 MAP_HOST_KAFKA=0.0.0.0 MAP_HOST_ZOOKEEPER=0.0.0.0 +MAP_HOST_NATS=0.0.0.0 MAP_HOST_AZURITE=0.0.0.0 MAP_HOST_GRAFANA=0.0.0.0 MAP_HOST_JAEGER=0.0.0.0 diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 5fc1564e680..2fe43c32f47 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -153,6 +153,9 @@ jobs: - name: Run Pulsar service run: DOCKER_SERVICES=pulsar make docker-compose-up + - name: Run NATS service + run: DOCKER_SERVICES=nats make docker-compose-up + - name: Install Rust run: rustup toolchain install working-directory: ./quickwit diff --git a/LICENSE-3rdparty.csv b/LICENSE-3rdparty.csv index b3c307aa803..20339d32348 100644 --- a/LICENSE-3rdparty.csv +++ b/LICENSE-3rdparty.csv @@ -48,6 +48,7 @@ async-channel,https://github.com/smol-rs/async-channel,Apache-2.0 OR MIT,Stjepan async-compression,https://github.com/Nullus157/async-compression,MIT OR Apache-2.0,"Wim Looman , Allen Bui " async-io,https://github.com/smol-rs/async-io,Apache-2.0 OR MIT,Stjepan Glavina async-lock,https://github.com/smol-rs/async-lock,Apache-2.0 OR MIT,Stjepan Glavina +async-nats,https://github.com/nats-io/nats.rs,Apache-2.0,"Tomasz Pietrek , Casper Beyer " async-process,https://github.com/smol-rs/async-process,Apache-2.0 OR MIT,Stjepan Glavina async-recursion,https://github.com/dcchut/async-recursion,MIT OR Apache-2.0,Robert Usher <266585+dcchut@users.noreply.github.com> async-signal,https://github.com/smol-rs/async-signal,Apache-2.0 OR MIT,John Nunley @@ -746,9 +747,11 @@ serde_derive,https://github.com/serde-rs/serde,MIT OR Apache-2.0,"Erick Tryzelaa serde_derive_internals,https://github.com/serde-rs/serde,MIT OR Apache-2.0,"Erick Tryzelaar , David Tolnay " serde_json,https://github.com/serde-rs/json,MIT OR Apache-2.0,"Erick Tryzelaar , David Tolnay " serde_json_borrow,https://github.com/PSeitz/serde_json_borrow,MIT,Pascal Seitz +serde_nanos,https://github.com/caspervonb/serde_nanos,MIT OR Apache-2.0,Casper Beyer serde_path_to_error,https://github.com/dtolnay/path-to-error,MIT OR Apache-2.0,David Tolnay serde_plain,https://github.com/mitsuhiko/serde-plain,MIT OR Apache-2.0,Armin Ronacher serde_qs,https://github.com/samscott89/serde_qs,MIT OR Apache-2.0,Sam Scott +serde_repr,https://github.com/dtolnay/serde-repr,MIT OR Apache-2.0,David Tolnay serde_spanned,https://github.com/toml-rs/toml,MIT OR Apache-2.0,The serde_spanned Authors serde_tokenstream,https://github.com/oxidecomputer/serde_tokenstream,Apache-2.0,The serde_tokenstream Authors serde_urlencoded,https://github.com/nox/serde_urlencoded,MIT OR Apache-2.0,Anthony Ramine @@ -856,6 +859,7 @@ tokio-retry2,https://github.com/naomijub/tokio-retry,MIT,"Julia Naomi tokio-util,https://github.com/tokio-rs/tokio,MIT,Tokio Contributors +tokio-websockets,https://github.com/Gelbpunkt/tokio-websockets,MIT,The tokio-websockets Authors toml,https://github.com/toml-rs/toml,MIT OR Apache-2.0,The toml Authors toml_datetime,https://github.com/toml-rs/toml,MIT OR Apache-2.0,The toml_datetime Authors toml_edit,https://github.com/toml-rs/toml,MIT OR Apache-2.0,The toml_edit Authors @@ -881,6 +885,7 @@ tracing-serde,https://github.com/tokio-rs/tracing,MIT,Tokio Contributors , David Barsky , Tokio Contributors " triomphe,https://github.com/Manishearth/triomphe,MIT OR Apache-2.0,"Manish Goregaokar , The Servo Project Developers" try-lock,https://github.com/seanmonstar/try-lock,MIT,Sean McArthur +tryhard,https://github.com/EmbarkStudios/tryhard,MIT OR Apache-2.0,Embark ttl_cache,https://github.com/stusmall/ttl_cache,MIT OR Apache-2.0,Stu Small twox-hash,https://github.com/shepmaster/twox-hash,MIT,Jake Goulding typeid,https://github.com/dtolnay/typeid,MIT OR Apache-2.0,David Tolnay diff --git a/config/tutorials/stackoverflow/nats-source.yaml b/config/tutorials/stackoverflow/nats-source.yaml new file mode 100644 index 00000000000..fac37d4cbd2 --- /dev/null +++ b/config/tutorials/stackoverflow/nats-source.yaml @@ -0,0 +1,11 @@ +# +# NATS source config file. +# +version: 0.8 +source_id: nats-source +source_type: nats +params: + uris: + - nats://localhost:4222 + stream: stackoverflow + consumer: quickwit-consumer diff --git a/config/tutorials/stackoverflow/send_messages_to_nats.py b/config/tutorials/stackoverflow/send_messages_to_nats.py new file mode 100644 index 00000000000..7cefd4e3f1f --- /dev/null +++ b/config/tutorials/stackoverflow/send_messages_to_nats.py @@ -0,0 +1,19 @@ +import asyncio + +import nats + + +async def main(): + client = await nats.connect("nats://localhost:4222") + jetstream = client.jetstream() + + with open("stackoverflow.posts.transformed-10000.json", encoding="utf8") as file: + for i, line in enumerate(file): + await jetstream.publish("stackoverflow.posts", line.strip().encode("utf-8")) + if i % 1000 == 0: + print(f"{i}/10000 messages sent.") + + await client.close() + + +asyncio.run(main()) diff --git a/docker-compose.yml b/docker-compose.yml index 24f5ed29b47..9198d13aba6 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -127,6 +127,22 @@ services: timeout: 10s retries: 100 + nats: + # 2.10 is the oldest version supporting multiple consumer filter subjects. + image: nats:${NATS_VERSION:-2.10-alpine} + container_name: nats + command: --jetstream --http_port 8222 + ports: + - "${MAP_HOST_NATS:-127.0.0.1}:4222:4222" + profiles: + - all + - nats + healthcheck: + test: ["CMD", "wget", "-q", "-O-", "http://localhost:8222/healthz"] + interval: 1s + timeout: 5s + retries: 100 + azurite: image: mcr.microsoft.com/azure-storage/azurite:${AZURITE_VERSION:-3.24.0} container_name: azurite diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 30894d3faad..41f5ad26787 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -183,6 +183,51 @@ EOF quickwit source create --index my-index --source-config source-config.yaml ``` +### NATS source + +A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-concepts/jetstream) stream through a durable consumer. Each message carries one payload in the source's [input format](#input-format): a single JSON object (`json`, the default), a plain text document (`plain_text`), or an OTLP export request whose log records or spans are each indexed as a separate document (`otlp_*` formats). Payloads must not exceed 1 MiB. + +A tutorial is available [here](/docs/ingest-data/nats.md). + +The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. + +Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged — with server-confirmed acknowledgments — before the pipeline stops, so nothing is indexed twice; messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. `max_ack_pending` bounds the messages in flight across all pipelines of the consumer, and therefore the aggregate throughput. + +**Scaling** + +Multiple indexing pipelines can share the consumer: NATS load-balances the messages across them, so scaling up or down is a plain `num_pipelines` update, and the control plane places the pipelines across the indexers of the cluster. + +**Monitoring** + +Being durable, the consumer is observable through NATS's own monitoring (`nats consumer info`, exporters): `num_pending` is the indexing lag and `num_ack_pending` the in-flight window, both available even while the pipelines are down. The source also reports `num_pending_acks` (messages indexed but whose split is not published yet) in its observable state. + +**NATS source parameters** + +| Property | Description | Default value | +| --- | --- | --- | +| `uris` | List of NATS server URIs (e.g. `nats://localhost:4222`). | required | +| `stream` | Name of the JetStream stream to consume. | required | +| `consumer` | Name of the pre-provisioned durable consumer to bind to. | required | +| `tls` | TLS options: `ca_certificates_path` (PEM file whose root certificates are trusted in addition to the system ones), and `client_certificate_path` + `client_key_path` (PEM files, set together) for mutual TLS. TLS itself is enabled by connecting to `tls://` URIs. The files are read by the indexer nodes when the connection is established. | optional | +| `authentication` | Authentication parameters: either `user_password` (with `user` and `password`) or `token`. | optional | + +*Adding a NATS source to an index with the [CLI](../reference/cli.md#source)* + +```bash +cat << EOF > source-config.yaml +version: 0.8 +source_id: my-nats-source +source_type: nats +num_pipelines: 2 +params: + uris: + - nats://localhost:4222 + stream: my-stream + consumer: my-consumer +EOF +./quickwit source create --index my-index --source-config source-config.yaml +``` + ### Pulsar source A Puslar source reads data from one or several Pulsar topics. Each message in topic(s) must hold a JSON object. diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md new file mode 100644 index 00000000000..2dc3f78cb42 --- /dev/null +++ b/docs/ingest-data/nats.md @@ -0,0 +1,247 @@ +--- +title: NATS +description: A short tutorial describing how to set up Quickwit to ingest data from NATS JetStream in a few minutes +tags: [nats, integration] +icon_url: /img/tutorials/nats.svg +sidebar_position: 4 +--- + +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +In this tutorial, we will describe how to set up Quickwit to ingest data from [NATS JetStream](https://docs.nats.io/nats-concepts/jetstream) in a few minutes. First, we will create an index and configure a NATS source. Then, we will create a JetStream stream and load some events from the [Stack Overflow dataset](https://www.kaggle.com/stackoverflow/stacksample) into it. Finally, we will execute some searches. + +## Prerequisites + +You will need the following to complete this tutorial: +- A local running [Quickwit instance](/docs/get-started/installation.md) +- A local running NATS server (2.10+) with JetStream enabled +- The [NATS CLI](https://github.com/nats-io/natscli) + +### Quickwit setup + +[Download](/docs/get-started/installation.md) Quickwit and start a server. Then open a new terminal to execute CLI commands with the same binary. + +```bash +./quickwit run +``` + +Test that the cluster is running: + +```bash +./quickwit index list +``` + +### NATS setup + + + + + +Download the [NATS server](https://docs.nats.io/running-a-nats-service/introduction/installation) and start it with JetStream enabled: + +```bash +nats-server --jetstream +``` + + + + + +```bash +docker run -it -p 4222:4222 nats:2.10 --jetstream +``` + +See the details on the [official documentation](https://docs.nats.io/running-a-nats-service/nats_docker). + + + + + +## Prepare Quickwit + +First, let's create a new index. Here is the index config and doc mapping corresponding to the schema of Stack Overflow posts: + +```yaml title="index-config.yaml" +# +# Index config file for Stack Overflow dataset. +# +version: 0.7 + +index_id: stackoverflow + +doc_mapping: + field_mappings: + - name: user + type: text + fast: true + tokenizer: raw + - name: tags + type: array + fast: true + tokenizer: raw + - name: type + type: text + fast: true + tokenizer: raw + - name: title + type: text + tokenizer: default + record: position + stored: true + - name: body + type: text + tokenizer: default + record: position + stored: true + - name: questionId + type: u64 + - name: answerId + type: u64 + - name: acceptedAnswerId + type: u64 + - name: creationDate + type: datetime + fast: true + input_formats: + - rfc3339 + fast_precision: seconds + timestamp_field: creationDate + +search_settings: + default_search_fields: [title, body] + +indexing_settings: + commit_timeout_secs: 10 +``` + +Execute these Bash commands to download the index config and create the `stackoverflow` index. + +```bash +# Download stackoverflow index config. +wget -O stackoverflow.yaml https://raw.githubusercontent.com/quickwit-oss/quickwit/main/config/tutorials/stackoverflow/index-config.yaml + +# Create index. +./quickwit index create --index-config stackoverflow.yaml +``` + +## Create a JetStream stream and a durable consumer + +The NATS source consumes a JetStream stream through a durable consumer, so messages must be published on subjects captured by a stream, and the consumer must be provisioned before creating the source — Quickwit only ever fetches it, and never creates, updates, nor deletes it. Let's create both: + +```bash +nats stream add stackoverflow --subjects "stackoverflow.posts" --defaults +nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack explicit --wait 5m --defaults +``` + +:::info + +The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. See the [NATS source reference](../configuration/source-config.md#nats-source) for details. + +::: + +## Create the NATS source + +A NATS source just needs to define the server URIs, the stream, and the durable consumer to bind to. + +```yaml title="nats-source.yaml" +# +# NATS source config file. +# +version: 0.8 +source_id: nats-source +source_type: nats +params: + uris: + - nats://localhost:4222 + stream: stackoverflow + consumer: quickwit-consumer +``` + +Run these commands to download the source config file and create the source. + +```bash +# Download NATS source config. +wget -O stackoverflow-nats-source.yaml https://raw.githubusercontent.com/quickwit-oss/quickwit/main/config/tutorials/stackoverflow/nats-source.yaml + +# Create source. +./quickwit source create --index stackoverflow --source-config stackoverflow-nats-source.yaml +``` + +As soon as the NATS source is created, the Quickwit control plane will ask an indexer to start a new indexing pipeline. You will see logs like below on the indexer: + +```bash +INFO spawn_pipeline{index=stackoverflow gen=0}: quickwit_indexing::source::nats_source: starting NATS source index_id=stackoverflow source_id=nats-source stream=stackoverflow consumer_name=quickwit-consumer +``` + +The consumer's indexing lag and ack floor are visible at any time — even while the pipeline is down — with `nats consumer info stackoverflow quickwit-consumer`. + +## Populate the stream + +To populate the stream, we will use a python script: + +```python title=send_messages_to_nats.py +import asyncio + +import nats + + +async def main(): + client = await nats.connect("nats://localhost:4222") + jetstream = client.jetstream() + + with open("stackoverflow.posts.transformed-10000.json", encoding="utf8") as file: + for i, line in enumerate(file): + await jetstream.publish("stackoverflow.posts", line.strip().encode("utf-8")) + if i % 1000 == 0: + print(f"{i}/10000 messages sent.") + + await client.close() + + +asyncio.run(main()) +``` + +Install the [python client](https://github.com/nats-io/nats.py) locally and run the script: + +```bash +# Download the first 10_000 Stackoverflow posts articles. +curl -O https://quickwit-datasets-public.s3.amazonaws.com/stackoverflow.posts.transformed-10000.json + +# Install nats python client. +pip3 install nats-py +wget https://raw.githubusercontent.com/quickwit-oss/quickwit/main/config/tutorials/stackoverflow/send_messages_to_nats.py +python3 send_messages_to_nats.py +``` + +## Time to search! + +You can run this command to inspect the properties of the index and check the current number of published splits and documents: + +```bash +# Display some general information about the index. +./quickwit index describe --index stackoverflow +``` + +You will notably see the number of published documents. + +You are now ready to execute some queries. + +```bash +curl 'http://localhost:7280/api/v1/stackoverflow/search?query=search+AND+engine' +``` + +If your Quickwit server is local, you can access the results through the Quickwit UI on [localhost:7280](http://localhost:7280/ui/search?query=&index_id=stackoverflow&max_hits=10). + +## Tear down resources (optional) + +Let's delete the files and resources created for the purpose of this tutorial. + +```bash +# Delete quickwit index. +./quickwit index delete --index stackoverflow --yes +# Delete NATS stream. +nats stream rm -f stackoverflow +``` + +This concludes the tutorial. If you have any questions regarding Quickwit or encounter any issues, don't hesitate to ask a [question](https://github.com/quickwit-oss/quickwit/discussions) or open an [issue](https://github.com/quickwit-oss/quickwit/issues) on [GitHub](https://github.com/quickwit-oss/quickwit) or contact us directly on [Discord](https://discord.com/invite/MT27AG5EVE). diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index caede75b90f..18957554214 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -582,6 +582,40 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-nats" +version = "0.50.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d83a251fa1a4c9d0fe6e816b7acd60549e473e08d14f27a1d992c2675abff05f" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-util", + "memchr", + "pin-project", + "portable-atomic", + "rand 0.10.2", + "regex", + "ring", + "rustls-native-certs", + "rustls-pki-types", + "rustls-webpki 0.103.13", + "serde", + "serde_json", + "serde_nanos", + "serde_repr", + "thiserror 2.0.18", + "time", + "tokio", + "tokio-rustls 0.26.4", + "tokio-stream", + "tokio-util", + "tokio-websockets", + "tracing", + "tryhard", + "url", +] + [[package]] name = "async-process" version = "2.5.0" @@ -8777,6 +8811,7 @@ dependencies = [ "arc-swap", "arrow", "async-compression", + "async-nats", "async-trait", "aws-sdk-kinesis", "aws-sdk-sqs", @@ -8796,6 +8831,8 @@ dependencies = [ "mockall", "oneshot 0.2.1", "openssl", + "opentelemetry", + "opentelemetry_sdk", "parquet", "percent-encoding", "proptest", @@ -8832,6 +8869,7 @@ dependencies = [ "time", "tokio", "tracing", + "tracing-opentelemetry", "ulid", "unicode-segmentation", "utoipa", @@ -10784,6 +10822,15 @@ dependencies = [ "serde_json", ] +[[package]] +name = "serde_nanos" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a93142f0367a4cc53ae0fead1bcda39e85beccfad3dcd717656cacab94b12985" +dependencies = [ + "serde", +] + [[package]] name = "serde_path_to_error" version = "0.1.20" @@ -10827,6 +10874,17 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_repr" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d3b1629de253c70a0508c3899572da79ca359fdab27c7920ff00406df418906" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] + [[package]] name = "serde_spanned" version = "1.1.1" @@ -11688,6 +11746,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -12257,6 +12326,27 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-websockets" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f591660438b3038dd04d16c938271c79e7e06260ad2ea2885a4861bfb238605d" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-sink", + "http 1.4.2", + "httparse", + "rand 0.8.6", + "ring", + "rustls-pki-types", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", + "webpki-roots 0.26.11", +] + [[package]] name = "toml" version = "1.1.2+spec-1.1.0" @@ -12626,6 +12716,16 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tryhard" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fe58ebd5edd976e0fe0f8a14d2a04b7c81ef153ea9a54eebc42e67c2c23b4e5" +dependencies = [ + "pin-project-lite", + "tokio", +] + [[package]] name = "ttl_cache" version = "0.5.1" diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml index 8904ed1dd3c..84cbbd71c88 100644 --- a/quickwit/Cargo.toml +++ b/quickwit/Cargo.toml @@ -104,6 +104,11 @@ arc-swap = "1.9" arrow = { version = "58", default-features = false, features = ["ipc"] } assert-json-diff = "2" async-compression = { version = "0.4", features = ["tokio", "gzip"] } +async-nats = { version = "0.50", default-features = false, features = [ + "jetstream", + "ring", + "server_2_10", +] } async-speed-limit = "0.4" async-trait = "0.1" atomic_float = "1.1" diff --git a/quickwit/quickwit-cli/Cargo.toml b/quickwit/quickwit-cli/Cargo.toml index 01857eb221f..6e115ac4023 100644 --- a/quickwit/quickwit-cli/Cargo.toml +++ b/quickwit/quickwit-cli/Cargo.toml @@ -108,6 +108,7 @@ release-feature-set = [ "pprof", "quickwit-indexing/kafka", "quickwit-indexing/kinesis", + "quickwit-indexing/nats", "quickwit-indexing/pulsar", "quickwit-indexing/sqs", "quickwit-indexing/vrl", @@ -122,6 +123,7 @@ release-feature-vendored-set = [ "pprof", "quickwit-indexing/kinesis", "quickwit-indexing/metrics", + "quickwit-indexing/nats", "quickwit-indexing/pulsar", "quickwit-indexing/sqs", "quickwit-indexing/vrl", @@ -136,6 +138,7 @@ release-macos-feature-vendored-set = [ "openssl-support", "quickwit-indexing/kinesis", "quickwit-indexing/metrics", + "quickwit-indexing/nats", "quickwit-indexing/pulsar", "quickwit-indexing/sqs", "quickwit-indexing/vrl", diff --git a/quickwit/quickwit-common/src/metrics.rs b/quickwit/quickwit-common/src/metrics.rs index 7244c3816e4..aa8f6e97adf 100644 --- a/quickwit/quickwit-common/src/metrics.rs +++ b/quickwit/quickwit-common/src/metrics.rs @@ -88,6 +88,9 @@ pub static IN_FLIGHT_KAFKA_SOURCE: LazyGauge = pub static IN_FLIGHT_KINESIS_SOURCE: LazyGauge = LazyGauge::new(|| in_flight_data_gauge("kinesis_source")); +pub static IN_FLIGHT_NATS_SOURCE: LazyGauge = + LazyGauge::new(|| in_flight_data_gauge("nats_source")); + pub static IN_FLIGHT_PUBSUB_SOURCE: LazyGauge = LazyGauge::new(|| in_flight_data_gauge("pubsub_source")); diff --git a/quickwit/quickwit-config/src/lib.rs b/quickwit/quickwit-config/src/lib.rs index 904e65f8180..1a06d42d2b9 100644 --- a/quickwit/quickwit-config/src/lib.rs +++ b/quickwit/quickwit-config/src/lib.rs @@ -61,9 +61,10 @@ use source_config::FileSourceParamsForSerde; pub use source_config::{ CLI_SOURCE_ID, FileSourceMessageType, FileSourceNotification, FileSourceParams, FileSourceSqs, INGEST_API_SOURCE_ID, INGEST_V2_SOURCE_ID, KafkaSourceParams, KinesisSourceParams, - PubSubSourceParams, PulsarSourceAuth, PulsarSourceParams, RegionOrEndpoint, SourceConfig, - SourceInputFormat, SourceParams, TransformConfig, VecSourceParams, VoidSourceParams, - load_source_config_from_user_config, load_source_config_update, + NatsSourceAuth, NatsSourceParams, NatsSourceTls, PubSubSourceParams, PulsarSourceAuth, + PulsarSourceParams, RegionOrEndpoint, SourceConfig, SourceInputFormat, SourceParams, + TransformConfig, VecSourceParams, VoidSourceParams, load_source_config_from_user_config, + load_source_config_update, }; use tracing::warn; @@ -116,6 +117,9 @@ pub fn disable_ingest_v1() -> bool { KafkaSourceParams, KinesisSourceParams, MergePolicyConfig, + NatsSourceAuth, + NatsSourceParams, + NatsSourceTls, ParquetIndexingConfig, ParquetMergePolicyConfig, PubSubSourceParams, diff --git a/quickwit/quickwit-config/src/source_config/mod.rs b/quickwit/quickwit-config/src/source_config/mod.rs index 81811e9f47c..9283f2a1dd5 100644 --- a/quickwit/quickwit-config/src/source_config/mod.rs +++ b/quickwit/quickwit-config/src/source_config/mod.rs @@ -86,6 +86,7 @@ impl SourceConfig { SourceParams::Kafka(params) => serde_json::to_value(params), SourceParams::Kinesis(params) => serde_json::to_value(params), SourceParams::Pulsar(params) => serde_json::to_value(params), + SourceParams::Nats(params) => serde_json::to_value(params), SourceParams::Stdin => serde_json::to_value(()), SourceParams::Vec(params) => serde_json::to_value(params), SourceParams::Void(params) => serde_json::to_value(params), @@ -231,6 +232,7 @@ pub enum SourceParams { #[serde(rename = "pubsub")] PubSub(PubSubSourceParams), Pulsar(PulsarSourceParams), + Nats(NatsSourceParams), Stdin, Vec(VecSourceParams), Void(VoidSourceParams), @@ -263,6 +265,7 @@ impl SourceParams { SourceParams::Kinesis(_) => SourceType::Kinesis, SourceParams::PubSub(_) => SourceType::PubSub, SourceParams::Pulsar(_) => SourceType::Pulsar, + SourceParams::Nats(_) => SourceType::Nats, SourceParams::Stdin => SourceType::Stdin, SourceParams::Vec(_) => SourceType::Vec, SourceParams::Void(_) => SourceType::Void, @@ -287,6 +290,7 @@ impl SourceParams { (SourceParams::Pulsar(current), SourceParams::Pulsar(new)) => { current.validate_update(new) } + (SourceParams::Nats(current), SourceParams::Nats(new)) => current.validate_update(new), (current, new) if current.source_type() != new.source_type() => Err(anyhow::anyhow!( "source type cannot be changed, current type {}", current.source_type(), @@ -620,6 +624,88 @@ fn default_consumer_name() -> String { "quickwit".to_string() } +#[derive( + Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize, utoipa::ToSchema, +)] +#[serde(deny_unknown_fields)] +pub struct NatsSourceParams { + /// The URIs of the NATS servers to connect to (e.g. `nats://localhost:4222`). + pub uris: Vec, + /// Name of the NATS JetStream stream that the source consumes. + pub stream: String, + /// Name of the durable consumer to bind to. The consumer is provisioned + /// externally. It must use the explicit ack policy: the source acknowledges + /// each message once the split containing it is published. + pub consumer: String, + /// TLS options: custom root CA and client certificates (mutual TLS). TLS + /// itself is enabled by connecting to `tls://` URIs. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tls: Option, + // Serde yaml has some specific behaviour when deserializing + // enums (see https://github.com/dtolnay/serde-yaml/issues/342) + // and requires explicitly stating `default` in order to make the parameter + // optional on the yaml config. + #[serde(default, with = "serde_yaml::with::singleton_map")] + /// Authentication for NATS. + pub authentication: Option, +} + +impl NatsSourceParams { + pub fn validate(&self) -> anyhow::Result<()> { + ensure!( + !self.uris.is_empty(), + "`uris` must contain at least one NATS server URI" + ); + ensure!(!self.stream.is_empty(), "`stream` must be set"); + ensure!(!self.consumer.is_empty(), "`consumer` must be set"); + if let Some(tls) = &self.tls { + ensure!( + tls.client_certificate_path.is_some() == tls.client_key_path.is_some(), + "`tls.client_certificate_path` and `tls.client_key_path` must be set together" + ); + } + Ok(()) + } + + fn validate_update(&self, _other: &Self) -> anyhow::Result<()> { + // The resume point is the consumer's ack floor, not the metastore + // checkpoint: re-binding to another stream or consumer is a plain + // reconfiguration and cannot corrupt any tracked progress. + Ok(()) + } +} + +/// TLS options for the NATS connection. Certificates are PEM files loaded by +/// the indexer when the connection is established, so the paths must exist on +/// the indexer nodes; validation only checks the options' consistency. +#[derive( + Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize, utoipa::ToSchema, +)] +#[serde(deny_unknown_fields)] +pub struct NatsSourceTls { + /// Path to a PEM file whose root certificates are trusted in addition to + /// the system ones (e.g. a private CA). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ca_certificates_path: Option, + /// Path to the client certificate PEM file, for mutual TLS. Requires + /// `client_key_path`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_certificate_path: Option, + /// Path to the client private key PEM file, for mutual TLS. Requires + /// `client_certificate_path`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_key_path: Option, +} + +#[derive( + Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize, utoipa::ToSchema, +)] +#[serde(rename_all = "snake_case")] +pub enum NatsSourceAuth { + UserPassword { user: String, password: String }, + Token(String), +} + #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, utoipa::ToSchema)] #[serde(deny_unknown_fields)] pub struct TransformConfig { @@ -1341,6 +1427,116 @@ mod tests { } } + #[test] + fn test_nats_source_params_deserialization() { + let base_params = NatsSourceParams { + uris: vec!["nats://localhost:4222".to_string()], + stream: "my-stream".to_string(), + consumer: "my-consumer".to_string(), + tls: None, + authentication: None, + }; + { + let yaml = r#" + uris: + - nats://localhost:4222 + stream: my-stream + consumer: my-consumer + "#; + let params = serde_yaml::from_str::(yaml).unwrap(); + assert_eq!(params, base_params); + params.validate().unwrap(); + } + { + let yaml = r#" + uris: + - tls://localhost:4222 + stream: my-stream + consumer: my-consumer + tls: + ca_certificates_path: /etc/ssl/private-ca.pem + client_certificate_path: /etc/ssl/client.pem + client_key_path: /etc/ssl/client.key + authentication: + token: my-token + "#; + let params = serde_yaml::from_str::(yaml).unwrap(); + assert_eq!( + params, + NatsSourceParams { + uris: vec!["tls://localhost:4222".to_string()], + tls: Some(NatsSourceTls { + ca_certificates_path: Some("/etc/ssl/private-ca.pem".to_string()), + client_certificate_path: Some("/etc/ssl/client.pem".to_string()), + client_key_path: Some("/etc/ssl/client.key".to_string()), + }), + authentication: Some(NatsSourceAuth::Token("my-token".to_string())), + ..base_params.clone() + } + ); + params.validate().unwrap(); + } + { + let yaml = r#" + uris: + - nats://localhost:4222 + stream: my-stream + "#; + serde_yaml::from_str::(yaml) + .expect_err("parameters should error on missing consumer"); + } + { + let params = NatsSourceParams { + uris: Vec::new(), + ..base_params.clone() + }; + params + .validate() + .expect_err("validation should reject empty uris"); + } + { + let params = NatsSourceParams { + consumer: String::new(), + ..base_params.clone() + }; + params + .validate() + .expect_err("validation should reject an empty consumer name"); + } + { + let params = NatsSourceParams { + tls: Some(NatsSourceTls { + ca_certificates_path: None, + client_certificate_path: Some("/etc/ssl/client.pem".to_string()), + client_key_path: None, + }), + ..base_params.clone() + }; + params + .validate() + .expect_err("validation should reject a client certificate without its key"); + } + } + + #[test] + fn test_nats_source_supports_multiple_pipelines() { + let source_config_yaml = r#" + version: 0.8 + source_id: my-nats-source + num_pipelines: 2 + source_type: nats + params: + uris: + - nats://localhost:4222 + stream: my-stream + consumer: my-consumer + "#; + let source_config = + load_source_config_from_user_config(ConfigFormat::Yaml, source_config_yaml.as_bytes()) + .unwrap(); + assert_eq!(source_config.num_pipelines.get(), 2); + } + #[cfg(feature = "vrl")] #[tokio::test] async fn test_load_ingest_api_source_config() { diff --git a/quickwit/quickwit-config/src/source_config/serialize.rs b/quickwit/quickwit-config/src/source_config/serialize.rs index 3f580ee8fa3..25e92fe0a28 100644 --- a/quickwit/quickwit-config/src/source_config/serialize.rs +++ b/quickwit/quickwit-config/src/source_config/serialize.rs @@ -116,6 +116,7 @@ impl SourceConfigForSerialization { | SourceParams::Pulsar(_) => { // TODO consider any validation opportunity } + SourceParams::Nats(params) => params.validate()?, SourceParams::PubSub(_) | SourceParams::Ingest | SourceParams::IngestApi @@ -127,10 +128,11 @@ impl SourceConfigForSerialization { SourceParams::PubSub(_) | SourceParams::Kafka(_) | SourceParams::File(FileSourceParams::Notifications(_)) => {} + SourceParams::Nats(_) => {} _ => { if self.num_pipelines > 1 { bail!( - "Quickwit currently supports multiple pipelines only for GCP PubSub or Kafka sources. open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources. open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" ); } } diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs index a656cdcfc5f..40ffe42145a 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -272,6 +272,7 @@ fn get_sources_to_schedule( } SourceParams::Kafka(_) | SourceParams::Kinesis(_) + | SourceParams::Nats(_) | SourceParams::PubSub(_) | SourceParams::Pulsar(_) | SourceParams::File(FileSourceParams::Notifications(_)) => { diff --git a/quickwit/quickwit-indexing/Cargo.toml b/quickwit/quickwit-indexing/Cargo.toml index 23d81a3e58c..172dfbf0133 100644 --- a/quickwit/quickwit-indexing/Cargo.toml +++ b/quickwit/quickwit-indexing/Cargo.toml @@ -14,6 +14,7 @@ license.workspace = true anyhow = { workspace = true } arc-swap = { workspace = true } async-compression = { workspace = true } +async-nats = { workspace = true, optional = true } async-trait = { workspace = true } aws-sdk-kinesis = { workspace = true, optional = true } aws-sdk-sqs = { workspace = true, optional = true } @@ -31,6 +32,7 @@ itertools = { workspace = true } libz-sys = { workspace = true, optional = true } oneshot = { workspace = true } openssl = { workspace = true, optional = true } +opentelemetry = { workspace = true, optional = true } percent-encoding = { workspace = true } pulsar = { workspace = true, optional = true } quickwit-query = { workspace = true } @@ -45,6 +47,7 @@ thiserror = { workspace = true } time = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } +tracing-opentelemetry = { workspace = true, optional = true } ulid = { workspace = true } utoipa = { workspace = true } vrl = { workspace = true, optional = true } @@ -83,6 +86,8 @@ kinesis = [ "quickwit-aws/kinesis", ] kinesis-localstack-tests = [] +nats = ["dep:async-nats", "dep:opentelemetry", "dep:tracing-opentelemetry"] +nats-broker-tests = ["nats"] pulsar = ["dep:pulsar"] pulsar-broker-tests = [] queue-sources = [] @@ -117,6 +122,7 @@ ci-test = [] bytes = { workspace = true } criterion = { workspace = true, features = ["async_tokio"] } mockall = { workspace = true } +opentelemetry_sdk = { workspace = true } parquet = { workspace = true } proptest = { workspace = true } prost = { workspace = true } diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index e23f44dc732..d86bf764dec 100644 --- a/quickwit/quickwit-indexing/src/source/mod.rs +++ b/quickwit/quickwit-indexing/src/source/mod.rs @@ -62,6 +62,8 @@ mod ingest_api_source; mod kafka_source; #[cfg(feature = "kinesis")] mod kinesis; +#[cfg(feature = "nats")] +mod nats_source; #[cfg(feature = "pulsar")] mod pulsar_source; #[cfg(feature = "queue-sources")] @@ -87,6 +89,8 @@ pub use gcp_pubsub_source::{GcpPubSubSource, GcpPubSubSourceFactory}; pub use kafka_source::{KafkaSource, KafkaSourceFactory}; #[cfg(feature = "kinesis")] pub use kinesis::kinesis_source::{KinesisSource, KinesisSourceFactory}; +#[cfg(feature = "nats")] +pub use nats_source::{NatsSource, NatsSourceFactory}; #[cfg(feature = "pulsar")] pub use pulsar_source::{PulsarSource, PulsarSourceFactory}; #[cfg(feature = "sqs")] @@ -94,8 +98,8 @@ pub use queue_sources::sqs_queue; use quickwit_actors::{Actor, ActorContext, ActorExitStatus, Handler}; use quickwit_common::metrics::{ IN_FLIGHT_FILE_SOURCE, IN_FLIGHT_INGEST_SOURCE, IN_FLIGHT_KAFKA_SOURCE, - IN_FLIGHT_KINESIS_SOURCE, IN_FLIGHT_OTHER_SOURCE, IN_FLIGHT_PUBSUB_SOURCE, - IN_FLIGHT_PULSAR_SOURCE, + IN_FLIGHT_KINESIS_SOURCE, IN_FLIGHT_NATS_SOURCE, IN_FLIGHT_OTHER_SOURCE, + IN_FLIGHT_PUBSUB_SOURCE, IN_FLIGHT_PULSAR_SOURCE, }; use quickwit_common::pubsub::EventBroker; use quickwit_common::runtimes::RuntimeType; @@ -434,10 +438,8 @@ impl Handler for SourceActor { /// messages, which are then redelivered and indexed again after the broker's /// ack timeout. Checkpoint-based sources resume from their last checkpoint, so /// a drain buys them nothing. -pub(crate) fn source_needs_drain(_source_params: &SourceParams) -> bool { - // No built-in source opts in yet: acknowledgment-based sources add their - // variant here as they land. - false +pub(crate) fn source_needs_drain(source_params: &SourceParams) -> bool { + matches!(source_params, SourceParams::Nats(_)) } #[async_trait] @@ -489,6 +491,8 @@ pub fn quickwit_supported_sources() -> &'static SourceLoader { source_factory.add_source(SourceType::Kafka, KafkaSourceFactory); #[cfg(feature = "kinesis")] source_factory.add_source(SourceType::Kinesis, KinesisSourceFactory); + #[cfg(feature = "nats")] + source_factory.add_source(SourceType::Nats, NatsSourceFactory); #[cfg(feature = "pulsar")] source_factory.add_source(SourceType::Pulsar, PulsarSourceFactory); source_factory.add_source(SourceType::Stdin, StdinSourceFactory); @@ -556,6 +560,17 @@ pub async fn check_source_connectivity( Ok(()) } } + #[allow(unused_variables)] + SourceParams::Nats(params) => { + #[cfg(not(feature = "nats"))] + anyhow::bail!("Quickwit was compiled without the `nats` feature"); + + #[cfg(feature = "nats")] + { + nats_source::check_connectivity(params).await?; + Ok(()) + } + } _ => Ok(()), } } @@ -608,6 +623,7 @@ impl BatchBuilder { SourceType::IngestV2 => &IN_FLIGHT_INGEST_SOURCE, SourceType::Kafka => &IN_FLIGHT_KAFKA_SOURCE, SourceType::Kinesis => &IN_FLIGHT_KINESIS_SOURCE, + SourceType::Nats => &IN_FLIGHT_NATS_SOURCE, SourceType::PubSub => &IN_FLIGHT_PUBSUB_SOURCE, SourceType::Pulsar => &IN_FLIGHT_PULSAR_SOURCE, _ => &IN_FLIGHT_OTHER_SOURCE, @@ -707,7 +723,11 @@ mod tests { #[cfg(all( test, - any(feature = "kafka-broker-tests", feature = "sqs-localstack-tests") + any( + feature = "kafka-broker-tests", + feature = "nats-broker-tests", + feature = "sqs-localstack-tests" + ) ))] pub fn with_metastore(mut self, metastore: MetastoreServiceClient) -> Self { self.metastore_opt = Some(metastore); @@ -865,6 +885,7 @@ mod tests { any( feature = "sqs-localstack-tests", feature = "kafka-broker-tests", + feature = "nats-broker-tests", feature = "pulsar-broker-tests" ) ))] diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs new file mode 100644 index 00000000000..aeb0acc66fd --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -0,0 +1,1017 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! A source consuming a NATS JetStream stream through a pre-provisioned +//! durable consumer. +//! +//! The consumer is only ever fetched — never created, updated, nor deleted — +//! so its lifecycle, subject filters, deliver policy, and ack tuning +//! (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The +//! consumer must use the explicit ack policy: each message is acknowledged — +//! with a server-confirmed ack — once the split containing it is published. +//! On a planned teardown (node shutdown, pipeline reassignment) the pipeline +//! is drained first: the in-flight messages are published and acknowledged, +//! and the prefetched but unprocessed ones are negatively acknowledged for +//! immediate redelivery — so delivery is exactly-once. A crash degrades it to +//! at-least-once: the unacknowledged messages — including ones whose split +//! ends up published — are redelivered after the consumer's `ack_wait` and +//! indexed again. The consumer's ack floor is the resume point; the metastore +//! checkpoint only carries synthetic positions used to correlate published +//! batches with the acknowledgments they release. +//! +//! Several indexing pipelines can share the consumer: NATS load-balances the +//! messages across them, so scaling is a plain `num_pipelines` update. Being +//! durable, the consumer is observable through NATS's own monitoring even +//! while the pipelines are down. +//! +//! When a message carries a W3C `traceparent` header, the source stitches the +//! processing of the message into the publisher's distributed trace. + +use std::collections::BTreeMap; +use std::fmt; +use std::path::PathBuf; +use std::str::FromStr; +use std::time::{Duration, Instant}; + +use anyhow::{Context as _, anyhow, bail, ensure}; +use async_nats::header::{HeaderMap, HeaderName}; +use async_nats::jetstream::consumer::pull::{ + MessagesError, MessagesErrorKind, Stream as DurableMessageStream, +}; +use async_nats::jetstream::consumer::{AckPolicy, PullConsumer}; +use async_nats::jetstream::message::AckKind; +use async_nats::{ConnectOptions, Subject, jetstream}; +use async_trait::async_trait; +use futures::{FutureExt, StreamExt}; +use opentelemetry::propagation::Extractor; +use opentelemetry::trace::TraceContextExt; +use opentelemetry::{Context as OtelContext, global}; +use quickwit_actors::ActorExitStatus; +use quickwit_common::rand::append_random_suffix; +use quickwit_config::{NatsSourceAuth, NatsSourceParams}; +use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpoint}; +use quickwit_proto::metastore::SourceType; +use quickwit_proto::types::Position; +use serde_json::{Value as JsonValue, json}; +use tokio::time; +use tracing::{Span, debug, info, warn}; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +use crate::source::{ + BATCH_NUM_BYTES_LIMIT, BatchBuilder, EMIT_BATCHES_TIMEOUT, Source, SourceContext, + SourceRuntime, SourceSink, TypedSourceFactory, +}; + +pub struct NatsSourceFactory; + +#[async_trait] +impl TypedSourceFactory for NatsSourceFactory { + type Source = NatsSource; + type Params = NatsSourceParams; + + async fn typed_create_source( + source_runtime: SourceRuntime, + source_params: NatsSourceParams, + ) -> anyhow::Result { + NatsSource::try_new(source_runtime, source_params).await + } +} + +#[derive(Default, Debug)] +pub struct NatsSourceState { + /// Number of bytes processed by the source. + pub num_bytes_processed: u64, + /// Number of messages processed by the source (including invalid messages). + pub num_messages_processed: u64, + /// Number of invalid messages, i.e., that were empty. + pub num_invalid_messages: u64, +} + +/// Source flavor binding to a pre-provisioned durable consumer. +/// +/// The consumer is only ever fetched — never created, updated, nor deleted — +/// so its lifecycle, subject filters, deliver policy, and ack tuning +/// (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. Each +/// message is acknowledged once the split containing it is published +/// ([`Source::suggest_truncate`]); see the module documentation for the +/// delivery semantics. +/// +/// The metastore checkpoint only carries synthetic positions (a per-pipeline +/// partition and a delivery counter) used to correlate published batches with +/// the acknowledgments they release; the resume point is the consumer's ack +/// floor, not the checkpoint. +/// +/// The source acknowledges inline in [`Source::suggest_truncate`], so once a +/// drain reports completion (see [`Source::is_drained`]) every ack has been +/// confirmed by the server. If the source actor is instead killed with batches +/// still in the pipeline, their `suggest_truncate` notifications are lost and +/// the unacknowledged messages — including ones whose split ends up published — +/// are redelivered after the consumer's `ack_wait`. +pub struct NatsSource { + source_runtime: SourceRuntime, + source_params: NatsSourceParams, + nats_client: async_nats::Client, + message_stream: DurableMessageStream, + consumer_name: String, + partition_id: PartitionId, + delivery_counter: u64, + /// Ack subjects of the messages delivered but not published yet, keyed by + /// delivery counter. Bounded by the consumer's `max_ack_pending`: the + /// server stops delivering when too many messages are unacknowledged. + pending_acks: BTreeMap, + state: NatsSourceState, +} + +impl fmt::Debug for NatsSource { + fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter + .debug_struct("NatsSource") + .field("index_uid", self.source_runtime.index_uid()) + .field("source_id", &self.source_runtime.source_id()) + .field("stream", &self.source_params.stream) + .field("consumer_name", &self.consumer_name) + .finish() + } +} + +impl NatsSource { + pub async fn try_new( + source_runtime: SourceRuntime, + source_params: NatsSourceParams, + ) -> anyhow::Result { + let consumer_name = source_params.consumer.clone(); + + info!( + index_id=%source_runtime.index_id(), + source_id=%source_runtime.source_id(), + stream=%source_params.stream, + %consumer_name, + "starting NATS source" + ); + + let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; + let message_stream = consumer + .messages() + .await + .context("failed to subscribe to NATS consumer messages")?; + + // A fresh partition per pipeline: durable-mode positions are delivery + // counters local to this instance, and pipelines sharing the consumer + // must not collide on a partition. + let partition_id = + PartitionId::from(append_random_suffix(&format!("nats-{consumer_name}"))); + + Ok(NatsSource { + source_runtime, + source_params, + nats_client, + message_stream, + consumer_name, + partition_id, + delivery_counter: 0, + pending_acks: BTreeMap::new(), + state: NatsSourceState::default(), + }) + } + + fn process_message( + &mut self, + message: jetstream::Message, + batch: &mut BatchBuilder, + ) -> anyhow::Result<()> { + let stream_sequence = message + .info() + .map_err(|error| anyhow!("failed to parse NATS message metadata: {error}"))? + .stream_sequence; + let _span_guard = remote_parented_span(&message, stream_sequence, &self.source_runtime) + .map(Span::entered); + let Some(ack_subject) = message.message.reply else { + bail!("NATS message carries no reply subject to acknowledge it on"); + }; + let doc = message.message.payload; + let num_bytes = doc.len() as u64; + + if doc.is_empty() { + warn!("message received from NATS was empty"); + self.state.num_invalid_messages += 1; + } else { + batch.add_doc(doc); + } + let from_position = if self.delivery_counter == 0 { + Position::Beginning + } else { + Position::offset(self.delivery_counter) + }; + self.delivery_counter += 1; + batch + .checkpoint_delta + .record_partition_delta( + self.partition_id.clone(), + from_position, + Position::offset(self.delivery_counter), + ) + .context("failed to record partition delta")?; + self.pending_acks.insert(self.delivery_counter, ack_subject); + + self.state.num_bytes_processed += num_bytes; + self.state.num_messages_processed += 1; + + Ok(()) + } +} + +#[async_trait] +impl Source for NatsSource { + async fn emit_batches( + &mut self, + source_sink: &SourceSink, + ctx: &SourceContext, + ) -> Result { + let now = Instant::now(); + let mut batch_builder = BatchBuilder::new(SourceType::Nats); + let deadline = time::sleep(*EMIT_BATCHES_TIMEOUT); + tokio::pin!(deadline); + let mut wait_before_next_batch = Duration::default(); + + loop { + tokio::select! { + message_res_opt = self.message_stream.next() => { + let message_res = message_res_opt + .ok_or_else(|| ActorExitStatus::from(anyhow!("NATS message stream ended unexpectedly")))?; + let message = match message_res { + Ok(message) => message, + Err(error) if is_transient_stream_error(&error) => { + warn!(%error, "transient NATS message stream error"); + wait_before_next_batch = TRANSIENT_STREAM_ERROR_BACKOFF; + break; + } + Err(error) => { + return Err(ActorExitStatus::from(anyhow!( + "failed to pull message from NATS consumer: {error}" + ))); + } + }; + self.process_message(message, &mut batch_builder).map_err(ActorExitStatus::from)?; + + if batch_builder.num_bytes >= BATCH_NUM_BYTES_LIMIT { + break; + } + } + _ = &mut deadline => { + break; + } + } + ctx.record_progress(); + } + + if !batch_builder.checkpoint_delta.is_empty() { + debug!( + num_docs=%batch_builder.docs.len(), + num_bytes=%batch_builder.num_bytes, + num_millis=%now.elapsed().as_millis(), + "sending doc batch to indexer" + ); + let message = batch_builder.build(); + source_sink.send_raw_doc_batch(message, ctx).await?; + } + Ok(wait_before_next_batch) + } + + async fn suggest_truncate( + &mut self, + checkpoint: SourceCheckpoint, + ctx: &SourceContext, + ) -> anyhow::Result<()> { + let Some(position) = checkpoint.position_for_partition(&self.partition_id) else { + return Ok(()); + }; + let Some(published_up_to) = position.as_u64() else { + return Ok(()); + }; + ctx.protect_future(ack_up_to( + &self.nats_client, + &mut self.pending_acks, + published_up_to, + )) + .await; + Ok(()) + } + + fn is_drained(&self) -> bool { + // Every delivered message is pending from `process_message` until the + // server confirmed its ack: an empty map means everything delivered so + // far has been published and acknowledged. + self.pending_acks.is_empty() + } + + fn name(&self) -> String { + format!("{self:?}") + } + + async fn finalize( + &mut self, + _exit_status: &ActorExitStatus, + _ctx: &SourceContext, + ) -> anyhow::Result<()> { + // The durable consumer itself is left untouched. Any ack still pending + // at this point was not settled by a drain: its message is redelivered + // after the consumer's `ack_wait`. + // + // Messages the client prefetched but that were never processed are + // negatively acknowledged so the server redelivers them right after + // this connection is gone, instead of after `ack_wait`. Only what is + // already buffered is NAK'd: the loop stops at the first pending poll, + // so a refill in flight cannot keep it alive. + while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { + let Ok(message) = message_res else { + continue; + }; + let Some(ack_subject) = message.message.reply else { + continue; + }; + let _ = self + .nats_client + .publish(ack_subject, AckKind::Nak(Some(NAK_REDELIVERY_DELAY)).into()) + .await; + } + let _ = self.nats_client.drain().await; + Ok(()) + } + + fn observable_state(&self) -> JsonValue { + let num_pending_acks = self.pending_acks.len(); + json!({ + "index_id": self.source_runtime.index_id(), + "source_id": self.source_runtime.source_id(), + "stream": self.source_params.stream, + "consumer_name": self.consumer_name, + "num_bytes_processed": self.state.num_bytes_processed, + "num_messages_processed": self.state.num_messages_processed, + "num_invalid_messages": self.state.num_invalid_messages, + "num_pending_acks": num_pending_acks, + }) + } +} + +/// Bounds the concurrent server-confirmed ack requests: enough to hide the +/// round-trip latency without flooding the connection. +const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; + +/// Pause before pulling again after a transient stream error, to avoid a hot +/// error loop while the connection recovers. `SuggestTruncate` is still +/// processed in the meantime. +const TRANSIENT_STREAM_ERROR_BACKOFF: Duration = Duration::from_secs(1); + +/// Delay before the server redelivers a negatively acknowledged message: long +/// enough for this connection to be gone, so the redelivery goes to a +/// surviving pipeline instead of bouncing back into this one's outstanding +/// pull request. +const NAK_REDELIVERY_DELAY: Duration = Duration::from_secs(1); + +/// Errors the pull stream recovers from on its own: the subscription stays +/// usable and polling simply resumes. `MissingHeartbeat` in particular also +/// fires when the stream was not polled for a while (e.g. under downstream +/// backpressure), not only when the server is unreachable. +fn is_transient_stream_error(error: &MessagesError) -> bool { + matches!( + error.kind(), + MessagesErrorKind::MissingHeartbeat + | MessagesErrorKind::Pull + | MessagesErrorKind::NoResponders + ) +} + +/// Sends server-confirmed acknowledgments ("double acks") for the messages up +/// to `published_up_to`. +/// +/// Entries only leave the pending map once the server confirmed their ack, so +/// `is_drained` never reports progress that could still be lost: a +/// fire-and-forget ack can be dropped by a reconnection even after a +/// successful client-side flush. A retried ack for a message the server +/// already saw acknowledged is simply confirmed again. Whatever could not be +/// acknowledged is redelivered after the consumer's `ack_wait`. +async fn ack_up_to( + nats_client: &async_nats::Client, + pending_acks: &mut BTreeMap, + published_up_to: u64, +) { + let acks: Vec<(u64, Subject)> = pending_acks + .range(..=published_up_to) + .map(|(delivery_counter, ack_subject)| (*delivery_counter, ack_subject.clone())) + .collect(); + if acks.is_empty() { + return; + } + let mut ack_results = futures::stream::iter(acks) + .map(|(delivery_counter, ack_subject)| async move { + // Mirrors `jetstream::Message::double_ack()`: a request instead of + // a plain publish, so the server's reply confirms the ack. + let ack_result = nats_client.request(ack_subject, AckKind::Ack.into()).await; + (delivery_counter, ack_result) + }) + .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); + let mut num_acks = 0usize; + let mut num_failed_acks = 0usize; + let mut last_ack_error = None; + while let Some((delivery_counter, ack_result)) = ack_results.next().await { + match ack_result { + Ok(_ack_reply) => { + pending_acks.remove(&delivery_counter); + num_acks += 1; + } + Err(error) => { + num_failed_acks += 1; + last_ack_error = Some(error); + } + } + } + if let Some(error) = last_ack_error { + warn!(%error, num_failed_acks, "failed to ack NATS messages"); + } + debug!(num_acks, "acked published messages"); +} + +/// Fetches the pre-provisioned durable consumer. Read-only by contract: the +/// consumer is never created nor updated. +async fn fetch_durable_consumer( + jetstream_stream: &jetstream::stream::Stream, + consumer_name: &str, +) -> anyhow::Result { + let consumer: PullConsumer = jetstream_stream + .get_consumer(consumer_name) + .await + .map_err(|error| anyhow!("failed to find NATS consumer `{consumer_name}`: {error}"))?; + let ack_policy = consumer.cached_info().config.ack_policy; + ensure!( + ack_policy == AckPolicy::Explicit, + "NATS consumer `{consumer_name}` must use the explicit ack policy, got `{ack_policy:?}`" + ); + Ok(consumer) +} + +/// W3C trace context extraction from NATS message headers, following the +/// pattern of `quickwit_common::tracing_utils` for gRPC metadata. +struct NatsHeaderExtractor<'a>(&'a HeaderMap); + +impl Extractor for NatsHeaderExtractor<'_> { + fn get(&self, key: &str) -> Option<&str> { + let header_name = HeaderName::from_str(key).ok()?; + self.0.get(header_name).map(|value| value.as_str()) + } + + fn keys(&self) -> Vec<&str> { + self.0.iter().map(|(key, _)| key.as_ref()).collect() + } +} + +fn extract_remote_context(headers: &HeaderMap) -> OtelContext { + global::get_text_map_propagator(|propagator| propagator.extract(&NatsHeaderExtractor(headers))) +} + +/// Builds a span parented on the publisher's trace when the message carries +/// a W3C `traceparent` header, stitching the processing of the message into +/// the publisher's distributed trace. Messages without a propagated context +/// cost nothing: no span is created. +fn remote_parented_span( + message: &jetstream::Message, + stream_sequence: u64, + source_runtime: &SourceRuntime, +) -> Option { + let headers = message.headers.as_ref()?; + let parent_context = extract_remote_context(headers); + if !parent_context.span().span_context().is_valid() { + return None; + } + let span = tracing::info_span!( + "process_nats_message", + index_id = %source_runtime.index_id(), + source_id = %source_runtime.source_id(), + subject = %message.subject, + stream_sequence, + ); + let _ = span.set_parent(parent_context); + Some(span) +} + +async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result { + let mut connect_options = ConnectOptions::new(); + match params.authentication.clone() { + None => {} + Some(NatsSourceAuth::UserPassword { user, password }) => { + connect_options = connect_options.user_and_password(user, password); + } + Some(NatsSourceAuth::Token(token)) => { + connect_options = connect_options.token(token); + } + } + if let Some(tls) = ¶ms.tls { + if let Some(ca_certificates_path) = &tls.ca_certificates_path { + connect_options = + connect_options.add_root_certificates(PathBuf::from(ca_certificates_path)); + } + if let (Some(certificate_path), Some(key_path)) = + (&tls.client_certificate_path, &tls.client_key_path) + { + connect_options = connect_options + .add_client_certificate(PathBuf::from(certificate_path), PathBuf::from(key_path)); + } + } + let client = async_nats::connect_with_options(¶ms.uris, connect_options) + .await + .with_context(|| { + format!( + "failed to connect to NATS servers `{}`", + params.uris.join(", ") + ) + })?; + Ok(client) +} + +/// Connects to the NATS servers and binds to the JetStream stream's +/// pre-provisioned durable consumer. +async fn connect_and_fetch_consumer( + params: &NatsSourceParams, +) -> anyhow::Result<(async_nats::Client, PullConsumer)> { + let nats_client = connect_nats(params).await?; + let jetstream_ctx = jetstream::new(nats_client.clone()); + let jetstream_stream = jetstream_ctx + .get_stream(¶ms.stream) + .await + .with_context(|| format!("failed to find NATS JetStream stream `{}`", params.stream))?; + let consumer = fetch_durable_consumer(&jetstream_stream, ¶ms.consumer).await?; + Ok((nats_client, consumer)) +} + +/// Checks whether we can connect to the NATS servers and find the JetStream +/// stream and the pre-provisioned consumer. +pub(crate) async fn check_connectivity(params: &NatsSourceParams) -> anyhow::Result<()> { + connect_and_fetch_consumer(params).await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_extract_trace_context_from_headers() { + use opentelemetry::propagation::TextMapPropagator; + use opentelemetry_sdk::propagation::TraceContextPropagator; + + // The global propagator is not installed in tests, so the extractor + // is exercised against an explicit W3C propagator. + let propagator = TraceContextPropagator::new(); + + let mut headers = HeaderMap::new(); + headers.insert( + "traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ); + let span_context = propagator + .extract(&NatsHeaderExtractor(&headers)) + .span() + .span_context() + .clone(); + assert!(span_context.is_valid()); + assert_eq!( + span_context.trace_id().to_string(), + "4bf92f3577b34da6a3ce929d0e0e4736" + ); + assert_eq!(span_context.span_id().to_string(), "00f067aa0ba902b7"); + + let empty_headers = HeaderMap::new(); + let span_context = propagator + .extract(&NatsHeaderExtractor(&empty_headers)) + .span() + .span_context() + .clone(); + assert!(!span_context.is_valid()); + } +} + +#[cfg(all(test, feature = "nats-broker-tests"))] +mod nats_broker_tests { + use std::num::NonZeroUsize; + use std::ops::Range; + use std::sync::Arc; + + use bytes::Bytes; + use quickwit_actors::{ActorHandle, Inbox, Universe}; + use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; + use quickwit_metastore::metastore_for_test; + use quickwit_proto::metastore::MetastoreServiceClient; + use quickwit_proto::types::IndexUid; + + use super::*; + use crate::actors::DocProcessor; + use crate::models::RawDocBatch; + use crate::source::test_setup_helper::setup_index; + use crate::source::tests::SourceRuntimeBuilder; + use crate::source::{SourceActor, SuggestTruncate, quickwit_supported_sources}; + + static NATS_URI: &str = "nats://localhost:4222"; + + async fn setup_nats_stream(stream_name: &str) -> jetstream::Context { + let client = async_nats::connect(NATS_URI).await.unwrap(); + let jetstream_ctx = jetstream::new(client); + jetstream_ctx + .create_stream(jetstream::stream::Config { + name: stream_name.to_string(), + subjects: vec![format!("{stream_name}.>")], + ..Default::default() + }) + .await + .unwrap(); + jetstream_ctx + } + + /// Publishes one JSON doc per ID on the subject and waits for each + /// publish ack, so stream sequences are assigned in `ids` order. Messages + /// carry a W3C `traceparent` header to exercise the trace propagation + /// path. + async fn publish_docs( + jetstream_ctx: &jetstream::Context, + subject: &str, + ids: Range, + ) -> Vec { + let mut headers = HeaderMap::new(); + headers.insert( + "traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ); + let mut docs = Vec::with_capacity(ids.len()); + for id in ids { + let doc = json!({ "id": id, "subject": subject }).to_string(); + jetstream_ctx + .publish_with_headers(subject.to_string(), headers.clone(), doc.clone().into()) + .await + .unwrap() + .await + .unwrap(); + docs.push(doc); + } + docs + } + + async fn create_source_actor( + universe: &Universe, + metastore: MetastoreServiceClient, + index_uid: IndexUid, + source_config: SourceConfig, + ) -> (ActorHandle, Inbox) { + let source_runtime = SourceRuntimeBuilder::new(index_uid, source_config) + .with_metastore(metastore) + .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_handle, doc_processor_inbox) + } + + async fn wait_for_processed_messages( + source_handle: &ActorHandle, + num_expected: u64, + ) { + loop { + let observation = source_handle.observe().await; + let num_messages_processed = observation + .state + .get("num_messages_processed") + .unwrap() + .as_u64() + .unwrap(); + if num_messages_processed >= num_expected { + return; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + fn merge_doc_batches(batches: Vec) -> RawDocBatch { + let mut merged_batch = RawDocBatch::default(); + for batch in batches { + merged_batch.docs.extend(batch.docs); + merged_batch + .checkpoint_delta + .extend(batch.checkpoint_delta) + .unwrap(); + } + merged_batch.docs.sort(); + merged_batch + } + + fn get_durable_source_config(stream: &str, consumer: &str) -> SourceConfig { + let source_id = append_random_suffix("test-nats-source--durable-source"); + SourceConfig { + source_id, + num_pipelines: NonZeroUsize::MIN, + enabled: true, + source_params: SourceParams::Nats(NatsSourceParams { + uris: vec![NATS_URI.to_string()], + stream: stream.to_string(), + consumer: consumer.to_string(), + tls: None, + authentication: None, + }), + transform_config: None, + input_format: SourceInputFormat::Json, + } + } + + /// Provisions the durable consumer the way an operator would: the source + /// itself only ever fetches it. + async fn provision_durable_consumer( + jetstream_ctx: &jetstream::Context, + stream: &str, + consumer_name: &str, + ) { + jetstream_ctx + .create_consumer_on_stream( + jetstream::consumer::pull::Config { + name: Some(consumer_name.to_string()), + durable_name: Some(consumer_name.to_string()), + ack_policy: AckPolicy::Explicit, + // Long enough for the tests to never hit a redelivery. + ack_wait: Duration::from_secs(300), + ..Default::default() + }, + stream, + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_durable_mode_ingestion_and_ack() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-ack-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let subject = format!("{stream}.logs"); + let expected_docs = publish_docs(&jetstream_ctx, &subject, 0..10).await; + + let index_id = append_random_suffix("test-nats-source--durable--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle, doc_processor_inbox) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + + wait_for_processed_messages(&source_handle, 10).await; + + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + let batch = merge_doc_batches(batches); + assert_eq!(batch.docs, expected_docs); + // Positions are synthetic delivery counters on a per-pipeline + // partition, not stream sequences. + assert_eq!(batch.checkpoint_delta.num_partitions(), 1); + + // A `SuggestTruncate` simulates the split publication notification: + // it must release the acks of the published messages. + let checkpoint = batch.checkpoint_delta.get_source_checkpoint(); + source_handle + .mailbox() + .send_message(SuggestTruncate(checkpoint)) + .await + .unwrap(); + + let mut consumer: PullConsumer = jetstream_ctx + .get_consumer_from_stream(consumer_name, stream.as_str()) + .await + .unwrap(); + let mut acked = false; + for _ in 0..100 { + let consumer_info = consumer.info().await.unwrap(); + if consumer_info.num_ack_pending == 0 { + assert_eq!(consumer_info.ack_floor.stream_sequence, 10); + acked = true; + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!(acked, "messages should be acked after truncation"); + + source_handle.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + /// End-to-end drain: a full indexing pipeline is asked to drain, which + /// must publish the in-flight batches and flush their acks BEFORE the + /// drain replies — the exactly-once guarantee on planned teardowns. + #[tokio::test] + async fn test_durable_mode_graceful_drain_acks_before_teardown() { + use quickwit_actors::Universe; + use quickwit_common::temp_dir::TempDirectory; + use quickwit_config::IndexingSettings; + use quickwit_doc_mapper::default_doc_mapper_for_test; + use quickwit_ingest::IngesterPool; + use quickwit_proto::indexing::IndexingPipelineId; + use quickwit_proto::types::{NodeId, PipelineUid}; + use quickwit_storage::{RamStorage, StorageResolver}; + + use crate::actors::pipeline_shared::DrainPipeline; + use crate::merge_policy::default_merge_policy; + use crate::{IndexingPipeline, IndexingPipelineParams, IndexingSplitStore}; + + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-drain--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-drain-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let index_id = append_random_suffix("test-nats-source--durable-drain--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + // Docs valid for the test doc mapper, so the pipeline indexes them. + let subject = format!("{stream}.logs"); + for id in 0..10 { + let doc = json!({"timestamp": 1_700_000_000 + id, "body": format!("drain test {id}")}) + .to_string(); + jetstream_ctx + .publish(subject.clone(), doc.into()) + .await + .unwrap() + .await + .unwrap(); + } + + let pipeline_id = IndexingPipelineId { + node_id: NodeId::from_str("test-node"), + index_uid, + source_id: source_config.source_id.clone(), + pipeline_uid: PipelineUid::for_test(0u128), + }; + let storage = Arc::new(RamStorage::default()); + let split_store = IndexingSplitStore::create_without_local_store_for_test(storage.clone()); + let (merge_planner_mailbox, _merge_planner_inbox) = universe.create_test_mailbox(); + let pipeline_params = IndexingPipelineParams { + pipeline_id, + doc_mapper: Arc::new(default_doc_mapper_for_test()), + source_config, + source_storage_resolver: StorageResolver::for_test(), + indexing_directory: TempDirectory::for_test(), + indexing_settings: IndexingSettings::for_test(), + fingerprinter_opt: None, + ingester_pool: IngesterPool::default(), + metastore, + queues_dir_path: PathBuf::from("./queues"), + storage, + split_store, + merge_policy: default_merge_policy(), + retention_policy: None, + max_concurrent_split_uploads_index: 4, + max_concurrent_split_uploads_merge: 5, + cooperative_indexing_permits: None, + merge_planner_mailbox_opt: Some(merge_planner_mailbox), + event_broker: Default::default(), + params_fingerprint: 42u64, + }; + let pipeline = IndexingPipeline::new(pipeline_params); + let (pipeline_mailbox, pipeline_handle) = universe.spawn_builder().spawn(pipeline); + + // Wait for the source to have delivered all the messages. The commit + // timeout has not elapsed: without the drain, nothing would have been + // published nor acked yet. + loop { + let observation = pipeline_handle.observe().await; + if observation.num_docs >= 10 { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + + pipeline_mailbox.send_message(DrainPipeline).await.unwrap(); + // The pipeline exits on its own once the in-flight batches are + // published and their acks flushed. + let (exit_status, _statistics) = pipeline_handle.join().await; + assert!(matches!(exit_status, ActorExitStatus::Success)); + + // The pipeline exit implies the acks are already flushed: no polling, + // no grace period. + let mut consumer: PullConsumer = jetstream_ctx + .get_consumer_from_stream(consumer_name, stream.as_str()) + .await + .unwrap(); + let consumer_info = consumer.info().await.unwrap(); + assert_eq!(consumer_info.num_ack_pending, 0); + assert_eq!(consumer_info.ack_floor.stream_sequence, 10); + + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + #[tokio::test] + async fn test_durable_mode_load_balancing() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-lb--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-lb-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let index_id = append_random_suffix("test-nats-source--durable-lb--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle_1, doc_processor_inbox_1) = create_source_actor( + &universe, + metastore.clone(), + index_uid.clone(), + source_config.clone(), + ) + .await; + let (source_handle_2, doc_processor_inbox_2) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + + let subject = format!("{stream}.logs"); + let mut expected_docs = publish_docs(&jetstream_ctx, &subject, 0..20).await; + expected_docs.sort(); + + // Work-queue delivery splits the messages between the two pipelines + // in some arbitrary way; together they must cover all of them + // exactly once. + loop { + let num_processed_1 = source_handle_1 + .observe() + .await + .state + .get("num_messages_processed") + .unwrap() + .as_u64() + .unwrap(); + let num_processed_2 = source_handle_2 + .observe() + .await + .state + .get("num_messages_processed") + .unwrap() + .as_u64() + .unwrap(); + if num_processed_1 + num_processed_2 >= 20 { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + source_handle_1.quit().await; + source_handle_2.quit().await; + + let mut all_docs: Vec = Vec::new(); + for batch in doc_processor_inbox_1 + .drain_for_test_typed::() + .into_iter() + .chain(doc_processor_inbox_2.drain_for_test_typed::()) + { + all_docs.extend(batch.docs); + } + all_docs.sort(); + assert_eq!(all_docs, expected_docs); + + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + #[tokio::test] + async fn test_durable_mode_missing_consumer() { + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-missing--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + + let index_id = append_random_suffix("test-nats-source--durable-missing--index"); + let source_config = get_durable_source_config(&stream, "does-not-exist"); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let source_runtime = SourceRuntimeBuilder::new(index_uid, source_config) + .with_metastore(metastore) + .build(); + let load_source_result = quickwit_supported_sources() + .load_source(source_runtime) + .await; + assert!( + load_source_result.is_err(), + "binding to a missing durable consumer should fail" + ); + + jetstream_ctx.delete_stream(&stream).await.unwrap(); + } +} diff --git a/quickwit/quickwit-metastore/src/metastore/mod.rs b/quickwit/quickwit-metastore/src/metastore/mod.rs index c12b8b11b61..98b2e14864e 100644 --- a/quickwit/quickwit-metastore/src/metastore/mod.rs +++ b/quickwit/quickwit-metastore/src/metastore/mod.rs @@ -1513,6 +1513,7 @@ fn use_shard_api(params: &SourceParams) -> bool { SourceParams::Kinesis(_) => false, SourceParams::PubSub(_) => false, SourceParams::Pulsar(_) => false, + SourceParams::Nats(_) => false, SourceParams::Stdin => panic!("stdin cannot be checkpointed"), SourceParams::Vec(_) => false, SourceParams::Void(_) => false, From 137846c641e2315f41d9b6b795e12bfeeac5a958 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 18:15:29 +0200 Subject: [PATCH 03/41] feat: indexing pipeline drain Drain every pipeline on teardown (sources without acknowledgment state exit as soon as their flush batch is pushed), share the drain state machine between both pipeline flavors, and fix the drain accounting: DrainAllPipelines now also waits for pipelines a plan change already detached, merge pipelines survive their draining indexing pipelines, and a paused draining pipeline is no longer reaped as exited. Co-Authored-By: Claude Fable 5 --- .../src/actors/indexing_pipeline.rs | 35 ++++----- .../src/actors/indexing_service.rs | 76 +++++++++++-------- .../parquet_pipeline/indexing_service_impl.rs | 3 - .../src/actors/parquet_pipeline/pipeline.rs | 35 ++++----- .../src/actors/pipeline_shared.rs | 46 ++++++++--- quickwit/quickwit-indexing/src/source/mod.rs | 12 --- 6 files changed, 111 insertions(+), 96 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs index 6b4203c54f2..6199c8b46d6 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs @@ -39,8 +39,8 @@ use tracing::{debug, error, info, instrument, warn}; use super::{DocProcessor, IndexSerializer, Indexer, MergePlanner, Packager}; use crate::SplitsUpdateMailbox; use crate::actors::pipeline_shared::{ - DRAIN_GRACE_MARGIN, DrainPipeline, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, - SuperviseLoop, wait_duration_before_retry, + DrainPipeline, DrainState, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, SuperviseLoop, + wait_duration_before_retry, }; use crate::actors::sequencer::Sequencer; use crate::actors::uploader::UploaderType; @@ -50,7 +50,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, Drain, SourceActor, SourceRuntime, quickwit_supported_sources, + AssignShards, Assignment, SourceActor, SourceRuntime, quickwit_supported_sources, }; use crate::split_store::IndexingSplitStore; @@ -86,10 +86,7 @@ pub struct IndexingPipeline { handles_opt: Option, // Killswitch used for the actors in the pipeline. This is not the supervisor killswitch. kill_switch: KillSwitch, - // Set when the pipeline is draining (see `DrainPipeline`): the pipeline is - // shutting down and must never respawn; past the deadline it gives up and - // kills whatever is left. - drain_deadline_opt: Option, + drain_state: DrainState, // The set of shard is something that can change dynamically without necessarily // requiring a respawn of the pipeline. @@ -142,7 +139,7 @@ impl IndexingPipeline { previous_generations_statistics: Default::default(), handles_opt: None, kill_switch: KillSwitch::default(), - drain_deadline_opt: None, + drain_state: DrainState::default(), statistics: IndexingStatistics { params_fingerprint, ..Default::default() @@ -270,7 +267,7 @@ impl IndexingPipeline { Health::Healthy => {} Health::FailureOrUnhealthy => { self.terminate().await; - if self.drain_deadline_opt.is_some() { + if self.drain_state.is_draining() { // A draining pipeline is never respawned: whatever could // not be settled is redelivered (see `DrainPipeline`). warn!( @@ -474,9 +471,7 @@ impl Handler for IndexingPipeline { ) -> Result<(), ActorExitStatus> { self.perform_observe(ctx); self.perform_health_check(ctx).await?; - if let Some(drain_deadline) = self.drain_deadline_opt - && Instant::now() >= drain_deadline - { + if self.drain_state.is_deadline_exceeded() { warn!( pipeline_id=?self.params.pipeline_id, "indexing pipeline could not be drained before the deadline" @@ -498,7 +493,7 @@ impl Handler for IndexingPipeline { spawn: Spawn, ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { - if self.drain_deadline_opt.is_some() { + if self.drain_state.is_draining() { // A draining pipeline is shutting down: never respawn it. return Ok(()); } @@ -535,19 +530,19 @@ impl Handler for IndexingPipeline { _drain: DrainPipeline, _ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { - if self.drain_deadline_opt.is_some() { + if self.drain_state.is_draining() { return Ok(()); } let Some(handles) = &self.handles_opt else { // Nothing is running, so nothing is in flight. return Err(ActorExitStatus::Success); }; - self.drain_deadline_opt = Some( - Instant::now() + self.params.indexing_settings.commit_timeout() + DRAIN_GRACE_MARGIN, - ); - // If the source is already dead, the health check terminates the - // draining pipeline instead of respawning it. - let _ = handles.source_mailbox.send_message(Drain).await; + self.drain_state + .start( + &handles.source_mailbox, + self.params.indexing_settings.commit_timeout(), + ) + .await; Ok(()) } } diff --git a/quickwit/quickwit-indexing/src/actors/indexing_service.rs b/quickwit/quickwit-indexing/src/actors/indexing_service.rs index 9be098a9928..80bf916dc60 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_service.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_service.rs @@ -62,7 +62,7 @@ use super::pipeline_shared::{ActorPipeline, PipelineHandle}; use super::{FinishPendingMergesAndShutdownPipeline, MergePlanner, MergeSchedulerService}; use crate::docs_clustering::Fingerprinter; use crate::models::{DetachIndexingPipeline, DetachMergePipeline, ObservePipeline, SpawnPipeline}; -use crate::source::{AssignShards, Assignment, source_needs_drain}; +use crate::source::{AssignShards, Assignment}; use crate::split_store::IndexingSplitCache; use crate::{IndexingPipeline, IndexingPipelineParams, IndexingSplitStore, IndexingStatistics}; @@ -421,7 +421,6 @@ impl IndexingService { ); } - let needs_drain = source_needs_drain(&source_config.source_params); let pipeline_params = IndexingPipelineParams { pipeline_id: indexing_pipeline_id.clone(), metastore: self.metastore.clone(), @@ -450,7 +449,6 @@ impl IndexingService { pipeline_id: indexing_pipeline_id, mailbox, handle, - needs_drain, })) } @@ -561,6 +559,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.state().is_exit(); + } + self.draining_pipelines.iter().any(|pipeline_handle| { + pipeline_handle.indexing_pipeline_id().pipeline_uid == pipeline_uid + && !pipeline_handle.state().is_exit() + }) + } + async fn handle_supervise(&mut self) -> Result<(), ActorExitStatus> { self.indexing_pipelines .retain(|pipeline_uid, pipeline_handle| { @@ -584,18 +594,14 @@ impl IndexingService { }); // Draining pipelines exit on their own (see `DrainPipeline`): reap the // detached ones and complete the `DrainAllPipelines` replies whose - // pipelines are all gone. + // pipelines are all gone. A paused pipeline is still alive, hence the + // `is_exit` checks rather than `is_running`. self.draining_pipelines - .retain(|pipeline_handle| pipeline_handle.state().is_running()); + .retain(|pipeline_handle| !pipeline_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.indexing_pipelines - .get(pipeline_uid) - .map(|pipeline_handle| pipeline_handle.state().is_running()) - .unwrap_or(false) - }); + pipeline_uids.retain(|pipeline_uid| self.is_pipeline_alive(*pipeline_uid)); if pipeline_uids.is_empty() { reply(()); } else { @@ -604,9 +610,12 @@ impl IndexingService { } } + // 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.indexing_pipeline_id().merge_pipeline_id()) .collect(); @@ -654,6 +663,7 @@ impl IndexingService { let parquet_index_uids_to_retain: HashSet = self .indexing_pipelines .values() + .chain(self.draining_pipelines.iter()) .filter(|h| { quickwit_common::is_parquet_pipeline_index( &h.indexing_pipeline_id().index_uid.index_id, @@ -1016,23 +1026,21 @@ impl IndexingService { } } } - // Acknowledgment-based sources are drained before their teardown so - // their in-flight messages are published and acknowledged instead of - // redelivered (see `DrainPipeline`). A drain can take up to the commit - // timeout, so it is not awaited here: the draining pipeline exits on - // its own and the supervise loop reaps its handle. + // Pipelines are drained before their teardown so their in-flight + // batches are flushed, published, and settled (acknowledged for + // acknowledgment-based sources) instead of dropped (see + // `DrainPipeline`). A drain can take up to the commit timeout, so it + // is not awaited here: the draining pipeline exits on its own and the + // supervise loop reaps its handle. for pipeline_handle in detached_pipeline_handles { - if pipeline_handle.needs_drain() { - pipeline_handle.start_drain().await; - self.draining_pipelines.push(pipeline_handle); - } else { - // Killing the pipeline ensures that all the pipeline actors will stop. - pipeline_handle.kill().await; - } + pipeline_handle.start_drain().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. + // indexes. The GC only removes queues of deleted indexes: a pipeline + // still draining such an index fails its publish and exits without + // respawning. if should_gc_ingest_api_queues && let Err(error) = self.run_ingest_api_queues_gc().await { warn!( %error, @@ -1126,11 +1134,12 @@ impl IndexingService { } } -/// Drains every running indexing pipeline whose source needs it (see -/// `DrainPipeline`): 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. +/// Drains every running indexing pipeline (see `DrainPipeline`): 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 — including ones a +/// plan change already detached and started draining — have all exited, +/// without blocking the service. #[derive(Debug)] pub struct DrainAllPipelines; @@ -1146,9 +1155,12 @@ impl DeferableReplyHandler for IndexingService { ) -> Result<(), ActorExitStatus> { let mut pipeline_uids: Vec = Vec::new(); for (pipeline_uid, pipeline_handle) in &self.indexing_pipelines { - if pipeline_handle.needs_drain() { - pipeline_handle.start_drain().await; - pipeline_uids.push(*pipeline_uid); + pipeline_handle.start_drain().await; + pipeline_uids.push(*pipeline_uid); + } + for pipeline_handle in &self.draining_pipelines { + if !pipeline_handle.state().is_exit() { + pipeline_uids.push(pipeline_handle.indexing_pipeline_id().pipeline_uid); } } if pipeline_uids.is_empty() { diff --git a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/indexing_service_impl.rs b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/indexing_service_impl.rs index 2a658ad10d2..857bf399838 100644 --- a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/indexing_service_impl.rs +++ b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/indexing_service_impl.rs @@ -26,7 +26,6 @@ use quickwit_proto::indexing::{IndexingError, IndexingPipelineId}; use crate::actors::pipeline_shared::ActorPipeline; use crate::actors::{MetricsPipeline, MetricsPipelineParams}; -use crate::source::source_needs_drain; use crate::{BoxedPipelineHandle, IndexingService}; impl IndexingService { @@ -87,7 +86,6 @@ impl IndexingService { ctx, )?; - let needs_drain = source_needs_drain(&source_config.source_params); let pipeline_params = MetricsPipelineParams { pipeline_id: indexing_pipeline_id.clone(), metastore: self.metastore.clone(), @@ -113,7 +111,6 @@ impl IndexingService { pipeline_id: indexing_pipeline_id, mailbox, handle, - needs_drain, })) } diff --git a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs index eaed82b6119..48a4685a655 100644 --- a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs @@ -46,15 +46,15 @@ use tracing::{debug, error, info, instrument, warn}; use super::{ParquetDocProcessor, ParquetIndexer, ParquetPackager, ParquetUploader}; use crate::actors::pipeline_shared::{ - DRAIN_GRACE_MARGIN, DrainPipeline, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, - SuperviseLoop, wait_duration_before_retry, + DrainPipeline, DrainState, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, SuperviseLoop, + wait_duration_before_retry, }; use crate::actors::sequencer::Sequencer; use crate::actors::{Publisher, UploaderType}; use crate::metrics::INDEXING_PIPELINES; use crate::models::{IndexingStatistics, SharedPublishToken}; use crate::source::{ - AssignShards, Assignment, Drain, SourceActor, SourceRuntime, quickwit_supported_sources, + AssignShards, Assignment, SourceActor, SourceRuntime, quickwit_supported_sources, }; struct MetricsPipelineHandles { @@ -114,10 +114,7 @@ pub struct MetricsPipeline { statistics: IndexingStatistics, handles_opt: Option, kill_switch: KillSwitch, - // Set when the pipeline is draining (see `DrainPipeline`): the pipeline is - // shutting down and must never respawn; past the deadline it gives up and - // kills whatever is left. - drain_deadline_opt: Option, + drain_state: DrainState, shard_ids: BTreeSet, // Id of the last indexing plan assigned to this pipeline. Kept here, like `shard_ids`, so it // can be re-sent to the source on respawn; the source adopts it as its publish token. @@ -167,7 +164,7 @@ impl MetricsPipeline { previous_generations_statistics: Default::default(), handles_opt: None, kill_switch: KillSwitch::default(), - drain_deadline_opt: None, + drain_state: DrainState::default(), statistics: IndexingStatistics { params_fingerprint, ..Default::default() @@ -293,7 +290,7 @@ impl MetricsPipeline { Health::Healthy => {} Health::FailureOrUnhealthy => { self.terminate().await; - if self.drain_deadline_opt.is_some() { + if self.drain_state.is_draining() { // A draining pipeline is never respawned: whatever could // not be settled is redelivered (see `DrainPipeline`). warn!( @@ -514,9 +511,7 @@ impl Handler for MetricsPipeline { ) -> Result<(), ActorExitStatus> { self.perform_observe(ctx); self.perform_health_check(ctx).await?; - if let Some(drain_deadline) = self.drain_deadline_opt - && Instant::now() >= drain_deadline - { + if self.drain_state.is_deadline_exceeded() { warn!( pipeline_id=?self.params.pipeline_id, "metrics pipeline could not be drained before the deadline" @@ -538,7 +533,7 @@ impl Handler for MetricsPipeline { spawn: Spawn, ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { - if self.drain_deadline_opt.is_some() { + if self.drain_state.is_draining() { // A draining pipeline is shutting down: never respawn it. return Ok(()); } @@ -575,19 +570,19 @@ impl Handler for MetricsPipeline { _drain: DrainPipeline, _ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { - if self.drain_deadline_opt.is_some() { + if self.drain_state.is_draining() { return Ok(()); } let Some(handles) = &self.handles_opt else { // Nothing is running, so nothing is in flight. return Err(ActorExitStatus::Success); }; - self.drain_deadline_opt = Some( - Instant::now() + self.params.indexing_settings.commit_timeout() + DRAIN_GRACE_MARGIN, - ); - // If the source is already dead, the health check terminates the - // draining pipeline instead of respawning it. - let _ = handles.source_mailbox.send_message(Drain).await; + self.drain_state + .start( + &handles.source_mailbox, + self.params.indexing_settings.commit_timeout(), + ) + .await; Ok(()) } } diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index 61a40a7df76..33c5a5384cc 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -65,10 +65,46 @@ pub struct DrainPipeline; /// covering the upload and publication of the final splits. pub(crate) const DRAIN_GRACE_MARGIN: Duration = Duration::from_secs(30); +/// Drain progress of a pipeline supervisor (see [`DrainPipeline`]): both +/// pipeline flavors share the deadline bookkeeping so they cannot diverge. +#[derive(Default)] +pub(crate) struct DrainState { + deadline_opt: Option, +} + +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() + } + + /// Records the drain deadline and asks the source to drain. If the source + /// is already dead, the pipeline health check terminates the draining + /// pipeline instead of respawning it. + pub(crate) async fn start( + &mut self, + source_mailbox: &Mailbox, + commit_timeout: Duration, + ) { + self.deadline_opt = Some(Instant::now() + commit_timeout + DRAIN_GRACE_MARGIN); + let _ = source_mailbox.send_message(Drain).await; + } + + /// 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) + } +} + // --------------------------------------------------------------------------- // Pipeline trait — type-erased handle for any indexing pipeline actor // --------------------------------------------------------------------------- +use std::time::Instant; + use async_trait::async_trait; use quickwit_actors::{ Actor, ActorExitStatus, ActorHandle, ActorState, DeferableReplyHandler, Health, Mailbox, @@ -77,7 +113,7 @@ use quickwit_actors::{ use quickwit_proto::indexing::IndexingPipelineId; use crate::models::IndexingStatistics; -use crate::source::AssignShards; +use crate::source::{AssignShards, Drain, SourceActor}; /// Trait that abstracts over the concrete pipeline actor type /// (`IndexingPipeline` or `MetricsPipeline`). This allows `PipelineHandle` @@ -92,9 +128,6 @@ pub trait PipelineHandle: Send + Sync { async fn send_assign_shards(&self, message: AssignShards) -> Result<(), SendError>; /// See [`DrainPipeline`]. Fire-and-forget: the pipeline exits on its own. async fn start_drain(&self); - /// Whether a teardown must drain before killing - /// (see [`crate::source::source_needs_drain`]). - fn needs_drain(&self) -> bool; async fn observe(&self) -> Observation; async fn join(self: Box) -> (ActorExitStatus, IndexingStatistics); async fn quit(self: Box) -> (ActorExitStatus, IndexingStatistics); @@ -107,7 +140,6 @@ pub(crate) struct ActorPipeline> pub pipeline_id: IndexingPipelineId, pub mailbox: Mailbox, pub handle: ActorHandle, - pub needs_drain: bool, } #[async_trait] @@ -147,10 +179,6 @@ where A: Actor let _ = self.mailbox.send_message(DrainPipeline).await; } - fn needs_drain(&self) -> bool { - self.needs_drain - } - async fn observe(&self) -> Observation { self.handle.observe().await } diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index e23f44dc732..266ffbe05aa 100644 --- a/quickwit/quickwit-indexing/src/source/mod.rs +++ b/quickwit/quickwit-indexing/src/source/mod.rs @@ -428,18 +428,6 @@ impl Handler for SourceActor { } } -/// Whether tearing down a pipeline for this source must go through a drain -/// ([`Drain`]) rather than a plain kill. Only acknowledgment-based -/// sources benefit: killing them loses the acknowledgments of their in-flight -/// messages, which are then redelivered and indexed again after the broker's -/// ack timeout. Checkpoint-based sources resume from their last checkpoint, so -/// a drain buys them nothing. -pub(crate) fn source_needs_drain(_source_params: &SourceParams) -> bool { - // No built-in source opts in yet: acknowledgment-based sources add their - // variant here as they land. - false -} - #[async_trait] impl Handler for SourceActor { type Reply = (); From ef62d741da62a91f65e12604fd9a0950e0e2b1f8 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 18:28:36 +0200 Subject: [PATCH 04/41] feat: indexing pipeline drain Bound the node-shutdown drain wait with a configurable indexer shutdown_drain_timeout (QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT, 300s by default) so a drain can no longer outlive the deployment's grace period. Co-Authored-By: Claude Fable 5 --- docs/configuration/node-config.md | 1 + .../quickwit-config/src/node_config/mod.rs | 24 +++++++++++++++ .../src/node_config/serialize.rs | 1 + quickwit/quickwit-serve/src/lib.rs | 30 ++++++++++++++----- 4 files changed, 48 insertions(+), 8 deletions(-) diff --git a/docs/configuration/node-config.md b/docs/configuration/node-config.md index cbe3f8b0022..52f8c925059 100644 --- a/docs/configuration/node-config.md +++ b/docs/configuration/node-config.md @@ -244,6 +244,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` | Maximum amount of time to wait for the indexing pipelines to drain gracefully on shutdown before giving up. It must exceed the largest `commit_timeout_secs` plus the time to upload and publish the final splits, and the deployment's shutdown grace period (e.g. `terminationGracePeriodSeconds` on Kubernetes) must exceed it in turn. Can be overridden with the `QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT` environment variable. | `300s` | Example: diff --git a/quickwit/quickwit-config/src/node_config/mod.rs b/quickwit/quickwit-config/src/node_config/mod.rs index 5e02be1af39..94c274d3f28 100644 --- a/quickwit/quickwit-config/src/node_config/mod.rs +++ b/quickwit/quickwit-config/src/node_config/mod.rs @@ -225,9 +225,31 @@ pub struct IndexerConfig { /// can't handle mixed prefix lengths. #[serde(default = "IndexerConfig::default_parquet_merge_use_streaming_engine")] pub parquet_merge_use_streaming_engine: bool, + /// Maximum amount of time to wait for the indexing pipelines to drain + /// gracefully on shutdown before giving up. Must exceed the largest commit + /// timeout plus the time to upload and publish the final splits, and the + /// deployment's shutdown grace period must exceed it in turn. Can be + /// overridden with the `QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT` environment + /// variable. + #[serde(default = "IndexerConfig::default_shutdown_drain_timeout")] + 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 } @@ -280,6 +302,7 @@ impl IndexerConfig { max_merge_write_throughput: None, merge_concurrency: NonZeroUsize::new(3).unwrap(), parquet_merge_use_streaming_engine: Self::default_parquet_merge_use_streaming_engine(), + shutdown_drain_timeout: Self::default_shutdown_drain_timeout(), }; Ok(indexer_config) } @@ -297,6 +320,7 @@ impl Default for IndexerConfig { merge_concurrency: Self::default_merge_concurrency(), max_merge_write_throughput: None, parquet_merge_use_streaming_engine: Self::default_parquet_merge_use_streaming_engine(), + 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 b69c354ed58..0c3c588a51d 100644 --- a/quickwit/quickwit-config/src/node_config/serialize.rs +++ b/quickwit/quickwit-config/src/node_config/serialize.rs @@ -869,6 +869,7 @@ mod tests { enable_cooperative_indexing: false, max_merge_write_throughput: Some(ByteSize::mb(100)), parquet_merge_use_streaming_engine: true, + shutdown_drain_timeout: IndexerConfig::default_shutdown_drain_timeout(), } ); assert_eq!( diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index 8f5973d3423..8c3a0225dfb 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -509,6 +509,7 @@ async fn shutdown_signal_handler( shutdown_signal: BoxFutureInfaillible<()>, universe: Universe, indexing_service_opt: Option>, + indexer_shutdown_drain_timeout: Duration, ingester_opt: Option, ingester_decommission_timeout: Duration, compactor_supervisor_opt: Option>, @@ -538,14 +539,25 @@ async fn shutdown_signal_handler( if let Err(error) = compactor_result { error!("failed to decommission compactor gracefully: {:?}", error); } - // Drain the indexing pipelines only after the ingester decommission: it - // waits for the persisted shard data to be indexed, which requires the - // pipelines to still be consuming. Draining publishes and settles the - // in-flight batches before the universe teardown drops them. - if let Some(indexing_service) = &indexing_service_opt - && let Err(error) = indexing_service.ask(DrainAllPipelines).await - { - error!("failed to drain indexing pipelines 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; @@ -920,6 +932,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 { @@ -1044,6 +1057,7 @@ pub async fn serve_quickwit( shutdown_signal, universe, quickwit_services.indexing_service_opt.clone(), + indexer_shutdown_drain_timeout, ingester_opt, ingester_decommission_timeout, compactor_supervisor_opt, From 3b175978ceefc0bbc57f2a0bea26c7b83cac70a0 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 18:34:51 +0200 Subject: [PATCH 05/41] feat: indexing pipeline drain --- quickwit/quickwit-config/src/node_config/mod.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/quickwit/quickwit-config/src/node_config/mod.rs b/quickwit/quickwit-config/src/node_config/mod.rs index 94c274d3f28..c93cc622b1a 100644 --- a/quickwit/quickwit-config/src/node_config/mod.rs +++ b/quickwit/quickwit-config/src/node_config/mod.rs @@ -232,7 +232,7 @@ pub struct IndexerConfig { /// overridden with the `QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT` environment /// variable. #[serde(default = "IndexerConfig::default_shutdown_drain_timeout")] - shutdown_drain_timeout: HumanDuration, + pub shutdown_drain_timeout: HumanDuration, } impl IndexerConfig { @@ -609,7 +609,9 @@ impl CacheConfig { fn deserialize_with_default<'de, D, const DEFAULT_CAPACITY: u64>( deserializer: D, ) -> Result - where D: Deserializer<'de> { + where + D: Deserializer<'de>, + { use serde_with::{DeserializeAs, FromInto, PickFirst, Same}; let mut cache_config: CacheConfig = From 430bde4df3a8f0f17762ce03e539890ee2bb36f3 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 18:40:33 +0200 Subject: [PATCH 06/41] feat: indexing pipeline drain --- quickwit/quickwit-config/src/node_config/mod.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/quickwit/quickwit-config/src/node_config/mod.rs b/quickwit/quickwit-config/src/node_config/mod.rs index c93cc622b1a..148bcdfbfda 100644 --- a/quickwit/quickwit-config/src/node_config/mod.rs +++ b/quickwit/quickwit-config/src/node_config/mod.rs @@ -609,9 +609,7 @@ impl CacheConfig { fn deserialize_with_default<'de, D, const DEFAULT_CAPACITY: u64>( deserializer: D, ) -> Result - where - D: Deserializer<'de>, - { + where D: Deserializer<'de> { use serde_with::{DeserializeAs, FromInto, PickFirst, Same}; let mut cache_config: CacheConfig = From 6b0adc1340fc91ab814ea145cd76b6c265ab5c65 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 19:02:12 +0200 Subject: [PATCH 07/41] feat: indexing pipeline drain --- .../src/actors/indexing_pipeline.rs | 11 +++++++++++ .../src/actors/parquet_pipeline/pipeline.rs | 11 +++++++++++ .../src/actors/pipeline_shared.rs | 17 +++++++++++------ quickwit/quickwit-indexing/src/source/mod.rs | 9 +++++++++ 4 files changed, 42 insertions(+), 6 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs index 6199c8b46d6..52c73dc7923 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs @@ -58,6 +58,9 @@ use crate::split_store::IndexingSplitStore; struct IndexingPipelineHandles { source_mailbox: Mailbox, source_handle: ActorHandle, + /// Captured from [`crate::source::Source::should_be_drained`] at spawn: + /// a teardown drains the pipeline when true, kills it otherwise. + source_should_be_drained: bool, doc_processor: ActorHandle, indexer: ActorHandle, index_serializer: ActorHandle, @@ -417,6 +420,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() @@ -435,6 +439,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, @@ -537,6 +542,12 @@ impl Handler for IndexingPipeline { // Nothing is running, so nothing is in flight. return Err(ActorExitStatus::Success); }; + if !handles.source_should_be_drained { + // The source settles nothing on a drain: tear the pipeline down + // right away, exactly as a kill would. + self.terminate().await; + return Err(ActorExitStatus::Success); + } self.drain_state .start( &handles.source_mailbox, diff --git a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs index 48a4685a655..9465f612788 100644 --- a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs @@ -60,6 +60,9 @@ use crate::source::{ struct MetricsPipelineHandles { source_mailbox: Mailbox, source_handle: ActorHandle, + /// Captured from [`crate::source::Source::should_be_drained`] at spawn: + /// a teardown drains the pipeline when true, kills it otherwise. + source_should_be_drained: bool, doc_processor: ActorHandle, indexer: ActorHandle, packager: ActorHandle, @@ -459,6 +462,7 @@ impl MetricsPipeline { 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() @@ -476,6 +480,7 @@ impl MetricsPipeline { self.handles_opt = Some(MetricsPipelineHandles { source_mailbox, source_handle, + source_should_be_drained, doc_processor: doc_processor_handle, indexer: indexer_handle, packager: packager_handle, @@ -577,6 +582,12 @@ impl Handler for MetricsPipeline { // Nothing is running, so nothing is in flight. return Err(ActorExitStatus::Success); }; + if !handles.source_should_be_drained { + // The source settles nothing on a drain: tear the pipeline down + // right away, exactly as a kill would. + self.terminate().await; + return Err(ActorExitStatus::Success); + } self.drain_state .start( &handles.source_mailbox, diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index 33c5a5384cc..b01e82e1aa0 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -52,12 +52,17 @@ pub(crate) struct Spawn { /// 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 for acknowledgment-based sources), the source — and with it -/// the whole pipeline — exits with success on its own. Past the drain deadline -/// (commit timeout + [`DRAIN_GRACE_MARGIN`]), the pipeline gives up, kills its -/// actors, and still exits: delivery degrades to at-least-once. A draining -/// pipeline is never respawned. Fire-and-forget: watch the pipeline actor's -/// state to know when it is done. +/// (acknowledged), the source — and with it the whole pipeline — exits with +/// success on its own. Past the drain deadline (commit timeout + +/// [`DRAIN_GRACE_MARGIN`]), the pipeline gives up, kills its actors, and +/// still exits: delivery degrades to at-least-once. A draining pipeline is +/// never respawned. Fire-and-forget: watch the pipeline actor's state to know +/// when it is done. +/// +/// Only sources opting in through +/// [`crate::source::Source::should_be_drained`] actually drain: for the +/// others the pipeline kills its actors and exits immediately, keeping the +/// exact teardown semantics they had before drains existed. #[derive(Debug)] pub struct DrainPipeline; diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index 266ffbe05aa..27d8958d145 100644 --- a/quickwit/quickwit-indexing/src/source/mod.rs +++ b/quickwit/quickwit-indexing/src/source/mod.rs @@ -296,6 +296,15 @@ 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. From 2ef145e39c04dfb7ebf98b7b55a15ecffa10af91 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 19:25:13 +0200 Subject: [PATCH 08/41] feat: indexing pipeline drain --- .../src/actors/indexing_pipeline.rs | 2 -- .../src/actors/indexing_service.rs | 24 +++++++++---------- 2 files changed, 11 insertions(+), 15 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs index 52c73dc7923..f4ee53c0d66 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs @@ -58,8 +58,6 @@ use crate::split_store::IndexingSplitStore; struct IndexingPipelineHandles { source_mailbox: Mailbox, source_handle: ActorHandle, - /// Captured from [`crate::source::Source::should_be_drained`] at spawn: - /// a teardown drains the pipeline when true, kills it otherwise. source_should_be_drained: bool, doc_processor: ActorHandle, indexer: ActorHandle, diff --git a/quickwit/quickwit-indexing/src/actors/indexing_service.rs b/quickwit/quickwit-indexing/src/actors/indexing_service.rs index 80bf916dc60..bd773328f2b 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_service.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_service.rs @@ -113,7 +113,7 @@ pub struct IndexingService { pub(crate) storage_resolver: StorageResolver, indexing_pipelines: HashMap, /// Detached pipelines draining before their teardown. They exit on their - /// own (see `DrainPipeline`) and are reaped by the supervise loop. + /// 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. @@ -592,12 +592,12 @@ impl IndexingService { } } }); - // Draining pipelines exit on their own (see `DrainPipeline`): reap the + // Draining pipelines exit on their own, reap the // detached ones and complete the `DrainAllPipelines` replies whose - // pipelines are all gone. A paused pipeline is still alive, hence the - // `is_exit` checks rather than `is_running`. + // pipelines are all gone. self.draining_pipelines .retain(|pipeline_handle| !pipeline_handle.state().is_exit()); + // Are we waiting for a global drain? 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 { @@ -1026,12 +1026,9 @@ impl IndexingService { } } } - // Pipelines are drained before their teardown so their in-flight - // batches are flushed, published, and settled (acknowledged for - // acknowledgment-based sources) instead of dropped (see - // `DrainPipeline`). A drain can take up to the commit timeout, so it - // is not awaited here: the draining pipeline exits on its own and the - // supervise loop reaps its handle. + + // 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 { pipeline_handle.start_drain().await; self.draining_pipelines.push(pipeline_handle); @@ -1134,11 +1131,10 @@ impl IndexingService { } } -/// Drains every running indexing pipeline (see `DrainPipeline`): it is meant +/// 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 — including ones a -/// plan change already detached and started draining — have all exited, +/// them. The reply is deferred until the drained pipelines have all exited, /// without blocking the service. #[derive(Debug)] pub struct DrainAllPipelines; @@ -1158,6 +1154,8 @@ impl DeferableReplyHandler for IndexingService { pipeline_handle.start_drain().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.state().is_exit() { pipeline_uids.push(pipeline_handle.indexing_pipeline_id().pipeline_uid); From 88d1242e92522bbe8672d47e37322c018861350a Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 19:37:54 +0200 Subject: [PATCH 09/41] feat: indexing pipeline drain --- .../src/actors/indexing_service.rs | 4 +--- .../src/actors/pipeline_shared.rs | 14 ++++++-------- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/indexing_service.rs b/quickwit/quickwit-indexing/src/actors/indexing_service.rs index bd773328f2b..1f34e8f23e5 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_service.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_service.rs @@ -1035,9 +1035,7 @@ impl IndexingService { } // 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. The GC only removes queues of deleted indexes: a pipeline - // still draining such an index fails its publish and exits without - // respawning. + // indexes. if should_gc_ingest_api_queues && let Err(error) = self.run_ingest_api_queues_gc().await { warn!( %error, diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index b01e82e1aa0..31f003449de 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -52,17 +52,14 @@ pub(crate) struct Spawn { /// 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 — and with it the whole pipeline — exits with -/// success on its own. Past the drain deadline (commit timeout + +/// (acknowledged), the source exits on its own. Past the drain deadline (commit timeout + /// [`DRAIN_GRACE_MARGIN`]), the pipeline gives up, kills its actors, and /// still exits: delivery degrades to at-least-once. A draining pipeline is -/// never respawned. Fire-and-forget: watch the pipeline actor's state to know -/// when it is done. +/// 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, keeping the -/// exact teardown semantics they had before drains existed. +/// others the pipeline kills its actors and exits immediately. #[derive(Debug)] pub struct DrainPipeline; @@ -149,9 +146,10 @@ pub(crate) struct ActorPipeline> #[async_trait] impl PipelineHandle for ActorPipeline -where A: Actor +where + A: Actor + DeferableReplyHandler - + DeferableReplyHandler + + DeferableReplyHandler, { fn indexing_pipeline_id(&self) -> &IndexingPipelineId { &self.pipeline_id From b4b2b6054844d297e76f0c998e99be13071247ce Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 19:44:36 +0200 Subject: [PATCH 10/41] feat: indexing pipeline drain --- quickwit/quickwit-indexing/src/actors/indexing_service.rs | 6 +++--- quickwit/quickwit-indexing/src/actors/pipeline_shared.rs | 5 ++--- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/indexing_service.rs b/quickwit/quickwit-indexing/src/actors/indexing_service.rs index 1f34e8f23e5..c90998d2fec 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_service.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_service.rs @@ -592,12 +592,12 @@ impl IndexingService { } } }); + // Draining pipelines exit on their own, reap the - // detached ones and complete the `DrainAllPipelines` replies whose - // pipelines are all gone. + // detached ones and complete the `DrainAllPipelines` replies if all + // of their pipelines are gone. self.draining_pipelines .retain(|pipeline_handle| !pipeline_handle.state().is_exit()); - // Are we waiting for a global drain? 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 { diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index 31f003449de..bf1c8de43ef 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -146,10 +146,9 @@ pub(crate) struct ActorPipeline> #[async_trait] impl PipelineHandle for ActorPipeline -where - A: Actor +where A: Actor + DeferableReplyHandler - + DeferableReplyHandler, + + DeferableReplyHandler { fn indexing_pipeline_id(&self) -> &IndexingPipelineId { &self.pipeline_id From 8e311e7ffcf62c8f46eb0005e3d5d15fd400e90e Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 20:25:22 +0200 Subject: [PATCH 11/41] feat: indexing pipeline drain --- .../src/actors/indexing_pipeline.rs | 39 +++++------ .../src/actors/parquet_pipeline/pipeline.rs | 39 +++++------ .../src/actors/pipeline_shared.rs | 65 ++++++++++++------- 3 files changed, 79 insertions(+), 64 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs index f4ee53c0d66..d7f444e40c2 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs @@ -39,8 +39,8 @@ use tracing::{debug, error, info, instrument, warn}; use super::{DocProcessor, IndexSerializer, Indexer, MergePlanner, Packager}; use crate::SplitsUpdateMailbox; use crate::actors::pipeline_shared::{ - DrainPipeline, DrainState, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, SuperviseLoop, - wait_duration_before_retry, + DrainAction, DrainPipeline, DrainState, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, + SuperviseLoop, wait_duration_before_retry, }; use crate::actors::sequencer::Sequencer; use crate::actors::uploader::UploaderType; @@ -533,26 +533,23 @@ impl Handler for IndexingPipeline { _drain: DrainPipeline, _ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { - if self.drain_state.is_draining() { - return Ok(()); - } - let Some(handles) = &self.handles_opt else { - // Nothing is running, so nothing is in flight. - return Err(ActorExitStatus::Success); - }; - if !handles.source_should_be_drained { - // The source settles nothing on a drain: tear the pipeline down - // right away, exactly as a kill would. - self.terminate().await; - return Err(ActorExitStatus::Success); + let running_source_opt = self + .handles_opt + .as_ref() + .map(|handles| (&handles.source_mailbox, handles.source_should_be_drained)); + let commit_timeout = self.params.indexing_settings.commit_timeout(); + match self + .drain_state + .on_drain_request(running_source_opt, commit_timeout) + .await + { + DrainAction::KeepRunning => Ok(()), + DrainAction::Exit => Err(ActorExitStatus::Success), + DrainAction::TerminateAndExit => { + self.terminate().await; + Err(ActorExitStatus::Success) + } } - self.drain_state - .start( - &handles.source_mailbox, - self.params.indexing_settings.commit_timeout(), - ) - .await; - Ok(()) } } diff --git a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs index 9465f612788..cd50cec98af 100644 --- a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs @@ -46,8 +46,8 @@ use tracing::{debug, error, info, instrument, warn}; use super::{ParquetDocProcessor, ParquetIndexer, ParquetPackager, ParquetUploader}; use crate::actors::pipeline_shared::{ - DrainPipeline, DrainState, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, SuperviseLoop, - wait_duration_before_retry, + DrainAction, DrainPipeline, DrainState, SPAWN_PIPELINE_SEMAPHORE, SUPERVISE_INTERVAL, Spawn, + SuperviseLoop, wait_duration_before_retry, }; use crate::actors::sequencer::Sequencer; use crate::actors::{Publisher, UploaderType}; @@ -575,26 +575,23 @@ impl Handler for MetricsPipeline { _drain: DrainPipeline, _ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { - if self.drain_state.is_draining() { - return Ok(()); - } - let Some(handles) = &self.handles_opt else { - // Nothing is running, so nothing is in flight. - return Err(ActorExitStatus::Success); - }; - if !handles.source_should_be_drained { - // The source settles nothing on a drain: tear the pipeline down - // right away, exactly as a kill would. - self.terminate().await; - return Err(ActorExitStatus::Success); + let running_source_opt = self + .handles_opt + .as_ref() + .map(|handles| (&handles.source_mailbox, handles.source_should_be_drained)); + let commit_timeout = self.params.indexing_settings.commit_timeout(); + match self + .drain_state + .on_drain_request(running_source_opt, commit_timeout) + .await + { + DrainAction::KeepRunning => Ok(()), + DrainAction::Exit => Err(ActorExitStatus::Success), + DrainAction::TerminateAndExit => { + self.terminate().await; + Err(ActorExitStatus::Success) + } } - self.drain_state - .start( - &handles.source_mailbox, - self.params.indexing_settings.commit_timeout(), - ) - .await; - Ok(()) } } diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index bf1c8de43ef..a012ae35e83 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -14,10 +14,19 @@ //! Shared infrastructure for indexing pipeline supervisors (logs and metrics). -use std::time::Duration; +use std::time::{Duration, Instant}; +use async_trait::async_trait; +use quickwit_actors::{ + Actor, ActorExitStatus, ActorHandle, ActorState, DeferableReplyHandler, Health, Mailbox, + Observation, SendError, Supervisable, +}; +use quickwit_proto::indexing::IndexingPipelineId; use tokio::sync::Semaphore; +use crate::models::IndexingStatistics; +use crate::source::{AssignShards, Drain, SourceActor}; + pub(crate) const SUPERVISE_INTERVAL: Duration = Duration::from_secs(1); const MAX_RETRY_DELAY: Duration = Duration::from_mins(10); @@ -68,28 +77,51 @@ pub struct DrainPipeline; pub(crate) const DRAIN_GRACE_MARGIN: Duration = Duration::from_secs(30); /// Drain progress of a pipeline supervisor (see [`DrainPipeline`]): both -/// pipeline flavors share the deadline bookkeeping so they cannot diverge. +/// pipeline flavors share the deadline bookkeeping and the drain-request +/// handling so they cannot diverge. #[derive(Default)] pub(crate) struct DrainState { deadline_opt: Option, } +/// Follow-up required from the pipeline supervisor after +/// [`DrainState::on_drain_request`]. Killing the actors is only possible from +/// the supervisor itself, so it is returned as data instead of being executed +/// here. +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() } - /// Records the drain deadline and asks the source to drain. If the source - /// is already dead, the pipeline health check terminates the draining - /// pipeline instead of respawning it. - pub(crate) async fn start( + pub(crate) async fn on_drain_request( &mut self, - source_mailbox: &Mailbox, + running_source_opt: Option<(&Mailbox, bool)>, commit_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 is no worse than a kill. + return DrainAction::TerminateAndExit; + } self.deadline_opt = Some(Instant::now() + commit_timeout + DRAIN_GRACE_MARGIN); let _ = source_mailbox.send_message(Drain).await; + DrainAction::KeepRunning } /// Whether the draining pipeline ran out of time to settle: the caller @@ -105,18 +137,6 @@ impl DrainState { // Pipeline trait — type-erased handle for any indexing pipeline actor // --------------------------------------------------------------------------- -use std::time::Instant; - -use async_trait::async_trait; -use quickwit_actors::{ - Actor, ActorExitStatus, ActorHandle, ActorState, DeferableReplyHandler, Health, Mailbox, - Observation, SendError, Supervisable, -}; -use quickwit_proto::indexing::IndexingPipelineId; - -use crate::models::IndexingStatistics; -use crate::source::{AssignShards, Drain, SourceActor}; - /// Trait that abstracts over the concrete pipeline actor type /// (`IndexingPipeline` or `MetricsPipeline`). This allows `PipelineHandle` /// to hold a single `Box`. @@ -146,9 +166,10 @@ pub(crate) struct ActorPipeline> #[async_trait] impl PipelineHandle for ActorPipeline -where A: Actor +where + A: Actor + DeferableReplyHandler - + DeferableReplyHandler + + DeferableReplyHandler, { fn indexing_pipeline_id(&self) -> &IndexingPipelineId { &self.pipeline_id From c007b105d85c305ec67e9bef3bc5f16fa3131844 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 20:27:46 +0200 Subject: [PATCH 12/41] feat: indexing pipeline drain --- quickwit/quickwit-indexing/src/actors/pipeline_shared.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index a012ae35e83..97d86c6c733 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -84,10 +84,6 @@ pub(crate) struct DrainState { deadline_opt: Option, } -/// Follow-up required from the pipeline supervisor after -/// [`DrainState::on_drain_request`]. Killing the actors is only possible from -/// the supervisor itself, so it is returned as data instead of being executed -/// here. pub(crate) enum DrainAction { /// The drain is now (or was already) in progress: keep supervising. KeepRunning, @@ -116,7 +112,7 @@ impl DrainState { }; if !source_should_be_drained { // The source settles nothing on a drain: tearing the pipeline down - // right away is no worse than a kill. + // right away. return DrainAction::TerminateAndExit; } self.deadline_opt = Some(Instant::now() + commit_timeout + DRAIN_GRACE_MARGIN); From d6588f32ab328a0e4513618ce73880b24ebd8840 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 20:33:42 +0200 Subject: [PATCH 13/41] feat: indexing pipeline drain --- docs/configuration/node-config.md | 2 +- .../quickwit-config/src/node_config/mod.rs | 10 ++---- .../src/actors/indexing_pipeline.rs | 5 ++- .../src/actors/indexing_service.rs | 7 ++-- .../src/actors/parquet_pipeline/pipeline.rs | 5 ++- .../src/actors/pipeline_shared.rs | 33 +++++++++---------- 6 files changed, 29 insertions(+), 33 deletions(-) diff --git a/docs/configuration/node-config.md b/docs/configuration/node-config.md index 52f8c925059..89dc2b27211 100644 --- a/docs/configuration/node-config.md +++ b/docs/configuration/node-config.md @@ -244,7 +244,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` | Maximum amount of time to wait for the indexing pipelines to drain gracefully on shutdown before giving up. It must exceed the largest `commit_timeout_secs` plus the time to upload and publish the final splits, and the deployment's shutdown grace period (e.g. `terminationGracePeriodSeconds` on Kubernetes) must exceed it in turn. Can be overridden with the `QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT` environment variable. | `300s` | +| `shutdown_drain_timeout` | Time budget granted to each indexing pipeline to drain gracefully before giving up, on node shutdown and when the control plane tears down a pipeline. It must exceed the largest `commit_timeout_secs` plus the time to upload and publish the final splits, and the deployment's shutdown grace period (e.g. `terminationGracePeriodSeconds` on Kubernetes) must exceed it in turn. Can be overridden with the `QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT` environment variable. | `300s` | Example: diff --git a/quickwit/quickwit-config/src/node_config/mod.rs b/quickwit/quickwit-config/src/node_config/mod.rs index 148bcdfbfda..975d85faead 100644 --- a/quickwit/quickwit-config/src/node_config/mod.rs +++ b/quickwit/quickwit-config/src/node_config/mod.rs @@ -225,12 +225,6 @@ pub struct IndexerConfig { /// can't handle mixed prefix lengths. #[serde(default = "IndexerConfig::default_parquet_merge_use_streaming_engine")] pub parquet_merge_use_streaming_engine: bool, - /// Maximum amount of time to wait for the indexing pipelines to drain - /// gracefully on shutdown before giving up. Must exceed the largest commit - /// timeout plus the time to upload and publish the final splits, and the - /// deployment's shutdown grace period must exceed it in turn. Can be - /// overridden with the `QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT` environment - /// variable. #[serde(default = "IndexerConfig::default_shutdown_drain_timeout")] pub shutdown_drain_timeout: HumanDuration, } @@ -609,7 +603,9 @@ impl CacheConfig { fn deserialize_with_default<'de, D, const DEFAULT_CAPACITY: u64>( deserializer: D, ) -> Result - where D: Deserializer<'de> { + where + D: Deserializer<'de>, + { use serde_with::{DeserializeAs, FromInto, PickFirst, Same}; let mut cache_config: CacheConfig = diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs index d7f444e40c2..09fd0700639 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs @@ -530,17 +530,16 @@ impl Handler for IndexingPipeline { async fn handle( &mut self, - _drain: DrainPipeline, + 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)); - let commit_timeout = self.params.indexing_settings.commit_timeout(); match self .drain_state - .on_drain_request(running_source_opt, commit_timeout) + .on_drain_request(running_source_opt, drain.drain_timeout) .await { DrainAction::KeepRunning => Ok(()), diff --git a/quickwit/quickwit-indexing/src/actors/indexing_service.rs b/quickwit/quickwit-indexing/src/actors/indexing_service.rs index c90998d2fec..4a1241958df 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; +use std::time::Duration; use anyhow::Context; use async_trait::async_trait; @@ -118,6 +119,7 @@ pub struct IndexingService { /// Pending `DrainAllPipelines` replies, completed by the supervise loop /// once the tracked pipelines have exited. drain_all_waiters: Vec, + drain_timeout: Duration, latest_indexing_plan_id: IndexingPlanId, counters: IndexingServiceCounters, pub(crate) max_concurrent_split_uploads: usize, @@ -190,6 +192,7 @@ impl IndexingService { indexing_pipelines: Default::default(), draining_pipelines: Vec::new(), drain_all_waiters: Vec::new(), + drain_timeout: indexer_config.shutdown_drain_timeout(), latest_indexing_plan_id: String::new(), counters: Default::default(), max_concurrent_split_uploads: indexer_config.max_concurrent_split_uploads, @@ -1030,7 +1033,7 @@ 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 { - pipeline_handle.start_drain().await; + pipeline_handle.start_drain(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 @@ -1149,7 +1152,7 @@ impl DeferableReplyHandler for IndexingService { ) -> Result<(), ActorExitStatus> { let mut pipeline_uids: Vec = Vec::new(); for (pipeline_uid, pipeline_handle) in &self.indexing_pipelines { - pipeline_handle.start_drain().await; + pipeline_handle.start_drain(self.drain_timeout).await; pipeline_uids.push(*pipeline_uid); } // Some pipelines were potentially already draining when shutting down, e.g. in case diff --git a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs index cd50cec98af..aa8f75d6eaf 100644 --- a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs @@ -572,17 +572,16 @@ impl Handler for MetricsPipeline { async fn handle( &mut self, - _drain: DrainPipeline, + 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)); - let commit_timeout = self.params.indexing_settings.commit_timeout(); match self .drain_state - .on_drain_request(running_source_opt, commit_timeout) + .on_drain_request(running_source_opt, drain.drain_timeout) .await { DrainAction::KeepRunning => Ok(()), diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index 97d86c6c733..d95ed27862c 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -61,20 +61,17 @@ pub(crate) struct Spawn { /// 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 the drain deadline (commit timeout + -/// [`DRAIN_GRACE_MARGIN`]), the pipeline gives up, kills its actors, and -/// still exits: delivery degrades to at-least-once. A draining pipeline is -/// never respawned. +/// (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; - -/// Extra time granted to a draining pipeline on top of the commit timeout, -/// covering the upload and publication of the final splits. -pub(crate) const DRAIN_GRACE_MARGIN: Duration = Duration::from_secs(30); +pub struct DrainPipeline { + pub drain_timeout: Duration, +} /// Drain progress of a pipeline supervisor (see [`DrainPipeline`]): both /// pipeline flavors share the deadline bookkeeping and the drain-request @@ -102,7 +99,7 @@ impl DrainState { pub(crate) async fn on_drain_request( &mut self, running_source_opt: Option<(&Mailbox, bool)>, - commit_timeout: Duration, + drain_timeout: Duration, ) -> DrainAction { if self.is_draining() { return DrainAction::KeepRunning; @@ -115,7 +112,7 @@ impl DrainState { // right away. return DrainAction::TerminateAndExit; } - self.deadline_opt = Some(Instant::now() + commit_timeout + DRAIN_GRACE_MARGIN); + self.deadline_opt = Some(Instant::now() + drain_timeout); let _ = source_mailbox.send_message(Drain).await; DrainAction::KeepRunning } @@ -145,7 +142,7 @@ pub trait PipelineHandle: Send + Sync { fn check_health(&self, check_for_progress: bool) -> Health; async fn send_assign_shards(&self, message: AssignShards) -> Result<(), SendError>; /// See [`DrainPipeline`]. Fire-and-forget: the pipeline exits on its own. - async fn start_drain(&self); + async fn start_drain(&self, drain_timeout: Duration); async fn observe(&self) -> Observation; async fn join(self: Box) -> (ActorExitStatus, IndexingStatistics); async fn quit(self: Box) -> (ActorExitStatus, IndexingStatistics); @@ -162,10 +159,9 @@ pub(crate) struct ActorPipeline> #[async_trait] impl PipelineHandle for ActorPipeline -where - A: Actor +where A: Actor + DeferableReplyHandler - + DeferableReplyHandler, + + DeferableReplyHandler { fn indexing_pipeline_id(&self) -> &IndexingPipelineId { &self.pipeline_id @@ -192,10 +188,13 @@ where Ok(()) } - async fn start_drain(&self) { + async fn start_drain(&self, drain_timeout: Duration) { // The pipeline mailbox is unbounded, so this should not block; a send // error means the pipeline is already dead, hence already settled. - let _ = self.mailbox.send_message(DrainPipeline).await; + let _ = self + .mailbox + .send_message(DrainPipeline { drain_timeout }) + .await; } async fn observe(&self) -> Observation { From 946b518c90840499b33022b3299183400d59fa23 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 20:34:16 +0200 Subject: [PATCH 14/41] feat: indexing pipeline drain --- quickwit/quickwit-config/src/node_config/mod.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/quickwit/quickwit-config/src/node_config/mod.rs b/quickwit/quickwit-config/src/node_config/mod.rs index 975d85faead..191a74e9a54 100644 --- a/quickwit/quickwit-config/src/node_config/mod.rs +++ b/quickwit/quickwit-config/src/node_config/mod.rs @@ -603,9 +603,7 @@ impl CacheConfig { fn deserialize_with_default<'de, D, const DEFAULT_CAPACITY: u64>( deserializer: D, ) -> Result - where - D: Deserializer<'de>, - { + where D: Deserializer<'de> { use serde_with::{DeserializeAs, FromInto, PickFirst, Same}; let mut cache_config: CacheConfig = From b615e9fee099737db3439d69d2c5a961af994655 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 20:34:35 +0200 Subject: [PATCH 15/41] feat: indexing pipeline drain --- .../quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs index aa8f75d6eaf..ef06e2eef30 100644 --- a/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs @@ -60,8 +60,6 @@ use crate::source::{ struct MetricsPipelineHandles { source_mailbox: Mailbox, source_handle: ActorHandle, - /// Captured from [`crate::source::Source::should_be_drained`] at spawn: - /// a teardown drains the pipeline when true, kills it otherwise. source_should_be_drained: bool, doc_processor: ActorHandle, indexer: ActorHandle, From cd9b6ccde02f6d711b1ed289475674a3163e193a Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 20:43:59 +0200 Subject: [PATCH 16/41] feat: indexing pipeline drain --- quickwit/quickwit-indexing/src/actors/pipeline_shared.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index d95ed27862c..977857b7b60 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -73,9 +73,7 @@ pub struct DrainPipeline { pub drain_timeout: Duration, } -/// Drain progress of a pipeline supervisor (see [`DrainPipeline`]): both -/// pipeline flavors share the deadline bookkeeping and the drain-request -/// handling so they cannot diverge. +/// Drain progress of a pipeline supervisor (see [`DrainPipeline`]). #[derive(Default)] pub(crate) struct DrainState { deadline_opt: Option, @@ -141,7 +139,6 @@ pub trait PipelineHandle: Send + Sync { fn last_observation(&self) -> IndexingStatistics; fn check_health(&self, check_for_progress: bool) -> Health; async fn send_assign_shards(&self, message: AssignShards) -> Result<(), SendError>; - /// See [`DrainPipeline`]. Fire-and-forget: the pipeline exits on its own. async fn start_drain(&self, drain_timeout: Duration); async fn observe(&self) -> Observation; async fn join(self: Box) -> (ActorExitStatus, IndexingStatistics); From a23df941fc3640f472c81ff74730379a4c41d539 Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 11:15:47 +0200 Subject: [PATCH 17/41] fix(indexers): disable new pipeline spawning when draining --- .../src/actors/indexing_service.rs | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/quickwit/quickwit-indexing/src/actors/indexing_service.rs b/quickwit/quickwit-indexing/src/actors/indexing_service.rs index 4a1241958df..0da46b1b393 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_service.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_service.rs @@ -42,6 +42,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, @@ -120,6 +121,9 @@ pub struct IndexingService { /// 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, pub(crate) max_concurrent_split_uploads: usize, @@ -193,6 +197,7 @@ impl IndexingService { 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, @@ -257,6 +262,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(), @@ -860,6 +873,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!( @@ -1150,6 +1172,8 @@ impl DeferableReplyHandler for IndexingService { 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 { pipeline_handle.start_drain(self.drain_timeout).await; From 90fd1a10e8de275c9ef61e5f5f0a839e6378291a Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 11:34:22 +0200 Subject: [PATCH 18/41] fix(indexers): drive the drain supervise loop at 1s rather than 30s) --- .../src/actors/indexing_service.rs | 117 +++++++++++++++++- 1 file changed, 114 insertions(+), 3 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/indexing_service.rs b/quickwit/quickwit-indexing/src/actors/indexing_service.rs index 0da46b1b393..0b58828e35d 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_service.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_service.rs @@ -60,7 +60,7 @@ use tokio::sync::Semaphore; use tracing::{debug, error, info, warn}; use super::merge_pipeline::{MergePipeline, MergePipelineParams}; -use super::pipeline_shared::{ActorPipeline, PipelineHandle}; +use super::pipeline_shared::{ActorPipeline, PipelineHandle, SUPERVISE_INTERVAL}; use super::{FinishPendingMergesAndShutdownPipeline, MergePlanner, MergeSchedulerService}; use crate::docs_clustering::Fingerprinter; use crate::models::{DetachIndexingPipeline, DetachMergePipeline, ObservePipeline, SpawnPipeline}; @@ -1170,7 +1170,7 @@ impl DeferableReplyHandler for IndexingService { &mut self, _msg: DrainAllPipelines, reply: impl FnOnce(()) + Send + Sync + 'static, - _ctx: &ActorContext, + ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { self.draining = true; @@ -1194,8 +1194,40 @@ impl DeferableReplyHandler for IndexingService { 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(()) } } @@ -1336,7 +1368,7 @@ mod tests { use std::path::Path; use std::time::Duration; - use quickwit_actors::{HEARTBEAT, Health, ObservationType, Universe}; + use quickwit_actors::{AskError, HEARTBEAT, Health, ObservationType, Universe}; use quickwit_cluster::{ChitchatTransport, create_cluster_for_test}; use quickwit_common::ServiceStream; use quickwit_common::rand::append_random_suffix; @@ -1496,6 +1528,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(); From 94a28c266e88d8693e60a473f37f7f9afd143e1e Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 11:44:46 +0200 Subject: [PATCH 19/41] fix(doc): rephrase the indexer shutdown drain timeout description --- docs/configuration/node-config.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/configuration/node-config.md b/docs/configuration/node-config.md index 89dc2b27211..2295caccb7b 100644 --- a/docs/configuration/node-config.md +++ b/docs/configuration/node-config.md @@ -244,7 +244,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 giving up, on node shutdown and when the control plane tears down a pipeline. It must exceed the largest `commit_timeout_secs` plus the time to upload and publish the final splits, and the deployment's shutdown grace period (e.g. `terminationGracePeriodSeconds` on Kubernetes) must exceed it in turn. Can be overridden with the `QW_INDEXER_SHUTDOWN_DRAIN_TIMEOUT` environment variable. | `300s` | +| `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` | Example: From 169e7a8ed75721c45d71e9f324ec90513773b2d6 Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 12:03:17 +0200 Subject: [PATCH 20/41] fix(doc): wait for DrainPipeline so that non-draining pipeline are killed synchronously to match previous semantics --- .../src/actors/pipeline_shared.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index 977857b7b60..2575d84fb0a 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -156,9 +156,10 @@ pub(crate) struct ActorPipeline> #[async_trait] impl PipelineHandle for ActorPipeline -where A: Actor +where + A: Actor + DeferableReplyHandler - + DeferableReplyHandler + + DeferableReplyHandler, { fn indexing_pipeline_id(&self) -> &IndexingPipelineId { &self.pipeline_id @@ -186,12 +187,11 @@ where A: Actor } async fn start_drain(&self, drain_timeout: Duration) { - // The pipeline mailbox is unbounded, so this should not block; a send - // error means the pipeline is already dead, hence already settled. - let _ = self - .mailbox - .send_message(DrainPipeline { drain_timeout }) - .await; + // 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 _ = self.mailbox.ask(DrainPipeline { drain_timeout }).await; } async fn observe(&self) -> Observation { From f77e519aa269fc57b1bef856d305595909020f23 Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 12:48:41 +0200 Subject: [PATCH 21/41] fix(doc): wait for DrainPipeline so that non-draining pipeline are killed synchronously to match previous semantics --- quickwit/quickwit-indexing/src/actors/pipeline_shared.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs index 2575d84fb0a..12ee34f6fb2 100644 --- a/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs +++ b/quickwit/quickwit-indexing/src/actors/pipeline_shared.rs @@ -156,10 +156,9 @@ pub(crate) struct ActorPipeline> #[async_trait] impl PipelineHandle for ActorPipeline -where - A: Actor +where A: Actor + DeferableReplyHandler - + DeferableReplyHandler, + + DeferableReplyHandler { fn indexing_pipeline_id(&self) -> &IndexingPipelineId { &self.pipeline_id From 39a89c8bf51d0b386790feffb125bd11a22a895d Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 14:52:20 +0200 Subject: [PATCH 22/41] fix: config --- docs/configuration/source-config.md | 2 +- .../src/source/nats_source.rs | 51 +++++-------------- 2 files changed, 15 insertions(+), 38 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 41f5ad26787..82d32c05393 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -191,7 +191,7 @@ A tutorial is available [here](/docs/ingest-data/nats.md). The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. -Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged — with server-confirmed acknowledgments — before the pipeline stops, so nothing is indexed twice; messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. `max_ack_pending` bounds the messages in flight across all pipelines of the consumer, and therefore the aggregate throughput. +Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. `max_ack_pending` bounds the messages in flight across all pipelines of the consumer, and therefore the aggregate throughput. **Scaling** diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index aeb0acc66fd..4cd165c49bf 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -15,20 +15,9 @@ //! A source consuming a NATS JetStream stream through a pre-provisioned //! durable consumer. //! -//! The consumer is only ever fetched — never created, updated, nor deleted — -//! so its lifecycle, subject filters, deliver policy, and ack tuning -//! (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The -//! consumer must use the explicit ack policy: each message is acknowledged — -//! with a server-confirmed ack — once the split containing it is published. -//! On a planned teardown (node shutdown, pipeline reassignment) the pipeline -//! is drained first: the in-flight messages are published and acknowledged, -//! and the prefetched but unprocessed ones are negatively acknowledged for -//! immediate redelivery — so delivery is exactly-once. A crash degrades it to -//! at-least-once: the unacknowledged messages — including ones whose split -//! ends up published — are redelivered after the consumer's `ack_wait` and -//! indexed again. The consumer's ack floor is the resume point; the metastore -//! checkpoint only carries synthetic positions used to correlate published -//! batches with the acknowledgments they release. +//! The consumer is only ever fetched so its lifecycle, subject filters, +//! deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to +//! whoever provisioned it. //! //! Several indexing pipelines can share the consumer: NATS load-balances the //! messages across them, so scaling is a plain `num_pipelines` update. Being @@ -98,26 +87,6 @@ pub struct NatsSourceState { pub num_invalid_messages: u64, } -/// Source flavor binding to a pre-provisioned durable consumer. -/// -/// The consumer is only ever fetched — never created, updated, nor deleted — -/// so its lifecycle, subject filters, deliver policy, and ack tuning -/// (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. Each -/// message is acknowledged once the split containing it is published -/// ([`Source::suggest_truncate`]); see the module documentation for the -/// delivery semantics. -/// -/// The metastore checkpoint only carries synthetic positions (a per-pipeline -/// partition and a delivery counter) used to correlate published batches with -/// the acknowledgments they release; the resume point is the consumer's ack -/// floor, not the checkpoint. -/// -/// The source acknowledges inline in [`Source::suggest_truncate`], so once a -/// drain reports completion (see [`Source::is_drained`]) every ack has been -/// confirmed by the server. If the source actor is instead killed with batches -/// still in the pipeline, their `suggest_truncate` notifications are lost and -/// the unacknowledged messages — including ones whose split ends up published — -/// are redelivered after the consumer's `ack_wait`. pub struct NatsSource { source_runtime: SourceRuntime, source_params: NatsSourceParams, @@ -308,6 +277,10 @@ impl Source for NatsSource { Ok(()) } + fn should_be_drained(&self) -> bool { + true + } + fn is_drained(&self) -> bool { // Every delivered message is pending from `process_message` until the // server confirmed its ack: an empty map means everything delivered so @@ -442,8 +415,7 @@ async fn ack_up_to( debug!(num_acks, "acked published messages"); } -/// Fetches the pre-provisioned durable consumer. Read-only by contract: the -/// consumer is never created nor updated. +/// Fetches the pre-provisioned durable consumer. async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, @@ -901,7 +873,12 @@ mod nats_broker_tests { tokio::time::sleep(Duration::from_millis(100)).await; } - pipeline_mailbox.send_message(DrainPipeline).await.unwrap(); + pipeline_mailbox + .send_message(DrainPipeline { + drain_timeout: Duration::from_secs(60), + }) + .await + .unwrap(); // The pipeline exits on its own once the in-flight batches are // published and their acks flushed. let (exit_status, _statistics) = pipeline_handle.join().await; From 3b5e971749118857ac74496297fa3b596a2675ab Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 14:59:27 +0200 Subject: [PATCH 23/41] fix: config --- .../src/source/nats_source.rs | 23 ++----------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 4cd165c49bf..9c243e70be0 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -297,15 +297,8 @@ impl Source for NatsSource { _exit_status: &ActorExitStatus, _ctx: &SourceContext, ) -> anyhow::Result<()> { - // The durable consumer itself is left untouched. Any ack still pending - // at this point was not settled by a drain: its message is redelivered - // after the consumer's `ack_wait`. - // - // Messages the client prefetched but that were never processed are - // negatively acknowledged so the server redelivers them right after - // this connection is gone, instead of after `ack_wait`. Only what is - // already buffered is NAK'd: the loop stops at the first pending poll, - // so a refill in flight cannot keep it alive. + // Any ack still pending at this point was not completed by a drain: + // its message is redelivered after the consumer's `ack_wait`. while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { let Ok(message) = message_res else { continue; @@ -367,13 +360,6 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { /// Sends server-confirmed acknowledgments ("double acks") for the messages up /// to `published_up_to`. -/// -/// Entries only leave the pending map once the server confirmed their ack, so -/// `is_drained` never reports progress that could still be lost: a -/// fire-and-forget ack can be dropped by a reconnection even after a -/// successful client-side flush. A retried ack for a message the server -/// already saw acknowledged is simply confirmed again. Whatever could not be -/// acknowledged is redelivered after the consumer's `ack_wait`. async fn ack_up_to( nats_client: &async_nats::Client, pending_acks: &mut BTreeMap, @@ -415,7 +401,6 @@ async fn ack_up_to( debug!(num_acks, "acked published messages"); } -/// Fetches the pre-provisioned durable consumer. async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, @@ -510,8 +495,6 @@ async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result anyhow::Result<(async_nats::Client, PullConsumer)> { @@ -525,8 +508,6 @@ async fn connect_and_fetch_consumer( Ok((nats_client, consumer)) } -/// Checks whether we can connect to the NATS servers and find the JetStream -/// stream and the pre-provisioned consumer. pub(crate) async fn check_connectivity(params: &NatsSourceParams) -> anyhow::Result<()> { connect_and_fetch_consumer(params).await?; Ok(()) From 2d9c42cc9de9d20dce3441aa984d26917049729e Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 15:31:16 +0200 Subject: [PATCH 24/41] feat: add tests --- .../src/source/nats_source.rs | 352 +++++++++++++++++- 1 file changed, 332 insertions(+), 20 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 9c243e70be0..0cb0c3cd10b 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -518,7 +518,7 @@ mod tests { use super::*; #[test] - fn test_extract_trace_context_from_headers() { + fn extract_trace_context_from_headers() { use opentelemetry::propagation::TextMapPropagator; use opentelemetry_sdk::propagation::TraceContextPropagator; @@ -551,6 +551,32 @@ mod tests { .clone(); assert!(!span_context.is_valid()); } + + #[test] + fn transient_stream_error_classification() { + let transient_kinds = [ + MessagesErrorKind::MissingHeartbeat, + MessagesErrorKind::Pull, + MessagesErrorKind::NoResponders, + ]; + for kind in transient_kinds { + assert!( + is_transient_stream_error(&MessagesError::new(kind)), + "`{kind:?}` should be retried, not kill the pipeline" + ); + } + let terminal_kinds = [ + MessagesErrorKind::ConsumerDeleted, + MessagesErrorKind::PushBasedConsumer, + MessagesErrorKind::Other, + ]; + for kind in terminal_kinds { + assert!( + !is_transient_stream_error(&MessagesError::new(kind)), + "`{kind:?}` should kill the pipeline" + ); + } + } } #[cfg(all(test, feature = "nats-broker-tests"))] @@ -562,6 +588,7 @@ mod nats_broker_tests { use bytes::Bytes; use quickwit_actors::{ActorHandle, Inbox, Universe}; use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; + use quickwit_metastore::checkpoint::SourceCheckpointDelta; use quickwit_metastore::metastore_for_test; use quickwit_proto::metastore::MetastoreServiceClient; use quickwit_proto::types::IndexUid; @@ -636,10 +663,15 @@ mod nats_broker_tests { (source_handle, doc_processor_inbox) } + /// Waits until the source reports at least `num_expected` processed + /// messages, panicking past `timeout` — kept tight where the delivery + /// delay itself is the assertion (NAK vs `ack_wait` redelivery). async fn wait_for_processed_messages( source_handle: &ActorHandle, num_expected: u64, + timeout: Duration, ) { + let deadline = Instant::now() + timeout; loop { let observation = source_handle.observe().await; let num_messages_processed = observation @@ -651,8 +683,34 @@ mod nats_broker_tests { if num_messages_processed >= num_expected { return; } + assert!( + Instant::now() < deadline, + "source did not process {num_expected} messages within {timeout:?}" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + /// Waits until the consumer reports the expected ack state, panicking + /// after ~10s. + async fn wait_for_consumer_ack_state( + consumer: &mut PullConsumer, + expected_num_ack_pending: usize, + expected_ack_floor: u64, + ) { + for _ in 0..100 { + let consumer_info = consumer.info().await.unwrap(); + if consumer_info.num_ack_pending == expected_num_ack_pending + && consumer_info.ack_floor.stream_sequence == expected_ack_floor + { + return; + } tokio::time::sleep(Duration::from_millis(100)).await; } + panic!( + "consumer never reached {expected_num_ack_pending} pending acks with ack floor \ + {expected_ack_floor}" + ); } fn merge_doc_batches(batches: Vec) -> RawDocBatch { @@ -688,10 +746,11 @@ mod nats_broker_tests { /// Provisions the durable consumer the way an operator would: the source /// itself only ever fetches it. - async fn provision_durable_consumer( + async fn provision_durable_consumer_with_ack_wait( jetstream_ctx: &jetstream::Context, stream: &str, consumer_name: &str, + ack_wait: Duration, ) { jetstream_ctx .create_consumer_on_stream( @@ -699,8 +758,7 @@ mod nats_broker_tests { name: Some(consumer_name.to_string()), durable_name: Some(consumer_name.to_string()), ack_policy: AckPolicy::Explicit, - // Long enough for the tests to never hit a redelivery. - ack_wait: Duration::from_secs(300), + ack_wait, ..Default::default() }, stream, @@ -709,8 +767,23 @@ mod nats_broker_tests { .unwrap(); } + async fn provision_durable_consumer( + jetstream_ctx: &jetstream::Context, + stream: &str, + consumer_name: &str, + ) { + // An `ack_wait` long enough for the test to never hit a redelivery. + provision_durable_consumer_with_ack_wait( + jetstream_ctx, + stream, + consumer_name, + Duration::from_secs(300), + ) + .await; + } + #[tokio::test] - async fn test_durable_mode_ingestion_and_ack() { + async fn durable_mode_ingestion_and_ack() { let universe = Universe::with_accelerated_time(); let metastore = metastore_for_test(); let stream = append_random_suffix("test-nats-source--durable--stream"); @@ -728,7 +801,7 @@ mod nats_broker_tests { let (source_handle, doc_processor_inbox) = create_source_actor(&universe, metastore, index_uid, source_config).await; - wait_for_processed_messages(&source_handle, 10).await; + wait_for_processed_messages(&source_handle, 10, Duration::from_secs(60)).await; let batches: Vec = doc_processor_inbox.drain_for_test_typed(); let batch = merge_doc_batches(batches); @@ -750,17 +823,7 @@ mod nats_broker_tests { .get_consumer_from_stream(consumer_name, stream.as_str()) .await .unwrap(); - let mut acked = false; - for _ in 0..100 { - let consumer_info = consumer.info().await.unwrap(); - if consumer_info.num_ack_pending == 0 { - assert_eq!(consumer_info.ack_floor.stream_sequence, 10); - acked = true; - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - assert!(acked, "messages should be acked after truncation"); + wait_for_consumer_ack_state(&mut consumer, 0, 10).await; source_handle.quit().await; jetstream_ctx.delete_stream(&stream).await.unwrap(); @@ -771,7 +834,7 @@ mod nats_broker_tests { /// must publish the in-flight batches and flush their acks BEFORE the /// drain replies — the exactly-once guarantee on planned teardowns. #[tokio::test] - async fn test_durable_mode_graceful_drain_acks_before_teardown() { + async fn durable_mode_graceful_drain_acks_before_teardown() { use quickwit_actors::Universe; use quickwit_common::temp_dir::TempDirectory; use quickwit_config::IndexingSettings; @@ -880,7 +943,7 @@ mod nats_broker_tests { } #[tokio::test] - async fn test_durable_mode_load_balancing() { + async fn durable_mode_load_balancing() { let universe = Universe::with_accelerated_time(); let metastore = metastore_for_test(); let stream = append_random_suffix("test-nats-source--durable-lb--stream"); @@ -950,7 +1013,256 @@ mod nats_broker_tests { } #[tokio::test] - async fn test_durable_mode_missing_consumer() { + async fn durable_mode_crash_redelivery_after_ack_wait() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-crash--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-crash-consumer"; + provision_durable_consumer_with_ack_wait( + &jetstream_ctx, + &stream, + consumer_name, + Duration::from_secs(3), + ) + .await; + + let subject = format!("{stream}.logs"); + let expected_docs = publish_docs(&jetstream_ctx, &subject, 0..10).await; + + let index_id = append_random_suffix("test-nats-source--durable-crash--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle_1, _doc_processor_inbox_1) = create_source_actor( + &universe, + metastore.clone(), + index_uid.clone(), + source_config.clone(), + ) + .await; + wait_for_processed_messages(&source_handle_1, 10, Duration::from_secs(60)).await; + // No `SuggestTruncate` was sent: quitting here loses the acks of the + // processed messages, like a crash would. + source_handle_1.quit().await; + + let (source_handle_2, doc_processor_inbox_2) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + wait_for_processed_messages(&source_handle_2, 10, Duration::from_secs(30)).await; + + let batches: Vec = doc_processor_inbox_2.drain_for_test_typed(); + // A message can be redelivered more than once if the test straddles + // several `ack_wait` windows. + let mut redelivered_docs = merge_doc_batches(batches).docs; + redelivered_docs.dedup(); + assert_eq!(redelivered_docs, expected_docs); + + source_handle_2.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + #[tokio::test] + async fn durable_mode_nak_redelivers_prefetched_messages_promptly() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-nak--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-nak-consumer"; + // The `ack_wait` is much longer than the test: only a NAK can explain + // a prompt redelivery. + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let index_id = append_random_suffix("test-nats-source--durable-nak--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + // The source is driven by hand rather than through a `SourceActor`: + // an idle actor keeps polling, so it would process the messages + // instead of leaving them prefetched. + let source_runtime = SourceRuntimeBuilder::new(index_uid.clone(), source_config.clone()) + .with_metastore(metastore.clone()) + .build(); + let mut source = quickwit_supported_sources() + .load_source(source_runtime) + .await + .unwrap(); + let (source_mailbox, _source_inbox) = universe.create_test_mailbox::(); + let (doc_processor_mailbox, _doc_processor_inbox) = + universe.create_test_mailbox::(); + let source_sink = SourceSink::from(doc_processor_mailbox); + let (observable_state_tx, _observable_state_rx) = + tokio::sync::watch::channel(JsonValue::Null); + let ctx: SourceContext = + quickwit_actors::ActorContext::for_test(&universe, source_mailbox, observable_state_tx); + + // A first empty emit issues the pull request: the messages published + // next are prefetched into the client buffer but never processed. + source.emit_batches(&source_sink, &ctx).await.unwrap(); + let subject = format!("{stream}.logs"); + let expected_docs = publish_docs(&jetstream_ctx, &subject, 0..10).await; + tokio::time::sleep(Duration::from_secs(1)).await; + assert_eq!( + source + .observable_state() + .get("num_messages_processed") + .unwrap(), + &json!(0), + "the messages must be prefetched, not processed" + ); + + source.finalize(&ActorExitStatus::Quit, &ctx).await.unwrap(); + drop(source); + + let (source_handle_2, doc_processor_inbox_2) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + wait_for_processed_messages(&source_handle_2, 10, Duration::from_secs(20)).await; + + let batches: Vec = doc_processor_inbox_2.drain_for_test_typed(); + assert_eq!(merge_doc_batches(batches).docs, expected_docs); + + source_handle_2.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + /// A truncate covering only part of the delivered messages must ack + /// exactly up to its position and leave the rest pending. + #[tokio::test] + async fn durable_mode_partial_truncate_acks_up_to_position() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-partial--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-partial-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let subject = format!("{stream}.logs"); + publish_docs(&jetstream_ctx, &subject, 0..10).await; + + let index_id = append_random_suffix("test-nats-source--durable-partial--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle, doc_processor_inbox) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + wait_for_processed_messages(&source_handle, 10, Duration::from_secs(60)).await; + + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + let batch = merge_doc_batches(batches); + let partition_id = batch.checkpoint_delta.partitions().next().unwrap().clone(); + + // Sequential publishes on a single pipeline: delivery counters map + // one-to-one to stream sequences, so truncating up to position 5 must + // ack stream sequences 1 to 5. + let partial_checkpoint = SourceCheckpointDelta::from_partition_delta( + partition_id, + Position::Beginning, + Position::offset(5u64), + ) + .unwrap() + .get_source_checkpoint(); + source_handle + .mailbox() + .send_message(SuggestTruncate(partial_checkpoint)) + .await + .unwrap(); + + let mut consumer: PullConsumer = jetstream_ctx + .get_consumer_from_stream(consumer_name, stream.as_str()) + .await + .unwrap(); + wait_for_consumer_ack_state(&mut consumer, 5, 5).await; + let observation = source_handle.observe().await; + assert_eq!( + observation.state.get("num_pending_acks").unwrap(), + &json!(5) + ); + + // The remainder is released by the next truncate. + let full_checkpoint = batch.checkpoint_delta.get_source_checkpoint(); + source_handle + .mailbox() + .send_message(SuggestTruncate(full_checkpoint)) + .await + .unwrap(); + wait_for_consumer_ack_state(&mut consumer, 0, 10).await; + + source_handle.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + /// An empty payload is counted invalid and skipped from the batch, but + /// its message must still be acknowledged on truncate: a poison message + /// must not wedge the consumer nor a drain. + #[tokio::test] + async fn durable_mode_empty_message_acked_and_counted_invalid() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-empty--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-empty-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let subject = format!("{stream}.logs"); + let valid_docs = [ + json!({ "id": 0, "subject": subject }).to_string(), + json!({ "id": 2, "subject": subject }).to_string(), + ]; + for payload in [ + Bytes::from(valid_docs[0].clone()), + Bytes::new(), + Bytes::from(valid_docs[1].clone()), + ] { + jetstream_ctx + .publish(subject.clone(), payload) + .await + .unwrap() + .await + .unwrap(); + } + + let index_id = append_random_suffix("test-nats-source--durable-empty--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle, doc_processor_inbox) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + wait_for_processed_messages(&source_handle, 3, Duration::from_secs(60)).await; + + let observation = source_handle.observe().await; + assert_eq!( + observation.state.get("num_invalid_messages").unwrap(), + &json!(1) + ); + + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + let batch = merge_doc_batches(batches); + assert_eq!(batch.docs, valid_docs); + + // The checkpoint covers the empty message too, and truncating acks it + // along with the valid ones. + source_handle + .mailbox() + .send_message(SuggestTruncate( + batch.checkpoint_delta.get_source_checkpoint(), + )) + .await + .unwrap(); + let mut consumer: PullConsumer = jetstream_ctx + .get_consumer_from_stream(consumer_name, stream.as_str()) + .await + .unwrap(); + wait_for_consumer_ack_state(&mut consumer, 0, 3).await; + + source_handle.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + #[tokio::test] + async fn durable_mode_missing_consumer() { let metastore = metastore_for_test(); let stream = append_random_suffix("test-nats-source--durable-missing--stream"); let jetstream_ctx = setup_nats_stream(&stream).await; From 8ac7c8d6f7a2f4580c0d062fd80394cd56935bb0 Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 15:59:59 +0200 Subject: [PATCH 25/41] feat: add tests --- quickwit/quickwit-config/src/source_config/serialize.rs | 3 ++- quickwit/quickwit-serve/src/index_api/rest_handler.rs | 3 +-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/quickwit/quickwit-config/src/source_config/serialize.rs b/quickwit/quickwit-config/src/source_config/serialize.rs index 25e92fe0a28..3b412a9ab19 100644 --- a/quickwit/quickwit-config/src/source_config/serialize.rs +++ b/quickwit/quickwit-config/src/source_config/serialize.rs @@ -132,7 +132,8 @@ impl SourceConfigForSerialization { _ => { if self.num_pipelines > 1 { bail!( - "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources. open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources.\ + open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" ); } } diff --git a/quickwit/quickwit-serve/src/index_api/rest_handler.rs b/quickwit/quickwit-serve/src/index_api/rest_handler.rs index bc4c8d9b105..09cd6eb09ea 100644 --- a/quickwit/quickwit-serve/src/index_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/index_api/rest_handler.rs @@ -1100,8 +1100,7 @@ mod tests { assert_eq!(resp.status(), 400); let body = std::str::from_utf8(resp.body()).unwrap(); assert!(body.contains( - "Quickwit currently supports multiple pipelines only for GCP PubSub or Kafka \ - sources" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources" )); } { From 928f4a9a5c02d9d1b870b1ee54387eb5a6d63346 Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 16:34:46 +0200 Subject: [PATCH 26/41] fix: fmt --- quickwit/quickwit-serve/src/index_api/rest_handler.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/quickwit/quickwit-serve/src/index_api/rest_handler.rs b/quickwit/quickwit-serve/src/index_api/rest_handler.rs index 09cd6eb09ea..1324508ff68 100644 --- a/quickwit/quickwit-serve/src/index_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/index_api/rest_handler.rs @@ -1100,7 +1100,8 @@ mod tests { assert_eq!(resp.status(), 400); let body = std::str::from_utf8(resp.body()).unwrap(); assert!(body.contains( - "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or \ + NATS sources" )); } { From a19e18f8439d769bb170642ff0ce8129f91cdf39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Tue, 8 Sep 2026 14:11:18 +0000 Subject: [PATCH 27/41] fix: nats --- .../src/source/nats_source.rs | 124 +++++++++++++++--- 1 file changed, 105 insertions(+), 19 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 0cb0c3cd10b..61f2cfbc598 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -126,11 +126,22 @@ impl NatsSource { source_id=%source_runtime.source_id(), stream=%source_params.stream, %consumer_name, + pull_max_bytes_per_batch=%pull_max_bytes_per_batch(), + pull_max_messages_per_batch=%pull_max_messages_per_batch(), "starting NATS source" ); let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; let message_stream = consumer + .stream() + .max_messages_per_batch(pull_max_messages_per_batch()) + .max_bytes_per_batch(pull_max_bytes_per_batch()) + // `stream()` defaults the heartbeat to zero, i.e. disabled, unlike + // `messages()`. It has to be set back, or a stalled subscription is + // never reported and `is_transient_stream_error` has nothing to act + // on. + .heartbeat(PULL_IDLE_HEARTBEAT) + .expires(PULL_EXPIRES) .messages() .await .context("failed to subscribe to NATS consumer messages")?; @@ -299,19 +310,40 @@ impl Source for NatsSource { ) -> anyhow::Result<()> { // Any ack still pending at this point was not completed by a drain: // its message is redelivered after the consumer's `ack_wait`. + let mut num_naks = 0usize; while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { - let Ok(message) = message_res else { - continue; + let message = match message_res { + Ok(message) => message, + Err(error) => { + warn!(%error, "failed to pull a prefetched NATS message to negatively acknowledge it"); + continue; + } }; let Some(ack_subject) = message.message.reply else { + warn!("prefetched NATS message carries no reply subject to negatively acknowledge it on"); continue; }; - let _ = self + if let Err(error) = self .nats_client .publish(ack_subject, AckKind::Nak(Some(NAK_REDELIVERY_DELAY)).into()) - .await; + .await + { + warn!(%error, "failed to negatively acknowledge a prefetched NATS message"); + continue; + } + num_naks += 1; + } + // The negative acknowledgments above are plain publishes, and `drain` + // returns as soon as its command is queued rather than once the + // connection has been written out. Without this flush they can be lost + // with the connection, and their messages then wait a full `ack_wait` + // instead of being redelivered to a surviving pipeline immediately. + if let Err(error) = self.nats_client.flush().await { + warn!(%error, num_naks, "failed to flush NATS negative acknowledgments"); + } + if let Err(error) = self.nats_client.drain().await { + warn!(%error, "failed to drain the NATS connection"); } - let _ = self.nats_client.drain().await; Ok(()) } @@ -332,7 +364,55 @@ impl Source for NatsSource { /// Bounds the concurrent server-confirmed ack requests: enough to hide the /// round-trip latency without flooding the connection. -const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; +/// Caps a pull batch in bytes, which the message-count cap alone cannot do. +/// +/// The server buffers a pull response on the connection's outbound queue, and a +/// connection that exceeds `max_pending` (64 MiB by default) is declared a slow +/// consumer and closed. Messages already written to it are lost while still +/// counted as delivered, so they only come back after `ack_wait` expires: with +/// 1 MiB messages, an uncapped batch of 200 asks the server for 200 MiB and the +/// source stalls for a whole `ack_wait` at a time. +/// +/// The value is bounded on both sides. It must stay well below the server's +/// `max_pending`, since several pull requests can be outstanding at once, and it +/// must stay *above* the server's `max_payload`: the server answers a pull whose +/// `max_bytes` is smaller than the next message with `409 Message Size Exceeds +/// MaxBytes` and delivers nothing, which turns an intermittent stall into a +/// permanent one. 10 MiB clears the 8 MiB `max_payload` ceiling a server can be +/// configured with in practice, and leaves a six-fold margin against the default +/// `max_pending`. +/// Overridable while the value is being tuned against real message sizes. +const DEFAULT_PULL_MAX_BYTES_PER_BATCH: usize = 10 * 1024 * 1024; + +fn pull_max_bytes_per_batch() -> usize { + quickwit_common::get_from_env_cached!( + usize, + "QW_NATS_PULL_MAX_BYTES_PER_BATCH", + DEFAULT_PULL_MAX_BYTES_PER_BATCH, + false + ) +} + +/// Messages per pull batch. The byte cap above is what actually bounds a batch; +/// this only binds for small messages, where it caps a pull far below the byte +/// budget (200 messages of 1 KiB is under 200 KiB against a 10 MiB budget) and +/// costs a round-trip per 200 messages. Left high so the byte cap is the single +/// thing deciding batch size, and overridable for the same reason as above. +const DEFAULT_PULL_MAX_MESSAGES_PER_BATCH: usize = 100_000; + +fn pull_max_messages_per_batch() -> usize { + quickwit_common::get_from_env_cached!( + usize, + "QW_NATS_PULL_MAX_MESSAGES_PER_BATCH", + DEFAULT_PULL_MAX_MESSAGES_PER_BATCH, + false + ) +} + +/// Both kept at the `Consumer::messages` defaults, so switching to the builder +/// changes the batch bounds and nothing else. +const PULL_IDLE_HEARTBEAT: Duration = Duration::from_secs(15); +const PULL_EXPIRES: Duration = Duration::from_secs(30); /// Pause before pulling again after a transient stream error, to avoid a hot /// error loop while the connection recovers. `SuggestTruncate` is still @@ -358,8 +438,20 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { ) } -/// Sends server-confirmed acknowledgments ("double acks") for the messages up -/// to `published_up_to`. +/// Acknowledges the messages up to `published_up_to`. +/// +/// The acknowledgments are plain publishes, not requests: a request per message +/// would put a server round-trip on the source actor's own task, which is the +/// task that also pulls, and the actor handles one message at a time. What makes +/// that safe is that the connection is only ever torn down through `finalize`, +/// which flushes before draining it, so an acknowledgment queued here is written +/// out before the connection goes away. +/// +/// The consequence for [`Source::is_drained`] is worth being explicit about: a +/// message leaves `pending_acks` once its acknowledgment is queued on the +/// connection rather than once the server has answered for it. A drain therefore +/// declares itself finished on "queued and about to be flushed" instead of +/// "server has it", and it is `finalize`'s flush that closes that gap. async fn ack_up_to( nats_client: &async_nats::Client, pending_acks: &mut BTreeMap, @@ -372,20 +464,12 @@ async fn ack_up_to( if acks.is_empty() { return; } - let mut ack_results = futures::stream::iter(acks) - .map(|(delivery_counter, ack_subject)| async move { - // Mirrors `jetstream::Message::double_ack()`: a request instead of - // a plain publish, so the server's reply confirms the ack. - let ack_result = nats_client.request(ack_subject, AckKind::Ack.into()).await; - (delivery_counter, ack_result) - }) - .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); let mut num_acks = 0usize; let mut num_failed_acks = 0usize; let mut last_ack_error = None; - while let Some((delivery_counter, ack_result)) = ack_results.next().await { - match ack_result { - Ok(_ack_reply) => { + for (delivery_counter, ack_subject) in acks { + match nats_client.publish(ack_subject, AckKind::Ack.into()).await { + Ok(()) => { pending_acks.remove(&delivery_counter); num_acks += 1; } @@ -396,6 +480,8 @@ async fn ack_up_to( } } if let Some(error) = last_ack_error { + // The messages stay in `pending_acks`, so the next truncate retries them + // and a drain keeps waiting rather than declaring itself finished. warn!(%error, num_failed_acks, "failed to ack NATS messages"); } debug!(num_acks, "acked published messages"); From 945370252b8c6ce94a63505703de8a3b53b710f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Tue, 8 Sep 2026 14:39:30 +0000 Subject: [PATCH 28/41] fix: nat source --- .../src/source/nats_source.rs | 45 +++++++++++-------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 61f2cfbc598..4612f04260a 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -438,20 +438,23 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { ) } -/// Acknowledges the messages up to `published_up_to`. -/// -/// The acknowledgments are plain publishes, not requests: a request per message -/// would put a server round-trip on the source actor's own task, which is the -/// task that also pulls, and the actor handles one message at a time. What makes -/// that safe is that the connection is only ever torn down through `finalize`, -/// which flushes before draining it, so an acknowledgment queued here is written -/// out before the connection goes away. +/// Bounds the concurrent server-confirmed ack requests: enough to hide the +/// round-trip latency without flooding the connection. Measured against 1 KiB +/// messages: 1024 is indistinguishable from 128, and 8192 is 20 % slower. +const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; + +/// Sends server-confirmed acknowledgments ("double acks") for the messages up +/// to `published_up_to`. /// -/// The consequence for [`Source::is_drained`] is worth being explicit about: a -/// message leaves `pending_acks` once its acknowledgment is queued on the -/// connection rather than once the server has answered for it. A drain therefore -/// declares itself finished on "queued and about to be flushed" instead of -/// "server has it", and it is `finalize`'s flush that closes that gap. +/// The confirmation is not for error reporting: it is what holds a message in +/// `pending_acks` long enough for [`Source::is_drained`] to still be false when +/// `Drain` arrives. `Drain` queues its force-commit batch and then checks +/// `is_drained` immediately, so a source that reports itself drained at that +/// instant exits before the in-flight documents are published, and they are +/// redelivered as duplicates. Replacing these requests with plain publishes was +/// measured at 37 % less broker CPU and 40 % more throughput at four pipelines, +/// and it broke exactly-once on `num_pipelines` reassignment (+91 581 duplicates +/// on a 2 000 000 document corpus), which is why they are requests. async fn ack_up_to( nats_client: &async_nats::Client, pending_acks: &mut BTreeMap, @@ -464,12 +467,20 @@ async fn ack_up_to( if acks.is_empty() { return; } + let mut ack_results = futures::stream::iter(acks) + .map(|(delivery_counter, ack_subject)| async move { + // Mirrors `jetstream::Message::double_ack()`, but through the + // client's muxed inbox rather than a subscription per message. + let ack_result = nats_client.request(ack_subject, AckKind::Ack.into()).await; + (delivery_counter, ack_result) + }) + .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); let mut num_acks = 0usize; let mut num_failed_acks = 0usize; let mut last_ack_error = None; - for (delivery_counter, ack_subject) in acks { - match nats_client.publish(ack_subject, AckKind::Ack.into()).await { - Ok(()) => { + while let Some((delivery_counter, ack_result)) = ack_results.next().await { + match ack_result { + Ok(_ack_reply) => { pending_acks.remove(&delivery_counter); num_acks += 1; } @@ -480,8 +491,6 @@ async fn ack_up_to( } } if let Some(error) = last_ack_error { - // The messages stay in `pending_acks`, so the next truncate retries them - // and a drain keeps waiting rather than declaring itself finished. warn!(%error, num_failed_acks, "failed to ack NATS messages"); } debug!(num_acks, "acked published messages"); From 17f226cca353c6b6859c955deb3bb65b2628b4c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Tue, 8 Sep 2026 15:37:55 +0000 Subject: [PATCH 29/41] feat: nats source --- docs/configuration/source-config.md | 33 +++- docs/ingest-data/nats.md | 6 +- .../src/source/nats_source.rs | 145 ++++++++++++++++-- 3 files changed, 161 insertions(+), 23 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 82d32c05393..a9dd816ba10 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -189,15 +189,36 @@ A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-conce A tutorial is available [here](/docs/ingest-data/nats.md). -The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. +The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use either the **explicit** or the **all** ack policy: the source acknowledges messages once the split containing them is published, and the consumer's ack floor is the resume point. Which of the two to pick is a throughput decision, covered in [Ack policy](#ack-policy) below. -Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. `max_ack_pending` bounds the messages in flight across all pipelines of the consumer, and therefore the aggregate throughput. +Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. -**Scaling** +#### Ack policy -Multiple indexing pipelines can share the consumer: NATS load-balances the messages across them, so scaling up or down is a plain `num_pipelines` update, and the control plane places the pipelines across the indexers of the cluster. +The ack policy decides how many acknowledgments a corpus costs, and at small message sizes that is what bounds the source's throughput. JetStream delivery is cheap; per-message acknowledgment is not. -**Monitoring** +- **`explicit`** — one acknowledgment per message. Required whenever several pipelines share the consumer, because acknowledging a message under the `all` policy would also acknowledge the messages a sibling pipeline is still indexing. +- **`all`** — one acknowledgment covers every message up to it, so the source sends a single acknowledgment per commit instead of one per message, the same shape as a Kafka offset commit. Quickwit rejects a source configured with `num_pipelines` greater than 1 on such a consumer, since it is only sound when the consumer has a **single puller**. + +Prefer `all`, and get parallelism from partitioning rather than from several pipelines on one consumer. On a 16 vCPU host indexing 1 KiB messages, moving one pipeline from `explicit` to `all` cut broker CPU by 56 % and total CPU per indexed GiB by 31 %; the gain is larger the smaller the messages, and negligible above a few hundred KiB, where the message rate is low enough that acknowledgments no longer matter. + +#### Scaling + +There are two ways to spread a stream over several pipelines, and they do not perform alike. + +**One stream per pipeline** (recommended). Partition the data into one stream per pipeline — the equivalent of Kafka partitions — give each stream its own consumer, and add one Quickwit source with `num_pipelines: 1` per stream. Each consumer then has a single puller, which is what makes the `all` ack policy usable, and no consumer reads a message it does not own. This is the only shape that scales: on the host above, sixteen streams indexed the same 1 KiB corpus 2.7 times faster than sixteen pipelines sharing one consumer. + +**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. This requires the `explicit` policy, and the resulting per-message acknowledgments cap the aggregate throughput regardless of how many pipelines are added — adding pipelines beyond a handful buys close to nothing at small message sizes. + +:::warning + +Do not partition a single stream by giving each pipeline a subject-filtered consumer. Unless a subject's messages happen to be contiguous in the stream, every consumer scans past the messages it does not own, and the broker saturates: sixteen subject-filtered consumers over sixteen interleaved subjects were 5.5 times *slower* than one shared consumer, with the broker pinned and the indexers idle. + +::: + +`max_ack_pending` bounds the messages in flight and therefore the achievable throughput. Leave it unlimited (`-1`) unless there is a reason to bound the source's memory: nothing is acknowledged until a split is published, so a whole commit window is always ack-pending, and a low value stalls the pipelines waiting on commits. Under the `all` policy the window is acknowledged by a single message, so the bound has to cover the entire window rather than a sliding fraction of it. + +#### Monitoring Being durable, the consumer is observable through NATS's own monitoring (`nats consumer info`, exporters): `num_pending` is the indexing lag and `num_ack_pending` the in-flight window, both available even while the pipelines are down. The source also reports `num_pending_acks` (messages indexed but whose split is not published yet) in its observable state. @@ -261,7 +282,7 @@ EOF ## Number of pipelines -The `num_pipelines` parameter is only available for distributed sources like Kafka, GCP PubSub, and Pulsar. +The `num_pipelines` parameter is only available for distributed sources like Kafka, GCP PubSub, NATS, and Pulsar. It defines the number of pipelines to run on a cluster for the source. The actual placement of these pipelines on the different indexer will be decided by the control plane. diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md index 2dc3f78cb42..39f7d8a44d9 100644 --- a/docs/ingest-data/nats.md +++ b/docs/ingest-data/nats.md @@ -131,12 +131,14 @@ The NATS source consumes a JetStream stream through a durable consumer, so messa ```bash nats stream add stackoverflow --subjects "stackoverflow.posts" --defaults -nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack explicit --wait 5m --defaults +nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack all --max-pending=-1 --wait 5m --defaults ``` :::info -The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. See the [NATS source reference](../configuration/source-config.md#nats-source) for details. +The consumer must use the `all` or the `explicit` ack policy: the source acknowledges messages once they are durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. + +`--ack all` is used here because this source runs a single pipeline. One acknowledgment then covers every message up to it, so a commit costs one acknowledgment instead of one per message — the same shape as a Kafka offset commit, and the difference between the two is what bounds throughput on small messages. It is only safe with a single puller, so a consumer shared by several pipelines needs `--ack explicit` instead; see [Ack policy](../configuration/source-config.md#ack-policy) and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. ::: diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 4612f04260a..e35e8e58d17 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -20,9 +20,14 @@ //! whoever provisioned it. //! //! Several indexing pipelines can share the consumer: NATS load-balances the -//! messages across them, so scaling is a plain `num_pipelines` update. Being -//! durable, the consumer is observable through NATS's own monitoring even -//! while the pipelines are down. +//! messages across them, so scaling is a plain `num_pipelines` update. That +//! shape requires `AckPolicy::Explicit`, which costs one acknowledgment per +//! message and is what bounds throughput on small messages; a consumer with a +//! single puller can use `AckPolicy::All` instead and spend one acknowledgment +//! per commit. See `ack_up_to` and `docs/configuration/source-config.md`. +//! +//! Being durable, the consumer is observable through NATS's own monitoring +//! even while the pipelines are down. //! //! When a message carries a W3C `traceparent` header, the source stitches the //! processing of the message into the publisher's distributed trace. @@ -95,13 +100,28 @@ pub struct NatsSource { consumer_name: String, partition_id: PartitionId, delivery_counter: u64, - /// Ack subjects of the messages delivered but not published yet, keyed by - /// delivery counter. Bounded by the consumer's `max_ack_pending`: the - /// server stops delivering when too many messages are unacknowledged. - pending_acks: BTreeMap, + /// Messages delivered but not published yet, keyed by delivery counter. + /// Bounded by the consumer's `max_ack_pending`: the server stops delivering + /// when too many messages are unacknowledged. + pending_acks: BTreeMap, + /// How the consumer wants its messages acknowledged, read off the consumer + /// at startup. `All` collapses a whole checkpoint range into a single ack, + /// `Explicit` costs one ack per message. See `ack_up_to`. + ack_policy: AckPolicy, state: NatsSourceState, } +/// A delivered message waiting to be published, and what it takes to ack it. +/// +/// The stream sequence is kept because `AckPolicy::All` acknowledges *by stream +/// sequence*: acking one message acks every message at or below its sequence. +/// Delivery order and stream order only coincide while nothing is redelivered, +/// so the sequence has to be carried rather than inferred from the map's key. +struct PendingAck { + ack_subject: Subject, + stream_sequence: u64, +} + impl fmt::Debug for NatsSource { fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { formatter @@ -131,7 +151,12 @@ impl NatsSource { "starting NATS source" ); - let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; + let (nats_client, consumer) = connect_and_fetch_consumer( + &source_params, + source_runtime.source_config.num_pipelines.get(), + ) + .await?; + let ack_policy = consumer.cached_info().config.ack_policy; let message_stream = consumer .stream() .max_messages_per_batch(pull_max_messages_per_batch()) @@ -161,6 +186,7 @@ impl NatsSource { partition_id, delivery_counter: 0, pending_acks: BTreeMap::new(), + ack_policy, state: NatsSourceState::default(), }) } @@ -202,7 +228,13 @@ impl NatsSource { Position::offset(self.delivery_counter), ) .context("failed to record partition delta")?; - self.pending_acks.insert(self.delivery_counter, ack_subject); + self.pending_acks.insert( + self.delivery_counter, + PendingAck { + ack_subject, + stream_sequence, + }, + ); self.state.num_bytes_processed += num_bytes; self.state.num_messages_processed += 1; @@ -283,6 +315,7 @@ impl Source for NatsSource { &self.nats_client, &mut self.pending_acks, published_up_to, + self.ack_policy, )) .await; Ok(()) @@ -457,12 +490,19 @@ const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; /// on a 2 000 000 document corpus), which is why they are requests. async fn ack_up_to( nats_client: &async_nats::Client, - pending_acks: &mut BTreeMap, + pending_acks: &mut BTreeMap, published_up_to: u64, + ack_policy: AckPolicy, ) { + if ack_policy == AckPolicy::All { + ack_all_up_to(nats_client, pending_acks, published_up_to).await; + return; + } let acks: Vec<(u64, Subject)> = pending_acks .range(..=published_up_to) - .map(|(delivery_counter, ack_subject)| (*delivery_counter, ack_subject.clone())) + .map(|(delivery_counter, pending_ack)| { + (*delivery_counter, pending_ack.ack_subject.clone()) + }) .collect(); if acks.is_empty() { return; @@ -496,9 +536,69 @@ async fn ack_up_to( debug!(num_acks, "acked published messages"); } +/// Acks the whole published range with a single request. +/// +/// `AckPolicy::All` makes the server treat one ack as an ack of every message at +/// or below that message's stream sequence, so the ack of the highest published +/// sequence stands in for the entire range. That is the same shape as a Kafka +/// offset commit, and it takes the per-message ack off the hot path: a broker +/// that serves ~400 000 messages/s of pure delivery only serves ~60 000/s once +/// every message needs its own acknowledgment. +/// +/// The hidden contract is that the consumer must have a single puller. Acking by +/// sequence also acks the messages a sibling pipeline holds but has not +/// published, so `fetch_durable_consumer` refuses this policy for a source with +/// more than one pipeline. +/// +/// Nothing is dropped from `pending_acks` unless the server confirmed the ack, +/// so a failure leaves the range pending and the next `suggest_truncate` retries +/// it. `is_drained` therefore still reports the truth. +async fn ack_all_up_to( + nats_client: &async_nats::Client, + pending_acks: &mut BTreeMap, + published_up_to: u64, +) { + // A redelivery is delivered late but carries its original, lower stream + // sequence, so delivery order and stream order can disagree. Acking the + // highest published sequence would then also ack a redelivered message that + // is still waiting to be published, and a crash would lose it. The highest + // sequence that is safe to ack is therefore the highest published one that + // is still below every unpublished one. + let first_unpublished_sequence = pending_acks + .range(published_up_to.saturating_add(1)..) + .map(|(_, pending_ack)| pending_ack.stream_sequence) + .min() + .unwrap_or(u64::MAX); + let Some((_, highest_safe_ack)) = pending_acks + .range(..=published_up_to) + .filter(|(_, pending_ack)| pending_ack.stream_sequence < first_unpublished_sequence) + .max_by_key(|(_, pending_ack)| pending_ack.stream_sequence) + else { + return; + }; + let ack_subject = highest_safe_ack.ack_subject.clone(); + let stream_sequence = highest_safe_ack.stream_sequence; + if let Err(error) = nats_client.request(ack_subject, AckKind::Ack.into()).await { + warn!(%error, stream_sequence, "failed to ack NATS messages up to a stream sequence"); + return; + } + // Only the messages the server actually acked are dropped. A published + // message above `stream_sequence` stays pending and is covered by the next + // call, which keeps `is_drained` honest either way. + let num_pending_before = pending_acks.len(); + pending_acks.retain(|delivery_counter, pending_ack| { + *delivery_counter > published_up_to || pending_ack.stream_sequence > stream_sequence + }); + debug!( + num_acked = num_pending_before - pending_acks.len(), + stream_sequence, "acked published messages up to a stream sequence" + ); +} + async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, + num_pipelines: usize, ) -> anyhow::Result { let consumer: PullConsumer = jetstream_stream .get_consumer(consumer_name) @@ -506,8 +606,19 @@ async fn fetch_durable_consumer( .map_err(|error| anyhow!("failed to find NATS consumer `{consumer_name}`: {error}"))?; let ack_policy = consumer.cached_info().config.ack_policy; ensure!( - ack_policy == AckPolicy::Explicit, - "NATS consumer `{consumer_name}` must use the explicit ack policy, got `{ack_policy:?}`" + matches!(ack_policy, AckPolicy::Explicit | AckPolicy::All), + "NATS consumer `{consumer_name}` must use the explicit or all ack policy, got \ + `{ack_policy:?}`" + ); + // `AckPolicy::All` acks by stream sequence, so one pipeline's ack would also + // ack the lower-sequenced messages its siblings are still indexing, and a + // crash would lose them. Sharing a consumer requires per-message acks. + ensure!( + ack_policy == AckPolicy::Explicit || num_pipelines == 1, + "NATS consumer `{consumer_name}` uses the all ack policy, which acknowledges every \ + message up to the acked one and is only safe with a single puller, but the source is \ + configured with {num_pipelines} pipelines: use the explicit ack policy, or one consumer \ + per pipeline" ); Ok(consumer) } @@ -592,6 +703,7 @@ async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result anyhow::Result<(async_nats::Client, PullConsumer)> { let nats_client = connect_nats(params).await?; let jetstream_ctx = jetstream::new(nats_client.clone()); @@ -599,12 +711,15 @@ async fn connect_and_fetch_consumer( .get_stream(¶ms.stream) .await .with_context(|| format!("failed to find NATS JetStream stream `{}`", params.stream))?; - let consumer = fetch_durable_consumer(&jetstream_stream, ¶ms.consumer).await?; + let consumer = + fetch_durable_consumer(&jetstream_stream, ¶ms.consumer, num_pipelines).await?; Ok((nats_client, consumer)) } pub(crate) async fn check_connectivity(params: &NatsSourceParams) -> anyhow::Result<()> { - connect_and_fetch_consumer(params).await?; + // Connectivity only: a single pipeline is the permissive case, so this never + // rejects a consumer the source itself would accept. + connect_and_fetch_consumer(params, 1).await?; Ok(()) } From 67839a1bbeae1fd158fe45c52743545b77e85518 Mon Sep 17 00:00:00 2001 From: tprevot Date: Tue, 8 Sep 2026 18:19:06 +0200 Subject: [PATCH 30/41] fix: comments --- .../src/source/nats_source.rs | 67 ++++++------------- 1 file changed, 20 insertions(+), 47 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index e35e8e58d17..dd86b7cd150 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -47,6 +47,7 @@ use async_nats::jetstream::consumer::{AckPolicy, PullConsumer}; use async_nats::jetstream::message::AckKind; use async_nats::{ConnectOptions, Subject, jetstream}; use async_trait::async_trait; +use bytesize::ByteSize; use futures::{FutureExt, StreamExt}; use opentelemetry::propagation::Extractor; use opentelemetry::trace::TraceContextExt; @@ -160,11 +161,9 @@ impl NatsSource { let message_stream = consumer .stream() .max_messages_per_batch(pull_max_messages_per_batch()) - .max_bytes_per_batch(pull_max_bytes_per_batch()) - // `stream()` defaults the heartbeat to zero, i.e. disabled, unlike - // `messages()`. It has to be set back, or a stalled subscription is - // never reported and `is_transient_stream_error` has nothing to act - // on. + .max_bytes_per_batch(pull_max_bytes_per_batch().0 as usize) + // default heartbeat that is disabled when using manual builder, + // required for transient conn errors. .heartbeat(PULL_IDLE_HEARTBEAT) .expires(PULL_EXPIRES) .messages() @@ -245,6 +244,7 @@ impl NatsSource { #[async_trait] impl Source for NatsSource { + #[tracing::instrument(skip(source_sink, ctx))] async fn emit_batches( &mut self, source_sink: &SourceSink, @@ -300,6 +300,7 @@ impl Source for NatsSource { Ok(wait_before_next_batch) } + #[tracing::instrument(skip(checkpoint, ctx))] async fn suggest_truncate( &mut self, checkpoint: SourceCheckpoint, @@ -353,7 +354,9 @@ impl Source for NatsSource { } }; let Some(ack_subject) = message.message.reply else { - warn!("prefetched NATS message carries no reply subject to negatively acknowledge it on"); + warn!( + "prefetched NATS message carries no reply subject to negatively acknowledge it on" + ); continue; }; if let Err(error) = self @@ -366,11 +369,7 @@ impl Source for NatsSource { } num_naks += 1; } - // The negative acknowledgments above are plain publishes, and `drain` - // returns as soon as its command is queued rather than once the - // connection has been written out. Without this flush they can be lost - // with the connection, and their messages then wait a full `ack_wait` - // instead of being redelivered to a surviving pipeline immediately. + if let Err(error) = self.nats_client.flush().await { warn!(%error, num_naks, "failed to flush NATS negative acknowledgments"); } @@ -397,7 +396,6 @@ impl Source for NatsSource { /// Bounds the concurrent server-confirmed ack requests: enough to hide the /// round-trip latency without flooding the connection. -/// Caps a pull batch in bytes, which the message-count cap alone cannot do. /// /// The server buffers a pull response on the connection's outbound queue, and a /// connection that exceeds `max_pending` (64 MiB by default) is declared a slow @@ -408,18 +406,12 @@ impl Source for NatsSource { /// /// The value is bounded on both sides. It must stay well below the server's /// `max_pending`, since several pull requests can be outstanding at once, and it -/// must stay *above* the server's `max_payload`: the server answers a pull whose -/// `max_bytes` is smaller than the next message with `409 Message Size Exceeds -/// MaxBytes` and delivers nothing, which turns an intermittent stall into a -/// permanent one. 10 MiB clears the 8 MiB `max_payload` ceiling a server can be -/// configured with in practice, and leaves a six-fold margin against the default -/// `max_pending`. -/// Overridable while the value is being tuned against real message sizes. -const DEFAULT_PULL_MAX_BYTES_PER_BATCH: usize = 10 * 1024 * 1024; - -fn pull_max_bytes_per_batch() -> usize { +/// must stay *above* the server's `max_payload`. +const DEFAULT_PULL_MAX_BYTES_PER_BATCH: ByteSize = ByteSize::mib(10); + +fn pull_max_bytes_per_batch() -> ByteSize { quickwit_common::get_from_env_cached!( - usize, + ByteSize, "QW_NATS_PULL_MAX_BYTES_PER_BATCH", DEFAULT_PULL_MAX_BYTES_PER_BATCH, false @@ -442,8 +434,8 @@ fn pull_max_messages_per_batch() -> usize { ) } -/// Both kept at the `Consumer::messages` defaults, so switching to the builder -/// changes the batch bounds and nothing else. +// Both kept at the `Consumer::messages` defaults, so switching to the builder +// changes the batch bounds and nothing else. const PULL_IDLE_HEARTBEAT: Duration = Duration::from_secs(15); const PULL_EXPIRES: Duration = Duration::from_secs(30); @@ -478,16 +470,6 @@ const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; /// Sends server-confirmed acknowledgments ("double acks") for the messages up /// to `published_up_to`. -/// -/// The confirmation is not for error reporting: it is what holds a message in -/// `pending_acks` long enough for [`Source::is_drained`] to still be false when -/// `Drain` arrives. `Drain` queues its force-commit batch and then checks -/// `is_drained` immediately, so a source that reports itself drained at that -/// instant exits before the in-flight documents are published, and they are -/// redelivered as duplicates. Replacing these requests with plain publishes was -/// measured at 37 % less broker CPU and 40 % more throughput at four pipelines, -/// and it broke exactly-once on `num_pipelines` reassignment (+91 581 duplicates -/// on a 2 000 000 document corpus), which is why they are requests. async fn ack_up_to( nats_client: &async_nats::Client, pending_acks: &mut BTreeMap, @@ -500,9 +482,7 @@ async fn ack_up_to( } let acks: Vec<(u64, Subject)> = pending_acks .range(..=published_up_to) - .map(|(delivery_counter, pending_ack)| { - (*delivery_counter, pending_ack.ack_subject.clone()) - }) + .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.ack_subject.clone())) .collect(); if acks.is_empty() { return; @@ -540,15 +520,8 @@ async fn ack_up_to( /// /// `AckPolicy::All` makes the server treat one ack as an ack of every message at /// or below that message's stream sequence, so the ack of the highest published -/// sequence stands in for the entire range. That is the same shape as a Kafka -/// offset commit, and it takes the per-message ack off the hot path: a broker -/// that serves ~400 000 messages/s of pure delivery only serves ~60 000/s once -/// every message needs its own acknowledgment. -/// -/// The hidden contract is that the consumer must have a single puller. Acking by -/// sequence also acks the messages a sibling pipeline holds but has not -/// published, so `fetch_durable_consumer` refuses this policy for a source with -/// more than one pipeline. +/// sequence stands in for the entire range. It requires that num_pipeline = 1, +/// otherwise this is unsound. /// /// Nothing is dropped from `pending_acks` unless the server confirmed the ack, /// so a failure leaves the range pending and the next `suggest_truncate` retries From 96efbf1a2459c96d9e7b5cd135b07404a6af40f1 Mon Sep 17 00:00:00 2001 From: tprevot Date: Tue, 8 Sep 2026 18:19:17 +0200 Subject: [PATCH 31/41] fix: comments --- quickwit/quickwit-indexing/src/source/nats_source.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index dd86b7cd150..03ac29d2868 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -355,7 +355,8 @@ impl Source for NatsSource { }; let Some(ack_subject) = message.message.reply else { warn!( - "prefetched NATS message carries no reply subject to negatively acknowledge it on" + "prefetched NATS message carries no reply subject to negatively acknowledge \ + it on" ); continue; }; From fcfcca902c01852e778e74b05ea66680fb527d9c Mon Sep 17 00:00:00 2001 From: tprevot Date: Tue, 8 Sep 2026 18:32:26 +0200 Subject: [PATCH 32/41] fix: remove the ack policy all --- .../src/source/nats_source.rs | 213 +++++------------- 1 file changed, 54 insertions(+), 159 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 03ac29d2868..c0b0dbb7b14 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -22,9 +22,7 @@ //! Several indexing pipelines can share the consumer: NATS load-balances the //! messages across them, so scaling is a plain `num_pipelines` update. That //! shape requires `AckPolicy::Explicit`, which costs one acknowledgment per -//! message and is what bounds throughput on small messages; a consumer with a -//! single puller can use `AckPolicy::All` instead and spend one acknowledgment -//! per commit. See `ack_up_to` and `docs/configuration/source-config.md`. +//! message and is what bounds throughput on small messages. //! //! Being durable, the consumer is observable through NATS's own monitoring //! even while the pipelines are down. @@ -104,23 +102,9 @@ pub struct NatsSource { /// Messages delivered but not published yet, keyed by delivery counter. /// Bounded by the consumer's `max_ack_pending`: the server stops delivering /// when too many messages are unacknowledged. - pending_acks: BTreeMap, - /// How the consumer wants its messages acknowledged, read off the consumer - /// at startup. `All` collapses a whole checkpoint range into a single ack, - /// `Explicit` costs one ack per message. See `ack_up_to`. - ack_policy: AckPolicy, - state: NatsSourceState, -} + pending_acks: BTreeMap, -/// A delivered message waiting to be published, and what it takes to ack it. -/// -/// The stream sequence is kept because `AckPolicy::All` acknowledges *by stream -/// sequence*: acking one message acks every message at or below its sequence. -/// Delivery order and stream order only coincide while nothing is redelivered, -/// so the sequence has to be carried rather than inferred from the map's key. -struct PendingAck { - ack_subject: Subject, - stream_sequence: u64, + state: NatsSourceState, } impl fmt::Debug for NatsSource { @@ -152,12 +136,7 @@ impl NatsSource { "starting NATS source" ); - let (nats_client, consumer) = connect_and_fetch_consumer( - &source_params, - source_runtime.source_config.num_pipelines.get(), - ) - .await?; - let ack_policy = consumer.cached_info().config.ack_policy; + let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; let message_stream = consumer .stream() .max_messages_per_batch(pull_max_messages_per_batch()) @@ -185,7 +164,6 @@ impl NatsSource { partition_id, delivery_counter: 0, pending_acks: BTreeMap::new(), - ack_policy, state: NatsSourceState::default(), }) } @@ -227,19 +205,57 @@ impl NatsSource { Position::offset(self.delivery_counter), ) .context("failed to record partition delta")?; - self.pending_acks.insert( - self.delivery_counter, - PendingAck { - ack_subject, - stream_sequence, - }, - ); + self.pending_acks.insert(self.delivery_counter, ack_subject); self.state.num_bytes_processed += num_bytes; self.state.num_messages_processed += 1; Ok(()) } + + /// Sends server-confirmed acknowledgments ("double acks") for the messages up + /// to `published_up_to`. + async fn ack_up_to(&mut self, published_up_to: u64) { + let acks: Vec<(u64, Subject)> = self + .pending_acks + .range(..=published_up_to) + .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.clone())) + .collect(); + if acks.is_empty() { + return; + } + let mut ack_results = futures::stream::iter(acks) + .map(async |(delivery_counter, ack_subject)| { + // Mirrors `jetstream::Message::double_ack()`, but through the + // client's muxed inbox rather than a subscription per message. + let ack_result = self + .nats_client + .request(ack_subject, AckKind::Ack.into()) + .await; + (delivery_counter, ack_result) + }) + .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); + + let mut num_acks = 0usize; + let mut num_failed_acks = 0usize; + let mut last_ack_error = None; + while let Some((delivery_counter, ack_result)) = ack_results.next().await { + match ack_result { + Ok(_ack_reply) => { + self.pending_acks.remove(&delivery_counter); + num_acks += 1; + } + Err(error) => { + num_failed_acks += 1; + last_ack_error = Some(error); + } + } + } + if let Some(error) = last_ack_error { + warn!(%error, num_failed_acks, "failed to ack NATS messages"); + } + debug!(num_acks, "acked published messages"); + } } #[async_trait] @@ -312,13 +328,7 @@ impl Source for NatsSource { let Some(published_up_to) = position.as_u64() else { return Ok(()); }; - ctx.protect_future(ack_up_to( - &self.nats_client, - &mut self.pending_acks, - published_up_to, - self.ack_policy, - )) - .await; + ctx.protect_future(self.ack_up_to(published_up_to)).await; Ok(()) } @@ -469,110 +479,9 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { /// messages: 1024 is indistinguishable from 128, and 8192 is 20 % slower. const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; -/// Sends server-confirmed acknowledgments ("double acks") for the messages up -/// to `published_up_to`. -async fn ack_up_to( - nats_client: &async_nats::Client, - pending_acks: &mut BTreeMap, - published_up_to: u64, - ack_policy: AckPolicy, -) { - if ack_policy == AckPolicy::All { - ack_all_up_to(nats_client, pending_acks, published_up_to).await; - return; - } - let acks: Vec<(u64, Subject)> = pending_acks - .range(..=published_up_to) - .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.ack_subject.clone())) - .collect(); - if acks.is_empty() { - return; - } - let mut ack_results = futures::stream::iter(acks) - .map(|(delivery_counter, ack_subject)| async move { - // Mirrors `jetstream::Message::double_ack()`, but through the - // client's muxed inbox rather than a subscription per message. - let ack_result = nats_client.request(ack_subject, AckKind::Ack.into()).await; - (delivery_counter, ack_result) - }) - .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); - let mut num_acks = 0usize; - let mut num_failed_acks = 0usize; - let mut last_ack_error = None; - while let Some((delivery_counter, ack_result)) = ack_results.next().await { - match ack_result { - Ok(_ack_reply) => { - pending_acks.remove(&delivery_counter); - num_acks += 1; - } - Err(error) => { - num_failed_acks += 1; - last_ack_error = Some(error); - } - } - } - if let Some(error) = last_ack_error { - warn!(%error, num_failed_acks, "failed to ack NATS messages"); - } - debug!(num_acks, "acked published messages"); -} - -/// Acks the whole published range with a single request. -/// -/// `AckPolicy::All` makes the server treat one ack as an ack of every message at -/// or below that message's stream sequence, so the ack of the highest published -/// sequence stands in for the entire range. It requires that num_pipeline = 1, -/// otherwise this is unsound. -/// -/// Nothing is dropped from `pending_acks` unless the server confirmed the ack, -/// so a failure leaves the range pending and the next `suggest_truncate` retries -/// it. `is_drained` therefore still reports the truth. -async fn ack_all_up_to( - nats_client: &async_nats::Client, - pending_acks: &mut BTreeMap, - published_up_to: u64, -) { - // A redelivery is delivered late but carries its original, lower stream - // sequence, so delivery order and stream order can disagree. Acking the - // highest published sequence would then also ack a redelivered message that - // is still waiting to be published, and a crash would lose it. The highest - // sequence that is safe to ack is therefore the highest published one that - // is still below every unpublished one. - let first_unpublished_sequence = pending_acks - .range(published_up_to.saturating_add(1)..) - .map(|(_, pending_ack)| pending_ack.stream_sequence) - .min() - .unwrap_or(u64::MAX); - let Some((_, highest_safe_ack)) = pending_acks - .range(..=published_up_to) - .filter(|(_, pending_ack)| pending_ack.stream_sequence < first_unpublished_sequence) - .max_by_key(|(_, pending_ack)| pending_ack.stream_sequence) - else { - return; - }; - let ack_subject = highest_safe_ack.ack_subject.clone(); - let stream_sequence = highest_safe_ack.stream_sequence; - if let Err(error) = nats_client.request(ack_subject, AckKind::Ack.into()).await { - warn!(%error, stream_sequence, "failed to ack NATS messages up to a stream sequence"); - return; - } - // Only the messages the server actually acked are dropped. A published - // message above `stream_sequence` stays pending and is covered by the next - // call, which keeps `is_drained` honest either way. - let num_pending_before = pending_acks.len(); - pending_acks.retain(|delivery_counter, pending_ack| { - *delivery_counter > published_up_to || pending_ack.stream_sequence > stream_sequence - }); - debug!( - num_acked = num_pending_before - pending_acks.len(), - stream_sequence, "acked published messages up to a stream sequence" - ); -} - async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, - num_pipelines: usize, ) -> anyhow::Result { let consumer: PullConsumer = jetstream_stream .get_consumer(consumer_name) @@ -580,20 +489,10 @@ async fn fetch_durable_consumer( .map_err(|error| anyhow!("failed to find NATS consumer `{consumer_name}`: {error}"))?; let ack_policy = consumer.cached_info().config.ack_policy; ensure!( - matches!(ack_policy, AckPolicy::Explicit | AckPolicy::All), - "NATS consumer `{consumer_name}` must use the explicit or all ack policy, got \ + ack_policy == AckPolicy::Explicit, + "NATS consumer `{consumer_name}` must use the explicit ack policy, got \ `{ack_policy:?}`" ); - // `AckPolicy::All` acks by stream sequence, so one pipeline's ack would also - // ack the lower-sequenced messages its siblings are still indexing, and a - // crash would lose them. Sharing a consumer requires per-message acks. - ensure!( - ack_policy == AckPolicy::Explicit || num_pipelines == 1, - "NATS consumer `{consumer_name}` uses the all ack policy, which acknowledges every \ - message up to the acked one and is only safe with a single puller, but the source is \ - configured with {num_pipelines} pipelines: use the explicit ack policy, or one consumer \ - per pipeline" - ); Ok(consumer) } @@ -677,7 +576,6 @@ async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result anyhow::Result<(async_nats::Client, PullConsumer)> { let nats_client = connect_nats(params).await?; let jetstream_ctx = jetstream::new(nats_client.clone()); @@ -685,15 +583,12 @@ async fn connect_and_fetch_consumer( .get_stream(¶ms.stream) .await .with_context(|| format!("failed to find NATS JetStream stream `{}`", params.stream))?; - let consumer = - fetch_durable_consumer(&jetstream_stream, ¶ms.consumer, num_pipelines).await?; + let consumer = fetch_durable_consumer(&jetstream_stream, ¶ms.consumer).await?; Ok((nats_client, consumer)) } pub(crate) async fn check_connectivity(params: &NatsSourceParams) -> anyhow::Result<()> { - // Connectivity only: a single pipeline is the permissive case, so this never - // rejects a consumer the source itself would accept. - connect_and_fetch_consumer(params, 1).await?; + connect_and_fetch_consumer(params).await?; Ok(()) } From 55b5e53ff2cda5a7fc2cb2276662c5c601578f5f Mon Sep 17 00:00:00 2001 From: tprevot Date: Tue, 8 Sep 2026 18:58:47 +0200 Subject: [PATCH 33/41] fix: remove the ack policy all --- quickwit/quickwit-indexing/src/source/nats_source.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index c0b0dbb7b14..f79c8ba68c3 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -24,9 +24,6 @@ //! shape requires `AckPolicy::Explicit`, which costs one acknowledgment per //! message and is what bounds throughput on small messages. //! -//! Being durable, the consumer is observable through NATS's own monitoring -//! even while the pipelines are down. -//! //! When a message carries a W3C `traceparent` header, the source stitches the //! processing of the message into the publisher's distributed trace. @@ -490,8 +487,7 @@ async fn fetch_durable_consumer( let ack_policy = consumer.cached_info().config.ack_policy; ensure!( ack_policy == AckPolicy::Explicit, - "NATS consumer `{consumer_name}` must use the explicit ack policy, got \ - `{ack_policy:?}`" + "NATS consumer `{consumer_name}` must use the explicit ack policy, got `{ack_policy:?}`" ); Ok(consumer) } From 3804790581b0606787111d6a114e61f60fc5a80d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Wed, 9 Sep 2026 09:26:31 +0000 Subject: [PATCH 34/41] feat: nats source --- docs/configuration/source-config.md | 19 ++++++++----------- docs/ingest-data/nats.md | 6 +++--- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index a9dd816ba10..7351e35927d 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -189,34 +189,31 @@ A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-conce A tutorial is available [here](/docs/ingest-data/nats.md). -The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use either the **explicit** or the **all** ack policy: the source acknowledges messages once the split containing them is published, and the consumer's ack floor is the resume point. Which of the two to pick is a throughput decision, covered in [Ack policy](#ack-policy) below. +The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. -#### Ack policy +#### Acknowledgment cost -The ack policy decides how many acknowledgments a corpus costs, and at small message sizes that is what bounds the source's throughput. JetStream delivery is cheap; per-message acknowledgment is not. +The source acknowledges every message individually and waits for the server to confirm each one — that confirmation is what tells the drain the pipeline is empty — so indexing a corpus costs one confirmed acknowledgment per document. At small message sizes that, rather than the indexer, is what bounds throughput. Measured on a 16 vCPU host across sixteen per-tenant consumers of 1 KiB documents: 72 MiB/s, with the broker spending 5.1 of its cores on the acknowledgment path against Quickwit's 4.7 on indexing, roughly ten times the broker CPU per document that an equivalent Kafka source spends on periodic offset commits. -- **`explicit`** — one acknowledgment per message. Required whenever several pipelines share the consumer, because acknowledging a message under the `all` policy would also acknowledge the messages a sibling pipeline is still indexing. -- **`all`** — one acknowledgment covers every message up to it, so the source sends a single acknowledgment per commit instead of one per message, the same shape as a Kafka offset commit. Quickwit rejects a source configured with `num_pipelines` greater than 1 on such a consumer, since it is only sound when the consumer has a **single puller**. - -Prefer `all`, and get parallelism from partitioning rather than from several pipelines on one consumer. On a 16 vCPU host indexing 1 KiB messages, moving one pipeline from `explicit` to `all` cut broker CPU by 56 % and total CPU per indexed GiB by 31 %; the gain is larger the smaller the messages, and negligible above a few hundred KiB, where the message rate is low enough that acknowledgments no longer matter. +The cost is per *message*, so it amortises as messages get larger and is negligible above a few hundred KiB, where the message rate is only a few hundred per second. Batching many log records into one message and letting an `otlp_*` input format expand them is the shape that gets that benefit. #### Scaling There are two ways to spread a stream over several pipelines, and they do not perform alike. -**One stream per pipeline** (recommended). Partition the data into one stream per pipeline — the equivalent of Kafka partitions — give each stream its own consumer, and add one Quickwit source with `num_pipelines: 1` per stream. Each consumer then has a single puller, which is what makes the `all` ack policy usable, and no consumer reads a message it does not own. This is the only shape that scales: on the host above, sixteen streams indexed the same 1 KiB corpus 2.7 times faster than sixteen pipelines sharing one consumer. +**One consumer per partition of the data** (recommended). Give each partition — a tenant, a subject, whatever the natural unit is — its own durable consumer with a single `filter_subject`, and add one Quickwit source with `num_pipelines: 1` per consumer. The partitions then scale and fail independently, and each gets its own checkpoint. Splitting the same corpus over one shared stream with a consumer per partition and over a stream per partition measured 21.28 s and 21.26 s, so partition whichever way suits the producers. Adding pipelines *within* a consumer does not help: sixteen consumers with two pipelines each measured 23.5 s against 21.3 s with one. -**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. This requires the `explicit` policy, and the resulting per-message acknowledgments cap the aggregate throughput regardless of how many pipelines are added — adding pipelines beyond a handful buys close to nothing at small message sizes. +**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. Because the acknowledgment cost is per message rather than per pipeline, adding pipelines to one consumer does not raise the aggregate throughput much at small message sizes: on the host above it was worth 1.15× between one and sixteen pipelines, against 1.4× for splitting the same work over sixteen consumers. :::warning -Do not partition a single stream by giving each pipeline a subject-filtered consumer. Unless a subject's messages happen to be contiguous in the stream, every consumer scans past the messages it does not own, and the broker saturates: sixteen subject-filtered consumers over sixteen interleaved subjects were 5.5 times *slower* than one shared consumer, with the broker pinned and the indexers idle. +Give each consumer **one** subject filter, not several. A consumer configured with `filter_subjects` (NATS 2.10's multi-filter) uses a different server path than one configured with a single `filter_subject`, and on a stream whose subjects are interleaved it is dramatically slower: sixteen multi-filter consumers over sixteen interleaved subjects took 7.5 times longer than the same sixteen consumers with one filter each, with the broker saturated and the indexers idle. ::: -`max_ack_pending` bounds the messages in flight and therefore the achievable throughput. Leave it unlimited (`-1`) unless there is a reason to bound the source's memory: nothing is acknowledged until a split is published, so a whole commit window is always ack-pending, and a low value stalls the pipelines waiting on commits. Under the `all` policy the window is acknowledged by a single message, so the bound has to cover the entire window rather than a sliding fraction of it. +`max_ack_pending` bounds the messages in flight and therefore the achievable throughput: nothing is acknowledged until a split is published, so a whole commit window is always ack-pending and the achievable rate is `max_ack_pending / commit_timeout` documents per second. Measured at four pipelines with a 10 s commit timeout, 1,000 gave 97 documents per second and 20,000 gave 2,115. Leave it unlimited (`-1`) unless there is a reason to bound the source's memory, and note the trade: unlimited also makes the replay window after a crash the whole delivered span above the ack floor rather than a bounded slice of it. #### Monitoring diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md index 39f7d8a44d9..3c63c766f92 100644 --- a/docs/ingest-data/nats.md +++ b/docs/ingest-data/nats.md @@ -131,14 +131,14 @@ The NATS source consumes a JetStream stream through a durable consumer, so messa ```bash nats stream add stackoverflow --subjects "stackoverflow.posts" --defaults -nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack all --max-pending=-1 --wait 5m --defaults +nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack explicit --max-pending=-1 --wait 5m --defaults ``` :::info -The consumer must use the `all` or the `explicit` ack policy: the source acknowledges messages once they are durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. +The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. -`--ack all` is used here because this source runs a single pipeline. One acknowledgment then covers every message up to it, so a commit costs one acknowledgment instead of one per message — the same shape as a Kafka offset commit, and the difference between the two is what bounds throughput on small messages. It is only safe with a single puller, so a consumer shared by several pipelines needs `--ack explicit` instead; see [Ack policy](../configuration/source-config.md#ack-policy) and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. +`--max-pending=-1` leaves the in-flight window unbounded: the source acknowledges only once a split is published, so a whole commit window is always ack-pending and a small bound throttles indexing. See the [NATS source reference](../configuration/source-config.md#nats-source) for the rest, and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. ::: From 2667e4c9de19ea4024c21c4f6dbf92008f01f920 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Wed, 9 Sep 2026 12:14:37 +0000 Subject: [PATCH 35/41] feat: nats docs --- docs/configuration/source-config.md | 77 ++++++++++++++++++++++++----- docs/ingest-data/nats.md | 13 +++-- 2 files changed, 75 insertions(+), 15 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 7351e35927d..1ff6a607f53 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -189,31 +189,84 @@ A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-conce A tutorial is available [here](/docs/ingest-data/nats.md). -The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. +The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning belong to whoever provisioned it. -Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. +Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. -#### Acknowledgment cost +#### Consumer invariants + +These are properties of the consumer, not of the source. Only the ack policy is enforced when the source is created; the rest are not checked, and getting them wrong looks like Quickwit being slow or duplicating rather than like a consumer misconfiguration. + +**`ack_policy` must be `explicit`.** The source acknowledges each message individually once the split containing it is published, and waits for the server to confirm the acknowledgment — the confirmation is what tells the drain that the pipeline is empty. The consumer's ack floor is therefore the resume point, and it is the only progress state that matters. Any other policy is rejected when the source is created. + +**`ack_wait` must exceed the end-to-end publish latency.** The timer starts at delivery and has to outlast three terms: + +``` +ack_wait > commit_timeout + split upload and publish + ack round trip +``` + +Shorter than their sum, and NATS redelivers messages that are still being indexed; they are then indexed twice. 5 minutes with a 60 s `commit_timeout` leaves a wide margin. + +It is also the *recovery* time from a lost delivery, so it should not be arbitrarily large either. If a message is delivered but never arrives — a dropped connection, a slow-consumer disconnect — nothing brings it back until the timer expires and the pipeline simply idles. The same injected fault idled a run for 278 s at `ack_wait=300s` and 23.2 s at `30s`. + +**`max_ack_pending` should be `-1`.** Nothing is acknowledged until a split is published, so a whole commit window is always ack-pending. That makes the setting a hard throughput cap rather than a safety valve: + +``` +achievable rate ≈ max_ack_pending / commit_timeout documents per second +``` + +Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing to a trickle here. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. -The source acknowledges every message individually and waits for the server to confirm each one — that confirmation is what tells the drain the pipeline is empty — so indexing a corpus costs one confirmed acknowledgment per document. At small message sizes that, rather than the indexer, is what bounds throughput. Measured on a 16 vCPU host across sixteen per-tenant consumers of 1 KiB documents: 72 MiB/s, with the broker spending 5.1 of its cores on the acknowledgment path against Quickwit's 4.7 on indexing, roughly ten times the broker CPU per document that an equivalent Kafka source spends on periodic offset commits. +**`deliver_policy`** is usually `all`, so a new consumer indexes the stream from the start. -The cost is per *message*, so it amortises as messages get larger and is negligible above a few hundred KiB, where the message rate is only a few hundred per second. Batching many log records into one message and letting an `otlp_*` input format expand them is the shape that gets that benefit. +**Give each consumer a single `filter_subject`.** A consumer configured with several filters through NATS 2.10's multi-filter `filter_subjects` appears to take a much slower path in the broker: on a stream whose subjects are interleaved, sixteen multi-filter consumers took 7.5 times longer than the same sixteen consumers with one filter each, with the broker saturated and the indexers idle. + +#### What can be tuned on the Quickwit side + +| Setting | Where | Effect | +| --- | --- | --- | +| `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | +| [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. It is also the length of the cooperative-indexing cycle. | +| [`docstore_compression_level`](index-config.md#indexing-settings) | index config | Defaults to zstd 8, which is what limits a single pipeline on small documents: the dedicated docstore thread saturates and back-pressures the indexer. Level 1 or 3 measured 21.1 s against 31.2 s on a 1.5 GiB corpus of 1 KiB documents. Not NATS-specific. | +| [`enable_cooperative_indexing`](node-config.md#indexer-configuration) | node config | See below. | +| `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. 10 MiB and 20 MiB measure identically. | +| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Does not bind at these message sizes — 200 and 100,000 measured identically — because the byte cap is what bounds a batch. | #### Scaling -There are two ways to spread a stream over several pipelines, and they do not perform alike. +**One consumer per partition of the data** (recommended). Give each partition — a tenant, a subject, whatever the natural unit is — its own durable consumer with a single `filter_subject`, and add one Quickwit source with `num_pipelines: 1` per consumer. The partitions then scale and fail independently, and each gets its own checkpoint. -**One consumer per partition of the data** (recommended). Give each partition — a tenant, a subject, whatever the natural unit is — its own durable consumer with a single `filter_subject`, and add one Quickwit source with `num_pipelines: 1` per consumer. The partitions then scale and fail independently, and each gets its own checkpoint. Splitting the same corpus over one shared stream with a consumer per partition and over a stream per partition measured 21.28 s and 21.26 s, so partition whichever way suits the producers. Adding pipelines *within* a consumer does not help: sixteen consumers with two pipelines each measured 23.5 s against 21.3 s with one. +**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. -**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. Because the acknowledgment cost is per message rather than per pipeline, adding pipelines to one consumer does not raise the aggregate throughput much at small message sizes: on the host above it was worth 1.15× between one and sixteen pipelines, against 1.4× for splitting the same work over sixteen consumers. +At 512 KiB per message the two shapes are equivalent — 52.1, 91.3 and 146.6 MiB/s at one, two and four units either way — and so is partitioning across separate streams instead of subjects. At 1 KiB they diverge: a shared consumer stops gaining past a couple of pipelines while per-tenant consumers keep scaling, because the acknowledgment cost is per message rather than per pipeline. The per-consumer shape is therefore the safe default: never worse, and better on small documents. -:::warning +#### Acknowledgment cost -Give each consumer **one** subject filter, not several. A consumer configured with `filter_subjects` (NATS 2.10's multi-filter) uses a different server path than one configured with a single `filter_subject`, and on a stream whose subjects are interleaved it is dramatically slower: sixteen multi-filter consumers over sixteen interleaved subjects took 7.5 times longer than the same sixteen consumers with one filter each, with the broker saturated and the indexers idle. +One confirmed acknowledgment per message is the source's dominant broker cost, and it is what decides whether the source or the indexer is the bottleneck. -::: +- At **512 KiB** per message it is invisible. The NATS and Kafka sources measured within a percent of each other across the whole pipeline sweep, and broker CPU stayed at 0.5–0.9 cores for both, which is the host's floor. Priced on the broker alone, a confirmed acknowledgment per message still leaves 1,138 MiB/s — several times what one indexer consumes. +- At **1 KiB** it is the ceiling. Sixteen per-tenant consumers reached 72 MiB/s while an equivalent Kafka source reached 150, and the broker spent 5.1 of the host's cores on acknowledgments against Quickwit's 4.7 on indexing — about ten times the broker CPU per document that a periodic offset commit costs. Priced on the broker alone, acknowledgment divides throughput by 6.5 at this size. + +The cost is per *message*, so the message rate is what matters. Batching many log records into one message and letting an `otlp_*` input format expand them into separate documents is what moves a workload from the second case to the first: one acknowledgment then covers the whole batch. + +#### Cooperative indexing + +[`enable_cooperative_indexing`](node-config.md#indexer-configuration) is a **node-level** switch, off by default. It cannot be set per index or per source: one semaphore per indexer, sized to the node's blocking-thread count, is shared by every pipeline that node runs. + +With it off, an indexer cuts a split when the commit timeout fires or the split hits its size limits, and every pipeline indexes whenever it has work. With it on, a pipeline first sleeps a phase offset derived from a hash of its pipeline id, spread over `[0, commit_timeout)`; it then takes a semaphore permit, cuts its split as soon as *its own mailbox drains* rather than waiting for the commit timeout, releases the permit, and sleeps out the rest of a `commit_timeout`-long cycle. So a bounded number of pipelines hold an `IndexWriter` at any instant, and their CPU, disk and network spikes are spread across the commit window instead of landing together. + +What it costs and buys, measured on a 4.75 GiB corpus with a 10 s commit timeout: + +| | throughput, coop off → on | first published split, off → on | +| --- | --- | --- | +| 512 KiB messages, 1 pipeline | 52.1 → 52.1 MiB/s | — | +| 512 KiB messages, 16 pipelines | 177.8 → 209.5 MiB/s | 13.2 s → 3.0 s | +| ~1 KiB messages, 1 pipeline | 33.6 → 24.2 MiB/s | — | +| ~1 KiB messages, 16 pipelines | 38.6 → 36.7 MiB/s | 13.2 s → 3.0 s | + +At 512 KiB it costs nothing measurable and still makes the first split searchable 4.4× sooner. On small documents it costs up to 28 % for a single busy pipeline, shrinking to a few percent by sixteen — the penalty is worst for one pipeline with a lot to do, which is the opposite of what the setting is for. Enable it when a node runs many pipelines that are each mostly idle, which is also when its bound on concurrent `IndexWriter`s is what keeps the node inside its memory budget. -`max_ack_pending` bounds the messages in flight and therefore the achievable throughput: nothing is acknowledged until a split is published, so a whole commit window is always ack-pending and the achievable rate is `max_ack_pending / commit_timeout` documents per second. Measured at four pipelines with a 10 s commit timeout, 1,000 gave 97 documents per second and 20,000 gave 2,115. Leave it unlimited (`-1`) unless there is a reason to bound the source's memory, and note the trade: unlimited also makes the replay window after a crash the whole delivered span above the ack floor rather than a bounded slice of it. +One caveat: the initial phase sleep is up to one full `commit_timeout`, so with a 60 s commit timeout a pipeline can sit idle for up to a minute after a restart or reassignment before it indexes anything. That is a ramp after each deploy rather than an ongoing cost; shorten the commit timeout if it matters. #### Monitoring diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md index 3c63c766f92..54b669faf58 100644 --- a/docs/ingest-data/nats.md +++ b/docs/ingest-data/nats.md @@ -136,9 +136,14 @@ nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack exp :::info -The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. +The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source. -`--max-pending=-1` leaves the in-flight window unbounded: the source acknowledges only once a split is published, so a whole commit window is always ack-pending and a small bound throttles indexing. See the [NATS source reference](../configuration/source-config.md#nats-source) for the rest, and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. +Two of those settings decide whether indexing works at all, and neither is checked when the source is created: + +- `--wait` (`ack_wait`) must outlast the whole path from delivery to acknowledgment — the commit timeout, plus the split upload and publish, plus the acknowledgment round trip. Shorter than that and NATS redelivers messages that are still being indexed, which duplicates them. `5m` against this tutorial's 10 s commit timeout is a wide margin. It is also the recovery time from a lost delivery, so it should not be arbitrarily large either. +- `--max-pending=-1` leaves the in-flight window unbounded. The source acknowledges only once a split is published, so a whole commit window is always ack-pending and a bounded value is a hard throughput cap of roughly `max_ack_pending / commit_timeout` documents per second. + +See [Consumer invariants](../configuration/source-config.md#consumer-invariants) for the details, and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. ::: @@ -242,8 +247,10 @@ Let's delete the files and resources created for the purpose of this tutorial. ```bash # Delete quickwit index. ./quickwit index delete --index stackoverflow --yes -# Delete NATS stream. +# Delete NATS stream, and with it the consumer it holds. nats stream rm -f stackoverflow ``` +Note that deleting a Quickwit source does **not** delete its consumer: the consumer is provisioned outside of Quickwit and outlives the source. Here the stream is deleted too, which removes it. Deleting only the source would leave the consumer behind, still holding its position. + This concludes the tutorial. If you have any questions regarding Quickwit or encounter any issues, don't hesitate to ask a [question](https://github.com/quickwit-oss/quickwit/discussions) or open an [issue](https://github.com/quickwit-oss/quickwit/issues) on [GitHub](https://github.com/quickwit-oss/quickwit) or contact us directly on [Discord](https://discord.com/invite/MT27AG5EVE). From 4fddd3ae7819664aaa0f7dec2a2c3ca2ab381853 Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 14:38:59 +0200 Subject: [PATCH 36/41] fix: source config --- docs/configuration/source-config.md | 45 ++++++----------------------- 1 file changed, 9 insertions(+), 36 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 1ff6a607f53..380e9095104 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -189,7 +189,7 @@ A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-conce A tutorial is available [here](/docs/ingest-data/nats.md). -The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning belong to whoever provisioned it. +The durable consumer is provisioned externally so its lifecycle, subject filters, deliver policy, and ack tuning belong to whoever provisioned it. Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. @@ -205,17 +205,17 @@ These are properties of the consumer, not of the source. Only the ack policy is ack_wait > commit_timeout + split upload and publish + ack round trip ``` -Shorter than their sum, and NATS redelivers messages that are still being indexed; they are then indexed twice. 5 minutes with a 60 s `commit_timeout` leaves a wide margin. +If `ack_wait` is shorter, NATS redelivers messages that are still being indexed; they are then indexed twice. 5 minutes with a 60 s `commit_timeout` leaves a wide margin. -It is also the *recovery* time from a lost delivery, so it should not be arbitrarily large either. If a message is delivered but never arrives — a dropped connection, a slow-consumer disconnect — nothing brings it back until the timer expires and the pipeline simply idles. The same injected fault idled a run for 278 s at `ack_wait=300s` and 23.2 s at `30s`. +It is also the *recovery* time from a lost delivery, so it should not be arbitrarily large either. If a message is delivered but never arrives (a dropped connection, a slow-consumer disconnect) nothing brings it back until the timer expires and the pipeline simply idles. -**`max_ack_pending` should be `-1`.** Nothing is acknowledged until a split is published, so a whole commit window is always ack-pending. That makes the setting a hard throughput cap rather than a safety valve: +**`max_ack_pending` should be `-1`.** Nothing is acknowledged until a split is published, so a whole commit window is always ack-pending. This setting is a hard throughput cap rather than a safety valve: ``` achievable rate ≈ max_ack_pending / commit_timeout documents per second ``` -Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing to a trickle here. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. +Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. **`deliver_policy`** is usually `all`, so a new consumer indexes the stream from the start. @@ -226,48 +226,21 @@ Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents | Setting | Where | Effect | | --- | --- | --- | | `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | -| [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. It is also the length of the cooperative-indexing cycle. | -| [`docstore_compression_level`](index-config.md#indexing-settings) | index config | Defaults to zstd 8, which is what limits a single pipeline on small documents: the dedicated docstore thread saturates and back-pressures the indexer. Level 1 or 3 measured 21.1 s against 31.2 s on a 1.5 GiB corpus of 1 KiB documents. Not NATS-specific. | -| [`enable_cooperative_indexing`](node-config.md#indexer-configuration) | node config | See below. | +| [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. | | `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. 10 MiB and 20 MiB measure identically. | -| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Does not bind at these message sizes — 200 and 100,000 measured identically — because the byte cap is what bounds a batch. | +| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Can be used along `QW_NATS_PULL_MAX_BYTES_PER_BATCH` to limit a batch if there is too much contention around acknowledgements | #### Scaling -**One consumer per partition of the data** (recommended). Give each partition — a tenant, a subject, whatever the natural unit is — its own durable consumer with a single `filter_subject`, and add one Quickwit source with `num_pipelines: 1` per consumer. The partitions then scale and fail independently, and each gets its own checkpoint. - -**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. - -At 512 KiB per message the two shapes are equivalent — 52.1, 91.3 and 146.6 MiB/s at one, two and four units either way — and so is partitioning across separate streams instead of subjects. At 1 KiB they diverge: a shared consumer stops gaining past a couple of pipelines while per-tenant consumers keep scaling, because the acknowledgment cost is per message rather than per pipeline. The per-consumer shape is therefore the safe default: never worse, and better on small documents. +Several pipelines can bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. #### Acknowledgment cost One confirmed acknowledgment per message is the source's dominant broker cost, and it is what decides whether the source or the indexer is the bottleneck. -- At **512 KiB** per message it is invisible. The NATS and Kafka sources measured within a percent of each other across the whole pipeline sweep, and broker CPU stayed at 0.5–0.9 cores for both, which is the host's floor. Priced on the broker alone, a confirmed acknowledgment per message still leaves 1,138 MiB/s — several times what one indexer consumes. +- At **512 KiB** per message it is invisible. The NATS and Kafka sources measured within a percent of each other across the whole pipeline sweep. - At **1 KiB** it is the ceiling. Sixteen per-tenant consumers reached 72 MiB/s while an equivalent Kafka source reached 150, and the broker spent 5.1 of the host's cores on acknowledgments against Quickwit's 4.7 on indexing — about ten times the broker CPU per document that a periodic offset commit costs. Priced on the broker alone, acknowledgment divides throughput by 6.5 at this size. -The cost is per *message*, so the message rate is what matters. Batching many log records into one message and letting an `otlp_*` input format expand them into separate documents is what moves a workload from the second case to the first: one acknowledgment then covers the whole batch. - -#### Cooperative indexing - -[`enable_cooperative_indexing`](node-config.md#indexer-configuration) is a **node-level** switch, off by default. It cannot be set per index or per source: one semaphore per indexer, sized to the node's blocking-thread count, is shared by every pipeline that node runs. - -With it off, an indexer cuts a split when the commit timeout fires or the split hits its size limits, and every pipeline indexes whenever it has work. With it on, a pipeline first sleeps a phase offset derived from a hash of its pipeline id, spread over `[0, commit_timeout)`; it then takes a semaphore permit, cuts its split as soon as *its own mailbox drains* rather than waiting for the commit timeout, releases the permit, and sleeps out the rest of a `commit_timeout`-long cycle. So a bounded number of pipelines hold an `IndexWriter` at any instant, and their CPU, disk and network spikes are spread across the commit window instead of landing together. - -What it costs and buys, measured on a 4.75 GiB corpus with a 10 s commit timeout: - -| | throughput, coop off → on | first published split, off → on | -| --- | --- | --- | -| 512 KiB messages, 1 pipeline | 52.1 → 52.1 MiB/s | — | -| 512 KiB messages, 16 pipelines | 177.8 → 209.5 MiB/s | 13.2 s → 3.0 s | -| ~1 KiB messages, 1 pipeline | 33.6 → 24.2 MiB/s | — | -| ~1 KiB messages, 16 pipelines | 38.6 → 36.7 MiB/s | 13.2 s → 3.0 s | - -At 512 KiB it costs nothing measurable and still makes the first split searchable 4.4× sooner. On small documents it costs up to 28 % for a single busy pipeline, shrinking to a few percent by sixteen — the penalty is worst for one pipeline with a lot to do, which is the opposite of what the setting is for. Enable it when a node runs many pipelines that are each mostly idle, which is also when its bound on concurrent `IndexWriter`s is what keeps the node inside its memory budget. - -One caveat: the initial phase sleep is up to one full `commit_timeout`, so with a 60 s commit timeout a pipeline can sit idle for up to a minute after a restart or reassignment before it indexes anything. That is a ramp after each deploy rather than an ongoing cost; shorten the commit timeout if it matters. - #### Monitoring Being durable, the consumer is observable through NATS's own monitoring (`nats consumer info`, exporters): `num_pending` is the indexing lag and `num_ack_pending` the in-flight window, both available even while the pipelines are down. The source also reports `num_pending_acks` (messages indexed but whose split is not published yet) in its observable state. From 29621c975173dfaaf3500a89069ea1840febf59d Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 16:52:11 +0200 Subject: [PATCH 37/41] feat(indexers): fixes --- docs/configuration/source-config.md | 18 +- docs/ingest-data/nats.md | 17 +- docs/overview/introduction.md | 2 +- docs/reference/rest-api.md | 4 +- quickwit/Cargo.lock | 1 - quickwit/quickwit-common/src/tracing_utils.rs | 51 ++- .../quickwit-config/src/source_config/mod.rs | 4 +- .../src/source_config/serialize.rs | 2 +- quickwit/quickwit-indexing/Cargo.toml | 5 +- .../src/source/nats_source.rs | 330 ++++++++++++++---- 10 files changed, 344 insertions(+), 90 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 380e9095104..de5251d90bf 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -23,7 +23,7 @@ The source ID is a string that uniquely identifies the source within an index. I ## Source type -The source type designates the kind of source being configured. As of version 0.5, available source types are `ingest-api`, `kafka`, `kinesis`, and `pulsar`. The `file` type is also supported but only for local ingestion from [the CLI](/docs/reference/cli.md#tool-local-ingest). +The source type designates the kind of source being configured. Available source types are `ingest-api`, `kafka`, `kinesis`, `nats`, `pubsub`, and `pulsar`. The `file` type is also supported but only for local ingestion from [the CLI](/docs/reference/cli.md#tool-local-ingest). ## Source parameters @@ -193,9 +193,11 @@ The durable consumer is provisioned externally so its lifecycle, subject filters Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. +The source never waits on unreachable NATS servers: acknowledgments and the negative acknowledgments sent at teardown are bounded by a 10 s timeout, unconfirmed acknowledgments stay pending and are retried at the next split publication, and a pipeline can still be stopped or reassigned during an outage. + #### Consumer invariants -These are properties of the consumer, not of the source. Only the ack policy is enforced when the source is created; the rest are not checked, and getting them wrong looks like Quickwit being slow or duplicating rather than like a consumer misconfiguration. +These are properties of the consumer, not of the source. Only the ack policy is enforced when the source is created. The pull limits and `max_deliver` are checked when a pipeline starts and only produce a warning in the indexer logs. The rest are not checked at all, and getting them wrong looks like Quickwit being slow or duplicating rather than like a consumer misconfiguration. **`ack_policy` must be `explicit`.** The source acknowledges each message individually once the split containing it is published, and waits for the server to confirm the acknowledgment — the confirmation is what tells the drain that the pipeline is empty. The consumer's ack floor is therefore the resume point, and it is the only progress state that matters. Any other policy is rejected when the source is created. @@ -217,6 +219,12 @@ achievable rate ≈ max_ack_pending / commit_timeout documents per second Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. +**`max_deliver` should stay unlimited (`-1`).** Redelivery is what recovers from a crash or a lost delivery. A finite value turns it into data loss: a message that exhausts it is never delivered again and is silently skipped. The source warns at start when it is finite. + +**`max_batch`, `max_bytes` and `max_expires` must not be below what the source pulls.** Each pull request asks for up to `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` messages and `QW_NATS_PULL_MAX_BYTES_PER_BATCH` bytes with a 30 s expiry. A consumer limit below any of those makes the server reject every pull request and the pipeline idles. The source warns at start. + +**The consumer must outlive the source.** Deleting it, or letting an `inactive_threshold` expire, while a source is bound to it does not stop the pipelines: the server answers their pull requests with "no responders", which is also what a JetStream outage produces, so the source retries indefinitely and warns. Recreate the consumer under the same name to resume, or disable the source. + **`deliver_policy`** is usually `all`, so a new consumer indexes the stream from the start. **Give each consumer a single `filter_subject`.** A consumer configured with several filters through NATS 2.10's multi-filter `filter_subjects` appears to take a much slower path in the broker: on a stream whose subjects are interleaved, sixteen multi-filter consumers took 7.5 times longer than the same sixteen consumers with one filter each, with the broker saturated and the indexers idle. @@ -227,8 +235,8 @@ Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents | --- | --- | --- | | `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | | [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. | -| `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. 10 MiB and 20 MiB measure identically. | -| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Can be used along `QW_NATS_PULL_MAX_BYTES_PER_BATCH` to limit a batch if there is too much contention around acknowledgements | +| `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. It must also stay above the server's `max_payload`, or larger messages are never delivered, and at or below the consumer's `max_bytes`; the source warns at start otherwise. 10 MiB and 20 MiB measure identically. | +| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Can be used along `QW_NATS_PULL_MAX_BYTES_PER_BATCH` to limit a batch if there is too much contention around acknowledgements. The client buffers eight such batches per subscription; messages beyond that are dropped client-side and come back after `ack_wait`. | #### Scaling @@ -252,7 +260,7 @@ Being durable, the consumer is observable through NATS's own monitoring (`nats c | `uris` | List of NATS server URIs (e.g. `nats://localhost:4222`). | required | | `stream` | Name of the JetStream stream to consume. | required | | `consumer` | Name of the pre-provisioned durable consumer to bind to. | required | -| `tls` | TLS options: `ca_certificates_path` (PEM file whose root certificates are trusted in addition to the system ones), and `client_certificate_path` + `client_key_path` (PEM files, set together) for mutual TLS. TLS itself is enabled by connecting to `tls://` URIs. The files are read by the indexer nodes when the connection is established. | optional | +| `tls` | TLS options: `ca_certificates_path` (PEM file whose root certificates are trusted instead of the system ones), and `client_certificate_path` + `client_key_path` (PEM files, set together) for mutual TLS. TLS itself is enabled by connecting to `tls://` URIs. The files are read when a connection is established: by the indexer nodes running the source, and by the node serving the source creation or update, which checks connectivity. | optional | | `authentication` | Authentication parameters: either `user_password` (with `user` and `password`) or `token`. | optional | *Adding a NATS source to an index with the [CLI](../reference/cli.md#source)* diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md index 54b669faf58..42e71e75691 100644 --- a/docs/ingest-data/nats.md +++ b/docs/ingest-data/nats.md @@ -131,17 +131,28 @@ The NATS source consumes a JetStream stream through a durable consumer, so messa ```bash nats stream add stackoverflow --subjects "stackoverflow.posts" --defaults -nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack explicit --max-pending=-1 --wait 5m --defaults +cat > quickwit-consumer.json < Context { + global::get_text_map_propagator(|propagator| propagator.extract(extractor)) +} + /// Extracts an OpenTelemetry context from incoming gRPC request metadata. /// Returns the empty context when no propagator is installed or no headers /// are present. pub fn extract_context(metadata: &MetadataMap) -> Context { - let extractor = MetadataExtractor(metadata); - global::get_text_map_propagator(|propagator| propagator.extract(&extractor)) + extract_context_from(&MetadataExtractor(metadata)) +} + +/// Extracts the context propagated through `extractor` (e.g. a W3C +/// `traceparent` header), or `None` when it carries no valid span context, so +/// callers can skip creating a span for it. +pub fn extract_remote_context(extractor: &impl Extractor) -> Option { + let context = extract_context_from(extractor); + if !context.span().span_context().is_valid() { + return None; + } + Some(context) +} + +/// Parents `span` on `parent_context`, typically one returned by +/// [`extract_remote_context`]. +pub fn set_span_parent(span: &Span, parent_context: Context) { + let _ = span.set_parent(parent_context); } /// Extracts a W3C trace context from incoming gRPC request metadata and @@ -87,7 +109,7 @@ pub fn extract_context(metadata: &MetadataMap) -> Context { /// caller's trace. pub fn set_current_span_parent_from_metadata(metadata: &MetadataMap) { let parent_context = extract_context(metadata); - let _ = Span::current().set_parent(parent_context); + set_span_parent(&Span::current(), parent_context); } /// Tonic interceptor that injects the active span's W3C trace context into @@ -149,6 +171,27 @@ mod tests { ); } + #[test] + fn extract_remote_context_requires_a_valid_span_context() { + // The helper reads the process-wide propagator, a no-op by default. + global::set_text_map_propagator(TraceContextPropagator::new()); + + let mut metadata = MetadataMap::new(); + assert!(extract_remote_context(&MetadataExtractor(&metadata)).is_none()); + + metadata.insert( + "traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + .parse() + .unwrap(), + ); + let context = extract_remote_context(&MetadataExtractor(&metadata)).unwrap(); + assert_eq!( + context.span().span_context().trace_id(), + known_span_context().trace_id() + ); + } + #[test] fn extract_returns_invalid_span_context_when_no_traceparent() { let propagator = TraceContextPropagator::new(); diff --git a/quickwit/quickwit-config/src/source_config/mod.rs b/quickwit/quickwit-config/src/source_config/mod.rs index 9283f2a1dd5..afcedc3c747 100644 --- a/quickwit/quickwit-config/src/source_config/mod.rs +++ b/quickwit/quickwit-config/src/source_config/mod.rs @@ -683,8 +683,8 @@ impl NatsSourceParams { )] #[serde(deny_unknown_fields)] pub struct NatsSourceTls { - /// Path to a PEM file whose root certificates are trusted in addition to - /// the system ones (e.g. a private CA). + /// Path to a PEM file whose root certificates are trusted instead of the + /// system ones (e.g. a private CA). #[serde(default, skip_serializing_if = "Option::is_none")] pub ca_certificates_path: Option, /// Path to the client certificate PEM file, for mutual TLS. Requires diff --git a/quickwit/quickwit-config/src/source_config/serialize.rs b/quickwit/quickwit-config/src/source_config/serialize.rs index 3b412a9ab19..e00402aafcd 100644 --- a/quickwit/quickwit-config/src/source_config/serialize.rs +++ b/quickwit/quickwit-config/src/source_config/serialize.rs @@ -132,7 +132,7 @@ impl SourceConfigForSerialization { _ => { if self.num_pipelines > 1 { bail!( - "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources.\ + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources. \ open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" ); } diff --git a/quickwit/quickwit-indexing/Cargo.toml b/quickwit/quickwit-indexing/Cargo.toml index 172dfbf0133..bf8aaaffc60 100644 --- a/quickwit/quickwit-indexing/Cargo.toml +++ b/quickwit/quickwit-indexing/Cargo.toml @@ -32,7 +32,6 @@ itertools = { workspace = true } libz-sys = { workspace = true, optional = true } oneshot = { workspace = true } openssl = { workspace = true, optional = true } -opentelemetry = { workspace = true, optional = true } percent-encoding = { workspace = true } pulsar = { workspace = true, optional = true } quickwit-query = { workspace = true } @@ -47,7 +46,6 @@ thiserror = { workspace = true } time = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } -tracing-opentelemetry = { workspace = true, optional = true } ulid = { workspace = true } utoipa = { workspace = true } vrl = { workspace = true, optional = true } @@ -86,7 +84,7 @@ kinesis = [ "quickwit-aws/kinesis", ] kinesis-localstack-tests = [] -nats = ["dep:async-nats", "dep:opentelemetry", "dep:tracing-opentelemetry"] +nats = ["dep:async-nats"] nats-broker-tests = ["nats"] pulsar = ["dep:pulsar"] pulsar-broker-tests = [] @@ -122,6 +120,7 @@ ci-test = [] bytes = { workspace = true } criterion = { workspace = true, features = ["async_tokio"] } mockall = { workspace = true } +opentelemetry = { workspace = true } opentelemetry_sdk = { workspace = true } parquet = { workspace = true } proptest = { workspace = true } diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index f79c8ba68c3..58091ba7dea 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -34,21 +34,19 @@ use std::str::FromStr; use std::time::{Duration, Instant}; use anyhow::{Context as _, anyhow, bail, ensure}; -use async_nats::header::{HeaderMap, HeaderName}; +use async_nats::connection::State; +use async_nats::header::HeaderMap; use async_nats::jetstream::consumer::pull::{ MessagesError, MessagesErrorKind, Stream as DurableMessageStream, }; use async_nats::jetstream::consumer::{AckPolicy, PullConsumer}; use async_nats::jetstream::message::AckKind; -use async_nats::{ConnectOptions, Subject, jetstream}; +use async_nats::{ConnectOptions, HeaderName, Subject, jetstream}; use async_trait::async_trait; use bytesize::ByteSize; use futures::{FutureExt, StreamExt}; -use opentelemetry::propagation::Extractor; -use opentelemetry::trace::TraceContextExt; -use opentelemetry::{Context as OtelContext, global}; use quickwit_actors::ActorExitStatus; -use quickwit_common::rand::append_random_suffix; +use quickwit_common::tracing_utils::{self, Extractor}; use quickwit_config::{NatsSourceAuth, NatsSourceParams}; use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpoint}; use quickwit_proto::metastore::SourceType; @@ -56,7 +54,6 @@ use quickwit_proto::types::Position; use serde_json::{Value as JsonValue, json}; use tokio::time; use tracing::{Span, debug, info, warn}; -use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::source::{ BATCH_NUM_BYTES_LIMIT, BatchBuilder, EMIT_BATCHES_TIMEOUT, Source, SourceContext, @@ -122,22 +119,30 @@ impl NatsSource { source_params: NatsSourceParams, ) -> anyhow::Result { let consumer_name = source_params.consumer.clone(); + let pull_max_bytes_per_batch = pull_max_bytes_per_batch(); info!( index_id=%source_runtime.index_id(), source_id=%source_runtime.source_id(), stream=%source_params.stream, %consumer_name, - pull_max_bytes_per_batch=%pull_max_bytes_per_batch(), + %pull_max_bytes_per_batch, pull_max_messages_per_batch=%pull_max_messages_per_batch(), "starting NATS source" ); let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; + warn_on_incompatible_pull_settings( + &nats_client, + &consumer.cached_info().config, + pull_max_bytes_per_batch, + pull_max_messages_per_batch(), + ); + let message_stream = consumer .stream() .max_messages_per_batch(pull_max_messages_per_batch()) - .max_bytes_per_batch(pull_max_bytes_per_batch().0 as usize) + .max_bytes_per_batch(pull_max_bytes_per_batch.0 as usize) // default heartbeat that is disabled when using manual builder, // required for transient conn errors. .heartbeat(PULL_IDLE_HEARTBEAT) @@ -146,11 +151,23 @@ impl NatsSource { .await .context("failed to subscribe to NATS consumer messages")?; - // A fresh partition per pipeline: durable-mode positions are delivery - // counters local to this instance, and pipelines sharing the consumer - // must not collide on a partition. - let partition_id = - PartitionId::from(append_random_suffix(&format!("nats-{consumer_name}"))); + // One partition per pipeline: positions are delivery counters local to + // this pipeline, and pipelines sharing the consumer must not collide on + // a partition. The pipeline uid is stable across respawns, so a respawn + // resumes the counter from the checkpoint instead of opening a new + // partition. + let partition_id = PartitionId::from(format!( + "nats-{consumer_name}-{}", + source_runtime.pipeline_uid() + )); + let checkpoint = source_runtime + .fetch_checkpoint() + .await + .context("failed to fetch the source checkpoint")?; + let delivery_counter = checkpoint + .position_for_partition(&partition_id) + .and_then(Position::as_u64) + .unwrap_or(0); Ok(NatsSource { source_runtime, @@ -159,7 +176,7 @@ impl NatsSource { message_stream, consumer_name, partition_id, - delivery_counter: 0, + delivery_counter, pending_acks: BTreeMap::new(), state: NatsSourceState::default(), }) @@ -221,14 +238,10 @@ impl NatsSource { if acks.is_empty() { return; } + let nats_client = &self.nats_client; let mut ack_results = futures::stream::iter(acks) .map(async |(delivery_counter, ack_subject)| { - // Mirrors `jetstream::Message::double_ack()`, but through the - // client's muxed inbox rather than a subscription per message. - let ack_result = self - .nats_client - .request(ack_subject, AckKind::Ack.into()) - .await; + let ack_result = send_ack(nats_client, ack_subject).await; (delivery_counter, ack_result) }) .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); @@ -238,7 +251,7 @@ impl NatsSource { let mut last_ack_error = None; while let Some((delivery_counter, ack_result)) = ack_results.next().await { match ack_result { - Ok(_ack_reply) => { + Ok(()) => { self.pending_acks.remove(&delivery_counter); num_acks += 1; } @@ -253,6 +266,38 @@ impl NatsSource { } debug!(num_acks, "acked published messages"); } + + /// Negatively acknowledges the messages the client prefetched, so a + /// surviving pipeline picks them up promptly instead of after `ack_wait`. + async fn nak_prefetched_messages(&mut self) { + let mut num_naks = 0usize; + while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { + let message = match message_res { + Ok(message) => message, + Err(error) => { + warn!(%error, "failed to pull a prefetched NATS message to negatively acknowledge it"); + continue; + } + }; + let Some(ack_subject) = message.message.reply else { + warn!( + "prefetched NATS message carries no reply subject to negatively acknowledge \ + it on" + ); + continue; + }; + if let Err(error) = self + .nats_client + .publish(ack_subject, AckKind::Nak(Some(NAK_REDELIVERY_DELAY)).into()) + .await + { + warn!(%error, "failed to negatively acknowledge a prefetched NATS message"); + continue; + } + num_naks += 1; + } + debug!(num_naks, "negatively acknowledged prefetched messages"); + } } #[async_trait] @@ -277,7 +322,7 @@ impl Source for NatsSource { let message = match message_res { Ok(message) => message, Err(error) if is_transient_stream_error(&error) => { - warn!(%error, "transient NATS message stream error"); + warn!(%error, "transient NATS message stream error, retrying"); wait_before_next_batch = TRANSIENT_STREAM_ERROR_BACKOFF; break; } @@ -351,38 +396,21 @@ impl Source for NatsSource { ) -> anyhow::Result<()> { // Any ack still pending at this point was not completed by a drain: // its message is redelivered after the consumer's `ack_wait`. - let mut num_naks = 0usize; - while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { - let message = match message_res { - Ok(message) => message, - Err(error) => { - warn!(%error, "failed to pull a prefetched NATS message to negatively acknowledge it"); - continue; - } - }; - let Some(ack_subject) = message.message.reply else { - warn!( - "prefetched NATS message carries no reply subject to negatively acknowledge \ - it on" - ); - continue; - }; - if let Err(error) = self - .nats_client - .publish(ack_subject, AckKind::Nak(Some(NAK_REDELIVERY_DELAY)).into()) - .await - { - warn!(%error, "failed to negatively acknowledge a prefetched NATS message"); - continue; + // + // Bounded for the same reason as `ACK_REQUEST_TIMEOUT`: a dead server + // must not hold the pipeline teardown, and node shutdown with it. Past + // the deadline, `ack_wait` redelivery takes over. + let teardown = async { + self.nak_prefetched_messages().await; + if let Err(error) = self.nats_client.flush().await { + warn!(%error, "failed to flush NATS negative acknowledgments"); } - num_naks += 1; - } - - if let Err(error) = self.nats_client.flush().await { - warn!(%error, num_naks, "failed to flush NATS negative acknowledgments"); - } - if let Err(error) = self.nats_client.drain().await { - warn!(%error, "failed to drain the NATS connection"); + if let Err(error) = self.nats_client.drain().await { + warn!(%error, "failed to drain the NATS connection"); + } + }; + if time::timeout(FINALIZE_TIMEOUT, teardown).await.is_err() { + warn!("timed out negatively acknowledging prefetched NATS messages"); } Ok(()) } @@ -402,6 +430,25 @@ impl Source for NatsSource { } } +/// Server-confirmed acknowledgment ("double ack"). Mirrors +/// `jetstream::Message::double_ack()`, but through the client's muxed inbox +/// rather than a subscription per message. +async fn send_ack(nats_client: &async_nats::Client, ack_subject: Subject) -> anyhow::Result<()> { + // Nothing drains the client's outbound queue while it reconnects, so a + // request issued now would only sit there. Unconfirmed acks stay pending + // and are retried by the next `SuggestTruncate`. + ensure!( + nats_client.connection_state() == State::Connected, + "NATS connection is down" + ); + let request = nats_client.request(ack_subject, AckKind::Ack.into()); + let ack_reply_res = time::timeout(ACK_REQUEST_TIMEOUT, request) + .await + .context("ack request timed out")?; + ack_reply_res.context("ack request failed")?; + Ok(()) +} + /// Bounds the concurrent server-confirmed ack requests: enough to hide the /// round-trip latency without flooding the connection. /// @@ -426,6 +473,70 @@ fn pull_max_bytes_per_batch() -> ByteSize { ) } +/// Warns about pull settings the server or the consumer would reject, or that +/// lose messages. The consumer's limits belong to whoever provisioned it, so +/// the source reports them rather than refusing to start. +fn warn_on_incompatible_pull_settings( + nats_client: &async_nats::Client, + consumer_config: &jetstream::consumer::Config, + pull_max_bytes_per_batch: ByteSize, + pull_max_messages_per_batch: usize, +) { + let pull_max_bytes = pull_max_bytes_per_batch.as_u64(); + if pull_max_bytes < BATCH_NUM_BYTES_LIMIT { + warn!( + %pull_max_bytes_per_batch, + indexing_batch_num_bytes=BATCH_NUM_BYTES_LIMIT, + "`QW_NATS_PULL_MAX_BYTES_PER_BATCH` is below the indexing batch size: every batch \ + costs several pull round trips" + ); + } + let max_payload = nats_client.server_info().max_payload as u64; + if max_payload > 0 && pull_max_bytes < max_payload { + warn!( + %pull_max_bytes_per_batch, + max_payload, + "`QW_NATS_PULL_MAX_BYTES_PER_BATCH` is below the server's `max_payload`: messages \ + larger than it can never be delivered" + ); + } + let consumer_max_bytes = consumer_config.max_bytes; + if consumer_max_bytes > 0 && pull_max_bytes > consumer_max_bytes as u64 { + warn!( + %pull_max_bytes_per_batch, + consumer_max_bytes, + "`QW_NATS_PULL_MAX_BYTES_PER_BATCH` exceeds the consumer's `max_bytes`: the server \ + rejects every pull request" + ); + } + let consumer_max_batch = consumer_config.max_batch; + if consumer_max_batch > 0 && pull_max_messages_per_batch as i64 > consumer_max_batch { + warn!( + pull_max_messages_per_batch, + consumer_max_batch, + "`QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` exceeds the consumer's `max_batch`: the server \ + rejects every pull request" + ); + } + let consumer_max_expires = consumer_config.max_expires; + if !consumer_max_expires.is_zero() && PULL_EXPIRES > consumer_max_expires { + warn!( + ?consumer_max_expires, + pull_expires=?PULL_EXPIRES, + "the consumer's `max_expires` is below the source's pull expiry: the server \ + rejects every pull request" + ); + } + let consumer_max_deliver = consumer_config.max_deliver; + if consumer_max_deliver > 0 { + warn!( + consumer_max_deliver, + "the consumer's `max_deliver` is finite: a message that exhausts it is never \ + redelivered, so a crash or a lost delivery can leave it unindexed" + ); + } +} + /// Messages per pull batch. The byte cap above is what actually bounds a batch; /// this only binds for small messages, where it caps a pull far below the byte /// budget (200 messages of 1 KiB is under 200 KiB against a 10 MiB budget) and @@ -442,6 +553,15 @@ fn pull_max_messages_per_batch() -> usize { ) } +/// Capacity of the client's per-subscription channel, in pull batches. The +/// connection handler drops a message that finds the channel full, and the +/// server already counts it as delivered, so it only comes back after +/// `ack_wait`. The client keeps up to 1.5 batches outstanding and re-pulls a +/// full batch every 35 s while the source is not polling (e.g. during a +/// drain), so the channel must hold several. Memory is bounded by the byte +/// cap rather than by this: eight batches are at most 80 MiB by default. +const SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES: usize = 8; + // Both kept at the `Consumer::messages` defaults, so switching to the builder // changes the batch bounds and nothing else. const PULL_IDLE_HEARTBEAT: Duration = Duration::from_secs(15); @@ -476,6 +596,12 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { /// messages: 1024 is indistinguishable from 128, and 8192 is 20 % slower. const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; +/// Bounds an ack request end to end, enqueueing included. +const ACK_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); + +/// Bounds the NAK sweep and connection drain at teardown. +const FINALIZE_TIMEOUT: Duration = Duration::from_secs(10); + async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, @@ -492,8 +618,8 @@ async fn fetch_durable_consumer( Ok(consumer) } -/// W3C trace context extraction from NATS message headers, following the -/// pattern of `quickwit_common::tracing_utils` for gRPC metadata. +/// Lets `quickwit_common::tracing_utils` read the W3C trace context from +/// NATS message headers. struct NatsHeaderExtractor<'a>(&'a HeaderMap); impl Extractor for NatsHeaderExtractor<'_> { @@ -507,10 +633,6 @@ impl Extractor for NatsHeaderExtractor<'_> { } } -fn extract_remote_context(headers: &HeaderMap) -> OtelContext { - global::get_text_map_propagator(|propagator| propagator.extract(&NatsHeaderExtractor(headers))) -} - /// Builds a span parented on the publisher's trace when the message carries /// a W3C `traceparent` header, stitching the processing of the message into /// the publisher's distributed trace. Messages without a propagated context @@ -521,10 +643,7 @@ fn remote_parented_span( source_runtime: &SourceRuntime, ) -> Option { let headers = message.headers.as_ref()?; - let parent_context = extract_remote_context(headers); - if !parent_context.span().span_context().is_valid() { - return None; - } + let parent_context = tracing_utils::extract_remote_context(&NatsHeaderExtractor(headers))?; let span = tracing::info_span!( "process_nats_message", index_id = %source_runtime.index_id(), @@ -532,12 +651,14 @@ fn remote_parented_span( subject = %message.subject, stream_sequence, ); - let _ = span.set_parent(parent_context); + tracing_utils::set_span_parent(&span, parent_context); Some(span) } async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result { - let mut connect_options = ConnectOptions::new(); + let subscription_capacity = + SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES * pull_max_messages_per_batch(); + let mut connect_options = ConnectOptions::new().subscription_capacity(subscription_capacity); match params.authentication.clone() { None => {} Some(NatsSourceAuth::UserPassword { user, password }) => { @@ -595,6 +716,7 @@ mod tests { #[test] fn extract_trace_context_from_headers() { use opentelemetry::propagation::TextMapPropagator; + use opentelemetry::trace::TraceContextExt; use opentelemetry_sdk::propagation::TraceContextPropagator; // The global propagator is not installed in tests, so the extractor @@ -618,6 +740,19 @@ mod tests { ); assert_eq!(span_context.span_id().to_string(), "00f067aa0ba902b7"); + // Go publishers canonicalize the header name. + let mut capitalized_headers = HeaderMap::new(); + capitalized_headers.insert( + "Traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ); + let span_context = propagator + .extract(&NatsHeaderExtractor(&capitalized_headers)) + .span() + .span_context() + .clone(); + assert!(span_context.is_valid()); + let empty_headers = HeaderMap::new(); let span_context = propagator .extract(&NatsHeaderExtractor(&empty_headers)) @@ -662,11 +797,12 @@ mod nats_broker_tests { use bytes::Bytes; use quickwit_actors::{ActorHandle, Inbox, Universe}; + use quickwit_common::rand::append_random_suffix; use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; - use quickwit_metastore::checkpoint::SourceCheckpointDelta; + use quickwit_metastore::checkpoint::{PartitionDelta, SourceCheckpointDelta}; use quickwit_metastore::metastore_for_test; use quickwit_proto::metastore::MetastoreServiceClient; - use quickwit_proto::types::IndexUid; + use quickwit_proto::types::{IndexUid, PipelineUid}; use super::*; use crate::actors::DocProcessor; @@ -1137,6 +1273,64 @@ mod nats_broker_tests { universe.assert_quit().await; } + #[tokio::test] + async fn durable_mode_respawn_resumes_partition_from_checkpoint() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-resume--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-resume-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let index_id = append_random_suffix("test-nats-source--durable-resume--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + // The partition of the pipeline the test runtime builds, as left by a + // previous incarnation that published 7 messages. + let partition_id = PartitionId::from(format!( + "nats-{consumer_name}-{}", + PipelineUid::for_test(0u128) + )); + let index_uid = setup_index( + metastore.clone(), + &index_id, + &source_config, + &[( + partition_id.clone(), + Position::Beginning, + Position::offset(7u64), + )], + ) + .await; + + let (source_handle, doc_processor_inbox) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + let subject = format!("{stream}.logs"); + publish_docs(&jetstream_ctx, &subject, 0..3).await; + wait_for_processed_messages(&source_handle, 3, Duration::from_secs(30)).await; + source_handle.quit().await; + + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + let checkpoint_delta = merge_doc_batches(batches).checkpoint_delta; + let partition_deltas: Vec<(PartitionId, PartitionDelta)> = + checkpoint_delta.iter().collect(); + // Same partition, counter continued from the checkpoint: the delta + // chains onto the published position instead of opening a new + // partition from the beginning. + assert_eq!( + partition_deltas, + vec![( + partition_id, + PartitionDelta { + from: Position::offset(7u64), + to: Position::offset(10u64), + } + )] + ); + + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + #[tokio::test] async fn durable_mode_nak_redelivers_prefetched_messages_promptly() { let universe = Universe::with_accelerated_time(); From 76ccf71c3d9889b9f50835363051e48de75c8f65 Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 17:01:50 +0200 Subject: [PATCH 38/41] feat(indexers): remove useless hb --- .../src/source/nats_source.rs | 35 +++---------------- 1 file changed, 5 insertions(+), 30 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 58091ba7dea..4f698e71287 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -143,9 +143,6 @@ impl NatsSource { .stream() .max_messages_per_batch(pull_max_messages_per_batch()) .max_bytes_per_batch(pull_max_bytes_per_batch.0 as usize) - // default heartbeat that is disabled when using manual builder, - // required for transient conn errors. - .heartbeat(PULL_IDLE_HEARTBEAT) .expires(PULL_EXPIRES) .messages() .await @@ -562,9 +559,8 @@ fn pull_max_messages_per_batch() -> usize { /// cap rather than by this: eight batches are at most 80 MiB by default. const SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES: usize = 8; -// Both kept at the `Consumer::messages` defaults, so switching to the builder -// changes the batch bounds and nothing else. -const PULL_IDLE_HEARTBEAT: Duration = Duration::from_secs(15); +/// Server-side lifetime of a pull request. The client re-pulls 5 s after it +/// elapses, which is also what recovers a pull request the server lost. const PULL_EXPIRES: Duration = Duration::from_secs(30); /// Pause before pulling again after a transient stream error, to avoid a hot @@ -579,15 +575,11 @@ const TRANSIENT_STREAM_ERROR_BACKOFF: Duration = Duration::from_secs(1); const NAK_REDELIVERY_DELAY: Duration = Duration::from_secs(1); /// Errors the pull stream recovers from on its own: the subscription stays -/// usable and polling simply resumes. `MissingHeartbeat` in particular also -/// fires when the stream was not polled for a while (e.g. under downstream -/// backpressure), not only when the server is unreachable. +/// usable and polling simply resumes. fn is_transient_stream_error(error: &MessagesError) -> bool { matches!( error.kind(), - MessagesErrorKind::MissingHeartbeat - | MessagesErrorKind::Pull - | MessagesErrorKind::NoResponders + MessagesErrorKind::Pull | MessagesErrorKind::NoResponders ) } @@ -740,19 +732,6 @@ mod tests { ); assert_eq!(span_context.span_id().to_string(), "00f067aa0ba902b7"); - // Go publishers canonicalize the header name. - let mut capitalized_headers = HeaderMap::new(); - capitalized_headers.insert( - "Traceparent", - "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", - ); - let span_context = propagator - .extract(&NatsHeaderExtractor(&capitalized_headers)) - .span() - .span_context() - .clone(); - assert!(span_context.is_valid()); - let empty_headers = HeaderMap::new(); let span_context = propagator .extract(&NatsHeaderExtractor(&empty_headers)) @@ -764,11 +743,7 @@ mod tests { #[test] fn transient_stream_error_classification() { - let transient_kinds = [ - MessagesErrorKind::MissingHeartbeat, - MessagesErrorKind::Pull, - MessagesErrorKind::NoResponders, - ]; + let transient_kinds = [MessagesErrorKind::Pull, MessagesErrorKind::NoResponders]; for kind in transient_kinds { assert!( is_transient_stream_error(&MessagesError::new(kind)), From 857a3fe72ce50ee784234ce549f0edc2f6a8d44b Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 22:17:11 +0200 Subject: [PATCH 39/41] fix: minor fixes --- .../src/source/nats_source.rs | 37 ++++++------------- 1 file changed, 11 insertions(+), 26 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 4f698e71287..01407691723 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -143,7 +143,6 @@ impl NatsSource { .stream() .max_messages_per_batch(pull_max_messages_per_batch()) .max_bytes_per_batch(pull_max_bytes_per_batch.0 as usize) - .expires(PULL_EXPIRES) .messages() .await .context("failed to subscribe to NATS consumer messages")?; @@ -266,8 +265,8 @@ impl NatsSource { /// Negatively acknowledges the messages the client prefetched, so a /// surviving pipeline picks them up promptly instead of after `ack_wait`. - async fn nak_prefetched_messages(&mut self) { - let mut num_naks = 0usize; + async fn nak_prefetched_messages(&mut self) -> usize { + let mut num_naks: usize = 0; while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { let message = match message_res { Ok(message) => message, @@ -293,7 +292,8 @@ impl NatsSource { } num_naks += 1; } - debug!(num_naks, "negatively acknowledged prefetched messages"); + + num_naks } } @@ -398,7 +398,13 @@ impl Source for NatsSource { // must not hold the pipeline teardown, and node shutdown with it. Past // the deadline, `ack_wait` redelivery takes over. let teardown = async { - self.nak_prefetched_messages().await; + let naks_count = self.nak_prefetched_messages().await; + info!( + naks_count, + "sending negative acks to nats, they will be redelivered in {}s", + ACK_REQUEST_TIMEOUT.as_secs() + ); + if let Err(error) = self.nats_client.flush().await { warn!(%error, "failed to flush NATS negative acknowledgments"); } @@ -515,23 +521,6 @@ fn warn_on_incompatible_pull_settings( rejects every pull request" ); } - let consumer_max_expires = consumer_config.max_expires; - if !consumer_max_expires.is_zero() && PULL_EXPIRES > consumer_max_expires { - warn!( - ?consumer_max_expires, - pull_expires=?PULL_EXPIRES, - "the consumer's `max_expires` is below the source's pull expiry: the server \ - rejects every pull request" - ); - } - let consumer_max_deliver = consumer_config.max_deliver; - if consumer_max_deliver > 0 { - warn!( - consumer_max_deliver, - "the consumer's `max_deliver` is finite: a message that exhausts it is never \ - redelivered, so a crash or a lost delivery can leave it unindexed" - ); - } } /// Messages per pull batch. The byte cap above is what actually bounds a batch; @@ -559,10 +548,6 @@ fn pull_max_messages_per_batch() -> usize { /// cap rather than by this: eight batches are at most 80 MiB by default. const SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES: usize = 8; -/// Server-side lifetime of a pull request. The client re-pulls 5 s after it -/// elapses, which is also what recovers a pull request the server lost. -const PULL_EXPIRES: Duration = Duration::from_secs(30); - /// Pause before pulling again after a transient stream error, to avoid a hot /// error loop while the connection recovers. `SuggestTruncate` is still /// processed in the meantime. From 3be262cb66ecbd7c6e02cb201edb443327c83548 Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 22:29:39 +0200 Subject: [PATCH 40/41] fix: minor fixes --- docker-compose.yml | 1 - docs/configuration/source-config.md | 6 +++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index 9198d13aba6..a2e88600caa 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -128,7 +128,6 @@ services: retries: 100 nats: - # 2.10 is the oldest version supporting multiple consumer filter subjects. image: nats:${NATS_VERSION:-2.10-alpine} container_name: nats command: --jetstream --http_port 8222 diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index de5251d90bf..7046843f3c8 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -219,9 +219,9 @@ achievable rate ≈ max_ack_pending / commit_timeout documents per second Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. -**`max_deliver` should stay unlimited (`-1`).** Redelivery is what recovers from a crash or a lost delivery. A finite value turns it into data loss: a message that exhausts it is never delivered again and is silently skipped. The source warns at start when it is finite. +**`max_deliver` should stay unlimited (`-1`).** Redelivery is what recovers from a crash or a lost delivery. A finite value turns it into data loss: a message that exhausts it is never delivered again and is silently skipped. -**`max_batch`, `max_bytes` and `max_expires` must not be below what the source pulls.** Each pull request asks for up to `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` messages and `QW_NATS_PULL_MAX_BYTES_PER_BATCH` bytes with a 30 s expiry. A consumer limit below any of those makes the server reject every pull request and the pipeline idles. The source warns at start. +**`max_batch`, `max_bytes` must not be below what the source pulls.** Each pull request asks for up to `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` messages and `QW_NATS_PULL_MAX_BYTES_PER_BATCH` bytes. A consumer limit below any of those makes the server reject every pull request and the pipeline idles. The source warns at start. **The consumer must outlive the source.** Deleting it, or letting an `inactive_threshold` expire, while a source is bound to it does not stop the pipelines: the server answers their pull requests with "no responders", which is also what a JetStream outage produces, so the source retries indefinitely and warns. Recreate the consumer under the same name to resume, or disable the source. @@ -233,7 +233,7 @@ Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents | Setting | Where | Effect | | --- | --- | --- | -| `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | +| `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages (1 KiB) the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | | [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. | | `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. It must also stay above the server's `max_payload`, or larger messages are never delivered, and at or below the consumer's `max_bytes`; the source warns at start otherwise. 10 MiB and 20 MiB measure identically. | | `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Can be used along `QW_NATS_PULL_MAX_BYTES_PER_BATCH` to limit a batch if there is too much contention around acknowledgements. The client buffers eight such batches per subscription; messages beyond that are dropped client-side and come back after `ack_wait`. | From f2d60fb505eda9594b18d810a988b162268c03e6 Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 23:11:51 +0200 Subject: [PATCH 41/41] feat: improve tracing --- docs/configuration/source-config.md | 2 + quickwit/quickwit-common/src/tracing_utils.rs | 16 ++++- .../src/source/nats_source.rs | 71 +++++++++++++++---- 3 files changed, 76 insertions(+), 13 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 7046843f3c8..44fcfc7981f 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -253,6 +253,8 @@ One confirmed acknowledgment per message is the source's dominant broker cost, a Being durable, the consumer is observable through NATS's own monitoring (`nats consumer info`, exporters): `num_pending` is the indexing lag and `num_ack_pending` the in-flight window, both available even while the pipelines are down. The source also reports `num_pending_acks` (messages indexed but whose split is not published yet) in its observable state. +When a message carries a W3C `traceparent` header and the indexers export their traces (see [distributed tracing](../distributed-tracing/plug-quickwit-to-jaeger.md)), the publisher's trace extends to Quickwit: a `process_nats_message` span covers the message's processing until it is acknowledged. + **NATS source parameters** | Property | Description | Default value | diff --git a/quickwit/quickwit-common/src/tracing_utils.rs b/quickwit/quickwit-common/src/tracing_utils.rs index c1711bb6b4e..4fb3774cfc6 100644 --- a/quickwit/quickwit-common/src/tracing_utils.rs +++ b/quickwit/quickwit-common/src/tracing_utils.rs @@ -20,10 +20,11 @@ use std::str::FromStr; +pub use opentelemetry::Context; +use opentelemetry::global; pub use opentelemetry::propagation::Extractor; use opentelemetry::propagation::Injector; use opentelemetry::trace::TraceContextExt; -use opentelemetry::{Context, global}; use tonic::Status; use tonic::metadata::{KeyAndValueRef, MetadataKey, MetadataMap, MetadataValue}; use tonic::service::Interceptor; @@ -102,6 +103,19 @@ pub fn set_span_parent(span: &Span, parent_context: Context) { let _ = span.set_parent(parent_context); } +/// The OpenTelemetry context of `span`, to parent or link spans created +/// later, possibly after `span` closed. +pub fn span_context(span: &Span) -> Context { + span.context() +} + +/// Links `span` to `linked_span`: a causal relation that is not a +/// parent/child one, e.g. a batch and each message it carries. No-op when +/// `linked_span` has no OpenTelemetry context. +pub fn link_span(span: &Span, linked_span: &Span) { + span.add_link(linked_span.context().span().span_context().clone()); +} + /// Extracts a W3C trace context from incoming gRPC request metadata and /// installs it as the parent of the currently-active tracing span. Use this /// at the entry of a gRPC handler that is itself wrapped in a diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 01407691723..eb9051b03d0 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -25,7 +25,9 @@ //! message and is what bounds throughput on small messages. //! //! When a message carries a W3C `traceparent` header, the source stitches the -//! processing of the message into the publisher's distributed trace. +//! processing and the acknowledgment of the message into the publisher's +//! distributed trace, and links the processing span to the batch that +//! carried the message. use std::collections::BTreeMap; use std::fmt; @@ -46,14 +48,14 @@ use async_trait::async_trait; use bytesize::ByteSize; use futures::{FutureExt, StreamExt}; use quickwit_actors::ActorExitStatus; -use quickwit_common::tracing_utils::{self, Extractor}; +use quickwit_common::tracing_utils::{self, Context as TraceContext, Extractor}; use quickwit_config::{NatsSourceAuth, NatsSourceParams}; use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpoint}; use quickwit_proto::metastore::SourceType; use quickwit_proto::types::Position; use serde_json::{Value as JsonValue, json}; use tokio::time; -use tracing::{Span, debug, info, warn}; +use tracing::{Instrument, Span, debug, info, warn}; use crate::source::{ BATCH_NUM_BYTES_LIMIT, BatchBuilder, EMIT_BATCHES_TIMEOUT, Source, SourceContext, @@ -96,11 +98,20 @@ pub struct NatsSource { /// Messages delivered but not published yet, keyed by delivery counter. /// Bounded by the consumer's `max_ack_pending`: the server stops delivering /// when too many messages are unacknowledged. - pending_acks: BTreeMap, + pending_acks: BTreeMap, state: NatsSourceState, } +/// A delivered message whose split is not published yet. +#[derive(Clone)] +struct PendingAck { + ack_subject: Subject, + /// Context of the message's processing span, when the publisher propagated + /// a trace: the ack is reported as its child. + trace_context: Option, +} + impl fmt::Debug for NatsSource { fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { formatter @@ -187,8 +198,11 @@ impl NatsSource { .info() .map_err(|error| anyhow!("failed to parse NATS message metadata: {error}"))? .stream_sequence; - let _span_guard = remote_parented_span(&message, stream_sequence, &self.source_runtime) - .map(Span::entered); + let batch_span = Span::current(); + let message_span = + remote_parented_span(&message, stream_sequence, &self.source_runtime, &batch_span); + let trace_context = message_span.as_ref().map(tracing_utils::span_context); + let _span_guard = message_span.map(Span::entered); let Some(ack_subject) = message.message.reply else { bail!("NATS message carries no reply subject to acknowledge it on"); }; @@ -215,7 +229,11 @@ impl NatsSource { Position::offset(self.delivery_counter), ) .context("failed to record partition delta")?; - self.pending_acks.insert(self.delivery_counter, ack_subject); + let pending_ack = PendingAck { + ack_subject, + trace_context, + }; + self.pending_acks.insert(self.delivery_counter, pending_ack); self.state.num_bytes_processed += num_bytes; self.state.num_messages_processed += 1; @@ -226,7 +244,7 @@ impl NatsSource { /// Sends server-confirmed acknowledgments ("double acks") for the messages up /// to `published_up_to`. async fn ack_up_to(&mut self, published_up_to: u64) { - let acks: Vec<(u64, Subject)> = self + let acks: Vec<(u64, PendingAck)> = self .pending_acks .range(..=published_up_to) .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.clone())) @@ -235,9 +253,17 @@ impl NatsSource { return; } let nats_client = &self.nats_client; + let source_runtime = &self.source_runtime; let mut ack_results = futures::stream::iter(acks) - .map(async |(delivery_counter, ack_subject)| { - let ack_result = send_ack(nats_client, ack_subject).await; + .map(async |(delivery_counter, pending_ack)| { + let ack_span = ack_span(pending_ack.trace_context, source_runtime); + let ack_result = send_ack(nats_client, pending_ack.ack_subject) + .instrument(ack_span.clone()) + .await; + if let Err(error) = &ack_result { + ack_span.record("otel.status_code", "ERROR"); + ack_span.record("otel.status_description", tracing::field::display(error)); + } (delivery_counter, ack_result) }) .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); @@ -612,12 +638,14 @@ impl Extractor for NatsHeaderExtractor<'_> { /// Builds a span parented on the publisher's trace when the message carries /// a W3C `traceparent` header, stitching the processing of the message into -/// the publisher's distributed trace. Messages without a propagated context -/// cost nothing: no span is created. +/// the publisher's distributed trace, and links it to `batch_span`, the +/// batch carrying the message. Messages without a propagated context cost +/// nothing: no span is created. fn remote_parented_span( message: &jetstream::Message, stream_sequence: u64, source_runtime: &SourceRuntime, + batch_span: &Span, ) -> Option { let headers = message.headers.as_ref()?; let parent_context = tracing_utils::extract_remote_context(&NatsHeaderExtractor(headers))?; @@ -629,9 +657,28 @@ fn remote_parented_span( stream_sequence, ); tracing_utils::set_span_parent(&span, parent_context); + tracing_utils::link_span(&span, batch_span); Some(span) } +/// Reports the ack of a message as a child of its processing span, so the +/// publisher's trace extends to durability. Disabled when the publisher +/// propagated no trace. +fn ack_span(trace_context: Option, source_runtime: &SourceRuntime) -> Span { + let Some(trace_context) = trace_context else { + return Span::none(); + }; + let span = tracing::info_span!( + "ack_nats_message", + index_id = %source_runtime.index_id(), + source_id = %source_runtime.source_id(), + otel.status_code = tracing::field::Empty, + otel.status_description = tracing::field::Empty, + ); + tracing_utils::set_span_parent(&span, trace_context); + span +} + async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result { let subscription_capacity = SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES * pull_max_messages_per_batch();