From 06d6adf4b7597fa316f3277c2aa60da27dac59fc Mon Sep 17 00:00:00 2001
From: Rob Ede <robjtede@icloud.com>
Date: Sat, 30 Jan 2021 09:57:04 +0000
Subject: [PATCH] util: use ReusableBoxFuture for PollSemaphore (#3463)

---
 tokio-util/Cargo.toml                 |  2 +-
 tokio-util/src/sync/poll_semaphore.rs | 30 +++++++++++++++------------
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/tokio-util/Cargo.toml b/tokio-util/Cargo.toml
index 6bfbdf6c4..9440e4ebb 100644
--- a/tokio-util/Cargo.toml
+++ b/tokio-util/Cargo.toml
@@ -47,12 +47,12 @@ futures-util = { version = "0.3.0", optional = true }
 log = "0.4"
 pin-project-lite = "0.2.0"
 slab = { version = "0.4.1", optional = true } # Backs `DelayQueue`
-async-stream = "0.3.0"
 
 [dev-dependencies]
 tokio = { version = "1.0.0", features = ["full"] }
 tokio-test = { version = "0.4.0" }
 
+async-stream = "0.3.0"
 futures = "0.3.0"
 futures-test = "0.3.5"
 
diff --git a/tokio-util/src/sync/poll_semaphore.rs b/tokio-util/src/sync/poll_semaphore.rs
index 6519bc663..d4594d03e 100644
--- a/tokio-util/src/sync/poll_semaphore.rs
+++ b/tokio-util/src/sync/poll_semaphore.rs
@@ -1,31 +1,28 @@
-use futures_core::Stream;
+use futures_core::{ready, Stream};
 use std::fmt;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
-use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore};
+
+use super::ReusableBoxFuture;
 
 /// A wrapper around [`Semaphore`] that provides a `poll_acquire` method.
 ///
 /// [`Semaphore`]: tokio::sync::Semaphore
 pub struct PollSemaphore {
     semaphore: Arc<Semaphore>,
-    inner: Pin<Box<dyn Stream<Item = OwnedSemaphorePermit> + Send + Sync>>,
+    permit_fut: ReusableBoxFuture<Result<OwnedSemaphorePermit, AcquireError>>,
 }
 
 impl PollSemaphore {
     /// Create a new `PollSemaphore`.
     pub fn new(semaphore: Arc<Semaphore>) -> Self {
+        let fut = Arc::clone(&semaphore).acquire_owned();
+
         Self {
-            semaphore: semaphore.clone(),
-            inner: Box::pin(async_stream::stream! {
-                loop {
-                    match semaphore.clone().acquire_owned().await {
-                        Ok(permit) => yield permit,
-                        Err(_closed) => break,
-                    }
-                }
-            }),
+            semaphore,
+            permit_fut: ReusableBoxFuture::new(fut),
         }
     }
 
@@ -58,7 +55,14 @@ impl PollSemaphore {
     /// the `Waker` from the `Context` passed to the most recent call is
     /// scheduled to receive a wakeup.
     pub fn poll_acquire(&mut self, cx: &mut Context<'_>) -> Poll<Option<OwnedSemaphorePermit>> {
-        self.inner.as_mut().poll_next(cx)
+        match ready!(self.permit_fut.poll(cx)) {
+            Ok(permit) => {
+                let next_fut = Arc::clone(&self.semaphore).acquire_owned();
+                self.permit_fut.set(next_fut);
+                Poll::Ready(Some(permit))
+            }
+            Err(_closed) => Poll::Ready(None),
+        }
     }
 }