diff --git a/lib/backend.py b/lib/backend.py index e01f7c4..c06e096 100644 --- a/lib/backend.py +++ b/lib/backend.py @@ -5,7 +5,7 @@ import base64 import subprocess import dataclasses import logging -from asyncio import wait, sleep, gather, Semaphore, FIRST_COMPLETED, create_task +from asyncio import sleep, gather, Semaphore from typing import Tuple, Awaitable, NoReturn, List, Union, Callable, Optional from functools import cached_property @@ -119,16 +119,6 @@ class Backend: return web.json_response(dict(error="invalid JSON"), status=422) workload = payload.count_workload() - async def wait_for_disconnection() -> None: - while request.transport and not request.transport.is_closing(): - await sleep(0.5) - - async def cancel_api_call_if_disconnected() -> web.Response: - await wait_for_disconnection() - log.debug(f"request with reqnum: {auth_data.reqnum} was canceled") - self.metrics._request_canceled(workload=workload, reqnum=auth_data.reqnum) - return web.Response(status=500) - async def make_request() -> Union[web.Response, web.StreamResponse]: log.debug(f"got request, {auth_data.reqnum}") self.metrics._request_start(workload=workload, reqnum=auth_data.reqnum) @@ -174,18 +164,16 @@ class Backend: return web.Response(status=401) try: - done, pending = await wait( - [ - create_task(make_request()), - create_task(cancel_api_call_if_disconnected()), - ], - return_when=FIRST_COMPLETED, - ) - [task.cancel() for task in pending] - return done.pop().result() + return await make_request() except Exception as e: log.debug(f"Exception in main handler loop {e}") return web.Response(status=500) + finally: + if request.task.cancelled(): + log.debug(f"request with reqnum: {auth_data.reqnum} was canceled") + self.metrics._request_canceled( + workload=workload, reqnum=auth_data.reqnum + ) async def __healthcheck(self): health_check_url = self.benchmark_handler.healthcheck_endpoint diff --git a/lib/server.py b/lib/server.py index b21c880..80e2959 100644 --- a/lib/server.py +++ b/lib/server.py @@ -27,7 +27,7 @@ def start_server(backend: Backend, routes: List[web.RouteDef], **kwargs): log.debug("starting server...") app = web.Application() app.add_routes(routes) - runner = web.AppRunner(app) + runner = web.AppRunner(app, handler_cancellation=True) await runner.setup() site = web.TCPSite( runner,