Track shared timers manually

Author: janbuchar

src/crawlee/crawlers/_abstract_http/_abstract_http_crawler.py

@@ -5,14 +5,12 @@
 from abc import ABC
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Generic
-from weakref import WeakKeyDictionary
 
 from more_itertools import partition
 from pydantic import ValidationError
 from typing_extensions import NotRequired, TypeVar
 
 from crawlee._request import Request, RequestOptions
-from crawlee._types import BasicCrawlingContext
 from crawlee._utils.docs import docs_group
 from crawlee._utils.time import SharedTimeout
 from crawlee._utils.urls import to_absolute_url_iterator
@@ -28,7 +26,7 @@
     from typing_extensions import Unpack
 
     from crawlee import RequestTransformAction
-    from crawlee._types import EnqueueLinksKwargs, ExtractLinksFunction
+    from crawlee._types import BasicCrawlingContext, EnqueueLinksKwargs, ExtractLinksFunction
 
     from ._abstract_http_parser import AbstractHttpParser
 
@@ -79,7 +77,7 @@ def __init__(
         self._parser = parser
         self._navigation_timeout = navigation_timeout or timedelta(minutes=1)
         self._pre_navigation_hooks: list[Callable[[BasicCrawlingContext], Awaitable[None]]] = []
-        self._shared_navigation_timeouts = WeakKeyDictionary[BasicCrawlingContext, SharedTimeout]()
+        self._shared_navigation_timeouts: dict[int, SharedTimeout] = {}
 
         if '_context_pipeline' not in kwargs:
             raise ValueError(
@@ -132,13 +130,17 @@ def _create_static_content_crawler_pipeline(self) -> ContextPipeline[ParsedHttpC
     async def _execute_pre_navigation_hooks(
         self, context: BasicCrawlingContext
     ) -> AsyncGenerator[BasicCrawlingContext, None]:
-        self._shared_navigation_timeouts[context] = SharedTimeout(self._navigation_timeout)
+        context_id = id(context)
+        self._shared_navigation_timeouts[context_id] = SharedTimeout(self._navigation_timeout)
 
-        for hook in self._pre_navigation_hooks:
-            async with self._shared_navigation_timeouts[context]:
-                await hook(context)
+        try:
+            for hook in self._pre_navigation_hooks:
+                async with self._shared_navigation_timeouts[context_id]:
+                    await hook(context)
 
-        yield context
+            yield context
+        finally:
+            self._shared_navigation_timeouts.pop(context_id, None)
 
     async def _parse_http_response(
         self, context: HttpCrawlingContext
@@ -240,7 +242,7 @@ async def _make_http_request(self, context: BasicCrawlingContext) -> AsyncGenera
         Yields:
             The original crawling context enhanced by HTTP response.
         """
-        async with self._shared_navigation_timeouts[context] as remaining_timeout:
+        async with self._shared_navigation_timeouts[id(context)] as remaining_timeout:
             result = await self._http_client.crawl(
                 request=context.request,
                 session=context.session,

src/crawlee/crawlers/_playwright/_playwright_crawler.py

@@ -6,7 +6,6 @@
 from datetime import timedelta
 from functools import partial
 from typing import TYPE_CHECKING, Any, Generic, Literal
-from weakref import WeakKeyDictionary
 
 import playwright.async_api
 from more_itertools import partition
@@ -149,7 +148,7 @@ def __init__(
         if configuration is not None:
             service_locator.set_configuration(configuration)
 
-        self._shared_navigation_timeouts = WeakKeyDictionary[BasicCrawlingContext, SharedTimeout]()
+        self._shared_navigation_timeouts: dict[int, SharedTimeout] = {}
 
         if browser_pool:
             # Raise an exception if browser_pool is provided together with other browser-related arguments.
@@ -241,14 +240,18 @@ async def _open_page(
             block_requests=partial(block_requests, page=crawlee_page.page),
         )
 
-        self._shared_navigation_timeouts[pre_navigation_context] = SharedTimeout(self._navigation_timeout)
+        context_id = id(pre_navigation_context)
+        self._shared_navigation_timeouts[context_id] = SharedTimeout(self._navigation_timeout)
 
-        async with browser_page_context(crawlee_page.page):
-            for hook in self._pre_navigation_hooks:
-                async with self._shared_navigation_timeouts[context]:
-                    await hook(pre_navigation_context)
+        try:
+            async with browser_page_context(crawlee_page.page):
+                for hook in self._pre_navigation_hooks:
+                    async with self._shared_navigation_timeouts[context_id]:
+                        await hook(pre_navigation_context)
 
-        yield pre_navigation_context
+            yield pre_navigation_context
+        finally:
+            self._shared_navigation_timeouts.pop(context_id, None)
 
     def _prepare_request_interceptor(
         self,
@@ -316,7 +319,7 @@ async def _navigate(
                 await context.page.route(context.request.url, route_handler)
 
             try:
-                async with self._shared_navigation_timeouts[context] as remaining_timeout:
+                async with self._shared_navigation_timeouts[id(context)] as remaining_timeout:
                     response = await context.page.goto(
                         context.request.url, timeout=remaining_timeout.total_seconds() * 1000
                     )