diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index c8411322314..d8db6417fbc 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -67,7 +67,7 @@ use crate::metrics::{ }; use crate::root::is_metadata_count_request_with_ast; use crate::search_permit_provider::{ - SearchPermit, SearchPermitFuture, compute_initial_memory_allocation, + BlockReasonHandle, SearchPermit, SearchPermitFuture, compute_initial_memory_allocation, }; use crate::service::{SearcherContext, deserialize_doc_mapper}; use crate::{QuickwitAggregations, SearchError}; @@ -2034,6 +2034,42 @@ pub async fn single_doc_mapping_leaf_search( Ok(leaf_search_response) } +/// Records the permit block reason onto the `waiting_for_leaf_search_split_semaphore` +/// span when the wait ends. +/// +/// Implemented as a drop guard (rather than reading after `.await`) so the reason is +/// recorded even when the wait future is cancelled before the permit is granted — the +/// long waits that hit a search deadline are exactly the ones that never get granted. +struct WaitBlockReasonRecorder { + wait_span: Span, + block_reason: BlockReasonHandle, + started_at: Instant, +} + +impl WaitBlockReasonRecorder { + fn new(wait_span: Span, block_reason: BlockReasonHandle) -> Self { + Self { + wait_span, + block_reason, + started_at: Instant::now(), + } + } +} + +impl Drop for WaitBlockReasonRecorder { + fn drop(&mut self) { + // Skip near-instant grants: a sub-millisecond wait isn't worth attributing and + // keeps the span uncluttered. + const MIN_WAIT_FOR_BLOCK_ATTRIBUTION: Duration = Duration::from_millis(1); + if self.started_at.elapsed() < MIN_WAIT_FOR_BLOCK_ATTRIBUTION { + return; + } + if let Some(block_reason) = self.block_reason.get() { + self.wait_span.record("blocked_on", block_reason.as_str()); + } + } +} + async fn run_local_search_tasks( local_search_tasks: Vec, index_storage: Arc, @@ -2056,7 +2092,18 @@ async fn run_local_search_tasks( split_id = split.split_id, num_docs = split.num_docs ); - let wait_span = info_span!(parent: &split_span, "waiting_for_leaf_search_split_semaphore"); + let wait_span = info_span!( + parent: &split_span, + "waiting_for_leaf_search_split_semaphore", + blocked_on = tracing::field::Empty + ); + // Records the block reason on `wait_span` when the wait ends — whether the permit + // is granted or the wait is cancelled on a search deadline (the important, long + // waits are exactly the ones that get cancelled before being granted). + let _block_reason_recorder = WaitBlockReasonRecorder::new( + wait_span.clone(), + search_permit_future.block_reason_handle(), + ); let leaf_split_search_permit = search_permit_future.instrument(wait_span).await; // We run simplify search request again: as we push split into the merge collector, diff --git a/quickwit/quickwit-search/src/search_permit_provider.rs b/quickwit/quickwit-search/src/search_permit_provider.rs index ff9b0368d68..6485089ccf1 100644 --- a/quickwit/quickwit-search/src/search_permit_provider.rs +++ b/quickwit/quickwit-search/src/search_permit_provider.rs @@ -17,7 +17,7 @@ use std::collections::binary_heap::PeekMut; use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; @@ -58,6 +58,66 @@ pub(crate) struct SplitSearchTaskMetadata { pub job_cost: usize, } +/// Why the head of the permit queue could not be granted. +/// +/// The actor writes it into the shared [`BlockReasonHandle`] each time it fails to serve +/// the head, so any waiter can read it whether its permit is eventually granted or the +/// wait is cancelled (e.g. on a search deadline). This is what lets us attribute the +/// acquisition latency to the binding resource. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PermitBlockReason { + /// All concurrent warmup/download slots were in use. + WarmupSlots, + /// The memory budget could not fit the request's estimated allocation. + MemoryBudget, +} + +impl PermitBlockReason { + pub fn as_str(&self) -> &'static str { + match self { + PermitBlockReason::WarmupSlots => "warmup_slots", + PermitBlockReason::MemoryBudget => "memory_budget", + } + } + + fn to_code(self) -> u8 { + match self { + PermitBlockReason::WarmupSlots => 1, + PermitBlockReason::MemoryBudget => 2, + } + } + + fn from_code(code: u8) -> Option { + match code { + 1 => Some(PermitBlockReason::WarmupSlots), + 2 => Some(PermitBlockReason::MemoryBudget), + _ => None, + } + } +} + +/// Shared cell recording the resource currently blocking the head of the permit queue. +/// +/// A single instance is shared by the actor (which writes it) and every +/// [`SearchPermitFuture`] (which reads it). Because permits are served in order, whatever +/// blocks the head effectively blocks every waiter, so a queue-level reason is correct for +/// a waiter at any position — including one cancelled deep in the queue before it ever +/// reaches the head. Uses an atomic so it can be read after the wait future is dropped on +/// cancellation. +#[derive(Clone, Default)] +pub struct BlockReasonHandle(Arc); + +impl BlockReasonHandle { + fn set(&self, reason: PermitBlockReason) { + self.0.store(reason.to_code(), Ordering::Relaxed); + } + + /// The resource currently blocking the queue, or `None` if it isn't blocked. + pub fn get(&self) -> Option { + PermitBlockReason::from_code(self.0.load(Ordering::Relaxed)) + } +} + pub enum SearchPermitMessage { RequestWithOffload { task_metadata: Vec, @@ -116,6 +176,7 @@ impl SearchPermitProvider { permits_requests: BinaryHeap::new(), total_memory_allocated: 0u64, total_job_cost: total_job_cost.clone(), + block_reason: BlockReasonHandle::default(), }; let actor_join_handle = Arc::new(tokio::spawn(actor.run())); Self { @@ -209,6 +270,9 @@ struct SearchPermitActor { /// Only the actor mutates it, so [`Ordering::Relaxed`] is sufficient. total_job_cost: Arc, permits_requests: BinaryHeap, + /// Shared with every [`SearchPermitFuture`]; set to the resource blocking the head of + /// the queue each time it can't be served. + block_reason: BlockReasonHandle, } struct SingleSplitPermitRequest { @@ -253,6 +317,7 @@ impl LeafPermitRequest { // `task_metadata` must not be empty. fn from_task_metadata( task_metadata: Vec, + block_reason: BlockReasonHandle, ) -> (Self, Vec) { assert!(!task_metadata.is_empty(), "task_metadata must not be empty"); // Stamped on every `SingleSplitPermitRequest` we're about to enqueue. @@ -272,7 +337,10 @@ impl LeafPermitRequest { job_cost: meta.job_cost, requested_at, }); - permits.push(SearchPermitFuture(rx)); + permits.push(SearchPermitFuture { + receiver: rx, + block_reason: block_reason.clone(), + }); } ( LeafPermitRequest { @@ -338,7 +406,7 @@ impl SearchPermitActor { SEARCHER_NODE_LOAD.set(new_load as f64); let (leaf_permit_request, permit_futures) = - LeafPermitRequest::from_task_metadata(task_metadata); + LeafPermitRequest::from_task_metadata(task_metadata, self.block_reason.clone()); self.permits_requests.push(leaf_permit_request); self.assign_available_permits(); let _ = permit_sender.send(permit_futures); @@ -374,12 +442,25 @@ impl SearchPermitActor { } fn pop_next_request_if_serviceable(&mut self) -> Option { + // Nothing is waiting, so there is no block to attribute (avoids leaving a stale + // reason set once the queue drains and slots/memory are exhausted). + if self.permits_requests.is_empty() { + return None; + } if self.num_warmup_slots_available == 0 { + self.block_reason.set(PermitBlockReason::WarmupSlots); return None; } - let available_memory = self + let available_memory = match self .total_memory_budget - .checked_sub(self.total_memory_allocated)?; + .checked_sub(self.total_memory_allocated) + { + Some(available_memory) => available_memory, + None => { + self.block_reason.set(PermitBlockReason::MemoryBudget); + return None; + } + }; let mut peeked = self.permits_requests.peek_mut()?; assert!( @@ -392,6 +473,8 @@ impl SearchPermitActor { } return Some(permit_request); } + // The head request needs more memory than is currently available. + self.block_reason.set(PermitBlockReason::MemoryBudget); None } @@ -510,13 +593,25 @@ impl Drop for SearchPermit { } } -pub struct SearchPermitFuture(oneshot::Receiver); +pub struct SearchPermitFuture { + receiver: oneshot::Receiver, + block_reason: BlockReasonHandle, +} + +impl SearchPermitFuture { + /// Handle to the reason this permit is blocked on. Readable at any time, including + /// after this future is dropped on cancellation, so the waiter can attribute the + /// wait even when the permit is never granted. + pub fn block_reason_handle(&self) -> BlockReasonHandle { + self.block_reason.clone() + } +} impl Future for SearchPermitFuture { type Output = SearchPermit; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let receiver = Pin::new(&mut self.get_mut().0); + let receiver = Pin::new(&mut self.get_mut().receiver); match receiver.poll(cx) { Poll::Ready(Ok(search_permit)) => Poll::Ready(search_permit), Poll::Ready(Err(_)) => panic!( @@ -837,6 +932,52 @@ mod tests { permits.push(try_get(next_blocked_permit_fut).await.unwrap()); } + #[tokio::test] + async fn test_permit_block_reason_warmup_slots() { + // A single warmup slot with ample memory: once a second request is queued, the + // queue can only be blocked on the warmup slot. The reason is queue-level and + // shared by every waiter's handle, readable without being granted a permit. + let permit_provider = SearchPermitProvider::new(1, ByteSize::mb(1000)); + let first_fut = permit_provider + .get_permits(make_splits(10, 1)) + .await + .into_iter() + .next() + .unwrap(); + // First was granted immediately; nothing is blocked yet. + assert_eq!(first_fut.block_reason_handle().get(), None); + let second_fut = permit_provider + .get_permits(make_splits(10, 1)) + .await + .into_iter() + .next() + .unwrap(); + // Second can't get the single slot: the queue is now blocked on warmup slots. + assert_eq!( + second_fut.block_reason_handle().get(), + Some(PermitBlockReason::WarmupSlots) + ); + // Draining still works: freeing the slot lets the second request through. + let first_permit = first_fut.await; + drop(first_permit); + let _second_permit = second_fut.await; + } + + #[tokio::test] + async fn test_permit_block_reason_memory_budget() { + // Ample slots but tight memory (100MB / 10MB = 10 permits): the 11th request can + // only be blocked on the memory budget. + let permit_provider = SearchPermitProvider::new(100, ByteSize::mb(100)); + let permit_futs = permit_provider.get_permits(make_splits(10, 11)).await; + assert_eq!(permit_futs.len(), 11); + // The 11th can't fit, so the queue is blocked on memory. The reason is queue-level + // and shared, so any waiter's handle reports it — even one deep in the queue. + assert_eq!( + permit_futs[10].block_reason_handle().get(), + Some(PermitBlockReason::MemoryBudget) + ); + } + #[tokio::test] async fn test_get_load_counts_queued_and_active() { let permit_provider = SearchPermitProvider::new(1, ByteSize::mb(100));