rename the method to get a permit and use it in all search requests

This commit is contained in:
Tamo 2024-03-26 17:28:03 +01:00
parent 3f23fbb46d
commit e433fd53e6
6 changed files with 47 additions and 26 deletions

View file

@ -1,18 +1,18 @@
use std::{sync::Arc, time::Duration};
use std::{num::NonZeroUsize, sync::Arc, time::Duration};
use meili_snap::snapshot;
use meilisearch::search_queue::SearchQueue;
#[actix_rt::test]
async fn search_queue_register() {
let queue = SearchQueue::new(4, 2);
let queue = SearchQueue::new(4, NonZeroUsize::new(2).unwrap());
// First, use all the cores
let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search())
let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit())
.await
.expect("I should get a permit straight away")
.unwrap();
let _permit2 = tokio::time::timeout(Duration::from_secs(1), queue.register_search())
let _permit2 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit())
.await
.expect("I should get a permit straight away")
.unwrap();
@ -20,7 +20,7 @@ async fn search_queue_register() {
// If we free one spot we should be able to register one new search
drop(permit1);
let permit3 = tokio::time::timeout(Duration::from_secs(1), queue.register_search())
let permit3 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit())
.await
.expect("I should get a permit straight away")
.unwrap();
@ -28,7 +28,7 @@ async fn search_queue_register() {
// And again
drop(permit3);
let _permit4 = tokio::time::timeout(Duration::from_secs(1), queue.register_search())
let _permit4 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit())
.await
.expect("I should get a permit straight away")
.unwrap();
@ -36,19 +36,19 @@ async fn search_queue_register() {
#[actix_rt::test]
async fn search_queue_wait_till_cores_available() {
let queue = Arc::new(SearchQueue::new(4, 1));
let queue = Arc::new(SearchQueue::new(4, NonZeroUsize::new(1).unwrap()));
// First, use all the cores
let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search())
let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit())
.await
.expect("I should get a permit straight away")
.unwrap();
let ret = tokio::time::timeout(Duration::from_secs(1), queue.register_search()).await;
let ret = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()).await;
assert!(ret.is_err(), "The capacity is full, we should not get a permit");
let q = queue.clone();
let task = tokio::task::spawn(async move { q.register_search().await });
let task = tokio::task::spawn(async move { q.try_get_search_permit().await });
// after dropping a permit the previous task should be able to finish
drop(permit1);
@ -60,20 +60,20 @@ async fn search_queue_wait_till_cores_available() {
#[actix_rt::test]
async fn search_queue_refuse_search_requests() {
let queue = Arc::new(SearchQueue::new(1, 1));
let queue = Arc::new(SearchQueue::new(1, NonZeroUsize::new(1).unwrap()));
// First, use the whole capacity of the
let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search())
let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit())
.await
.expect("I should get a permit straight away")
.unwrap();
let q = queue.clone();
let permit2 = tokio::task::spawn(async move { q.register_search().await });
let permit2 = tokio::task::spawn(async move { q.try_get_search_permit().await });
// Here the queue is full. By registering two new search requests the permit 2 and 3 should be thrown out
let q = queue.clone();
let _permit3 = tokio::task::spawn(async move { q.register_search().await });
let _permit3 = tokio::task::spawn(async move { q.try_get_search_permit().await });
let permit2 = tokio::time::timeout(Duration::from_secs(1), permit2)
.await
@ -85,14 +85,14 @@ async fn search_queue_refuse_search_requests() {
#[actix_rt::test]
async fn search_request_crashes_while_holding_permits() {
let queue = Arc::new(SearchQueue::new(1, 1));
let queue = Arc::new(SearchQueue::new(1, NonZeroUsize::new(1).unwrap()));
let (send, recv) = tokio::sync::oneshot::channel();
// This first request take a cpu
let q = queue.clone();
tokio::task::spawn(async move {
let _permit = q.register_search().await.unwrap();
let _permit = q.try_get_search_permit().await.unwrap();
recv.await.unwrap();
panic!("oops an unexpected crash happened")
});
@ -100,7 +100,7 @@ async fn search_request_crashes_while_holding_permits() {
// This second request waits in the queue till the first request finishes
let q = queue.clone();
let task = tokio::task::spawn(async move {
let _permit = q.register_search().await.unwrap();
let _permit = q.try_get_search_permit().await.unwrap();
});
// By sending something in the channel the request holding a CPU will panic and should lose its permit
@ -113,7 +113,7 @@ async fn search_request_crashes_while_holding_permits() {
.unwrap();
// I should even be able to take second permit here
let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search())
let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit())
.await
.expect("I should get a permit straight away")
.unwrap();