sync: add PollSemaphore::poll_acquire_many (#5137)

tokio-util/tests/poll_semaphore.rs

@@ -13,6 +13,14 @@ fn semaphore_poll(
     tokio_test::task::spawn(fut)
 }
 
+fn semaphore_poll_many(
+    sem: &mut PollSemaphore,
+    permits: u32,
+) -> tokio_test::task::Spawn<impl Future<Output = SemRet> + '_> {
+    let fut = futures::future::poll_fn(move |cx| sem.poll_acquire_many(cx, permits));
+    tokio_test::task::spawn(fut)
+}
+
 #[tokio::test]
 async fn it_works() {
     let sem = Arc::new(Semaphore::new(1));
@@ -34,3 +42,43 @@ async fn it_works() {
     assert!(semaphore_poll(&mut poll_sem).await.is_none());
     assert!(semaphore_poll(&mut poll_sem).await.is_none());
 }
+
+#[tokio::test]
+async fn can_acquire_many_permits() {
+    let sem = Arc::new(Semaphore::new(4));
+    let mut poll_sem = PollSemaphore::new(sem.clone());
+
+    let permit1 = semaphore_poll(&mut poll_sem).poll();
+    assert!(matches!(permit1, Poll::Ready(Some(_))));
+
+    let permit2 = semaphore_poll_many(&mut poll_sem, 2).poll();
+    assert!(matches!(permit2, Poll::Ready(Some(_))));
+
+    assert_eq!(sem.available_permits(), 1);
+
+    drop(permit2);
+
+    let mut permit4 = semaphore_poll_many(&mut poll_sem, 4);
+    assert!(permit4.poll().is_pending());
+
+    drop(permit1);
+
+    let permit4 = permit4.poll();
+    assert!(matches!(permit4, Poll::Ready(Some(_))));
+    assert_eq!(sem.available_permits(), 0);
+}
+
+#[tokio::test]
+async fn can_poll_different_amounts_of_permits() {
+    let sem = Arc::new(Semaphore::new(4));
+    let mut poll_sem = PollSemaphore::new(sem.clone());
+    assert!(semaphore_poll_many(&mut poll_sem, 5).poll().is_pending());
+    assert!(semaphore_poll_many(&mut poll_sem, 4).poll().is_ready());
+
+    let permit = sem.acquire_many(4).await.unwrap();
+    assert!(semaphore_poll_many(&mut poll_sem, 5).poll().is_pending());
+    assert!(semaphore_poll_many(&mut poll_sem, 4).poll().is_pending());
+    drop(permit);
+    assert!(semaphore_poll_many(&mut poll_sem, 5).poll().is_pending());
+    assert!(semaphore_poll_many(&mut poll_sem, 4).poll().is_ready());
+}