search: attribute permit-wait latency to the blocking resource

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<LocalSearchTask>,
     index_storage: Arc<dyn Storage + 'static>,
@@ -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,

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<Self> {
+        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<AtomicU8>);
+
+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> {
+        PermitBlockReason::from_code(self.0.load(Ordering::Relaxed))
+    }
+}
+
 pub enum SearchPermitMessage {
     RequestWithOffload {
         task_metadata: Vec<SplitSearchTaskMetadata>,
@@ -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<AtomicUsize>,
     permits_requests: BinaryHeap<LeafPermitRequest>,
+    /// 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<SplitSearchTaskMetadata>,
+        block_reason: BlockReasonHandle,
     ) -> (Self, Vec<SearchPermitFuture>) {
         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<SingleSplitPermitRequest> {
+        // 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<SearchPermit>);
+pub struct SearchPermitFuture {
+    receiver: oneshot::Receiver<SearchPermit>,
+    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<Self::Output> {
-        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));