add async thread helper class

vdusek · vdusek · commit 1321f82d8ee1 · 2025-02-10T14:52:16.000+01:00
diff --git a/src/apify/scrapy/_async_thread.py b/src/apify/scrapy/_async_thread.py
@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+from concurrent import futures
+from datetime import timedelta
+from logging import getLogger
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Coroutine
+
+logger = getLogger(__name__)
+
+
+class AsyncThread:
+    """Class for running an asyncio event loop in a separate thread.
+
+    This allows running asynchronous coroutines from synchronous code by executingthem on an event loop
+    that runs in its own dedicated thread.
+    """
+
+    def __init__(self) -> None:
+        self._eventloop = asyncio.new_event_loop()
+
+        # Start the event loop in a dedicated daemon thread.
+        self._thread = threading.Thread(
+            target=self._start_event_loop,
+            daemon=True,
+        )
+        self._thread.start()
+
+    def run_coro(
+        self,
+        coro: Coroutine,
+        timeout: timedelta = timedelta(seconds=60),
+    ) -> Any:
+        """Run a coroutine on an event loop running in a separate thread.
+
+        This method schedules the coroutine to run on the event loop and blocks until the coroutine completes
+        or the specified timeout is reached.
+
+        Args:
+            coro: The coroutine to run.
+            timeout: The maximum number of seconds to wait for the coroutine to finish.
+
+        Returns:
+            The result returned by the coroutine.
+
+        Raises:
+            TimeoutError: If the coroutine does not complete within the timeout.
+            Exception: Any exception raised during coroutine execution.
+        """
+        if not self._eventloop.is_running():
+            logger.warning('Event loop is not running! Ignoring coroutine execution.')
+            return None
+
+        # Submit the coroutine to the event loop running in the other thread.
+        future = asyncio.run_coroutine_threadsafe(coro, self._eventloop)
+        try:
+            # Wait for the coroutine's result until the specified timeout.
+            return future.result(timeout=timeout.total_seconds())
+        except futures.TimeoutError as exc:
+            logger.exception('Coroutine execution timed out.', exc_info=exc)
+            raise
+        except Exception as exc:
+            logger.exception('Coroutine execution raised an exception.', exc_info=exc)
+            raise
+
+    def close(self, timeout: timedelta = timedelta(seconds=60)) -> None:
+        """Close the event loop and its thread gracefully.
+
+        This method cancels all pending tasks, stops the event loop, and waits for the thread to exit.
+        If the thread does not exit within the given timeout, a forced shutdown is attempted.
+
+        Args:
+            timeout: The maximum number of seconds to wait for the event loop thread to exit.
+        """
+        if self._eventloop.is_running():
+            # Cancel all pending tasks in the event loop.
+            self.run_coro(self._shutdown_tasks())
+
+        # Schedule the event loop to stop.
+        self._eventloop.call_soon_threadsafe(self._eventloop.stop)
+
+        # Wait for the event loop thread to finish execution.
+        self._thread.join(timeout=timeout.total_seconds())
+
+        # If the thread is still running after the timeout, force a shutdown.
+        if self._thread.is_alive():
+            logger.warning('Event loop thread did not exit cleanly! Forcing shutdown...')
+            self._force_exit_event_loop()
+
+    def _start_event_loop(self) -> None:
+        """Set up and run the asyncio event loop in the dedicated thread."""
+        asyncio.set_event_loop(self._eventloop)
+        try:
+            self._eventloop.run_forever()
+        finally:
+            self._eventloop.close()
+            logger.debug('Asyncio event loop has been closed.')
+
+    async def _shutdown_tasks(self) -> None:
+        """Cancel all pending tasks in the event loop."""
+        # Retrieve all tasks for the event loop, excluding the current task.
+        tasks = [task for task in asyncio.all_tasks(self._eventloop) if task is not asyncio.current_task()]
+        if not tasks:
+            return
+
+        # Cancel each pending task.
+        for task in tasks:
+            task.cancel()
+
+        # Wait until all tasks have been cancelled or finished.
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+    def _force_exit_event_loop(self) -> None:
+        """Forcefully shut down the event loop and its thread."""
+        try:
+            logger.info('Forced shutdown of the event loop and its thread...')
+            self._eventloop.call_soon_threadsafe(self._eventloop.stop)
+            self._thread.join(timeout=5)
+        except Exception as exc:
+            logger.exception('Exception occurred during forced event loop shutdown.', exc_info=exc)
diff --git a/src/apify/scrapy/scheduler.py b/src/apify/scrapy/scheduler.py
@@ -1,100 +1,25 @@
 from __future__ import annotations
 
-import asyncio
-import threading
 import traceback
-from concurrent import futures
 from logging import getLogger
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from scrapy import Spider
 from scrapy.core.scheduler import BaseScheduler
 from scrapy.utils.reactor import is_asyncio_reactor_installed
 
+from ._async_thread import AsyncThread
+from .requests import to_apify_request, to_scrapy_request
 from apify import Configuration
 from apify.apify_storage_client import ApifyStorageClient
-from apify.scrapy.requests import to_apify_request, to_scrapy_request
 from apify.storages import RequestQueue
 
 if TYPE_CHECKING:
-    from collections.abc import Coroutine
-
     from scrapy.http.request import Request
     from twisted.internet.defer import Deferred
 
 logger = getLogger(__name__)
 
-_TIMEOUT = 60
-"""The timeout for waiting on asyncio coroutines to finish."""
-
-
-def _start_event_loop(eventloop: asyncio.AbstractEventLoop) -> None:
-    """Set and run the event loop until it is stopped.
-
-    Args:
-        eventloop: The asyncio event loop to run.
-    """
-    asyncio.set_event_loop(eventloop)
-    try:
-        eventloop.run_forever()
-    finally:
-        eventloop.close()
-        logger.debug('Asyncio event loop has been closed.')
-
-
-def _run_async_coro(eventloop: asyncio.AbstractEventLoop, coro: Coroutine) -> Any:
-    """Run a coroutine on the given loop in our separate thread, waiting for its result.
-
-    Args:
-        eventloop: The asyncio event loop to run the coroutine on.
-        coro: The coroutine to run.
-
-    Returns:
-        The result of the coroutine.
-    """
-    if not eventloop.is_running():
-        logger.warning('Event loop is not running! Ignoring coroutine execution.')
-        return None
-
-    future = asyncio.run_coroutine_threadsafe(coro, eventloop)
-    try:
-        return future.result(timeout=_TIMEOUT)
-    except futures.TimeoutError as exc:
-        logger.exception('Coroutine execution timed out.', exc_info=exc)
-        raise
-    except Exception as exc:
-        logger.exception('Coroutine execution raised an exception.', exc_info=exc)
-        raise
-
-
-async def _shutdown_async_tasks(eventloop: asyncio.AbstractEventLoop) -> None:
-    """Cancel and wait for all pending tasks on the current event loop.
-
-    Args:
-        eventloop: The asyncio event loop to cancel tasks on.
-    """
-    tasks = [task for task in asyncio.all_tasks(eventloop) if task is not asyncio.current_task()]
-    if not tasks:
-        return
-    for task in tasks:
-        task.cancel()
-    await asyncio.gather(*tasks, return_exceptions=True)
-
-
-def _force_exit_event_loop(eventloop: asyncio.AbstractEventLoop, thread: threading.Thread) -> None:
-    """Forcefully shut down the event loop and its thread.
-
-    Args:
-        eventloop: The asyncio event loop to stop.
-        thread: The thread running the event loop.
-    """
-    try:
-        logger.info('Forced shutdown of the event loop and its thread...')
-        eventloop.call_soon_threadsafe(eventloop.stop)
-        thread.join(timeout=5)
-    except Exception as exc:
-        logger.exception('Exception occurred during forced event loop shutdown.', exc_info=exc)
-
 
 class ApifyScheduler(BaseScheduler):
     """A Scrapy scheduler that uses the Apify request queue to manage requests.
@@ -112,10 +37,8 @@ def __init__(self) -> None:
         self._rq: RequestQueue | None = None
         self.spider: Spider | None = None
 
-        # Create a new event loop and run it in a separate thread.
-        self._eventloop = asyncio.new_event_loop()
-        self._thread = threading.Thread(target=lambda: _start_event_loop(self._eventloop), daemon=True)
-        self._thread.start()
+        # A thread with the asyncio event loop to run coroutines on.
+        self._async_thread = AsyncThread()
 
     def open(self, spider: Spider) -> Deferred[None] | None:
         """Open the scheduler.
@@ -133,7 +56,7 @@ async def open_rq() -> RequestQueue:
             return await RequestQueue.open()
 
         try:
-            self._rq = _run_async_coro(self._eventloop, open_rq())
+            self._rq = self._async_thread.run_coro(open_rq())
         except Exception:
             traceback.print_exc()
             raise
@@ -150,20 +73,7 @@ def close(self, reason: str) -> None:
         """
         logger.debug(f'Closing {self.__class__.__name__} due to {reason}...')
         try:
-            if self._eventloop.is_running():
-                # Cancel all pending tasks in the event loop.
-                _run_async_coro(self._eventloop, _shutdown_async_tasks(self._eventloop))
-
-            # Stop the event loop.
-            self._eventloop.call_soon_threadsafe(self._eventloop.stop)
-
-            # Wait for the event loop thread to exit.
-            self._thread.join(timeout=_TIMEOUT)
-
-            # If the thread is still alive, execute a forced shutdown.
-            if self._thread.is_alive():
-                logger.warning('Event loop thread did not exit cleanly! Forcing shutdown...')
-                _force_exit_event_loop(self._eventloop, self._thread)
+            self._async_thread.close()
 
         except KeyboardInterrupt:
             logger.warning('Shutdown interrupted by KeyboardInterrupt!')
@@ -184,7 +94,7 @@ def has_pending_requests(self) -> bool:
             raise TypeError('self._rq must be an instance of the RequestQueue class')
 
         try:
-            is_finished = _run_async_coro(self._eventloop, self._rq.is_finished())
+            is_finished = self._async_thread.run_coro(self._rq.is_finished())
         except Exception:
             traceback.print_exc()
             raise
@@ -217,7 +127,7 @@ def enqueue_request(self, request: Request) -> bool:
             raise TypeError('self._rq must be an instance of the RequestQueue class')
 
         try:
-            result = _run_async_coro(self._eventloop, self._rq.add_request(apify_request))
+            result = self._async_thread.run_coro(self._rq.add_request(apify_request))
         except Exception:
             traceback.print_exc()
             raise
@@ -236,7 +146,7 @@ def next_request(self) -> Request | None:
             raise TypeError('self._rq must be an instance of the RequestQueue class')
 
         try:
-            apify_request = _run_async_coro(self._eventloop, self._rq.fetch_next_request())
+            apify_request = self._async_thread.run_coro(self._rq.fetch_next_request())
         except Exception:
             traceback.print_exc()
             raise
@@ -248,10 +158,10 @@ def next_request(self) -> Request | None:
         if not isinstance(self.spider, Spider):
             raise TypeError('self.spider must be an instance of the Spider class')
 
-        # Let the request queue know that the request is being handled. Every request should be marked as handled,
-        # retrying is handled by the Scrapy's RetryMiddleware.
+        # Let the request queue know that the request is being handled. Every request should
+        # be marked as handled, retrying is handled by the Scrapy's RetryMiddleware.
         try:
-            _run_async_coro(self._eventloop, self._rq.mark_request_as_handled(apify_request))
+            self._async_thread.run_coro(self._rq.mark_request_as_handled(apify_request))
         except Exception:
             traceback.print_exc()
             raise
