fix: resolve core engine concurrency hangs

tests/test_orchestrator.py

@@ -58,7 +58,7 @@ def test_set_concurrency():
     original_max = orchestrator.MAX_CONCURRENT_REQUESTS
     set_concurrency(10)
     assert orchestrator.MAX_CONCURRENT_REQUESTS == 10
-    assert orchestrator._shared_executor._max_workers == 10
+    assert orchestrator._shared_executor._max_workers == max(10 * 2, 250)
 
     # restore
     set_concurrency(original_max)

user_scanner/core/email_orchestrator.py

@@ -2,7 +2,7 @@
 import httpx
 from pathlib import Path
 from types import ModuleType
-from typing import List, Optional, Set, Union
+from typing import List, Optional, Set, Union, Callable
 
 from colorama import Fore, Style
 
@@ -65,9 +65,12 @@ async def _async_worker(
     sem: asyncio.Semaphore,
     configs: ScanConfig,
     printed_cats: Optional[Set] = None,
+    on_start: Optional[Callable[[str], None]] = None,
 ) -> Result:
     async with sem:
         site_name = get_site_name(module)
+        if on_start:
+            on_start(site_name)
         func = get_scan_func(module)
         actual_cat = find_category(module) or "Email"
 
@@ -88,10 +91,13 @@ async def _async_worker(
 
         try:
             import inspect
+            module_timeout = (get_global_timeout() or 15.0) + 10.0
             if inspect.iscoroutinefunction(func):
-                result = await func(email)
+                result = await asyncio.wait_for(func(email), timeout=module_timeout)
             else:
-                result = await asyncio.to_thread(func, email)
+                result = await asyncio.wait_for(asyncio.to_thread(func, email), timeout=module_timeout)
+        except asyncio.TimeoutError:
+            result = Result.error(f"Module execution timed out after {module_timeout}s")
         except Exception as e:
             result = Result.error(e)
 
@@ -104,22 +110,10 @@ async def _run_batch(
     configs: ScanConfig,
     printed_cats: Optional[Set] = None,
 ) -> List[Result]:
-    sem = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
-    tasks = []
-    for module in modules:
-        tasks.append(
-            _async_worker(
-                module,
-                email,
-                sem,
-                configs,
-                printed_cats=printed_cats,
-            )
-        )
-
-    if not tasks:
+    if not modules:
         return []
 
+    sem = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
     results = []
 
     with Progress(
@@ -130,7 +124,25 @@ async def _run_batch(
         TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
         transient=True,
     ) as progress:
-        task_id = progress.add_task(f"[cyan]Scanning {email}...", total=len(tasks))
+        task_id = progress.add_task(f"[cyan]Scanning {email}...", total=len(modules))
+
+        def on_start_cb(site: str):
+            progress.update(task_id, description=f"[cyan]Scanning {email}... ({site})")
+
+        tasks = []
+        for module in modules:
+            t = asyncio.create_task(
+                _async_worker(
+                    module,
+                    email,
+                    sem,
+                    configs,
+                    printed_cats=printed_cats,
+                    on_start=on_start_cb,
+                )
+            )
+            t.add_done_callback(lambda t: progress.advance(task_id))
+            tasks.append(t)
 
         for coro in asyncio.as_completed(tasks):
             result = await coro
@@ -144,16 +156,18 @@ async def _run_batch(
 
             result.show(configs)
             results.append(result)
-            progress.advance(task_id)
 
     return results
 
 
 async def _run_email_module_batch_async(
     module: Union[ModuleType, List[ModuleType]], email: str, configs: ScanConfig
 ) -> List[Result]:
+    loop = asyncio.get_running_loop()
+    import concurrent.futures
+    loop.set_default_executor(concurrent.futures.ThreadPoolExecutor(max_workers=250))
     modules = [module] if isinstance(module, ModuleType) else list(module)
-    return await _run_batch(modules, email, configs)
+    return await _run_batch(modules, email, configs, printed_cats=set())
 
 
 def run_email_module_batch(
@@ -165,6 +179,9 @@ def run_email_module_batch(
 async def _run_email_category_batch_async(
     category_path: Path, email: str, configs: ScanConfig
 ) -> List[Result]:
+    loop = asyncio.get_running_loop()
+    import concurrent.futures
+    loop.set_default_executor(concurrent.futures.ThreadPoolExecutor(max_workers=250))
     cat_name = category_path.stem.capitalize()
     modules = load_modules(category_path)
     printed_cats = set()
@@ -187,31 +204,23 @@ def run_email_category_batch(
 
 
 async def _run_email_full_batch_async(email: str, configs: ScanConfig) -> List[Result]:
+    loop = asyncio.get_running_loop()
+    import concurrent.futures
+    loop.set_default_executor(concurrent.futures.ThreadPoolExecutor(max_workers=250))
     categories = load_categories(True, configs.no_nsfw)
     all_results = []
     printed_cats: Set[str] = set()
 
     # 1. Pre-spawn all tasks for all categories (global concurrency)
-    category_tasks = []
+    category_modules = []
     total_tasks = 0
+    sem = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
+
     for cat_name, cat_path in categories.items():
         display_name = cat_name.capitalize()
         modules = load_modules(cat_path)
-
-        sem = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
-        tasks = []
-        for module in modules:
-            tasks.append(
-                _async_worker(
-                    module,
-                    email,
-                    sem,
-                    configs,
-                    printed_cats=printed_cats,
-                )
-            )
-        category_tasks.append((display_name, tasks))
-        total_tasks += len(tasks)
+        category_modules.append((display_name, modules))
+        total_tasks += len(modules)
 
     # 2. Await tasks category by category to stream grouped output
     with Progress(
@@ -224,7 +233,28 @@ async def _run_email_full_batch_async(email: str, configs: ScanConfig) -> List[R
     ) as progress:
         task_id = progress.add_task(f"[cyan]Scanning {email}...", total=total_tasks)
 
-        for display_name, tasks in category_tasks:
+        def on_start_cb(site: str):
+            progress.update(task_id, description=f"[cyan]Scanning {email}... ({site})")
+
+        spawned_category_tasks = []
+        for display_name, modules in category_modules:
+            tasks = []
+            for module in modules:
+                t = asyncio.create_task(
+                    _async_worker(
+                        module,
+                        email,
+                        sem,
+                        configs,
+                        printed_cats=printed_cats,
+                        on_start=on_start_cb,
+                    )
+                )
+                t.add_done_callback(lambda t: progress.advance(task_id))
+                tasks.append(t)
+            spawned_category_tasks.append((display_name, tasks))
+
+        for display_name, tasks in spawned_category_tasks:
             if not tasks:
                 continue
 
@@ -244,7 +274,6 @@ async def _run_email_full_batch_async(email: str, configs: ScanConfig) -> List[R
 
                 result.show(configs)
                 all_results.append(result)
-                progress.advance(task_id)
 
     return all_results
 

user_scanner/core/impersonate.py

@@ -11,6 +11,7 @@
 DEFAULT_TIMEOUT = 15.0
 
 _sessions: dict[tuple, cffi.Session] = {}
+_key_locks: dict[tuple, threading.Lock] = {}
 _warmed: set[tuple] = set()
 _lock = threading.Lock()
 
@@ -93,12 +94,20 @@ def _get_warm_session(
                 proxies={"http": proxy, "https": proxy} if proxy else None,
             )
             _sessions[key] = session
-
-        if warmup_url and key not in _warmed:
-            # A blocked (403) warm-up still returns normally and sets the cookie;
-            # only a network error leaves the session unwarmed for a later retry.
-            session.get(warmup_url, timeout=_timeout())
-            _warmed.add(key)
+            _key_locks[key] = threading.Lock()
+
+        key_lock = _key_locks.get(key)
+        if key_lock is None:
+            key_lock = threading.Lock()
+            _key_locks[key] = key_lock
+
+    if warmup_url and key not in _warmed:
+        with key_lock:
+            if key not in _warmed:
+                # A blocked (403) warm-up still returns normally and sets the cookie;
+                # only a network error leaves the session unwarmed for a later retry.
+                session.get(warmup_url, timeout=_timeout())
+                _warmed.add(key)
 
     return session
 

user_scanner/core/orchestrator.py

@@ -25,23 +25,26 @@
 
 
 MAX_CONCURRENT_REQUESTS = 60
-_shared_executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_CONCURRENT_REQUESTS)
+_shared_executor = concurrent.futures.ThreadPoolExecutor(max_workers=max(MAX_CONCURRENT_REQUESTS * 2, 250))
 
 def set_concurrency(val: int):
     global MAX_CONCURRENT_REQUESTS, _shared_executor
     MAX_CONCURRENT_REQUESTS = val
-    _shared_executor = concurrent.futures.ThreadPoolExecutor(max_workers=val)
+    _shared_executor = concurrent.futures.ThreadPoolExecutor(max_workers=max(val * 2, 250))
 
 async def _async_worker(
     module: ModuleType,
     username: str,
     sem: asyncio.Semaphore,
     configs: ScanConfig,
     printed_cats: Optional[Set] = None,
-    cat_override: Optional[str] = None
+    cat_override: Optional[str] = None,
+    on_start: Optional[Callable[[str], None]] = None
 ) -> Result:
     async with sem:
         site_name = get_site_name(module)
+        if on_start:
+            on_start(site_name)
         func = get_scan_func(module)
         actual_cat = cat_override or find_category(module) or "Unknown"
 
@@ -58,11 +61,17 @@ async def _async_worker(
             return Result.skipped().update(**params)
 
         try:
+            module_timeout = (get_global_timeout() or 15.0) + 10.0
             if inspect.iscoroutinefunction(func):
-                result = await func(username)
+                result = await asyncio.wait_for(func(username), timeout=module_timeout)
             else:
                 loop = asyncio.get_running_loop()
-                result = await loop.run_in_executor(_shared_executor, func, username)
+                result = await asyncio.wait_for(
+                    loop.run_in_executor(_shared_executor, func, username),
+                    timeout=module_timeout
+                )
+        except asyncio.TimeoutError:
+            result = Result.error(f"Module execution timed out after {module_timeout}s")
         except Exception as e:
             result = Result.error(e)
 
@@ -80,14 +89,6 @@ async def _run_batch(
     if sem is None:
         sem = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
 
-    tasks = []
-    for module in modules:
-        tasks.append(
-            asyncio.create_task(
-                _async_worker(module, username, sem, configs, cat_override=cat_override)
-            )
-        )
-
     results = []
 
     with Progress(
@@ -98,7 +99,18 @@ async def _run_batch(
         TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
         transient=True,
     ) as progress:
-        task_id = progress.add_task(f"[cyan]Scanning {username}...", total=len(tasks))
+        task_id = progress.add_task(f"[cyan]Scanning {username}...", total=len(modules))
+
+        def on_start_cb(site: str):
+            progress.update(task_id, description=f"[cyan]Scanning {username}... ({site})")
+
+        tasks = []
+        for module in modules:
+            t = asyncio.create_task(
+                _async_worker(module, username, sem, configs, cat_override=cat_override, on_start=on_start_cb)
+            )
+            t.add_done_callback(lambda t: progress.advance(task_id))
+            tasks.append(t)
 
         for coro in asyncio.as_completed(tasks):
             result = await coro
@@ -112,7 +124,6 @@ async def _run_batch(
 
             result.show(configs)
             results.append(result)
-            progress.advance(task_id)
 
     return results
 
@@ -121,7 +132,7 @@ def run_user_module(
     module: Union[ModuleType, List[ModuleType]], username: str, configs: ScanConfig
 ) -> List[Result]:
     modules = [module] if isinstance(module, ModuleType) else list(module)
-    return asyncio.run(_run_batch(modules, username, configs))
+    return asyncio.run(_run_batch(modules, username, configs, printed_cats=set()))
 
 
 def run_user_category(
@@ -153,20 +164,13 @@ async def _run_user_full_async(username: str, configs: ScanConfig) -> List[Resul
     sem = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
 
     # 1. Pre-spawn all tasks for all categories (global concurrency)
-    category_tasks = []
+    category_modules = []
     total_tasks = 0
     for cat_name, cat_path in categories:
         display_name = cat_name.capitalize()
         modules = load_modules(cat_path)
-        tasks = []
-        for module in modules:
-            tasks.append(
-                asyncio.create_task(
-                    _async_worker(module, username, sem, configs, cat_override=display_name)
-                )
-            )
-        category_tasks.append((display_name, tasks))
-        total_tasks += len(tasks)
+        category_modules.append((display_name, modules))
+        total_tasks += len(modules)
 
     # 2. Await tasks category by category to stream grouped output
     with Progress(
@@ -179,7 +183,21 @@ async def _run_user_full_async(username: str, configs: ScanConfig) -> List[Resul
     ) as progress:
         task_id = progress.add_task(f"[cyan]Scanning {username}...", total=total_tasks)
 
-        for display_name, tasks in category_tasks:
+        def on_start_cb(site: str):
+            progress.update(task_id, description=f"[cyan]Scanning {username}... ({site})")
+
+        spawned_category_tasks = []
+        for display_name, modules in category_modules:
+            tasks = []
+            for module in modules:
+                t = asyncio.create_task(
+                    _async_worker(module, username, sem, configs, cat_override=display_name, on_start=on_start_cb)
+                )
+                t.add_done_callback(lambda t: progress.advance(task_id))
+                tasks.append(t)
+            spawned_category_tasks.append((display_name, tasks))
+
+        for display_name, tasks in spawned_category_tasks:
             if not tasks:
                 continue
 
@@ -198,7 +216,6 @@ async def _run_user_full_async(username: str, configs: ScanConfig) -> List[Resul
 
                 result.show(configs)
                 all_results.append(result)
-                progress.advance(task_id)
 
     return all_results