diff --git a/bbot/core/engine.py b/bbot/core/engine.py index 52d4b871a..cde424b6a 100644 --- a/bbot/core/engine.py +++ b/bbot/core/engine.py @@ -537,9 +537,20 @@ def new_child_task(self, client_id, coro): self.child_tasks[client_id] = {task} return task - async def finished_tasks(self, client_id): + async def finished_tasks(self, client_id, timeout=None): child_tasks = self.child_tasks.get(client_id, set()) - done, pending = await asyncio.wait(child_tasks, return_when=asyncio.FIRST_COMPLETED) + try: + done, pending = await asyncio.wait(child_tasks, return_when=asyncio.FIRST_COMPLETED, timeout=timeout) + except BaseException as e: + if isinstance(e, (TimeoutError, asyncio.TimeoutError)): + done = set() + self.log.warning(f"{self.name}: Timeout after {timeout:,} seconds in finished_tasks({child_tasks})") + for task in child_tasks: + task.cancel() + else: + self.log.error(f"{self.name}: Unhandled exception in finished_tasks({child_tasks}): {e}") + self.log.trace(traceback.format_exc()) + raise self.child_tasks[client_id] = pending return done diff --git a/bbot/core/helpers/dns/engine.py b/bbot/core/helpers/dns/engine.py index 91efca10d..16be33c35 100644 --- a/bbot/core/helpers/dns/engine.py +++ b/bbot/core/helpers/dns/engine.py @@ -361,7 +361,7 @@ def new_task(query): while tasks: # While there are tasks pending # Wait for the first task to complete - finished = await self.finished_tasks(client_id) + finished = await self.finished_tasks(client_id, timeout=120) for task in finished: results = task.result() @@ -388,7 +388,7 @@ def new_task(query, rdtype): while tasks: # While there are tasks pending # Wait for the first task to complete - finished = await self.finished_tasks(client_id) + finished = self.finished_tasks(client_id, timeout=120) for task in finished: answers, errors = task.result() diff --git a/bbot/core/helpers/web/engine.py b/bbot/core/helpers/web/engine.py index bc58057ed..30e037e6c 100644 --- a/bbot/core/helpers/web/engine.py +++ b/bbot/core/helpers/web/engine.py @@ -100,7 +100,7 @@ def new_task(): while tasks: # While there are tasks pending # Wait for the first task to complete - finished = await self.finished_tasks(client_id) + finished = await self.finished_tasks(client_id, timeout=120) for task in finished: response = task.result()