add timeout for resolve_raw_batch()

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
 

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()

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()