Enhancement: websocket heartbeat (#13739)

---------

Co-authored-by: shamoon <4887959+shamoon@users.noreply.github.com>

src/paperless/consumers.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import asyncio
+import contextlib
 import json
 from typing import TYPE_CHECKING
 
@@ -14,8 +16,13 @@ if TYPE_CHECKING:
     from documents.plugins.helpers import PermissionsData
     from documents.plugins.helpers import StatusUpdatePayload
 
+HEARTBEAT_INTERVAL = 30
+HEARTBEAT_MESSAGE = json.dumps({"type": "heartbeat"})
+
 
 class StatusConsumer(AsyncWebsocketConsumer):
+    heartbeat_task: asyncio.Task | None = None
+
     def _authenticated(self) -> bool:
         user: AbstractBaseUser | AnonymousUser | None = self.scope.get("user")
         return user is not None and user.is_authenticated
@@ -39,10 +46,27 @@ class StatusConsumer(AsyncWebsocketConsumer):
             return
         await self.channel_layer.group_add("status_updates", self.channel_name)
         await self.accept()
+        self._start_heartbeat()
 
     async def disconnect(self, code: int) -> None:
+        await self._stop_heartbeat()
         await self.channel_layer.group_discard("status_updates", self.channel_name)
 
+    def _start_heartbeat(self) -> None:
+        self.heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+
+    async def _stop_heartbeat(self) -> None:
+        if self.heartbeat_task is not None:
+            self.heartbeat_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self.heartbeat_task
+            self.heartbeat_task = None
+
+    async def _heartbeat_loop(self) -> None:
+        while True:
+            await asyncio.sleep(HEARTBEAT_INTERVAL)
+            await self.send(HEARTBEAT_MESSAGE)
+
     async def status_update(self, event: StatusUpdatePayload) -> None:
         if not self._authenticated():
             await self.close()

src/paperless/tests/test_websockets.py

@@ -189,6 +189,23 @@ class TestWebSockets:
 
         await communicator.disconnect()
 
+    @pytest.mark.anyio
+    async def test_heartbeat(self, mocker: MockerFixture) -> None:
+        mocker.patch(
+            "paperless.consumers.StatusConsumer._authenticated",
+            return_value=True,
+        )
+        mocker.patch("paperless.consumers.HEARTBEAT_INTERVAL", 0.01)
+
+        communicator = WebsocketCommunicator(application, "/ws/status/")
+        connected, _ = await communicator.connect()
+        assert connected
+
+        assert await communicator.receive_json_from() == {"type": "heartbeat"}
+        assert await communicator.receive_json_from() == {"type": "heartbeat"}
+
+        await communicator.disconnect()
+
     def test_manager_send_progress(self, mocker: MockerFixture) -> None:
         mock_group_send = mocker.patch(
             "channels.layers.InMemoryChannelLayer.group_send",