From d2ce52c21310c2abcf54fd8a55421c5cdc129411 Mon Sep 17 00:00:00 2001 From: Moritz Gunz Date: Mon, 22 Mar 2021 15:14:10 +0100 Subject: [PATCH 1/6] chore: Avoid allocation if PollSemaphore is unused --- tokio-util/src/sync/poll_semaphore.rs | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tokio-util/src/sync/poll_semaphore.rs b/tokio-util/src/sync/poll_semaphore.rs index 6b22b0d633f..24055fb8d16 100644 --- a/tokio-util/src/sync/poll_semaphore.rs +++ b/tokio-util/src/sync/poll_semaphore.rs @@ -12,17 +12,15 @@ use super::ReusableBoxFuture; /// [`Semaphore`]: tokio::sync::Semaphore pub struct PollSemaphore { semaphore: Arc, - permit_fut: ReusableBoxFuture>, + permit_fut: Option>>, } impl PollSemaphore { /// Create a new `PollSemaphore`. pub fn new(semaphore: Arc) -> Self { - let fut = Arc::clone(&semaphore).acquire_owned(); - Self { semaphore, - permit_fut: ReusableBoxFuture::new(fut), + permit_fut: None, } } @@ -55,10 +53,20 @@ 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> { - let result = ready!(self.permit_fut.poll(cx)); + let permit_future = match self.permit_fut.as_mut() { + Some(fut) => fut, + None => { + let next_fut = Arc::clone(&self.semaphore).acquire_owned(); + self.permit_fut = Some(ReusableBoxFuture::new(next_fut)); + + self.permit_fut.as_mut().unwrap() + } + }; + + let result = ready!(permit_future.poll(cx)); let next_fut = Arc::clone(&self.semaphore).acquire_owned(); - self.permit_fut.set(next_fut); + permit_future.set(next_fut); match result { Ok(permit) => Poll::Ready(Some(permit)), From 4cd3c5777bb4a64d7d904f9691d3a0ee46f97bae Mon Sep 17 00:00:00 2001 From: Moritz Gunz Date: Mon, 22 Mar 2021 15:43:48 +0100 Subject: [PATCH 2/6] chore: Add a fast-path for when a permit is immediately available --- tokio-util/src/sync/poll_semaphore.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tokio-util/src/sync/poll_semaphore.rs b/tokio-util/src/sync/poll_semaphore.rs index 24055fb8d16..88c6aa3da23 100644 --- a/tokio-util/src/sync/poll_semaphore.rs +++ b/tokio-util/src/sync/poll_semaphore.rs @@ -3,7 +3,7 @@ use std::fmt; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore}; +use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore, TryAcquireError}; use super::ReusableBoxFuture; @@ -56,6 +56,13 @@ impl PollSemaphore { let permit_future = match self.permit_fut.as_mut() { Some(fut) => fut, None => { + // avoid allocations completely if we can grab a permit immediately + match Arc::clone(&self.semaphore).try_acquire_owned() { + Ok(permit) => return Poll::Ready(Some(permit)), + Err(TryAcquireError::Closed) => return Poll::Ready(None), + Err(TryAcquireError::NoPermits) => {} + } + let next_fut = Arc::clone(&self.semaphore).acquire_owned(); self.permit_fut = Some(ReusableBoxFuture::new(next_fut)); From f9009f47a35d5e3503762d55fbaebd79de1a5859 Mon Sep 17 00:00:00 2001 From: Moritz Gunz Date: Mon, 22 Mar 2021 16:13:45 +0100 Subject: [PATCH 3/6] chore: Use .get_or_insert instead of .unwrap Co-authored-by: Alice Ryhl --- tokio-util/src/sync/poll_semaphore.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tokio-util/src/sync/poll_semaphore.rs b/tokio-util/src/sync/poll_semaphore.rs index 88c6aa3da23..b53f230dc1d 100644 --- a/tokio-util/src/sync/poll_semaphore.rs +++ b/tokio-util/src/sync/poll_semaphore.rs @@ -64,9 +64,7 @@ impl PollSemaphore { } let next_fut = Arc::clone(&self.semaphore).acquire_owned(); - self.permit_fut = Some(ReusableBoxFuture::new(next_fut)); - - self.permit_fut.as_mut().unwrap() + self.permit_fut.get_or_insert(ReusableBoxFuture::new(next_fut)); } }; From f75e2957426b7c17a5cfac054f28b0937eb6bccd Mon Sep 17 00:00:00 2001 From: Moritz Gunz Date: Mon, 22 Mar 2021 16:19:35 +0100 Subject: [PATCH 4/6] fix: Remove superfluous semicolon --- tokio-util/src/sync/poll_semaphore.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tokio-util/src/sync/poll_semaphore.rs b/tokio-util/src/sync/poll_semaphore.rs index b53f230dc1d..e18759e6157 100644 --- a/tokio-util/src/sync/poll_semaphore.rs +++ b/tokio-util/src/sync/poll_semaphore.rs @@ -64,7 +64,8 @@ impl PollSemaphore { } let next_fut = Arc::clone(&self.semaphore).acquire_owned(); - self.permit_fut.get_or_insert(ReusableBoxFuture::new(next_fut)); + self.permit_fut + .get_or_insert(ReusableBoxFuture::new(next_fut)) } }; From 80d245139a5b4d5ee91f9f6e0112d84f93055b56 Mon Sep 17 00:00:00 2001 From: Moritz Gunz Date: Mon, 22 Mar 2021 16:19:41 +0100 Subject: [PATCH 5/6] chore: Erase future when semaphore is closed --- tokio-util/src/sync/poll_semaphore.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tokio-util/src/sync/poll_semaphore.rs b/tokio-util/src/sync/poll_semaphore.rs index e18759e6157..c2ec20eae3e 100644 --- a/tokio-util/src/sync/poll_semaphore.rs +++ b/tokio-util/src/sync/poll_semaphore.rs @@ -59,7 +59,10 @@ impl PollSemaphore { // avoid allocations completely if we can grab a permit immediately match Arc::clone(&self.semaphore).try_acquire_owned() { Ok(permit) => return Poll::Ready(Some(permit)), - Err(TryAcquireError::Closed) => return Poll::Ready(None), + Err(TryAcquireError::Closed) => { + self.permit_fut = None; + return Poll::Ready(None); + } Err(TryAcquireError::NoPermits) => {} } From 3cdc068dd903d22236c7841da495aba58a95d0eb Mon Sep 17 00:00:00 2001 From: Moritz Gunz Date: Mon, 22 Mar 2021 16:26:11 +0100 Subject: [PATCH 6/6] fix: Erase future in the right match --- tokio-util/src/sync/poll_semaphore.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tokio-util/src/sync/poll_semaphore.rs b/tokio-util/src/sync/poll_semaphore.rs index c2ec20eae3e..a0a531c9171 100644 --- a/tokio-util/src/sync/poll_semaphore.rs +++ b/tokio-util/src/sync/poll_semaphore.rs @@ -59,10 +59,7 @@ impl PollSemaphore { // avoid allocations completely if we can grab a permit immediately match Arc::clone(&self.semaphore).try_acquire_owned() { Ok(permit) => return Poll::Ready(Some(permit)), - Err(TryAcquireError::Closed) => { - self.permit_fut = None; - return Poll::Ready(None); - } + Err(TryAcquireError::Closed) => return Poll::Ready(None), Err(TryAcquireError::NoPermits) => {} } @@ -79,7 +76,10 @@ impl PollSemaphore { match result { Ok(permit) => Poll::Ready(Some(permit)), - Err(_closed) => Poll::Ready(None), + Err(_closed) => { + self.permit_fut = None; + Poll::Ready(None) + } } } }