Fix VM logs websocket endpoint getting stuck

src/aleph/vm/orchestrator/views/operator.py

@@ -93,6 +93,7 @@ async def stream_logs(request: web.Request) -> web.StreamResponse:
                     logger.debug(message)
 
                     await ws.send_json({"type": log_type, "message": message})
+                    queue.task_done()
 
             finally:
                 await ws.close()

src/aleph/vm/utils/logs.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from collections.abc import Callable, Generator
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import TypedDict
 
 from systemd import journal
@@ -32,31 +32,47 @@
     For more information refer to the sd-journal(3) manpage
     and systemd.journal module documentation.
     """
-    r = journal.Reader()
-    r.add_match(SYSLOG_IDENTIFIER=stdout_identifier)
-    r.add_match(SYSLOG_IDENTIFIER=stderr_identifier)
-    queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
+    journal_reader = journal.Reader()
+    journal_reader.add_match(SYSLOG_IDENTIFIER=stdout_identifier)
+    journal_reader.add_match(SYSLOG_IDENTIFIER=stderr_identifier)
+    queue: asyncio.Queue = asyncio.Queue(maxsize=5)
+    tasks: list[asyncio.Task] = []
 
-    def _ready_for_read() -> None:
-        change_type = r.process()  # reset fd status
-        if change_type != journal.APPEND:
-            return
+    loop = asyncio.get_event_loop()
+
+    async def process_messages() -> None:
+        """Enqueue all the available log entries, wait if queue is full, then wait for new message via add_reader"""
+        # Remove reader so we don't get called again while processing
+        loop.remove_reader(journal_reader.fileno())
         entry: EntryDict
-        for entry in r:
+        for entry in journal_reader:
             log_type = "stdout" if entry["SYSLOG_IDENTIFIER"] == stdout_identifier else "stderr"
             msg = entry["MESSAGE"]
-            asyncio.create_task(queue.put((log_type, msg)))
+            # will wait if queue is full
+            await queue.put((log_type, msg))
+            journal_reader.process()  # reset fd status
+        journal_reader.process()  # reset fd status
+        # Call _ready_for_read read when entries are readable again, this is non-blocking
+        loop.add_reader(journal_reader.fileno(), _ready_for_read)
+
+    def _ready_for_read() -> None:
+        # wrapper around process_messages as add_reader don't take an async func
+        task = loop.create_task(process_messages(), name=f"process_messages-queue-{id(queue)}")
+        tasks.append(task)
+        task.add_done_callback(tasks.remove)
 
     if skip_past:
-        r.seek_tail()
+        # seek_tail doesn't work see https://github.com/systemd/systemd/issues/17662
+        journal_reader.seek_realtime(datetime.now() - timedelta(seconds=10))
 
-    loop = asyncio.get_event_loop()
-    loop.add_reader(r.fileno(), _ready_for_read)
+    _ready_for_read()
 
     def do_cancel():
-        logger.info(f"cancelling reader {r}")
-        loop.remove_reader(r.fileno())
-        r.close()
+        logger.info(f"cancelling queue and reader {journal_reader}")
+        loop.remove_reader(journal_reader.fileno())
+        for task in tasks:
+            task.cancel()
+        journal_reader.close()
 
     return queue, do_cancel
 

tests/supervisor/test_log.py

@@ -0,0 +1,15 @@
+from asyncio import QueueEmpty
+
+from aleph.vm.utils.logs import make_logs_queue
+
+
+def test_make_logs_queue():
+    stdout_identifier = "test_stdout"
+    stderr_identifier = "test_stderr"
+    queue, do_cancel = make_logs_queue(stdout_identifier, stderr_identifier)
+    import pytest
+
+    with pytest.raises(QueueEmpty):
+        while queue.get_nowait():
+            queue.task_done()
+    do_cancel()