Prevents single split searches from different leaf_search from interleaving #5509
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,213 @@ | ||
| // Copyright (C) 2024 Quickwit, Inc. | ||
| // | ||
| // Quickwit is offered under the AGPL v3.0 and as commercial software. | ||
| // For commercial licensing, contact us at hello@quickwit.io. | ||
| // | ||
| // AGPL: | ||
| // This program is free software: you can redistribute it and/or modify | ||
| // it under the terms of the GNU Affero General Public License as | ||
| // published by the Free Software Foundation, either version 3 of the | ||
| // License, or (at your option) any later version. | ||
| // | ||
| // This program is distributed in the hope that it will be useful, | ||
| // but WITHOUT ANY WARRANTY; without even the implied warranty of | ||
| // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the | ||
| // GNU Affero General Public License for more details. | ||
| // | ||
| // You should have received a copy of the GNU Affero General Public License | ||
| // along with this program. If not, see <http://www.gnu.org/licenses/>. | ||
|
|
||
| use std::collections::VecDeque; | ||
| use std::sync::{Arc, Mutex, Weak}; | ||
|
|
||
| use quickwit_common::metrics::GaugeGuard; | ||
| use tokio::sync::oneshot; | ||
|
|
||
| /// `SearchPermits` is a distributor of permits to perform single split | ||
| /// search operation. | ||
| /// | ||
| /// Requests are served in order. | ||
| #[derive(Clone)] | ||
| pub struct SearchPermits { | ||
| inner: Arc<Mutex<InnerSearchPermits>>, | ||
| } | ||
|
|
||
| impl SearchPermits { | ||
| pub fn new(num_permits: usize) -> SearchPermits { | ||
| SearchPermits { | ||
| inner: Arc::new(Mutex::new(InnerSearchPermits { | ||
| num_permits_available: num_permits, | ||
| permits_requests: VecDeque::new(), | ||
| })), | ||
| } | ||
| } | ||
|
|
||
| /// Returns a list of future permits in the form of a Receiver channel. | ||
| pub fn get_one_permit(&self) -> oneshot::Receiver<SearchPermit> { | ||
| let mut permits_lock = self.inner.lock().unwrap(); | ||
| permits_lock.get_permit(&self.inner) | ||
| } | ||
|
|
||
| /// Returns a list of future permits in the form of a Receiver channel. | ||
| /// | ||
| /// 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). | ||
| pub fn get_permits_futures(&self, num_permits: usize) -> Vec<oneshot::Receiver<SearchPermit>> { | ||
| let mut permits_lock = self.inner.lock().unwrap(); | ||
| permits_lock.get_permits(num_permits, &self.inner) | ||
| } | ||
| } | ||
|
|
||
| struct InnerSearchPermits { | ||
| num_permits_available: usize, | ||
| permits_requests: VecDeque<oneshot::Sender<SearchPermit>>, | ||
| } | ||
|
|
||
| impl InnerSearchPermits { | ||
| fn get_permits( | ||
| &mut self, | ||
| num_permits: usize, | ||
| inner: &Arc<Mutex<InnerSearchPermits>>, | ||
|
There was a problem hiding this comment. i think inner is a very confusing name, the InnerSearchPermits is already "the inner thing", so an There was a problem hiding this comment. I don't quite get the confusion/comment. I have not found a common practise on how to name the object of this pattern in Rust (in C++ What would be a better name? would inner_arc help a bit? |
||
| ) -> Vec<oneshot::Receiver<SearchPermit>> { | ||
| 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); | ||
| } | ||
| self.assign_available_permits(inner); | ||
| permits | ||
| } | ||
|
|
||
| fn get_permit( | ||
| &mut self, | ||
| inner: &Arc<Mutex<InnerSearchPermits>>, | ||
| ) -> oneshot::Receiver<SearchPermit> { | ||
| let (tx, rx) = oneshot::channel(); | ||
| self.permits_requests.push_back(tx); | ||
| self.assign_available_permits(inner); | ||
| rx | ||
| } | ||
|
|
||
| fn recycle_permit(&mut self, inner: &Arc<Mutex<InnerSearchPermits>>) { | ||
| self.num_permits_available += 1; | ||
| self.assign_available_permits(inner); | ||
| } | ||
|
|
||
| fn assign_available_permits(&mut self, inner: &Arc<Mutex<InnerSearchPermits>>) { | ||
| while self.num_permits_available > 0 { | ||
| let Some(sender) = self.permits_requests.pop_front() else { | ||
| break; | ||
| }; | ||
| let send_res = sender.send(SearchPermit { | ||
| _ongoing_gauge_guard: GaugeGuard::from_gauge( | ||
| &crate::SEARCH_METRICS.leaf_search_single_split_tasks_ongoing, | ||
| ), | ||
| inner_opt: Some(Arc::downgrade(inner)), | ||
| }); | ||
| match send_res { | ||
| Ok(()) => { | ||
| self.num_permits_available -= 1; | ||
| } | ||
| Err(mut search_permit) => { | ||
| search_permit.disable_drop(); | ||
| drop(search_permit); | ||
| } | ||
| } | ||
| } | ||
| crate::SEARCH_METRICS | ||
| .leaf_search_single_split_tasks_pending | ||
| .set(self.permits_requests.len() as i64); | ||
| } | ||
| } | ||
|
|
||
| pub struct SearchPermit { | ||
| _ongoing_gauge_guard: GaugeGuard<'static>, | ||
| inner_opt: Option<Weak<Mutex<InnerSearchPermits>>>, | ||
|
There was a problem hiding this comment. I don't believe there is need for Weak here: i find this type, There was a problem hiding this comment. I prefer the I'll add some comments. There was a problem hiding this comment. you ended up changing it to a bool. If you think an Option is better, you can put it back, just add a 1-2 line comment on the purpose of None There was a problem hiding this comment. Yeah, I thought after packing everything as |
||
| } | ||
|
|
||
| impl SearchPermit { | ||
| fn disable_drop(&mut self) { | ||
| self.inner_opt = None; | ||
| } | ||
| } | ||
|
|
||
| impl Drop for SearchPermit { | ||
| fn drop(&mut self) { | ||
| let Some(inner) = self.inner_opt.take() else { | ||
| return; | ||
| }; | ||
| let Some(inner) = inner.upgrade() else { | ||
| return; | ||
| }; | ||
| let mut inner_guard = inner.lock().unwrap(); | ||
| inner_guard.recycle_permit(&inner); | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
| use tokio::task::JoinSet; | ||
|
|
||
| #[tokio::test] | ||
| async fn test_search_permits_get_permits_future() { | ||
| // We test here that `get_permits_futures` does not interleave | ||
| let search_permits = super::SearchPermits::new(1); | ||
| let mut all_futures = Vec::new(); | ||
| let first_batch_of_permits = search_permits.get_permits_futures(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)), | ||
| ); | ||
|
|
||
| let second_batch_of_permits = search_permits.get_permits_futures(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)), | ||
| ); | ||
|
|
||
| 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 = super::SearchPermits::new(1); | ||
| let permit_rx = search_permits.get_one_permit(); | ||
| let permit_rx2 = search_permits.get_one_permit(); | ||
| drop(permit_rx2); | ||
| drop(permit_rx); | ||
| let _permit_rx = search_permits.get_one_permit(); | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this is a misnomer, it sounds like this is a akin to a tokio
SemaphorePermitwith num_permits > 1, but this is actually a semaphore, or a PermitProvider, or something of the likeThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah you are right.