Give indexing plan a time-based ID; use it to tiebreak publish tokens

Author: nadav-govari

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

@@ -35,6 +35,7 @@ use quickwit_proto::types::NodeId;
 use scheduling::{SourceToSchedule, SourceToScheduleType};
 use serde::Serialize;
 use tracing::{debug, info, warn};
+use ulid::Ulid;
 
 use crate::indexing_plan::PhysicalIndexingPlan;
 use crate::indexing_scheduler::change_tracker::{NotifyChangeOnDrop, RebuildNotifier};
@@ -414,11 +415,16 @@ impl IndexingScheduler {
     ) {
         debug!(new_physical_plan=?new_physical_plan, "apply physical indexing plan");
         APPLY_PLAN_TOTAL.inc();
+        // ULIDs are time-ordered, so a later application always yields a greater id. Indexers use
+        // it as the publish token for the shards they acquire. This guarantees a stale plan
+        // can never steal a shard from a more recent one.
+        let indexing_plan_id = Ulid::new().to_string();
         for (node_id, indexing_tasks) in new_physical_plan.indexing_tasks_per_indexer() {
             // We don't want to block on a slow indexer so we apply this change asynchronously
             // TODO not blocking is cool, but we need to make sure there is not accumulation
             // possible here.
             let notify_on_drop = notify_on_drop.clone();
+            let indexing_plan_id = indexing_plan_id.clone();
             tokio::spawn({
                 let indexer = indexers
                     .iter()
@@ -430,7 +436,10 @@ impl IndexingScheduler {
                     if let Err(error) = indexer
                         .client
                         .clone()
-                        .apply_indexing_plan(ApplyIndexingPlanRequest { indexing_tasks })
+                        .apply_indexing_plan(ApplyIndexingPlanRequest {
+                            indexing_tasks,
+                            indexing_plan_id,
+                        })
                         .await
                     {
                         warn!(

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

@@ -89,6 +89,9 @@ pub struct IndexingPipeline {
     // requiring a respawn of the pipeline.
     // We keep the list of shards here however, to reassign them after a respawn.
     shard_ids: BTreeSet<ShardId>,
+    // 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,
     _indexing_pipelines_gauge_guard: GaugeGuard,
 }
 
@@ -137,6 +140,7 @@ impl IndexingPipeline {
                 ..Default::default()
             },
             shard_ids: Default::default(),
+            indexing_plan_id: String::new(),
             _indexing_pipelines_gauge_guard: indexing_pipelines_gauge_guard,
         }
     }
@@ -402,6 +406,7 @@ impl IndexingPipeline {
             .spawn(actor_source);
         let assign_shards_message = AssignShards(Assignment {
             shard_ids: self.shard_ids.clone(),
+            indexing_plan_id: self.indexing_plan_id.clone(),
         });
         source_mailbox.send_message(assign_shards_message).await?;
 
@@ -496,6 +501,8 @@ impl Handler<AssignShards> for IndexingPipeline {
     ) -> Result<(), ActorExitStatus> {
         self.shard_ids
             .clone_from(&assign_shards_message.0.shard_ids);
+        self.indexing_plan_id
+            .clone_from(&assign_shards_message.0.indexing_plan_id);
         // If the pipeline is running, we forward the message to its source.
         // If it is not, it will be respawned soon, and the shards will be assigned afterward.
         if let Some(handles) = &self.handles_opt {

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

@@ -760,8 +760,8 @@ impl IndexingService {
     /// or not.
     ///
     /// If a pipeline actor has failed, this function just logs an error.
-    async fn assign_shards_to_pipelines(&mut self, tasks: &[IndexingTask]) {
-        for task in tasks {
+    async fn assign_shards_to_pipelines(&mut self, plan_request: &ApplyIndexingPlanRequest) {
+        for task in &plan_request.indexing_tasks {
             if task.shard_ids.is_empty() {
                 continue;
             }
@@ -771,6 +771,7 @@ impl IndexingService {
             };
             let assignment = Assignment {
                 shard_ids: task.shard_ids.iter().cloned().collect(),
+                indexing_plan_id: plan_request.indexing_plan_id.clone(),
             };
             let message = AssignShards(assignment);
 
@@ -785,9 +786,10 @@ impl IndexingService {
     /// - Starting the pipelines that are not running.
     async fn apply_indexing_plan(
         &mut self,
-        tasks: &[IndexingTask],
+        plan_request: ApplyIndexingPlanRequest,
         ctx: &ActorContext<Self>,
     ) -> Result<(), IndexingError> {
+        let tasks = &plan_request.indexing_tasks;
         let pipeline_diff = self.compute_pipeline_diff(tasks);
 
         if !pipeline_diff.pipelines_to_shutdown.is_empty() {
@@ -801,7 +803,7 @@ impl IndexingService {
                 .spawn_pipelines(&pipeline_diff.pipelines_to_spawn, ctx)
                 .await?;
         }
-        self.assign_shards_to_pipelines(tasks).await;
+        self.assign_shards_to_pipelines(&plan_request).await;
         self.update_chitchat_running_plan().await;
 
         if !spawn_pipeline_failures.is_empty() {
@@ -1135,7 +1137,7 @@ impl Handler<ApplyIndexingPlanRequest> for IndexingService {
         ctx: &ActorContext<Self>,
     ) -> Result<Self::Reply, ActorExitStatus> {
         Ok(self
-            .apply_indexing_plan(&plan_request.indexing_tasks, ctx)
+            .apply_indexing_plan(plan_request, ctx)
             .await
             .map(|_| ApplyIndexingPlanResponse {}))
     }
@@ -1465,7 +1467,10 @@ mod tests {
                 },
             ];
             indexing_service
-                .ask_for_res(ApplyIndexingPlanRequest { indexing_tasks })
+                .ask_for_res(ApplyIndexingPlanRequest {
+                    indexing_tasks,
+                    indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5FAV".to_string(),
+                })
                 .await
                 .unwrap();
             assert_eq!(
@@ -1531,6 +1536,7 @@ mod tests {
             indexing_service
                 .ask_for_res(ApplyIndexingPlanRequest {
                     indexing_tasks: indexing_tasks.clone(),
+                    indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5FAV".to_string(),
                 })
                 .await
                 .unwrap();
@@ -1587,6 +1593,7 @@ mod tests {
             indexing_service
                 .ask_for_res(ApplyIndexingPlanRequest {
                     indexing_tasks: indexing_tasks.clone(),
+                    indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5FAV".to_string(),
                 })
                 .await
                 .unwrap();
@@ -1646,6 +1653,7 @@ mod tests {
             indexing_service
                 .ask_for_res(ApplyIndexingPlanRequest {
                     indexing_tasks: indexing_tasks.clone(),
+                    indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5FAV".to_string(),
                 })
                 .await
                 .unwrap();
@@ -1665,6 +1673,7 @@ mod tests {
         indexing_service
             .ask_for_res(ApplyIndexingPlanRequest {
                 indexing_tasks: Vec::new(),
+                indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5FAV".to_string(),
             })
             .await
             .unwrap();
@@ -2072,6 +2081,7 @@ mod tests {
                         params_fingerprint: 0,
                     },
                 ],
+                indexing_plan_id: "01ARZ3NDEKTSV4RRFFQ69G5FAV".to_string(),
             })
             .await
             .unwrap();

quickwit/quickwit-indexing/src/actors/parquet_pipeline/pipeline.rs

@@ -114,6 +114,9 @@ pub struct MetricsPipeline {
     handles_opt: Option<MetricsPipelineHandles>,
     kill_switch: KillSwitch,
     shard_ids: BTreeSet<ShardId>,
+    // 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,
     _indexing_pipelines_gauge_guard: GaugeGuard,
 }
 
@@ -163,6 +166,7 @@ impl MetricsPipeline {
                 ..Default::default()
             },
             shard_ids: Default::default(),
+            indexing_plan_id: String::new(),
             _indexing_pipelines_gauge_guard: indexing_pipelines_gauge_guard,
         }
     }
@@ -447,6 +451,7 @@ impl MetricsPipeline {
             .spawn(actor_source);
         let assign_shards_message = AssignShards(Assignment {
             shard_ids: self.shard_ids.clone(),
+            indexing_plan_id: self.indexing_plan_id.clone(),
         });
         source_mailbox.send_message(assign_shards_message).await?;
 
@@ -539,6 +544,8 @@ impl Handler<AssignShards> for MetricsPipeline {
     ) -> Result<(), ActorExitStatus> {
         self.shard_ids
             .clone_from(&assign_shards_message.0.shard_ids);
+        self.indexing_plan_id
+            .clone_from(&assign_shards_message.0.indexing_plan_id);
         if let Some(handles) = &self.handles_opt {
             info!(
                 shard_ids=?assign_shards_message.0.shard_ids,