use timer to timeout functions

src/endpoints.py

@@ -15,14 +15,14 @@ from src.models import (
     LinkResponse,
     Solution,
 )
-from src.utils import CamoufoxDepType, get_camoufox, logger
+from src.utils import CamoufoxDepClass, TimeoutTimer, get_camoufox, logger
 
 warnings.filterwarnings("ignore", category=SyntaxWarning)
 
 
 router = APIRouter()
 
-CamoufoxDep = Annotated[CamoufoxDepType, Depends(get_camoufox)]
+CamoufoxDep = Annotated[CamoufoxDepClass, Depends(get_camoufox)]
 
 
 @router.get("/", include_in_schema=False)
@@ -54,20 +54,24 @@ async def read_item(request: LinkRequest, dep: CamoufoxDep) -> LinkResponse:
     """Handle POST requests."""
     start_time = int(time.time() * 1000)
 
-    request.url = request.url.replace('"', "").strip()
-    page_request = await dep.page.goto(request.url)
-    status = page_request.status if page_request else HTTPStatus.OK
+    timer = TimeoutTimer(duration=request.max_timeout)
 
-    if await dep.page.title() in CHALLENGE_TITLES:
-        logger.info("Challenge detected, attempting to solve...")
-        # Solve the captcha
-        remaining_timeout = (
-            request.max_timeout - (int(time.time()) * 1000 - start_time) / 1000
+    request.url = request.url.replace('"', "").strip()
+    try:
+        page_request = await dep.page.goto(
+            request.url, timeout=timer.remaining() * 1000
         )
-        logger.debug(
-            "Remaining timeout for solving the challenge: %d ms", remaining_timeout
+        status = page_request.status if page_request else HTTPStatus.OK
+        await dep.page.wait_for_load_state(
+            state="domcontentloaded", timeout=timer.remaining() * 1000
         )
-        try:
+        await dep.page.wait_for_load_state(
+            "networkidle", timeout=timer.remaining() * 1000
+        )
+
+        if await dep.page.title() in CHALLENGE_TITLES:
+            logger.info("Challenge detected, attempting to solve...")
+            # Solve the captcha
             await wait_for(
                 dep.solver.solve_captcha(  # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType]
                     captcha_container=dep.page,
@@ -75,16 +79,16 @@ async def read_item(request: LinkRequest, dep: CamoufoxDep) -> LinkResponse:
                     wait_checkbox_attempts=1,
                     wait_checkbox_delay=0.5,
                 ),
-                timeout=remaining_timeout,
+                timeout=timer.remaining(),
             )
             status = HTTPStatus.OK
-        except TimeoutError as e:
-            logger.error("Timed out while solving the challenge")
-            raise HTTPException(
-                status_code=408,
-                detail="Timed out while solving the challenge",
-            ) from e
-        logger.debug("Challenge solved successfully.")
+            logger.debug("Challenge solved successfully.")
+    except TimeoutError as e:
+        logger.error("Timed out while solving the challenge")
+        raise HTTPException(
+            status_code=408,
+            detail="Timed out while solving the challenge",
+        ) from e
 
     cookies = await dep.context.cookies()
 

src/middlewares.py

@@ -25,6 +25,6 @@ class LogRequest(BaseHTTPMiddleware):
         if response.status_code == HTTPStatus.OK:
             logger.info(f"Done {request_body.url} in {process_time:.2f}s")
         else:
-            logger.info(f"Failed {request_body.url} in {process_time:.2f}s")
+            logger.warning(f"Failed {request_body.url} in {process_time:.2f}s")
 
         return response

src/utils.py

@@ -1,4 +1,5 @@
 import logging
+import time
 from collections.abc import AsyncGenerator
 from typing import NamedTuple, cast
 
@@ -8,6 +9,7 @@ from playwright_captcha import (
     ClickSolver,
     FrameworkType,
 )
+from pydantic import BaseModel, Field
 
 from src.consts import (
     ADDON_PATH,
@@ -32,13 +34,22 @@ if len(logger.handlers) == 0:
     logger.addHandler(logging.StreamHandler())
 
 
-class CamoufoxDepType(NamedTuple):
+class TimeoutTimer(BaseModel):
+    duration: int  # in seconds
+    start_time: float = Field(default_factory=time.perf_counter)
+
+    def remaining(self) -> float:
+        """Get remaining time in seconds."""
+        return max(0, self.duration - (time.perf_counter() - self.start_time))
+
+
+class CamoufoxDepClass(NamedTuple):
     page: Page
     solver: ClickSolver
     context: BrowserContext
 
 
-async def get_camoufox() -> AsyncGenerator[CamoufoxDepType, None]:
+async def get_camoufox() -> AsyncGenerator[CamoufoxDepClass, None]:
     """Get Camoufox instance."""
     proxy_config = (
         {
@@ -72,4 +83,4 @@ async def get_camoufox() -> AsyncGenerator[CamoufoxDepType, None]:
             max_attempts=MAX_ATTEMPTS,
             attempt_delay=1,
         ) as solver:
-            yield CamoufoxDepType(page, solver, context)
+            yield CamoufoxDepClass(page, solver, context)