diff --git a/nvchecker/httpclient/aiohttp_httpclient.py b/nvchecker/httpclient/aiohttp_httpclient.py index 6af4b0f..e96ea0a 100644 --- a/nvchecker/httpclient/aiohttp_httpclient.py +++ b/nvchecker/httpclient/aiohttp_httpclient.py @@ -1,7 +1,6 @@ # MIT licensed # Copyright (c) 2013-2020 lilydjwg , et al. -import atexit import asyncio from typing import Optional, Dict @@ -16,16 +15,15 @@ logger = structlog.get_logger(logger_name=__name__) connector = aiohttp.TCPConnector(limit=20) class AiohttpSession(BaseSession): + session = None + def setup( self, concurreny: int = 20, timeout: int = 20, ) -> None: - self.session = aiohttp.ClientSession( - connector = aiohttp.TCPConnector(limit=concurreny), - timeout = aiohttp.ClientTimeout(total=timeout), - trust_env = True, - ) + self._concurreny = concurreny + self._timeout = timeout async def request_impl( self, url: str, *, @@ -38,6 +36,14 @@ class AiohttpSession(BaseSession): body = None, verify_cert: bool = True, ) -> Response: + if self.session is None: + # need to create in async context + self.session = aiohttp.ClientSession( + connector = aiohttp.TCPConnector(limit=self._concurreny), + timeout = aiohttp.ClientTimeout(total=self._timeout), + trust_env = True, + ) + kwargs = { 'headers': headers, 'params': params, @@ -77,9 +83,4 @@ class AiohttpSession(BaseSession): body = await res.content.read() return Response(res.headers, body) -@atexit.register -def cleanup(): - loop = asyncio.get_event_loop() - loop.run_until_complete(session.session.close()) - session = AiohttpSession() diff --git a/nvchecker/httpclient/httpx_httpclient.py b/nvchecker/httpclient/httpx_httpclient.py index 7609a9a..cacb915 100644 --- a/nvchecker/httpclient/httpx_httpclient.py +++ b/nvchecker/httpclient/httpx_httpclient.py @@ -1,7 +1,6 @@ # MIT licensed # Copyright (c) 2020 lilydjwg , et al. -import atexit from typing import Dict, Optional, Tuple import httpx @@ -75,10 +74,4 @@ class HttpxSession(BaseSession): await client.aclose() del self.clients -@atexit.register -def cleanup(): - import asyncio - loop = asyncio.get_event_loop() - loop.run_until_complete(session.aclose()) - session = HttpxSession()