fix: http timing module

Author: smypmsa

common/http_timing.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import time
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import aiohttp
 
@@ -12,34 +12,26 @@
 class HttpTimingCollector:
     """Utility class for measuring HTTP request timing with detailed breakdown."""
 
-    def __init__(self):
-        self.timing: Dict[str, float] = {}
+    def __init__(self) -> None:
+        self.timing: dict[str, float] = {}
 
     def create_trace_config(self) -> aiohttp.TraceConfig:
         """Create aiohttp trace configuration for detailed timing measurement."""
         trace_config = aiohttp.TraceConfig()
 
-        async def on_request_start(session, context, params):
+        async def on_request_start(_session: Any, _context: Any, _params: Any) -> None:
             self.timing["start"] = time.monotonic()
 
-        async def on_dns_resolvehost_start(session, context, params):
-            self.timing["dns_start"] = time.monotonic()
-
-        async def on_dns_resolvehost_end(session, context, params):
-            self.timing["dns_end"] = time.monotonic()
-
-        async def on_connection_create_start(session, context, params):
+        async def on_connection_create_start(_session: Any, _context: Any, _params: Any) -> None:
             self.timing["conn_start"] = time.monotonic()
 
-        async def on_connection_create_end(session, context, params):
+        async def on_connection_create_end(_session: Any, _context: Any, _params: Any) -> None:
             self.timing["conn_end"] = time.monotonic()
 
-        async def on_request_end(session, context, params):
+        async def on_request_end(_session: Any, _context: Any, _params: Any) -> None:
             self.timing["end"] = time.monotonic()
 
         trace_config.on_request_start.append(on_request_start)
-        trace_config.on_dns_resolvehost_start.append(on_dns_resolvehost_start)
-        trace_config.on_dns_resolvehost_end.append(on_dns_resolvehost_end)
         trace_config.on_connection_create_start.append(on_connection_create_start)
         trace_config.on_connection_create_end.append(on_connection_create_end)
         trace_config.on_request_end.append(on_request_end)
@@ -52,52 +44,75 @@ def get_connection_time(self) -> float:
             return self.timing["conn_end"] - self.timing["conn_start"]
         return 0.0
 
-    def get_dns_time(self) -> float:
-        """Get DNS resolution time in seconds."""
-        if "dns_start" in self.timing and "dns_end" in self.timing:
-            return self.timing["dns_end"] - self.timing["dns_start"]
+    def get_total_time(self) -> float:
+        """Get total request time in seconds."""
+        if "start" in self.timing and "end" in self.timing:
+            return self.timing["end"] - self.timing["start"]
         return 0.0
 
 
 async def measure_http_request_timing(
     session: aiohttp.ClientSession,
     method: str,
     url: str,
-    headers: Optional[Dict[str, str]] = None,
-    json_data: Optional[Dict[str, Any]] = None,
+    headers: Optional[dict[str, str]] = None,
+    json_data: Optional[dict[str, Any]] = None,
     exclude_connection_time: bool = True,
 ) -> tuple[float, aiohttp.ClientResponse]:
     """Measure HTTP request timing with retry logic and detailed breakdown.
 
+    Args:
+        exclude_connection_time: If True, exclude connection establishment time
+                               from the returned timing to measure pure API response time
+
     Returns:
         tuple: (response_time_seconds, response)
     """
-    response_time = 0.0
-    response = None
+    timing_collector = HttpTimingCollector()
+    trace_config: aiohttp.TraceConfig = timing_collector.create_trace_config()
 
-    for retry_count in range(MAX_RETRIES):
-        start_time = time.monotonic()
+    # Add timing trace to session temporarily
+    session._trace_configs.append(trace_config)
 
-        # Send request
-        if method.upper() == "POST":
-            response = await session.post(
-                url, headers=headers, json=json_data
-            )
-        else:
-            response = await session.get(url, headers=headers)
+    try:
+        response = None
+
+        for retry_count in range(MAX_RETRIES):
+            # Reset timing for each retry
+            timing_collector.timing.clear()
 
-        response_time = time.monotonic() - start_time
+            # Send request
+            if method.upper() == "POST":
+                response = await session.post(
+                    url, headers=headers, json=json_data
+                )
+            else:
+                response = await session.get(url, headers=headers)
 
-        # Handle rate limiting
-        if response.status == 429 and retry_count < MAX_RETRIES - 1:
-            wait_time = int(response.headers.get("Retry-After", 3))
-            await response.release()
-            await asyncio.sleep(wait_time)
-            continue
+            # Handle rate limiting
+            if response.status == 429 and retry_count < MAX_RETRIES - 1:
+                wait_time = int(response.headers.get("Retry-After", 3))
+                await response.release()
+                await asyncio.sleep(wait_time)
+                continue
 
-        break
+            break
+
+        if not response:
+            raise ValueError("No response received")
+
+        # Calculate response time based on exclusion setting
+        total_time: float = timing_collector.get_total_time()
+
+        if exclude_connection_time:
+            connection_time: float = timing_collector.get_connection_time()
+            response_time: float = max(0.0, total_time - connection_time)
+        else:
+            response_time = total_time
 
-    if not response:
-        raise ValueError("No response received")
+        return response_time, response
 
-    return response_time, response
+    finally:
+        # Remove trace config to avoid affecting other requests
+        if trace_config in session._trace_configs:
+            session._trace_configs.remove(trace_config)