bookmarkmgr: cronet: Recycle executor and threads

Author: lubo

bookmarkmgr/bookmarkmgr/__main__.py

@@ -3,6 +3,7 @@
 import argparse
 import asyncio
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
 from getpass import getpass
 import logging
 import signal
@@ -48,6 +49,13 @@ def parse(value: str) -> float | int | str:
 
 
 async def run_command(args: argparse.Namespace, raindrop_api_key: str) -> None:
+    # Cronet may use a lot of threads.
+    asyncio.get_running_loop().set_default_executor(
+        ThreadPoolExecutor(
+            sys.maxsize,
+        ),
+    )
+
     async with RaindropClient(raindrop_api_key) as raindrop_client:
         match args.command:
             case "export-collection":

bookmarkmgr/bookmarkmgr/cronet/managers/executor.py

@@ -1,11 +1,7 @@
 import asyncio
-from concurrent.futures import ThreadPoolExecutor
-from queue import Queue
-from typing import Any, Self, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from collections.abc import Awaitable
+from typing import Any, Self
 
+from bookmarkmgr.asyncio import ForgivingTaskGroup
 from bookmarkmgr.cronet._cronet import ffi, lib
 from bookmarkmgr.cronet.errors import NotContextManagerError
 from bookmarkmgr.cronet.types import Executor, Runnable
@@ -16,15 +12,21 @@ def _executor_execute(executor: Executor, runnable: Runnable) -> None:
     manager: ExecutorManager = ffi.from_handle(
         lib.Cronet_Executor_GetClientContext(executor),
     )
-    manager.enqueue_runnable(runnable)
+    manager.submit_runnable(runnable)
+
+
+def _process_runnable(runnable: Runnable) -> None:
+    try:
+        lib.Cronet_Runnable_Run(runnable)
+    finally:
+        lib.Cronet_Runnable_Destroy(runnable)
 
 
 class ExecutorManager:
     def __init__(self) -> None:
-        self._handle = ffi.new_handle(self)
-        self._queue: Queue[Runnable | None] = Queue()
-        self._worker: Awaitable[None] | None = None
         self._executor: Executor | None = None
+        self._handle = ffi.new_handle(self)
+        self._task_group: ForgivingTaskGroup | None = None
 
     async def __aenter__(self) -> Self:
         if self._executor is None:
@@ -33,10 +35,11 @@ async def __aenter__(self) -> Self:
             )
             lib.Cronet_Executor_SetClientContext(self._executor, self._handle)
 
-        self._processing_allowed = True
+        if self._task_group is None:
+            self._task_group = ForgivingTaskGroup()
+            await self._task_group.__aenter__()
 
-        if self._worker is None:
-            self._worker = asyncio.create_task(self._spawn_worker_thread())
+        self._processing_allowed = True
 
         return self
 
@@ -46,43 +49,36 @@ async def __aexit__(
         *args: Any,  # noqa: PYI036
         **kwargs: Any,
     ) -> None:
-        self.shutdown(process_pending=exc_type is None)
+        self._processing_allowed = False
 
         try:
-            if self._worker is not None:
-                await self._worker
+            if self._task_group is not None:
+                await self._task_group.__aexit__(exc_type, *args, **kwargs)
+                self._task_group = None
         finally:
             if self._executor is not None:
                 lib.Cronet_Executor_Destroy(self._executor)
                 self._executor = None
 
-            self._worker = None
+    def submit_runnable(self, runnable: Runnable) -> None:
+        if self._task_group is None:
+            message = "ExecutorManager has not been entered"
+            raise RuntimeError(message)
 
-    async def _spawn_worker_thread(self) -> None:
-        with ThreadPoolExecutor() as pool:
-            await asyncio.get_running_loop().run_in_executor(
-                pool,
-                self._worker_loop,
-            )
+        if not self._processing_allowed:
+            lib.Cronet_Runnable_Destroy(runnable)
+            return
 
-    def _worker_loop(self) -> None:
-        while (runnable := self._queue.get()) is not None:
-            try:
-                if self._processing_allowed:
-                    lib.Cronet_Runnable_Run(runnable)
-            finally:
-                lib.Cronet_Runnable_Destroy(runnable)
-
-    def enqueue_runnable(self, runnable: Runnable) -> None:
-        self._queue.put_nowait(runnable)
+        self._task_group.create_task(
+            asyncio.to_thread(
+                _process_runnable,
+                runnable,
+            ),
+        )
 
     @property
     def executor(self) -> Executor:
         if self._executor is None:
             raise NotContextManagerError
 
         return self._executor
-
-    def shutdown(self, *, process_pending: bool = True) -> None:
-        self._processing_allowed = process_pending
-        self._queue.put_nowait(None)

bookmarkmgr/bookmarkmgr/cronet/session.py

@@ -49,8 +49,10 @@
 class Session:
     def __init__(self) -> None:
         self._engine: Engine | None = None
+        self._executor_manager = ExecutorManager()
 
     async def __aenter__(self) -> Self:
+        await self._executor_manager.__aenter__()
         self._open()
 
         return self
@@ -60,7 +62,10 @@ async def __aexit__(
         *args: Any,  # noqa: PYI036
         **kwargs: Any,
     ) -> None:
-        self.close()
+        try:
+            await self._executor_manager.__aexit__(*args, **kwargs)
+        finally:
+            self._close()
 
     def _dispose_engine(self) -> None:
         if self._engine is None:
@@ -87,7 +92,7 @@ def _open(self) -> None:
             self._dispose_engine()
             raise
 
-    def close(self) -> None:
+    def _close(self) -> None:
         if self._engine is None:
             return
 
@@ -144,7 +149,6 @@ async def request(
                     **kwargs,
                 ),
             ) as callback_manager,
-            ExecutorManager() as executor_manager,
         ):
             lib.Cronet_UrlRequestParams_http_method_set(
                 parameters,
@@ -170,7 +174,7 @@ async def request(
                     url.encode(),
                     parameters,
                     callback_manager.callback,
-                    executor_manager.executor,
+                    self._executor_manager.executor,
                 ),
             )
 
@@ -317,8 +321,8 @@ async def _request(self, *args: Any, **kwargs: Any) -> Response:
         async with self._RateLimiterMixin_rate_limiter:
             return await super()._request(*args, **kwargs)
 
-    def close(self) -> None:
-        super().close()
+    def _close(self) -> None:
+        super()._close()
 
         self._RateLimiterMixin_rate_limiter.close()
 
@@ -357,8 +361,8 @@ async def _request(
         ):
             return await super()._request(method, url, *args, **kwargs)
 
-    def close(self) -> None:
-        super().close()
+    def _close(self) -> None:
+        super()._close()
 
         for rate_limiter in self.__rate_limiters.values():
             rate_limiter.close()