Skip to content

Commit

Permalink
Bring some clarifications and remove single permit getter
Browse files Browse the repository at this point in the history
  • Loading branch information
@rdettai
rdettai committed Nov 25, 2024
1 parent e1b7125 commit fd1c9a1
Show file tree
Hide file tree
Showing 2 changed files with 117 additions and 100 deletions.
13 changes: 6 additions & 7 deletions quickwit/quickwit-search/src/list_terms.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -328,18 +328,17 @@ pub async fn leaf_list_terms(
splits: &[SplitIdAndFooterOffsets],
) -> Result<LeafListTermsResponse, SearchError> {
info!(split_offsets = ?PrettySample::new(splits, 5));
let permits = searcher_context
.search_permit_provider
.get_permits(splits.len());
let leaf_search_single_split_futures: Vec<_> = splits
.iter()
.map(|split| {
.zip(permits.into_iter())
.map(|(split, search_permit_recv)| {
let index_storage_clone = index_storage.clone();
let searcher_context_clone = searcher_context.clone();
async move {
let _leaf_split_search_permit = searcher_context_clone
.search_permit_provider
.get_permit()
.await
.expect("Failed to acquire permit. This should never happen! Please, report on https://github.com/quickwit-oss/quickwit/issues.");

let _leaf_split_search_permit = search_permit_recv.await;
// TODO dedicated counter and timer?
crate::SEARCH_METRICS.leaf_searches_splits_total.inc();
let timer = crate::SEARCH_METRICS
Expand Down
204 changes: 111 additions & 93 deletions quickwit/quickwit-search/src/search_permit_provider.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,10 @@
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

use bytesize::ByteSize;
use quickwit_common::metrics::GaugeGuard;
Expand All @@ -28,17 +31,22 @@ use tracing::warn;
/// `SearchPermitProvider` is a distributor of permits to perform single split
/// search operation.
///
/// Requests are served in order.
/// - Two types of resources are managed: memory allocations and download slots.
/// - Requests are served in order.
#[derive(Clone)]
pub struct SearchPermitProvider {
inner_arc: Arc<Mutex<InnerSearchPermitProvider>>,
}

impl SearchPermitProvider {
pub fn new(num_permits: usize, memory_budget: ByteSize, initial_allocation: ByteSize) -> SearchPermitProvider {
pub fn new(
num_download_slots: usize,
memory_budget: ByteSize,
initial_allocation: ByteSize,
) -> SearchPermitProvider {
SearchPermitProvider {
inner_arc: Arc::new(Mutex::new(InnerSearchPermitProvider {
num_permits_available: num_permits,
num_download_slots_available: num_download_slots,
memory_budget: memory_budget.as_u64(),
permits_requests: VecDeque::new(),
memory_allocated: 0u64,
Expand All @@ -47,29 +55,20 @@ impl SearchPermitProvider {
}
}

/// Returns a future permit in the form of a oneshot Receiver channel.
///
/// At this point the permit is not acquired yet.
#[must_use]
pub fn get_permit(&self) -> oneshot::Receiver<SearchPermit> {
let mut permits_lock = self.inner_arc.lock().unwrap();
permits_lock.get_permit(&self.inner_arc)
}

/// Returns a list of future permits in the form of oneshot Receiver channels.
/// Returns a list of future permits in the form of awaitable futures.
///
/// The permits returned are guaranteed to be resolved in order.
/// In addition, the permits are guaranteed to be resolved before permits returned by
/// subsequent calls to this function (or `get_permit`).
#[must_use]
pub fn get_permits(&self, num_permits: usize) -> Vec<oneshot::Receiver<SearchPermit>> {
pub fn get_permits(&self, num_permits: usize) -> Vec<SearchPermitFuture> {
let mut permits_lock = self.inner_arc.lock().unwrap();
permits_lock.get_permits(num_permits, &self.inner_arc)
}
}

struct InnerSearchPermitProvider {
num_permits_available: usize,
num_download_slots_available: usize,

// Note it is possible for memory_allocated to exceed memory_budget temporarily,
// if and only if a split leaf search task ended up using more than `initial_allocation`.
Expand All @@ -82,33 +81,29 @@ struct InnerSearchPermitProvider {
}

impl InnerSearchPermitProvider {
fn get_permit(
&mut self,
inner_arc: &Arc<Mutex<InnerSearchPermitProvider>>,
) -> oneshot::Receiver<SearchPermit> {
let (tx, rx) = oneshot::channel();
self.permits_requests.push_back(tx);
self.assign_available_permits(inner_arc);
rx
}

fn get_permits(
&mut self,
num_permits: usize,
inner_arc: &Arc<Mutex<InnerSearchPermitProvider>>,
) -> Vec<oneshot::Receiver<SearchPermit>> {
) -> Vec<SearchPermitFuture> {
let mut permits = Vec::with_capacity(num_permits);
for _ in 0..num_permits {
let (tx, rx) = oneshot::channel();
self.permits_requests.push_back(tx);
permits.push(rx);
permits.push(SearchPermitFuture(rx));
}
self.assign_available_permits(inner_arc);
permits
}

/// Called each time a permit is requested or released
///
/// Calling lock on `inner_arc` inside this method will cause a deadlock as
/// `&mut self` and `inner_arc` reference the same instance.
fn assign_available_permits(&mut self, inner_arc: &Arc<Mutex<InnerSearchPermitProvider>>) {
while self.num_permits_available > 0 && self.memory_allocated + self.initial_allocation <= self.memory_budget {
while self.num_download_slots_available > 0
&& self.memory_allocated + self.initial_allocation <= self.memory_budget
{
let Some(permit_requester_tx) = self.permits_requests.pop_front() else {
break;
};
Expand All @@ -124,7 +119,7 @@ impl InnerSearchPermitProvider {
});
match send_res {
Ok(()) => {
self.num_permits_available -= 1;
self.num_download_slots_available -= 1;
self.memory_allocated += self.initial_allocation;
}
Err(search_permit) => {
Expand Down Expand Up @@ -153,14 +148,22 @@ impl SearchPermit {
/// We can then set the actual memory usage.
pub fn set_actual_memory_usage_and_release_permit_after(&mut self, new_memory_usage: u64) {
if new_memory_usage > self.memory_allocation {
warn!(memory_usage=new_memory_usage, memory_allocation=self.memory_allocation, "current leaf search is consuming more memory than the initial allocation");
warn!(
memory_usage = new_memory_usage,
memory_allocation = self.memory_allocation,
"current leaf search is consuming more memory than the initial allocation"
);
}
let mut inner_guard = self.inner_arc.lock().unwrap();
let delta = new_memory_usage as i64 - inner_guard.initial_allocation as i64;
inner_guard.memory_allocated += delta as u64;
inner_guard.num_permits_available += 1;
inner_guard.num_download_slots_available += 1;
if inner_guard.memory_allocated > inner_guard.memory_budget {
warn!(memory_allocated=inner_guard.memory_allocated, memory_budget=inner_guard.memory_budget, "memory allocated exceeds memory budget");
warn!(
memory_allocated = inner_guard.memory_allocated,
memory_budget = inner_guard.memory_budget,
"memory allocated exceeds memory budget"
);
}
self.memory_allocation = new_memory_usage;
inner_guard.assign_available_permits(&self.inner_arc);
Expand All @@ -182,77 +185,92 @@ impl Drop for SearchPermit {
}
let mut inner_guard = self.inner_arc.lock().unwrap();
if self.warmup_permit_held {
inner_guard.num_permits_available += 1;
inner_guard.num_download_slots_available += 1;
}
inner_guard.memory_allocated -= self.memory_allocation;
inner_guard.assign_available_permits(&self.inner_arc);
}
}

#[cfg(test)]
mod tests {
use tokio::task::JoinSet;
pub struct SearchPermitFuture(oneshot::Receiver<SearchPermit>);

use super::*;
impl Future for SearchPermitFuture {
type Output = Option<SearchPermit>;

#[tokio::test]
async fn test_search_permits_get_permits_future() {
// We test here that `get_permits_futures` does not interleave
let search_permits = SearchPermitProvider::new(1);
let mut all_futures = Vec::new();
let first_batch_of_permits = search_permits.get_permits(10);
assert_eq!(first_batch_of_permits.len(), 10);
all_futures.extend(
first_batch_of_permits
.into_iter()
.enumerate()
.map(move |(i, fut)| ((1, i), fut)),
);
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let receiver = Pin::new(&mut self.get_mut().0);
match receiver.poll(cx) {
Poll::Ready(Ok(search_permit)) => Poll::Ready(Some(search_permit)),
Poll::Ready(Err(_)) => panic!("Failed to acquire permit. This should never happen! Please, report on https://github.com/quickwit-oss/quickwit/issues."),
Poll::Pending => Poll::Pending,
}
}
}

let second_batch_of_permits = search_permits.get_permits(10);
assert_eq!(second_batch_of_permits.len(), 10);
all_futures.extend(
second_batch_of_permits
.into_iter()
.enumerate()
.map(move |(i, fut)| ((2, i), fut)),
);
// #[cfg(test)]
// mod tests {
// use tokio::task::JoinSet;

use rand::seq::SliceRandom;
// not super useful, considering what join set does, but still a tiny bit more sound.
all_futures.shuffle(&mut rand::thread_rng());
// use super::*;

let mut join_set = JoinSet::new();
for (res, fut) in all_futures {
join_set.spawn(async move {
let permit = fut.await;
(res, permit)
});
}
let mut ordered_result: Vec<(usize, usize)> = Vec::with_capacity(20);
while let Some(Ok(((batch_id, order), _permit))) = join_set.join_next().await {
ordered_result.push((batch_id, order));
}
// #[tokio::test]
// async fn test_search_permits_get_permits_future() {
// // We test here that `get_permits_futures` does not interleave
// let search_permits = SearchPermitProvider::new(1);
// let mut all_futures = Vec::new();
// let first_batch_of_permits = search_permits.get_permits(10);
// assert_eq!(first_batch_of_permits.len(), 10);
// all_futures.extend(
// first_batch_of_permits
// .into_iter()
// .enumerate()
// .map(move |(i, fut)| ((1, i), fut)),
// );

assert_eq!(ordered_result.len(), 20);
for (i, res) in ordered_result[0..10].iter().enumerate() {
assert_eq!(res, &(1, i));
}
for (i, res) in ordered_result[10..20].iter().enumerate() {
assert_eq!(res, &(2, i));
}
}
// let second_batch_of_permits = search_permits.get_permits(10);
// assert_eq!(second_batch_of_permits.len(), 10);
// all_futures.extend(
// second_batch_of_permits
// .into_iter()
// .enumerate()
// .map(move |(i, fut)| ((2, i), fut)),
// );

#[tokio::test]
async fn test_search_permits_receiver_race_condition() {
// Here we test that we don't have a problem if the Receiver is dropped.
// In particular, we want to check that there is not a race condition where drop attempts to
// lock the mutex.
let search_permits = SearchPermitProvider::new(1);
let permit_rx = search_permits.get_permit();
let permit_rx2 = search_permits.get_permit();
drop(permit_rx2);
drop(permit_rx);
let _permit_rx = search_permits.get_permit();
}
}
// use rand::seq::SliceRandom;
// // not super useful, considering what join set does, but still a tiny bit more sound.
// all_futures.shuffle(&mut rand::thread_rng());

// let mut join_set = JoinSet::new();
// for (res, fut) in all_futures {
// join_set.spawn(async move {
// let permit = fut.await;
// (res, permit)
// });
// }
// let mut ordered_result: Vec<(usize, usize)> = Vec::with_capacity(20);
// while let Some(Ok(((batch_id, order), _permit))) = join_set.join_next().await {
// ordered_result.push((batch_id, order));
// }

// assert_eq!(ordered_result.len(), 20);
// for (i, res) in ordered_result[0..10].iter().enumerate() {
// assert_eq!(res, &(1, i));
// }
// for (i, res) in ordered_result[10..20].iter().enumerate() {
// assert_eq!(res, &(2, i));
// }
// }

// #[tokio::test]
// async fn test_search_permits_receiver_race_condition() {
// // Here we test that we don't have a problem if the Receiver is dropped.
// // In particular, we want to check that there is not a race condition where drop attempts
// to // lock the mutex.
// let search_permits = SearchPermitProvider::new(1);
// let permit_rx = search_permits.get_permit();
// let permit_rx2 = search_permits.get_permit();
// drop(permit_rx2);
// drop(permit_rx);
// let _permit_rx = search_permits.get_permit();
// }
// }

0 comments on commit fd1c9a1

Please sign in to comment.