Skip to content

Commit

Permalink
Avoid to hang wait_closed on worker pool
Browse files Browse the repository at this point in the history
  • Loading branch information
mosquito committed Mar 9, 2024
1 parent ce5461d commit 5cd7ec6
Showing 1 changed file with 10 additions and 3 deletions.
13 changes: 10 additions & 3 deletions aiomisc/worker_pool.py
Expand Up @@ -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-")
Expand Down Expand Up @@ -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))

Expand Down

0 comments on commit 5cd7ec6

Please sign in to comment.