From 2ab46f1994b59f481f33c92766c0f49748f1a50d Mon Sep 17 00:00:00 2001
From: Stef <stbuwalda@gmail.com>
Date: Fri, 5 Sep 2025 18:05:50 +0200
Subject: [PATCH] Switched from manual tming to using trace request context.

---
 poll_services.py | 59 +++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 51 insertions(+), 8 deletions(-)

diff --git a/poll_services.py b/poll_services.py
index e4aae3c..525a320 100644
--- a/poll_services.py
+++ b/poll_services.py
@@ -5,21 +5,30 @@ import time
 from models import log, service
 from sqlalchemy.orm import sessionmaker
 from config import timeout as timeout_
+from typing import Optional
+from types import SimpleNamespace
 
 
-async def ping(client: aiohttp.ClientSession, s: service) -> int:
+async def ping(
+    client: aiohttp.ClientSession,
+    s: service,
+    ctx: Optional[SimpleNamespace] = None,
+) -> int:
+    ctx = ctx or SimpleNamespace()
     match s.ping_method:
         case 0:
             r = await client.head(
                 url=s.url,
                 ssl=True if s.public_access else False,
                 allow_redirects=True,
+                trace_request_ctx=ctx,  # type: ignore
             )
         case 1:
             r = await client.get(
                 url=s.url,
                 ssl=True if s.public_access else False,
                 allow_redirects=True,
+                trace_request_ctx=ctx,  # type: ignore
             )
         case _:
             raise Exception("UNKNOWN PING METHOD")
@@ -28,12 +37,14 @@ async def ping(client: aiohttp.ClientSession, s: service) -> int:
 
 async def check_service(client: aiohttp.ClientSession, s: service) -> log:
     try:
-        # TODO: Use aiohttp latency timing rather than timing it manually
-        before = time.perf_counter()
-        status = await ping(client=client, s=s)
-        after = time.perf_counter()
+        ctx = SimpleNamespace()
+        status = await ping(client=client, s=s, ctx=ctx)
+        print(status)
+        print(vars(ctx))
         if status == 200:
-            return log(service_id=s.id, ping=int((after - before) * 1000))
+            print("test")
+            print(ctx.duration_ms)
+            return log(service_id=s.id, ping=int(ctx.duration_ms))
         else:
             return log(service_id=s.id, ping=None)
     except aiohttp.ConnectionTimeoutError:
@@ -49,13 +60,45 @@ def start_async_loop():
     loop.run_forever()
 
 
+def setup_client() -> aiohttp.ClientSession:
+    timeout = aiohttp.client.ClientTimeout(total=timeout_ / 1000)
+    # Each request will get its own context
+    trace_config = aiohttp.TraceConfig()
+
+    async def on_start(
+        session: aiohttp.ClientSession,
+        context: SimpleNamespace,
+        params: aiohttp.TraceRequestStartParams,
+    ):
+        ctx = context.trace_request_ctx
+        ctx.start = time.perf_counter()  # store per-request
+
+    async def on_end(
+        session: aiohttp.ClientSession,
+        context: SimpleNamespace,
+        params: aiohttp.TraceRequestEndParams,
+    ):
+        ctx = context.trace_request_ctx
+        ctx.end = time.perf_counter()
+        ctx.duration_ms = int((ctx.end - ctx.start) * 1000)
+
+    trace_config.on_request_start.append(on_start)
+    trace_config.on_request_end.append(on_end)
+    client = aiohttp.ClientSession(
+        timeout=timeout, auto_decompress=False, trace_configs=[trace_config]
+    )
+    return client
+
+
 async def update_services(loop: asyncio.AbstractEventLoop):
     print("Starting service updates...")
     # Create new session
     with app.app_context():
         WorkerSession = sessionmaker(bind=db.engine)
-    timeout = aiohttp.client.ClientTimeout(total=timeout_ / 1000)
-    client = aiohttp.ClientSession(timeout=timeout, auto_decompress=False)
+
+    client = setup_client()
+
+    # Actual update loop
     while True:
         session = WorkerSession()
         sleeptask = asyncio.create_task(asyncio.sleep(timeout_ / 1000 + 1))