This commit is contained in:
Alexey
2026-10-30 22:34:37 +03:00
parent e7e588ebe5
commit 5a5009174a
5 changed files with 242 additions and 66 deletions
@@ -1,19 +1,28 @@
use super::super::ConnRegistry;
use super::super::pool_writer_security_tests::make_pool;
use super::*;
use std::future::{Future, poll_fn};
use std::task::Poll;
use super::super::pool_writer_security_tests::make_pool;
use super::super::ConnRegistry;
#[tokio::test]
async fn shutdown_cancels_registration_waiting_for_cleanup_capacity() {
let mut pool = make_pool().await;
Arc::get_mut(&mut Arc::get_mut(&mut pool).unwrap().routing).unwrap().registry =
Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 1));
Arc::get_mut(&mut Arc::get_mut(&mut pool).unwrap().routing)
.unwrap()
.registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 1));
let (lease, _receiver) = pool.register_connection().await.unwrap();
let mut waiting = Box::pin(pool.register_connection());
assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending));
assert!(matches!(
poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await,
Poll::Pending
));
pool.begin_shutdown();
assert!(tokio::time::timeout(Duration::from_secs(1), waiting).await.unwrap().is_none());
assert!(
tokio::time::timeout(Duration::from_secs(1), waiting)
.await
.unwrap()
.is_none()
);
lease.unregister().await;
assert!(pool.shutdown_until(Duration::from_secs(1)).await);
assert!(pool.register_connection().await.is_none());
@@ -13,16 +13,28 @@ async fn finite_cleanup_capacity_waits_until_dropped_route_is_received() {
}
assert_eq!(registry.routing.map.len(), 128);
assert_eq!(registry.cleanup_tx.capacity(), 0);
assert!(matches!(cleanup.try_recv(), Err(mpsc::error::TryRecvError::Empty)));
assert!(matches!(
cleanup.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
));
let mut waiting = Box::pin(registry.register_leased());
assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending));
assert!(matches!(
poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await,
Poll::Pending
));
let retired = clients.last().unwrap().0.conn_id();
drop(clients.pop().unwrap());
assert!(!registry.routing.map.contains_key(&retired));
assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending));
assert!(matches!(
poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await,
Poll::Pending
));
assert_eq!(cleanup.recv().await, Some(retired));
registry.unregister(retired).await;
let (lease, _) = tokio::time::timeout(Duration::from_secs(1), waiting).await.unwrap().unwrap();
let (lease, _) = tokio::time::timeout(Duration::from_secs(1), waiting)
.await
.unwrap()
.unwrap();
lease.unregister().await;
for (lease, _) in clients {
lease.unregister().await;
@@ -39,7 +51,10 @@ async fn cancelled_registration_leaks_neither_route_nor_cleanup_permit() {
let (lease, _receiver) = registry.register_leased().await.unwrap();
for _ in 0..64 {
let mut waiting = Box::pin(registry.register_leased());
assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending));
assert!(matches!(
poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await,
Poll::Pending
));
drop(waiting);
assert_eq!(registry.routing.map.len(), 1);
assert_eq!(registry.cleanup_tx.capacity(), 0);
@@ -48,7 +63,9 @@ async fn cancelled_registration_leaks_neither_route_nor_cleanup_permit() {
assert!(registry.routing.map.is_empty());
assert_eq!(registry.cleanup_tx.capacity(), 1);
let (lease, _) = tokio::time::timeout(Duration::from_secs(1), registry.register_leased())
.await.unwrap().unwrap();
.await
.unwrap()
.unwrap();
lease.unregister().await;
assert_eq!(registry.cleanup_tx.capacity(), 1);
}