Limit using load instead of shard number

Author: rdettai

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

@@ -17,7 +17,7 @@ use std::collections::BTreeMap;
 use std::collections::btree_map::Entry;
 
 use itertools::Itertools;
-use quickwit_proto::indexing::CpuCapacity;
+use quickwit_proto::indexing::{CpuCapacity, PIPELINE_FULL_CAPACITY};
 
 use super::scheduling_logic_model::*;
 use crate::indexing_scheduler::scheduling::inflate_node_capacities_if_necessary;
@@ -356,8 +356,11 @@ fn place_unassigned_shards_single_source(
 ) -> Result<(), NotEnoughCapacity> {
     let mut num_shards = source.num_shards;
     // To ensure that merges can keep up, try not to assign more than 3
-    // shards per indexer for a source (except if there aren't enough nodes)
-    let limit_num_shards_per_indexer_per_source = 3.max(num_shards.div_ceil(num_indexers as u32));
+    // pipelines per indexer for a source (except if there aren't enough nodes)
+    let target_limit_num_shards_per_indexer_per_source =
+        3 * PIPELINE_FULL_CAPACITY.cpu_millis() / source.load_per_shard.get();
+    let limit_num_shards_per_indexer_per_source = target_limit_num_shards_per_indexer_per_source
+        .max(num_shards.div_ceil(num_indexers as u32));
     while num_shards > 0 {
         let Some((indexer_ord, available_capacity)) = indexer_with_capacities.next() else {
             return Err(NotEnoughCapacity);
@@ -608,6 +611,21 @@ mod tests {
         problem.add_source(4, NonZeroU32::new(1_000).unwrap());
         problem.add_source(4, NonZeroU32::new(1_000).unwrap());
         problem.inc_affinity(0, 1);
+        problem.inc_affinity(1, 0);
+        let mut solution = problem.new_solution();
+        place_unassigned_shards_with_affinity(&problem, &mut solution);
+        assert_eq!(solution.indexer_assignments[0].num_shards(1), 4);
+        assert_eq!(solution.indexer_assignments[1].num_shards(0), 4);
+    }
+
+    #[test]
+    fn test_limit_placement_to_three_pipelines() {
+        let mut problem =
+            SchedulingProblem::with_indexer_cpu_capacities(vec![mcpu(16_000), mcpu(16_000)]);
+        let pipeline_full_capacity = NonZeroU32::new(PIPELINE_FULL_CAPACITY.cpu_millis()).unwrap();
+        problem.add_source(4, pipeline_full_capacity);
+        problem.add_source(4, pipeline_full_capacity);
+        problem.inc_affinity(0, 1);
         problem.inc_affinity(0, 1);
         problem.inc_affinity(0, 0);
         problem.inc_affinity(1, 0);