From 092d7682ec5624c1745ca8b0ce2c2ed909fc8303 Mon Sep 17 00:00:00 2001
From: Alexandru Matei <alexandru.matei@uipath.com>
Date: Thu, 11 Apr 2024 12:42:11 +0300
Subject: [PATCH] fix hang in wait_all_exit

self.wait() without await returns a Future, so the notified() object is not
created until we await on the returned future.  That means notify_waiters()
can be called before notified() is. This leads to notified() waiting
forever because notify_waiters is called only once, when the last waiter is
dropped.

notify_waiters() and notified() form a happens-before relationship.
There are two possible scenarios:

1. If notified() comes before notify_waiters() this means we can
   safely await on notified().

2. If notified() comes after notify_waiters() this means that what happened
   before it is visible in the notified() thread. Waiting on notified() at
   this point will block but we can check for waiters count, which is
   guaranteed to be 0 because it was set before notify_waiters() call.

Let's move notified() call before checking that the number of waiters
is 0.

Signed-off-by: Alexandru Matei <alexandru.matei@uipath.com>
---
 src/asynchronous/shutdown.rs | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/src/asynchronous/shutdown.rs b/src/asynchronous/shutdown.rs
index 6ea6a92c..4d883c35 100644
--- a/src/asynchronous/shutdown.rs
+++ b/src/asynchronous/shutdown.rs
@@ -144,24 +144,25 @@ impl Notifier {
     /// Wait for all [`Waiter`]s to drop.
     pub async fn wait_all_exit(&self) -> Result<(), Elapsed> {
         //debug_assert!(self.shared.is_shutdown());
-        if self.waiters() == 0 {
-            return Ok(());
-        }
-        let wait = self.wait();
-        if self.waiters() == 0 {
-            return Ok(());
-        }
-        wait.await
-    }
-
-    async fn wait(&self) -> Result<(), Elapsed> {
         if let Some(tm) = self.wait_time {
-            timeout(tm, self.shared.notify_exit.notified()).await
+            timeout(tm, self.wait()).await
         } else {
-            self.shared.notify_exit.notified().await;
+            self.wait().await;
             Ok(())
         }
     }
+
+    async fn wait(&self) {
+        while self.waiters() > 0 {
+            let notified = self.shared.notify_exit.notified();
+            if self.waiters() == 0 {
+                return;
+            }
+            notified.await;
+            // Some waiters could have been created in the meantime 
+            // by calling `subscribe`, loop again
+        }
+    }
 }
 
 impl Drop for Notifier {