Add test showing shard assignation fragmentation

Author: rdettai

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

@@ -999,6 +999,84 @@ mod tests {
         assert_eq!(source_1_on_node1 + source_1_on_node2, 4);
     }
 
+    #[test]
+    fn test_build_physical_plan_second_iteration() {
+        let indexer1 = "indexer1".to_string();
+        let indexer2 = "indexer2".to_string();
+        let indexer3 = "indexer3".to_string();
+        let mut sources = Vec::new();
+        for _ in 0..10 {
+            sources.push(SourceToSchedule {
+                source_uid: source_id(),
+                source_type: SourceToScheduleType::NonSharded {
+                    num_pipelines: 4,
+                    load_per_pipeline: NonZeroU32::new(MAX_LOAD_PER_PIPELINE.cpu_millis()).unwrap(),
+                },
+                params_fingerprint: 0,
+            });
+        }
+        let mut indexer_id_to_cpu_capacities = FnvHashMap::default();
+        indexer_id_to_cpu_capacities.insert(indexer1.clone(), mcpu(16_000));
+        indexer_id_to_cpu_capacities.insert(indexer2.clone(), mcpu(16_000));
+        indexer_id_to_cpu_capacities.insert(indexer3.clone(), mcpu(16_000));
+        let shard_locations = ShardLocations::default();
+        let indexing_plan = build_physical_indexing_plan(
+            &sources,
+            &indexer_id_to_cpu_capacities,
+            None,
+            &shard_locations,
+        );
+        assert_eq!(indexing_plan.indexing_tasks_per_indexer().len(), 3);
+
+        for source in &sources {
+            let pipelines_per_indexer_for_source = indexing_plan
+                .indexing_tasks_per_indexer()
+                .values()
+                .map(|tasks| {
+                    tasks
+                        .iter()
+                        .filter(|t| t.source_id == source.source_uid.source_id)
+                        .count()
+                })
+                .collect_vec();
+            assert!(pipelines_per_indexer_for_source.contains(&3));
+            assert!(pipelines_per_indexer_for_source.contains(&1));
+            assert!(pipelines_per_indexer_for_source.contains(&0));
+            assert_eq!(pipelines_per_indexer_for_source.iter().sum::<usize>(), 4);
+        }
+
+        for source in &mut sources {
+            if let SourceToScheduleType::NonSharded { num_pipelines, .. } = &mut source.source_type
+            {
+                *num_pipelines = 5;
+            }
+        }
+
+        let new_indexing_plan = build_physical_indexing_plan(
+            &sources,
+            &indexer_id_to_cpu_capacities,
+            Some(&indexing_plan),
+            &shard_locations,
+        );
+
+        for source in &sources {
+            let pipelines_per_indexer_for_source = new_indexing_plan
+                .indexing_tasks_per_indexer()
+                .values()
+                .map(|tasks| {
+                    tasks
+                        .iter()
+                        .filter(|t| t.source_id == source.source_uid.source_id)
+                        .count()
+                })
+                .collect_vec();
+            assert!(pipelines_per_indexer_for_source.contains(&3));
+            assert!(pipelines_per_indexer_for_source.contains(&2));
+            assert!(pipelines_per_indexer_for_source.contains(&0));
+            assert_eq!(pipelines_per_indexer_for_source.iter().sum::<usize>(), 5);
+        }
+    }
+
     fn make_indexing_tasks(
         source_uid: &SourceUid,
         shards: &[(PipelineUid, &[ShardId])],

quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs

@@ -620,7 +620,7 @@ mod tests {
     }
 
     #[test]
-    fn test_limit_placement_to_three_pipelines() {
+    fn test_placement_limit_with_affinity() {
         let mut problem =
             SchedulingProblem::with_indexer_cpu_capacities(vec![mcpu(16_000), mcpu(16_000)]);
         let max_load_per_pipeline = NonZeroU32::new(MAX_LOAD_PER_PIPELINE.cpu_millis()).unwrap();
@@ -635,6 +635,8 @@ mod tests {
         assert_eq!(solution.indexer_assignments[0].num_shards(1), 3);
         assert_eq!(solution.indexer_assignments[0].num_shards(0), 1);
         assert_eq!(solution.indexer_assignments[1].num_shards(0), 3);
+        // one shard was not placed because indexer 0 was full and it had no
+        // affinity with indexer 1
         assert_eq!(solution.indexer_assignments[1].num_shards(1), 0);
     }