Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 26 additions & 11 deletions common/metric_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from common.metric_config import MetricConfig, MetricLabelKey, MetricLabels
from common.metrics_handler import MetricsHandler

MAX_RETRIES = 3
MAX_RETRIES = 2


class WebSocketMetric(BaseMetric):
Expand Down Expand Up @@ -59,18 +59,29 @@ async def collect_metric(self) -> None:
"""Collects single WebSocket message."""
websocket = None

try:
async def _collect_ws_data():
nonlocal websocket
websocket = await self.connect()
await self.subscribe(websocket)
data = await self.listen_for_data(websocket)

if data is not None:
latency: int | float = self.process_data(data)
self.update_metric_value(latency)
self.mark_success()
return
return data
raise ValueError("No data in response")

try:
data = await asyncio.wait_for(
_collect_ws_data(),
timeout=self.config.timeout
)
latency: int | float = self.process_data(data)
self.update_metric_value(latency)
self.mark_success()

except asyncio.TimeoutError:
self.mark_failure()
self.handle_error(TimeoutError(f"WebSocket metric collection exceeded {self.config.timeout}s timeout"))

except Exception as e:
self.mark_failure()
self.handle_error(e)
Expand All @@ -97,13 +108,19 @@ def get_endpoint(self) -> str:

async def collect_metric(self) -> None:
try:
data = await self.fetch_data()
data = await asyncio.wait_for(
self.fetch_data(),
timeout=self.config.timeout
)
if data is not None:
latency: int | float = self.process_data(data)
self.update_metric_value(latency)
self.mark_success()
return
raise ValueError("No data in response")
except asyncio.TimeoutError:
self.mark_failure()
self.handle_error(TimeoutError(f"Metric collection exceeded {self.config.timeout}s timeout"))
except Exception as e:
self.mark_failure()
self.handle_error(e)
Expand Down Expand Up @@ -205,7 +222,6 @@ async def on_request_end(session, context, params):
trace_config.on_request_end.append(on_request_end)

async with aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=self.config.timeout),
trace_configs=[trace_config],
) as session:
response_time = 0.0
Expand All @@ -219,7 +235,7 @@ async def on_request_end(session, context, params):
response_time: float = time.monotonic() - start_time

if response.status == 429 and retry_count < MAX_RETRIES - 1:
wait_time = int(response.headers.get("Retry-After", 15))
wait_time = int(response.headers.get("Retry-After", 3))
await response.release()
await asyncio.sleep(wait_time)
continue
Expand Down Expand Up @@ -279,7 +295,6 @@ async def _send_request(
"Content-Type": "application/json",
},
json=self._base_request,
timeout=self.config.timeout, # type: ignore
)

def process_data(self, value: float) -> float:
Expand Down