From 5cd7ec6b877aa56ca26995f31ff0555d66eae24f Mon Sep 17 00:00:00 2001
From: Dmitry Orlov <me@mosquito.su>
Date: Sat, 9 Mar 2024 01:07:07 +0100
Subject: [PATCH] Avoid to hang wait_closed on worker pool

---
 aiomisc/worker_pool.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/aiomisc/worker_pool.py b/aiomisc/worker_pool.py
index e1cce0ff..e4280bd1 100644
--- a/aiomisc/worker_pool.py
+++ b/aiomisc/worker_pool.py
@@ -51,6 +51,8 @@ class WorkerPool:
     worker_ids: Tuple[bytes, ...]
     pids: Set[int]
 
+    SERVER_CLOSE_TIMEOUT = 1
+
     if hasattr(socket, "AF_UNIX"):
         def _create_socket(self) -> None:
             path = mktemp(suffix=".sock", prefix="worker-")
@@ -281,15 +283,20 @@ def __reject_futures(self) -> None:
 
     @shield
     async def close(self) -> None:
-        async with self.__closing_lock:
+        async with (self.__closing_lock):
             if self.__closing:
                 return
 
             self._kill_supervisor()
             self.__closing = True
             self.server.close()
-            await self.server.wait_closed()
-
+            await asyncio.gather(
+                asyncio.wait_for(
+                    self.server.wait_closed(),
+                    timeout=self.SERVER_CLOSE_TIMEOUT,
+                ),
+                return_exceptions=True,
+            )
             await cancel_tasks(tuple(self.__task_store))
             await cancel_tasks(tuple(self.__futures))