Separate CPU-bound thread pool from I/O executor

Author: liquidsec

bbot/core/helpers/helper.py

@@ -212,6 +212,8 @@ def loop(self):
             # during heavy scans where YARA, regex, DNS, and HTTP all compete for threads
             thread_pool_size = max(32, (os.cpu_count() or 1) * 4)
             self._loop.set_default_executor(ThreadPoolExecutor(max_workers=thread_pool_size))
+            # separate pool for CPU-bound work (YARA, regex) so it never queues behind I/O
+            self._cpu_executor = ThreadPoolExecutor(max_workers=max(8, os.cpu_count() or 4))
         return self._loop
 
     def run_in_executor(self, callback, *args, **kwargs):
@@ -225,6 +227,18 @@ def run_in_executor(self, callback, *args, **kwargs):
         callback = partial(callback, **kwargs)
         return self.loop.run_in_executor(None, callback, *args)
 
+    def run_in_executor_cpu(self, callback, *args, **kwargs):
+        """
+        Run short CPU-bound work that releases the GIL in a dedicated thread pool,
+        separate from I/O so it never queues behind long-running network calls.
+
+        Examples:
+            Execute callback:
+            >>> result = await self.helpers.run_in_executor_cpu(callback_fn, arg1, arg2)
+        """
+        callback = partial(callback, **kwargs)
+        return self.loop.run_in_executor(self._cpu_executor, callback, *args)
+
     def run_in_executor_mp(self, callback, *args, **kwargs):
         """
         Same as run_in_executor() except with a process pool executor

bbot/core/helpers/regex.py

@@ -29,19 +29,19 @@ def compile(self, *args, **kwargs):
 
     async def search(self, compiled_regex, *args, **kwargs):
         self.ensure_compiled_regex(compiled_regex)
-        return await self.parent_helper.run_in_executor(compiled_regex.search, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(compiled_regex.search, *args, **kwargs)
 
     async def match(self, compiled_regex, *args, **kwargs):
         self.ensure_compiled_regex(compiled_regex)
-        return await self.parent_helper.run_in_executor(compiled_regex.match, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(compiled_regex.match, *args, **kwargs)
 
     async def sub(self, compiled_regex, *args, **kwargs):
         self.ensure_compiled_regex(compiled_regex)
-        return await self.parent_helper.run_in_executor(compiled_regex.sub, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(compiled_regex.sub, *args, **kwargs)
 
     async def findall(self, compiled_regex, *args, **kwargs):
         self.ensure_compiled_regex(compiled_regex)
-        return await self.parent_helper.run_in_executor(compiled_regex.findall, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(compiled_regex.findall, *args, **kwargs)
 
     async def findall_multi(self, compiled_regexes, *args, threads=10, **kwargs):
         """
@@ -55,7 +55,7 @@ async def findall_multi(self, compiled_regexes, *args, threads=10, **kwargs):
         tasks = {}
 
         def new_task(regex_name, r):
-            task = self.parent_helper.run_in_executor(r.findall, *args, **kwargs)
+            task = self.parent_helper.run_in_executor_cpu(r.findall, *args, **kwargs)
             tasks[task] = regex_name
 
         compiled_regexes = dict(compiled_regexes)
@@ -77,15 +77,15 @@ def new_task(regex_name, r):
 
     async def finditer(self, compiled_regex, *args, **kwargs):
         self.ensure_compiled_regex(compiled_regex)
-        return await self.parent_helper.run_in_executor(self._finditer, compiled_regex, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(self._finditer, compiled_regex, *args, **kwargs)
 
     async def finditer_multi(self, compiled_regexes, *args, **kwargs):
         """
         Same as finditer() but with multiple regexes
         """
         for r in compiled_regexes:
             self.ensure_compiled_regex(r)
-        return await self.parent_helper.run_in_executor(self._finditer_multi, compiled_regexes, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(self._finditer_multi, compiled_regexes, *args, **kwargs)
 
     def _finditer_multi(self, compiled_regexes, *args, **kwargs):
         matches = []
@@ -98,16 +98,16 @@ def _finditer(self, compiled_regex, *args, **kwargs):
         return list(compiled_regex.finditer(*args, **kwargs))
 
     async def extract_params_html(self, *args, **kwargs):
-        return await self.parent_helper.run_in_executor(misc.extract_params_html, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(misc.extract_params_html, *args, **kwargs)
 
     async def extract_emails(self, *args, **kwargs):
-        return await self.parent_helper.run_in_executor(misc.extract_emails, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(misc.extract_emails, *args, **kwargs)
 
     async def search_dict_values(self, *args, **kwargs):
         def _search_dict_values(*_args, **_kwargs):
             return list(misc.search_dict_values(*_args, **_kwargs))
 
-        return await self.parent_helper.run_in_executor(_search_dict_values, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(_search_dict_values, *args, **kwargs)
 
     async def recursive_decode(self, *args, **kwargs):
-        return await self.parent_helper.run_in_executor(misc.recursive_decode, *args, **kwargs)
+        return await self.parent_helper.run_in_executor_cpu(misc.recursive_decode, *args, **kwargs)

bbot/core/helpers/yara_helper.py

@@ -40,7 +40,7 @@ async def match(self, compiled_rules, text):
         Given a compiled YARA rule and a body of text, return a list of strings that match the rule
         """
         matched_strings = []
-        matches = await self.parent_helper.run_in_executor(compiled_rules.match, data=text)
+        matches = await self.parent_helper.run_in_executor_cpu(compiled_rules.match, data=text)
         if matches:
             for match in matches:
                 for string_match in match.strings:

bbot/modules/dnsbrute_mutations.py

@@ -44,7 +44,7 @@ async def handle_event(self, event):
 
     async def get_parent_event(self, subdomain):
         start = time.time()
-        parent_host = await self.helpers.run_in_executor(self.helpers.closest_match, subdomain, self.parent_events)
+        parent_host = await self.helpers.run_in_executor_cpu(self.helpers.closest_match, subdomain, self.parent_events)
         elapsed = time.time() - start
         self.trace(f"{subdomain}: got closest match among {len(self.parent_events):,} parent events in {elapsed:.2f}s")
         return self.parent_events[parent_host]

bbot/modules/internal/excavate.py

@@ -1145,7 +1145,7 @@ async def search(self, data, event, content_type, discovery_context="HTTP respon
 
         for label, data_instance in data_items:
             # Your existing processing code
-            for result in await self.helpers.run_in_executor(self.yara_rules.match, data=f"{data_instance}"):
+            for result in await self.helpers.run_in_executor_cpu(self.yara_rules.match, data=f"{data_instance}"):
                 rule_name = result.rule
 
                 # Skip specific operations for 'parameter_extraction' rule on decoded_data

bbot/scanner/scanner.py

@@ -1193,7 +1193,7 @@ async def dns_yara_rules(self):
             if self.dns_yara_rules_uncompiled is not None:
                 import yara
 
-                self._dns_yara_rules = await self.helpers.run_in_executor(
+                self._dns_yara_rules = await self.helpers.run_in_executor_cpu(
                     yara.compile, source="\n".join(self.dns_yara_rules_uncompiled.values())
                 )
         return self._dns_yara_rules
@@ -1209,7 +1209,7 @@ async def extract_in_scope_hostnames(self, s):
         matches = set()
         dns_yara_rules = await self.dns_yara_rules()
         if dns_yara_rules is not None:
-            for match in await self.helpers.run_in_executor(dns_yara_rules.match, data=s):
+            for match in await self.helpers.run_in_executor_cpu(dns_yara_rules.match, data=s):
                 for string in match.strings:
                     for instance in string.instances:
                         matches.add(str(instance))