From 48a5b6220e3a0841535959a1f7113e87501dc1ea Mon Sep 17 00:00:00 2001
From: Taiki Endo <te316e89@gmail.com>
Date: Sat, 19 Feb 2022 19:26:11 +0900
Subject: [PATCH] Fix semaphore lifetime issue as well

---
 src/semaphore.rs   | 31 ++++++++++++++++++-------------
 tests/semaphore.rs |  9 +++++++++
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/src/semaphore.rs b/src/semaphore.rs
index 65768d7..094b061 100644
--- a/src/semaphore.rs
+++ b/src/semaphore.rs
@@ -1,3 +1,4 @@
+use std::future::Future;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
@@ -135,6 +136,21 @@ impl Semaphore {
         }
     }
 
+    async fn acquire_arc_impl(self: Arc<Self>) -> SemaphoreGuardArc {
+        let mut listener = None;
+
+        loop {
+            if let Some(guard) = self.try_acquire_arc() {
+                return guard;
+            }
+
+            match listener.take() {
+                None => listener = Some(self.event.listen()),
+                Some(l) => l.await,
+            }
+        }
+    }
+
     /// Waits for an owned permit for a concurrent operation.
     ///
     /// Returns a guard that releases the permit when dropped.
@@ -150,19 +166,8 @@ impl Semaphore {
     /// let guard = s.acquire_arc().await;
     /// # });
     /// ```
-    pub async fn acquire_arc(self: &Arc<Self>) -> SemaphoreGuardArc {
-        let mut listener = None;
-
-        loop {
-            if let Some(guard) = self.try_acquire_arc() {
-                return guard;
-            }
-
-            match listener.take() {
-                None => listener = Some(self.event.listen()),
-                Some(l) => l.await,
-            }
-        }
+    pub fn acquire_arc(self: &Arc<Self>) -> impl Future<Output = SemaphoreGuardArc> {
+        self.clone().acquire_arc_impl()
     }
 }
 
diff --git a/tests/semaphore.rs b/tests/semaphore.rs
index 8487dcc..2dc42d4 100644
--- a/tests/semaphore.rs
+++ b/tests/semaphore.rs
@@ -81,3 +81,12 @@ fn multi_resource() {
         rx1.recv().unwrap();
     });
 }
+
+#[test]
+fn lifetime() {
+    // Show that the future keeps the semaphore alive.
+    let _fut = {
+        let mutex = Arc::new(Semaphore::new(2));
+        mutex.acquire_arc()
+    };
+}