diff --git a/quickwit/quickwit-proto/build.rs b/quickwit/quickwit-proto/build.rs index 5d51c8cd3f9..b7d08264920 100644 --- a/quickwit/quickwit-proto/build.rs +++ b/quickwit/quickwit-proto/build.rs @@ -222,6 +222,7 @@ fn main() -> Result<(), Box> { 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"); diff --git a/quickwit/quickwit-proto/protos/quickwit/search.proto b/quickwit/quickwit-proto/protos/quickwit/search.proto index 18136f6b3f3..d266f889508 100644 --- a/quickwit/quickwit-proto/protos/quickwit/search.proto +++ b/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 { diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs index 1c1ed0b1b03..347892fb124 100644 --- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs +++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs @@ -211,6 +211,11 @@ pub struct SearchRequest { /// the raw IntermediateAggregationResults bytes instead. #[prost(bool, tag = "19")] pub skip_aggregation_finalization: bool, + /// Scheduling priority for leaf search execution. Negative values are allowed, + /// and lower values have higher priority. Callers that omit it get priority 0. + #[prost(int32, tag = "21")] + #[serde(default)] + pub priority: i32, } #[derive(serde::Serialize, serde::Deserialize, utoipa::ToSchema)] #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index 99ff06c312b..765c023255e 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -1893,7 +1893,7 @@ async fn schedule_search_tasks( ) -> ScheduleSearchTaskResult { let task_metadata: Vec = 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(); diff --git a/quickwit/quickwit-search/src/leaf_cache.rs b/quickwit/quickwit-search/src/leaf_cache.rs index 4cb9f37d6c5..39758b08956 100644 --- a/quickwit/quickwit-search/src/leaf_cache.rs +++ b/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()); + } } diff --git a/quickwit/quickwit-search/src/list_terms.rs b/quickwit/quickwit-search/src/list_terms.rs index 9c6e025ac92..03aa702a7ff 100644 --- a/quickwit/quickwit-search/src/list_terms.rs +++ b/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(); diff --git a/quickwit/quickwit-search/src/root.rs b/quickwit/quickwit-search/src/root.rs index 896a763f508..c31a848f8ea 100644 --- a/quickwit/quickwit-search/src/root.rs +++ b/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" ); diff --git a/quickwit/quickwit-search/src/search_permit_provider.rs b/quickwit/quickwit-search/src/search_permit_provider.rs index becc3ff6d18..954601cf1fc 100644 --- a/quickwit/quickwit-search/src/search_permit_provider.rs +++ b/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, @@ -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, } 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, ) -> (Self, Vec) { 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 { + make_splits_with_priority(memory_mb, count, 0) + } + + fn make_splits_with_priority( + memory_mb: u64, + count: usize, + priority: i32, + ) -> Vec { (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; diff --git a/quickwit/quickwit-serve/src/elasticsearch_api/rest_handler.rs b/quickwit/quickwit-serve/src/elasticsearch_api/rest_handler.rs index a5795de6b4e..cec582d934d 100644 --- a/quickwit/quickwit-serve/src/elasticsearch_api/rest_handler.rs +++ b/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, )) diff --git a/quickwit/quickwit-serve/src/search_api/rest_handler.rs b/quickwit/quickwit-serve/src/search_api/rest_handler.rs index b1400fa12c0..01ad87d78e9 100644 --- a/quickwit/quickwit-serve/src/search_api/rest_handler.rs +++ b/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) }