Switch to django-asgi-lifecycle instead of custom asgi lifecycle impl…

…ementation

api/Pipfile

@@ -23,6 +23,7 @@ aiohttp = "~=3.8"
 aws-requests-auth = "~=0.4"
 deepdiff = "~=6.4"
 Django = "~=4.2"
+django-asgi-lifespan = "~=0.1"
 django-cors-headers = "~=4.2"
 django-log-request-id = "~=2.0"
 django-oauth-toolkit = "~=2.3"

api/api/utils/aiohttp.py

@@ -3,8 +3,8 @@
 import weakref
 
 import aiohttp
-
-from conf.asgi import APPLICATION_LIFECYCLE
+import sentry_sdk
+from django_asgi_lifespan.signals import asgi_shutdown
 
 
 logger = logging.getLogger(__name__)
@@ -19,6 +19,24 @@
 ] = weakref.WeakKeyDictionary()
 
 
+@asgi_shutdown.connect
+async def _close_sessions(sender, **kwargs):
+    logger.debug("Closing aiohttp sessions on application shutdown")
+
+    closed_sessions = 0
+
+    while _SESSIONS:
+        loop, session = _SESSIONS.popitem()
+        try:
+            await session.close()
+            closed_sessions += 1
+        except BaseException as exc:
+            logger.error(exc)
+            sentry_sdk.capture_exception(exc)
+
+    logger.debug("Successfully closed %s session(s)", closed_sessions)
+
+
 async def get_aiohttp_session() -> aiohttp.ClientSession:
     """
     Safely retrieve a shared aiohttp session for the current event loop.
@@ -56,7 +74,6 @@ async def get_aiohttp_session() -> aiohttp.ClientSession:
 
         if create_session:
             session = aiohttp.ClientSession()
-            APPLICATION_LIFECYCLE.register_shutdown_handler(session.close)
             _SESSIONS[loop] = session
 
         return _SESSIONS[loop]

api/api/utils/check_dead_links/__init__.py

@@ -36,13 +36,13 @@ def _get_expiry(status, default):
 _timeout = aiohttp.ClientTimeout(total=2)
 
 
-async def _head(url: str, session: aiohttp.ClientSession) -> tuple[str, int]:
+async def _head(url: str) -> tuple[str, int]:
     try:
-        request = session.head(
+        session = await get_aiohttp_session()
+        response = await session.head(
             url, allow_redirects=False, headers=HEADERS, timeout=_timeout
         )
-        async with request as response:
-            return url, response.status
+        return url, response.status
     except (aiohttp.ClientError, asyncio.TimeoutError) as exception:
         _log_validation_failure(exception)
         return url, -1
@@ -52,8 +52,7 @@ async def _head(url: str, session: aiohttp.ClientSession) -> tuple[str, int]:
 @async_to_sync
 async def _make_head_requests(urls: list[str]) -> list[tuple[str, int]]:
     tasks = []
-    session = await get_aiohttp_session()
-    tasks = [asyncio.ensure_future(_head(url, session)) for url in urls]
+    tasks = [asyncio.ensure_future(_head(url)) for url in urls]
     responses = asyncio.gather(*tasks)
     await responses
     return responses.result()

api/conf/asgi.py

@@ -2,15 +2,14 @@
 
 from django.conf import settings
 from django.contrib.staticfiles.handlers import ASGIStaticFilesHandler
-from django.core.asgi import get_asgi_application
 
-from conf.lifecycle_handler import ASGILifecycleHandler
+from django_asgi_lifespan.asgi import get_asgi_application
 
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "conf.settings")
 
 
-application = APPLICATION_LIFECYCLE = ASGILifecycleHandler(get_asgi_application())
+application = get_asgi_application()
 
 
 if settings.ENVIRONMENT == "local":

api/test/conftest.py

@@ -1,19 +1,30 @@
 import pytest
 from asgiref.sync import async_to_sync
 
-from conf.asgi import APPLICATION_LIFECYCLE
+from conf.asgi import application
 
 
 @pytest.fixture(scope="session", autouse=True)
-def call_application_shutdown():
+def ensure_asgi_lifecycle():
     """
-    Call application shutdown during test session teardown.
+    Call application shutdown lifecycle event.
 
     This cannot be an async fixture because the scope is session
     and pytest-asynio's `event_loop` fixture, which is auto-used
     for async tests and fixtures, is function scoped, which is
     incomatible with session scoped fixtures. `async_to_sync` works
     fine here, so it's not a problem.
+
+    This cannot yet call the startup signal due to:
+    https://github.com/illagrenan/django-asgi-lifespan/pull/80
     """
+    scope = {"type": "lifespan"}
+
+    async def noop(*args, **kwargs):
+        ...
+
+    async def shutdown():
+        return {"type": "lifespan.shutdown"}
+
     yield
-    async_to_sync(APPLICATION_LIFECYCLE.shutdown)()
+    async_to_sync(application)(scope, shutdown, noop)

api/test/unit/utils/test_aiohttp.py

@@ -1,10 +1,8 @@
 import asyncio
 
 import pytest
-from asgiref.sync import async_to_sync
 
 from api.utils.aiohttp import get_aiohttp_session
-from conf.asgi import APPLICATION_LIFECYCLE
 
 
 @pytest.fixture(autouse=True)
@@ -18,7 +16,6 @@ def _get_new_loop():
 
     yield _get_new_loop
 
-    async_to_sync(APPLICATION_LIFECYCLE.shutdown)()
     for loop in loops:
         loop.close()
 
@@ -53,12 +50,3 @@ def test_multiple_loops_reuse_separate_sessions(get_new_loop):
 
     assert loop_1_session_1 is loop_1_session_2
     assert loop_2_session_1 is loop_2_session_2
-
-
-def test_registers_shutdown_cleanup(get_new_loop):
-    assert len(APPLICATION_LIFECYCLE._on_shutdown) == 0
-
-    loop = get_new_loop()
-    loop.run_until_complete(get_aiohttp_session())
-
-    assert len(APPLICATION_LIFECYCLE._on_shutdown) == 1

api/test/unit/utils/test_check_dead_links.py

@@ -1,8 +1,8 @@
 import asyncio
-from unittest import mock
 
 import pook
 import pytest
+from aiohttp.client import ClientSession
 
 from api.utils.check_dead_links import HEADERS, check_dead_links
 
@@ -30,7 +30,7 @@ def test_sends_user_agent():
         assert url in requested_urls
 
 
-def test_handles_timeout():
+def test_handles_timeout(monkeypatch):
     """
     Test that case where timeout occurs.
 
@@ -42,13 +42,11 @@ def test_handles_timeout():
     image_urls = [f"https://example.org/{i}" for i in range(len(results))]
     start_slice = 0
 
-    def raise_timeout_error(*args, **kwargs):
+    async def raise_timeout_error(*args, **kwargs):
         raise asyncio.TimeoutError()
 
-    with mock.patch(
-        "aiohttp.client.ClientSession._request", side_effect=raise_timeout_error
-    ):
-        check_dead_links(query_hash, start_slice, results, image_urls)
+    monkeypatch.setattr(ClientSession, "_request", raise_timeout_error)
+    check_dead_links(query_hash, start_slice, results, image_urls)
 
     # `check_dead_links` directly modifies the results list
     # if the results are timing out then they're considered dead and discarded