schedule leaf searches by (priority, remaining_splits) lexicographically

quickwit/quickwit-proto/build.rs

@@ -222,6 +222,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     prost_config
         .file_descriptor_set_path("src/codegen/quickwit/search_descriptor.bin")
         .protoc_arg("--experimental_allow_proto3_optional")
+        .field_attribute("SearchRequest.priority", "#[serde(default)]")
         // Box the large `LeafSearchResponse` variant so the oneof stays small
         // (the `Error` variant only carries a `String`).
         .boxed("LambdaSingleSplitResult.outcome.response");

quickwit/quickwit-proto/protos/quickwit/search.proto

@@ -273,6 +273,13 @@ message SearchRequest {
   // When true, skip finalization of aggregation results and return
   // the raw IntermediateAggregationResults bytes instead.
   bool skip_aggregation_finalization = 19;
+
+  // Field 20 is intentionally reserved for wire compatibility with a downstream fork.
+  reserved 20;
+
+  // Scheduling priority for leaf search execution. Negative values are allowed,
+  // and lower values have higher priority. Callers that omit it get priority 0.
+  int32 priority = 21;
 }
 
 enum CountHits {

quickwit/quickwit-search/src/leaf.rs

@@ -1893,7 +1893,7 @@ async fn schedule_search_tasks(
 ) -> ScheduleSearchTaskResult {
     let task_metadata: Vec<crate::search_permit_provider::SplitSearchTaskMetadata> = splits
         .iter()
-        .map(|(split, _)| {
+        .map(|(split, search_request)| {
             let memory_allocation = compute_initial_memory_allocation(
                 split,
                 searcher_context
@@ -1904,6 +1904,7 @@ async fn schedule_search_tasks(
             crate::search_permit_provider::SplitSearchTaskMetadata {
                 memory_allocation,
                 job_cost,
+                priority: search_request.priority,
             }
         })
         .collect();

quickwit/quickwit-search/src/leaf_cache.rs

@@ -107,6 +107,8 @@ impl CacheKey {
         // it doesn't matter whether or not we count all hits at the scale of a
         // single split: either we did process it and got everything, or we didn't.
         search_request.count_hits = CountHits::CountAll.into();
+        // Priority only affects scheduling, not search results.
+        search_request.priority = 0;
 
         CacheKey {
             split_id: split_info.split_id,
@@ -436,4 +438,30 @@ mod tests {
         assert!(cache.get(split_3.clone(), query_2).is_none());
         assert!(cache.get(split_3, query_2bis).is_some());
     }
+
+    #[test]
+    fn test_leaf_search_cache_ignores_priority() {
+        let cache = LeafSearchCache::new(&ByteSize::mb(64).into());
+        let split = SplitIdAndFooterOffsets {
+            split_id: "split".to_string(),
+            ..Default::default()
+        };
+        let high_priority_request = SearchRequest {
+            priority: -10,
+            ..Default::default()
+        };
+        let low_priority_request = SearchRequest {
+            priority: 10,
+            ..Default::default()
+        };
+
+        cache.put(
+            split.clone(),
+            high_priority_request,
+            LeafSearchResponse::default(),
+        );
+
+        // the two requests differ only in priority, so it should be a cache-hit
+        assert!(cache.get(split, low_priority_request).is_some());
+    }
 }

quickwit/quickwit-search/src/list_terms.rs

@@ -339,6 +339,7 @@ pub async fn leaf_list_terms(
             crate::search_permit_provider::SplitSearchTaskMetadata {
                 memory_allocation,
                 job_cost,
+                priority: 0,
             }
         })
         .collect();

quickwit/quickwit-search/src/root.rs

@@ -376,6 +376,7 @@ fn simplify_search_request_for_scroll_api(req: &SearchRequest) -> crate::Result<
         count_hits: quickwit_proto::search::CountHits::Underestimate as i32,
         ignore_missing_indexes: req.ignore_missing_indexes,
         skip_aggregation_finalization: false,
+        priority: req.priority,
     })
 }
 
@@ -1320,6 +1321,7 @@ pub async fn root_search(
         count_required = search_request.count_hits().as_str_name(),
         num_docs = num_docs,
         num_splits = num_splits,
+        priority = search_request.priority,
         "root_search"
     );
 

quickwit/quickwit-search/src/search_permit_provider.rs

@@ -34,11 +34,11 @@ use crate::metrics::{
 
 /// Distributor of permits to perform split search operation.
 ///
-/// Requests are served in order. Each permit initially reserves a slot for the
-/// warmup (limit concurrent downloads) and a pessimistic amount of memory. Once
-/// the warmup is completed, the actual memory usage is set and the warmup slot
-/// is released. Once the search is completed and the permit is dropped, the
-/// remaining memory is also released.
+/// Requests are served by priority, then by fewest remaining splits. Each permit initially
+/// reserves a slot for the warmup (limit concurrent downloads) and a pessimistic amount of
+/// memory. Once the warmup is completed, the actual memory usage is set and the warmup slot is
+/// released. Once the search is completed and the permit is dropped, the remaining memory is also
+/// released.
 #[derive(Clone)]
 pub struct SearchPermitProvider {
     message_sender: mpsc::UnboundedSender<SearchPermitMessage>,
@@ -56,6 +56,8 @@ pub(crate) struct SplitSearchTaskMetadata {
     /// Estimated cost of this task, in the same arbitrary unit as [`Job::cost()`].
     /// Used to report the current load of this node to the job placer.
     pub job_cost: usize,
+    /// Priority of the leaf request this split belongs to.
+    pub priority: i32,
 }
 
 pub enum SearchPermitMessage {
@@ -219,19 +221,23 @@ struct SingleSplitPermitRequest {
 }
 
 struct LeafPermitRequest {
+    /// Lower values have higher priority.
+    priority: i32,
     /// Single split permit requests for this leaf search.
     single_split_permit_requests: std::vec::IntoIter<SingleSplitPermitRequest>,
 }
 
 impl Ord for LeafPermitRequest {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        // we compare other with self and not the other way arround because we want a min-heap and
+        // we compare other with self and not the other way around because we want a min-heap and
         // Rust's is a max-heap
-        other
-            .single_split_permit_requests
-            .as_slice()
-            .len()
-            .cmp(&self.single_split_permit_requests.as_slice().len())
+        other.priority.cmp(&self.priority).then_with(|| {
+            other
+                .single_split_permit_requests
+                .as_slice()
+                .len()
+                .cmp(&self.single_split_permit_requests.as_slice().len())
+        })
     }
 }
 
@@ -255,6 +261,7 @@ impl LeafPermitRequest {
         task_metadata: Vec<SplitSearchTaskMetadata>,
     ) -> (Self, Vec<SearchPermitFuture>) {
         assert!(!task_metadata.is_empty(), "task_metadata must not be empty");
+        let priority = task_metadata[0].priority;
         // Stamped on every `SingleSplitPermitRequest` we're about to enqueue.
         // The actor will compute `requested_at.elapsed()` at grant time to
         // report the permit's acquisition latency.
@@ -276,6 +283,7 @@ impl LeafPermitRequest {
         }
         (
             LeafPermitRequest {
+                priority,
                 single_split_permit_requests: single_split_permit_requests.into_iter(),
             },
             permits,
@@ -524,14 +532,74 @@ mod tests {
     use super::*;
 
     fn make_splits(memory_mb: u64, count: usize) -> Vec<SplitSearchTaskMetadata> {
+        make_splits_with_priority(memory_mb, count, 0)
+    }
+
+    fn make_splits_with_priority(
+        memory_mb: u64,
+        count: usize,
+        priority: i32,
+    ) -> Vec<SplitSearchTaskMetadata> {
         (0..count)
             .map(|_| SplitSearchTaskMetadata {
                 memory_allocation: ByteSize::mb(memory_mb),
                 job_cost: 5,
+                priority,
             })
             .collect()
     }
 
+    #[tokio::test]
+    async fn test_search_permit_priority_precedes_remaining_splits() {
+        let permit_provider = SearchPermitProvider::new(1, ByteSize::mb(100));
+        let blocker = permit_provider
+            .get_permits(make_splits(10, 1))
+            .await
+            .pop()
+            .unwrap()
+            .await;
+
+        let negative_priority = permit_provider
+            .get_permits(make_splits_with_priority(10, 3, -10))
+            .await;
+        let default_priority = permit_provider.get_permits(make_splits(10, 1)).await;
+        let positive_priority = permit_provider
+            .get_permits(make_splits_with_priority(10, 1, 10))
+            .await;
+
+        let mut join_set = JoinSet::new();
+        for (request, permit_futures) in [
+            ("negative", negative_priority),
+            ("default", default_priority),
+            ("positive", positive_priority),
+        ] {
+            for (split_idx, permit_future) in permit_futures.into_iter().enumerate() {
+                join_set.spawn(async move {
+                    let permit = permit_future.await;
+                    (request, split_idx, permit)
+                });
+            }
+        }
+
+        drop(blocker);
+
+        let mut execution_order = Vec::new();
+        while let Some(result) = join_set.join_next().await {
+            let (request, split_idx, _permit) = result.unwrap();
+            execution_order.push((request, split_idx));
+        }
+        assert_eq!(
+            execution_order,
+            vec![
+                ("negative", 0),
+                ("negative", 1),
+                ("negative", 2),
+                ("default", 0),
+                ("positive", 0),
+            ]
+        );
+    }
+
     #[tokio::test]
     async fn test_search_permit_order() {
         let permit_provider = SearchPermitProvider::new(1, ByteSize::mb(100));
@@ -833,14 +901,17 @@ mod tests {
             SplitSearchTaskMetadata {
                 memory_allocation: ByteSize::mb(10),
                 job_cost: 7,
+                priority: 0,
             },
             SplitSearchTaskMetadata {
                 memory_allocation: ByteSize::mb(10),
                 job_cost: 3,
+                priority: 0,
             },
             SplitSearchTaskMetadata {
                 memory_allocation: ByteSize::mb(10),
                 job_cost: 5,
+                priority: 0,
             },
         ];
         let mut permit_futs = permit_provider.get_permits(splits).await;

quickwit/quickwit-serve/src/elasticsearch_api/rest_handler.rs

@@ -581,6 +581,7 @@ fn build_request_for_es_api(
             count_hits,
             ignore_missing_indexes,
             skip_aggregation_finalization: false,
+            ..Default::default()
         },
         has_doc_id_field,
     ))

quickwit/quickwit-serve/src/search_api/rest_handler.rs

@@ -266,6 +266,7 @@ pub fn search_request_from_api_request(
         count_hits: search_request.count_all.into(),
         ignore_missing_indexes: false,
         skip_aggregation_finalization: false,
+        ..Default::default()
     };
     Ok(search_request)
 }