From 7a021f954d12f0cace12b1e19852577dd54d7e16 Mon Sep 17 00:00:00 2001 From: liquidsec Date: Mon, 30 Sep 2024 15:32:15 -0400 Subject: [PATCH] more efficient task handling --- bbot/modules/baddns.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/bbot/modules/baddns.py b/bbot/modules/baddns.py index 4b9f977c2..84a5b6fed 100644 --- a/bbot/modules/baddns.py +++ b/bbot/modules/baddns.py @@ -71,13 +71,19 @@ async def handle_event(self, event): kwargs["raw_query_retry_wait"] = 0 module_instance = ModuleClass(event.data, **kwargs) + task = asyncio.create_task(module_instance.dispatch()) + tasks.append((module_instance, task)) - tasks.append((module_instance, asyncio.create_task(module_instance.dispatch()))) + async for completed_task in self.helpers.as_completed([task for _, task in tasks]): - for module_instance, task in tasks: + module_instance = next((m for m, t in tasks if t == completed_task), None) + try: + task_result = await completed_task + except Exception as e: + self.hugewarning(f"Task for {module_instance} raised an error: {e}") + task_result = None - complete_task = await task - if complete_task: + if task_result: results = module_instance.analyze() if results and len(results) > 0: for r in results: @@ -122,5 +128,4 @@ async def handle_event(self, event): tags=[f"baddns-{module_instance.name.lower()}"], context=f'{{module}}\'s "{r_dict["module"]}" module found {{event.type}}: {{event.data}}', ) - await module_instance.cleanup()