Re-work: keep one publish token, thread through indexing pipeline

Author: nadav-govari

quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs

@@ -420,7 +420,7 @@ impl IndexingScheduler {
     ) {
         debug!(new_physical_plan=?new_physical_plan, "apply physical indexing plan");
         APPLY_PLAN_TOTAL.inc();
-        // The indexing plan ID is a monotonically increased time based ID that's used as the
+        // The indexing plan ID is a monotonically increasing time based ID that's used as the
         // publish token for indexers, which ensures indexing plans and shard acquisition are always
         // informed by the most recent plan.
         let indexing_plan_id = Ulid::new().to_string();
@@ -446,12 +446,15 @@ impl IndexingScheduler {
                 IngesterStatus::Retiring | IngesterStatus::Decommissioning
             )
             .then_some(APPLY_INDEXING_PLAN_TIMEOUT);
+
             let notify_on_drop = notify_on_drop.clone();
             let indexing_plan_id = indexing_plan_id.clone();
             tokio::spawn(async move {
                 let client = indexer.client.clone();
-                let apply_plan_fut =
-                    client.apply_indexing_plan(ApplyIndexingPlanRequest { indexing_tasks,  indexing_plan_id });
+                let apply_plan_fut = client.apply_indexing_plan(ApplyIndexingPlanRequest {
+                    indexing_tasks,
+                    indexing_plan_id,
+                });
                 let apply_result = match apply_deadline {
                     Some(timeout) => tokio::time::timeout(timeout, apply_plan_fut).await,
                     None => Ok(apply_plan_fut.await),

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

@@ -41,9 +41,7 @@ use tokio::runtime::Handle;
 use super::vrl_processing::*;
 use crate::actors::Indexer;
 use crate::metrics::{PROCESSED_BYTES, PROCESSED_DOCS_TOTAL};
-use crate::models::{
-    NewPublishLock, NewPublishToken, ProcessedDoc, ProcessedDocBatch, PublishLock, RawDocBatch,
-};
+use crate::models::{NewPublishLock, ProcessedDoc, ProcessedDocBatch, PublishLock, RawDocBatch};
 
 const PLAIN_TEXT: &str = "plain_text";
 pub(super) struct JsonDoc {
@@ -607,20 +605,6 @@ impl Handler<NewPublishLock> for DocProcessor {
     }
 }
 
-#[async_trait]
-impl Handler<NewPublishToken> for DocProcessor {
-    type Reply = ();
-
-    async fn handle(
-        &mut self,
-        message: NewPublishToken,
-        ctx: &ActorContext<Self>,
-    ) -> Result<(), ActorExitStatus> {
-        ctx.send_message(&self.indexer_mailbox, message).await?;
-        Ok(())
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

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

@@ -89,7 +89,6 @@ impl Handler<IndexedSplitBatchBuilder> for IndexSerializer {
             splits,
             checkpoint_delta_opt: batch_builder.checkpoint_delta_opt,
             publish_lock: batch_builder.publish_lock,
-            publish_token_opt: batch_builder.publish_token_opt,
             merge_task_opt: None,
             batch_parent_span: batch_builder.batch_parent_span,
         };

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

@@ -38,7 +38,7 @@ use quickwit_proto::indexing::{IndexingPipelineId, PipelineMetrics};
 use quickwit_proto::metastore::{
     LastDeleteOpstampRequest, MetastoreService, MetastoreServiceClient,
 };
-use quickwit_proto::types::{DocMappingUid, PublishToken};
+use quickwit_proto::types::DocMappingUid;
 use quickwit_query::get_quickwit_fastfield_normalizer_manager;
 use serde::Serialize;
 use tantivy::schema::Schema;
@@ -55,7 +55,7 @@ use super::cooperative_indexing::{CooperativeIndexingCycle, CooperativeIndexingP
 use crate::metrics::SPLIT_BUILDERS;
 use crate::models::{
     CommitTrigger, EmptySplit, IndexedSplitBatchBuilder, IndexedSplitBuilder, NewPublishLock,
-    NewPublishToken, ProcessedDoc, ProcessedDocBatch, PublishLock,
+    ProcessedDoc, ProcessedDocBatch, PublishLock,
 };
 
 // Random partition ID used to gather partitions exceeding the maximum number of partitions.
@@ -93,7 +93,6 @@ struct IndexerState {
     indexing_directory: TempDirectory,
     indexing_settings: IndexingSettings,
     publish_lock: PublishLock,
-    publish_token_opt: Option<PublishToken>,
     schema: Schema,
     doc_mapping_uid: DocMappingUid,
     tokenizer_manager: TokenizerManager,
@@ -219,7 +218,6 @@ impl IndexerState {
             source_delta: SourceCheckpointDelta::default(),
         };
         let publish_lock = self.publish_lock.clone();
-        let publish_token_opt = self.publish_token_opt.clone();
 
         let split_builders_guard = GaugeGuard::new(&SPLIT_BUILDERS, 1.0);
 
@@ -231,7 +229,6 @@ impl IndexerState {
             other_indexed_split_opt: None,
             checkpoint_delta,
             publish_lock,
-            publish_token_opt,
             last_delete_opstamp,
             memory_usage: GaugeGuard::new(&IN_FLIGHT_INDEX_WRITER, 0.0),
             cooperative_indexing_period,
@@ -349,7 +346,6 @@ struct IndexingWorkbench {
 
     checkpoint_delta: IndexCheckpointDelta,
     publish_lock: PublishLock,
-    publish_token_opt: Option<PublishToken>,
     // On workbench creation, we fetch from the metastore the last delete task opstamp.
     // We use this value to set the `delete_opstamp` of the workbench splits.
     last_delete_opstamp: u64,
@@ -513,21 +509,6 @@ impl Handler<NewPublishLock> for Indexer {
     }
 }
 
-#[async_trait]
-impl Handler<NewPublishToken> for Indexer {
-    type Reply = ();
-
-    async fn handle(
-        &mut self,
-        message: NewPublishToken,
-        _ctx: &ActorContext<Self>,
-    ) -> Result<(), ActorExitStatus> {
-        let NewPublishToken(publish_token) = message;
-        self.indexer_state.publish_token_opt = Some(publish_token);
-        Ok(())
-    }
-}
-
 impl Indexer {
     pub fn new(
         pipeline_id: IndexingPipelineId,
@@ -563,7 +544,6 @@ impl Indexer {
                 indexing_directory,
                 indexing_settings,
                 publish_lock: PublishLock::default(),
-                publish_token_opt: None,
                 schema,
                 doc_mapping_uid: doc_mapper.doc_mapping_uid(),
                 tokenizer_manager: tokenizer_manager.tantivy_manager().clone(),
@@ -630,7 +610,6 @@ impl Indexer {
             other_indexed_split_opt,
             checkpoint_delta,
             publish_lock,
-            publish_token_opt,
             batch_parent_span,
             memory_usage,
             split_builders_guard,
@@ -656,7 +635,6 @@ impl Indexer {
                         index_uid: self.indexer_state.pipeline_id.index_uid.clone(),
                         checkpoint_delta,
                         publish_lock,
-                        publish_token_opt,
                         batch_parent_span,
                     },
                 )
@@ -680,7 +658,6 @@ impl Indexer {
                 splits,
                 checkpoint_delta_opt: Some(checkpoint_delta),
                 publish_lock,
-                publish_token_opt,
                 commit_trigger,
                 batch_parent_span,
                 memory_usage,

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

@@ -46,7 +46,7 @@ use crate::actors::uploader::UploaderType;
 use crate::actors::{Publisher, Uploader};
 use crate::merge_policy::MergePolicy;
 use crate::metrics::{ACTOR_NAME, BACKPRESSURE_MICROS, INDEXING_PIPELINES};
-use crate::models::IndexingStatistics;
+use crate::models::{IndexingStatistics, SharedPublishToken};
 use crate::source::{
     AssignShards, Assignment, SourceActor, SourceRuntime, quickwit_supported_sources,
 };
@@ -92,6 +92,7 @@ pub struct IndexingPipeline {
     // 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.
     indexing_plan_id: String,
+    publish_token: SharedPublishToken,
     _indexing_pipelines_gauge_guard: GaugeGuard,
 }
 
@@ -141,6 +142,7 @@ impl IndexingPipeline {
             },
             shard_ids: Default::default(),
             indexing_plan_id: String::new(),
+            publish_token: SharedPublishToken::default(),
             _indexing_pipelines_gauge_guard: indexing_pipelines_gauge_guard,
         }
     }
@@ -310,6 +312,7 @@ impl IndexingPipeline {
             self.params.metastore.clone(),
             Some(self.params.merge_planner_mailbox.clone()),
             Some(source_mailbox.clone()),
+            self.publish_token.clone(),
         );
         let (publisher_mailbox, publisher_handle) = ctx
             .spawn_actor()
@@ -394,6 +397,7 @@ impl IndexingPipeline {
             storage_resolver: self.params.source_storage_resolver.clone(),
             event_broker: self.params.event_broker.clone(),
             indexing_setting: self.params.indexing_settings.clone(),
+            publish_token: self.publish_token.clone(),
         };
         let source = ctx
             .protect_future(quickwit_supported_sources().load_source(source_runtime))

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

@@ -52,7 +52,7 @@ use quickwit_proto::metastore::{
     ListIndexesMetadataRequest, ListSplitsRequest, MetastoreResult, MetastoreService,
     MetastoreServiceClient,
 };
-use quickwit_proto::types::{IndexId, IndexUid, NodeId, PipelineUid, ShardId};
+use quickwit_proto::types::{IndexId, IndexUid, NodeId, PipelineUid, PublishToken, ShardId};
 use quickwit_storage::StorageResolver;
 use serde::{Deserialize, Serialize};
 use time::OffsetDateTime;
@@ -110,6 +110,7 @@ pub struct IndexingService {
     pub(crate) ingester_pool: IngesterPool,
     pub(crate) storage_resolver: StorageResolver,
     indexing_pipelines: HashMap<PipelineUid, BoxedPipelineHandle>,
+    latest_indexing_plan_id: PublishToken,
     counters: IndexingServiceCounters,
     local_split_store: Arc<IndexingSplitCache>,
     pub(crate) max_concurrent_split_uploads: usize,
@@ -182,6 +183,7 @@ impl IndexingService {
             storage_resolver,
             local_split_store: Arc::new(local_split_store),
             indexing_pipelines: Default::default(),
+            latest_indexing_plan_id: String::new(),
             counters: Default::default(),
             max_concurrent_split_uploads: indexer_config.max_concurrent_split_uploads,
             #[cfg(feature = "metrics")]
@@ -789,6 +791,17 @@ impl IndexingService {
         plan_request: ApplyIndexingPlanRequest,
         ctx: &ActorContext<Self>,
     ) -> Result<(), IndexingError> {
+        // Plan ids are ULIDs
+        if plan_request.indexing_plan_id < self.latest_indexing_plan_id {
+            info!(
+                dropped_plan_id = %plan_request.indexing_plan_id,
+                latest_plan_id = %self.latest_indexing_plan_id,
+                "ignoring stale indexing plan"
+            );
+            return Ok(());
+        }
+        self.latest_indexing_plan_id = plan_request.indexing_plan_id.clone();
+
         let pipeline_diff = self.compute_pipeline_diff(&plan_request.indexing_tasks);
 
         if !pipeline_diff.pipelines_to_shutdown.is_empty() {
@@ -1684,6 +1697,110 @@ mod tests {
         universe.assert_quit().await;
     }
 
+    #[tokio::test]
+    async fn test_indexing_service_drops_superseded_plan() {
+        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-plan-gate");
+        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-plan-gate--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.clone(), &source_config).unwrap();
+        metastore.add_source(add_source_request).await.unwrap();
+
+        let universe = Universe::new();
+        let temp_dir = tempfile::tempdir().unwrap();
+        let (indexing_service, indexing_service_handle) = spawn_indexing_service_for_test(
+            temp_dir.path(),
+            &universe,
+            metastore.clone(),
+            cluster.clone(),
+        )
+        .await;
+
+        let params_fingerprint =
+            indexing_pipeline_params_fingerprint(&index_config, &source_config);
+        let task = |pipeline_uid: u128| IndexingTask {
+            index_uid: Some(index_uid.clone()),
+            source_id: source_config.source_id.clone(),
+            shard_ids: Vec::new(),
+            pipeline_uid: Some(PipelineUid::for_test(pipeline_uid)),
+            params_fingerprint,
+        };
+
+        indexing_service
+            .ask_for_res(ApplyIndexingPlanRequest {
+                indexing_tasks: vec![task(0), task(1)],
+                indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5F50".to_string(),
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            indexing_service_handle
+                .observe()
+                .await
+                .num_running_pipelines,
+            2
+        );
+
+        // A superseded (older id) plan that would drop a pipeline is ignored.
+        indexing_service
+            .ask_for_res(ApplyIndexingPlanRequest {
+                indexing_tasks: vec![task(0)],
+                indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5F40".to_string(),
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            indexing_service_handle
+                .observe()
+                .await
+                .num_running_pipelines,
+            2
+        );
+
+        // A newer plan applies, dropping the second pipeline.
+        indexing_service
+            .ask_for_res(ApplyIndexingPlanRequest {
+                indexing_tasks: vec![task(0)],
+                indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5F60".to_string(),
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            indexing_service_handle
+                .observe()
+                .await
+                .num_running_pipelines,
+            1
+        );
+
+        indexing_service_handle.quit().await;
+        universe.assert_quit().await;
+    }
+
     #[tokio::test]
     async fn test_indexing_service_shutdown_merge_pipeline_when_no_indexing_pipeline() {
         quickwit_common::setup_logging_for_tests();