Merge pull request #2718 from langbot-app/fix/nonblocking-console-logging

fix(logging): prevent blocked stdout from freezing LangBot

src/langbot/pkg/core/boot.py

@@ -62,7 +62,6 @@ async def main(loop: asyncio.AbstractEventLoop):
         import signal
 
         def signal_handler(sig, frame):
-            print('[Signal] Program exit.')
             runtime_loop.call_soon_threadsafe(shutdown_requested.set)
 
         signal.signal(signal.SIGINT, signal_handler)

src/langbot/pkg/core/bootutils/log.py

@@ -1,7 +1,9 @@
 import logging
 import logging.handlers
 import os
+import queue
 import sys
+import threading
 import time
 
 import colorlog
@@ -23,6 +25,122 @@ LOG_FILE_BACKUP_COUNT = 5  # Keep 5 backup files (total ~50MB max)
 
 LOG_DIR = 'data/logs'
 
+CONSOLE_QUEUE_CAPACITY = 256
+CONSOLE_RECORD_MAX_BYTES = 16 * 1024
+CONSOLE_CLOSE_TIMEOUT = 0.2
+
+
+class NonBlockingConsoleHandler(logging.Handler):
+    """Isolate a stalled console from callers with at most 4 MiB of queued bytes.
+
+    Only formatted UTF-8 bytes cross the thread boundary, so queued records do
+    not retain arguments or exception frames. Overflow drops console copies;
+    other handlers still receive the original records.
+    """
+
+    def __init__(self, stream, *, capacity: int = CONSOLE_QUEUE_CAPACITY):
+        if capacity <= 0:
+            raise ValueError('Console queue capacity must be positive')
+        # Drain startup prints before the worker can fill the pipe. Otherwise
+        # interpreter finalization could block flushing the original stream.
+        stream.flush()
+        super().__init__()
+        self._queue: queue.Queue[bytes | None] = queue.Queue(maxsize=capacity)
+        # logging.shutdown holds Handler.lock while joining the worker.
+        self._state_lock = threading.Lock()
+        self._failed = threading.Event()
+        self.dropped_records = 0
+        self.write_error: str | None = None
+        # The worker owns this descriptor. Do not share a buffered Python stream
+        # lock with the main thread or interpreter finalization.
+        self._fd = os.dup(stream.fileno())
+        self._worker = threading.Thread(target=self._consume, name='langbot-console-log', daemon=True)
+        try:
+            self._worker.start()
+        except Exception:
+            os.close(self._fd)
+            super().close()
+            raise
+
+    def emit(self, record: logging.LogRecord) -> None:
+        with self._state_lock:
+            if self._closed or self._failed.is_set() or self._queue.full():
+                self.dropped_records += 1
+                return
+            try:
+                text = self.format(record)
+                payload = (text[:CONSOLE_RECORD_MAX_BYTES] + '\n').encode('utf-8', errors='backslashreplace')
+                if len(text) > CONSOLE_RECORD_MAX_BYTES or len(payload) > CONSOLE_RECORD_MAX_BYTES:
+                    marker = b'\n[console log truncated]\n'
+                    payload = (
+                        payload[: CONSOLE_RECORD_MAX_BYTES - len(marker)]
+                        .decode('utf-8', errors='ignore')
+                        .encode('utf-8')
+                        + marker
+                    )
+                self._queue.put_nowait(payload)
+            except queue.Full:
+                self.dropped_records += 1
+            except Exception:
+                # Logging's default handleError writes to stderr, which may
+                # share the same stalled container logging driver.
+                self.dropped_records += 1
+
+    def _write(self, payload: bytes) -> None:
+        remaining = memoryview(payload)
+        while remaining:
+            written = os.write(self._fd, remaining)
+            if written == 0:
+                raise OSError('Console write made no progress')
+            remaining = remaining[written:]
+
+    def _consume(self) -> None:
+        reported_drops = 0
+        try:
+            while (payload := self._queue.get()) is not None:
+                self._write(payload)
+                with self._state_lock:
+                    dropped = self.dropped_records
+                if dropped > reported_drops:
+                    self._write(
+                        f'[WARNING] Console logging dropped {dropped - reported_drops} records.\n'.encode('utf-8')
+                    )
+                    reported_drops = dropped
+        except OSError as exc:
+            with self._state_lock:
+                self.write_error = str(exc)
+                self._failed.set()
+                self.dropped_records += 1
+        finally:
+            os.close(self._fd)
+            with self._state_lock:
+                while True:
+                    try:
+                        if self._queue.get_nowait() is not None:
+                            self.dropped_records += 1
+                    except queue.Empty:
+                        break
+
+    def close(self) -> None:
+        with self._state_lock:
+            if self._closed:
+                return
+            super().close()
+            # A full queue must not make shutdown wait for the writer.
+            while True:
+                try:
+                    self._queue.put_nowait(None)
+                    break
+                except queue.Full:
+                    try:
+                        self._queue.get_nowait()
+                        self.dropped_records += 1
+                    except queue.Empty:
+                        # The writer consumed the last record after Full;
+                        # retry inserting the shutdown sentinel.
+                        continue
+        self._worker.join(timeout=CONSOLE_CLOSE_TIMEOUT)
+
 
 class DailyGroupedRotatingFileHandler(logging.handlers.RotatingFileHandler):
     """File handler that writes to ``data/logs/langbot-YYYY-MM-DD.log``.
@@ -94,6 +212,10 @@ async def init_logging(extra_handlers: list[logging.Handler] = None) -> logging.
 
     qcg_logger = logging.getLogger('langbot')
 
+    for handler in qcg_logger.handlers[:]:
+        qcg_logger.removeHandler(handler)
+        handler.close()
+
     qcg_logger.setLevel(level)
 
     color_formatter = colorlog.ColoredFormatter(
@@ -102,10 +224,7 @@ async def init_logging(extra_handlers: list[logging.Handler] = None) -> logging.
         log_colors=log_colors_config,
     )
 
-    stream_handler = logging.StreamHandler(sys.stdout)
-    # stream_handler.setLevel(level)
-    # stream_handler.setFormatter(color_formatter)
-    stream_handler.stream = open(sys.stdout.fileno(), mode='w', encoding='utf-8', buffering=1)
+    stream_handler = NonBlockingConsoleHandler(sys.stdout)
 
     # Rotate by size within a day and switch files when the date changes,
     # so long-running processes still produce a log file for the current day.

tests/unit_tests/core/test_boot.py

@@ -38,6 +38,8 @@ async def test_main_signal_handler_handles_sigint_before_app_created(monkeypatch
     captured_handler = {}
     print_exc = Mock()
     monkeypatch.setattr(boot.traceback, 'print_exc', print_exc)
+    # Signal handling must not touch a potentially blocked stdout stream.
+    monkeypatch.setattr(boot, 'print', Mock(side_effect=AssertionError('stdout must not be used')), raising=False)
 
     def fake_signal(sig, handler):
         captured_handler[sig] = handler

tests/unit_tests/core/test_console_logging.py

@@ -0,0 +1,259 @@
+"""Console backpressure must not stop request handling or process shutdown."""
+
+from __future__ import annotations
+
+import logging
+import os
+from pathlib import Path
+import subprocess
+import sys
+import textwrap
+import threading
+import time
+import weakref
+
+import pytest
+
+from langbot.pkg.core.bootutils import log as logmod
+from langbot.pkg.core.bootutils.log import NonBlockingConsoleHandler
+
+
+def _logger(handler):
+    logger = logging.Logger('console-test', logging.INFO)
+    logger.addHandler(handler)
+    return logger
+
+
+def test_console_truncation_preserves_utf8_and_other_handlers(tmp_path):
+    message = '界🙂' * logmod.CONSOLE_RECORD_MAX_BYTES
+    with (tmp_path / 'console').open('w+b') as stream:
+        handler = NonBlockingConsoleHandler(stream)
+        logger = _logger(handler)
+        other = logging.FileHandler(tmp_path / 'full.log', encoding='utf-8')
+        logger.addHandler(other)
+        try:
+            logger.info('%s', message)
+            handler.close()
+            assert not handler._worker.is_alive()
+            stream.seek(0)
+            payload = stream.read()
+            assert len(payload) <= logmod.CONSOLE_RECORD_MAX_BYTES
+            assert payload.decode('utf-8').endswith('[console log truncated]\n')
+            assert (tmp_path / 'full.log').read_text(encoding='utf-8') == message + '\n'
+        finally:
+            handler.close()
+            other.close()
+
+
+def test_close_with_full_queue_has_bounded_wait_and_reports_drops(tmp_path):
+    entered = threading.Event()
+    release = threading.Event()
+
+    class StalledConsole(NonBlockingConsoleHandler):
+        def _write(self, payload):
+            entered.set()
+            assert release.wait(5), 'test writer was not released'
+            super()._write(payload)
+
+    with (tmp_path / 'console').open('w+b') as stream:
+        handler = StalledConsole(stream, capacity=2)
+        logger = _logger(handler)
+        try:
+            logger.info('first')
+            assert entered.wait(2)
+            logger.info('second')
+            logger.info('third')
+            logger.info('overflow')
+            assert handler.dropped_records == 1
+            started = time.monotonic()
+            handler.close()
+            assert time.monotonic() - started < 1
+            logger.info('after close')
+            assert handler.dropped_records == 3
+        finally:
+            release.set()
+            handler.close()
+            handler._worker.join(timeout=2)
+        assert not handler._worker.is_alive()
+        stream.seek(0)
+        output = stream.read().decode('utf-8')
+        assert 'first\n' in output
+        assert 'third\n' in output
+        assert 'Console logging dropped 3 records' in output
+
+
+def test_broken_console_does_not_write_errors_to_stderr(tmp_path, capsys):
+    class BrokenConsole(NonBlockingConsoleHandler):
+        def _write(self, payload):
+            raise BrokenPipeError('reader exited')
+
+    with (tmp_path / 'console').open('w+b') as stream:
+        handler = BrokenConsole(stream)
+        logger = _logger(handler)
+        try:
+            logger.info('first')
+            assert handler._failed.wait(2)
+            logger.info('second')
+            handler._worker.join(timeout=2)
+            assert handler.dropped_records == 2
+            assert 'reader exited' in handler.write_error
+            assert handler._queue.empty()
+        finally:
+            handler.close()
+    assert capsys.readouterr().err == ''
+
+
+def test_logging_shutdown_drains_healthy_console_without_handler_lock_deadlock(tmp_path):
+    entered = threading.Event()
+    release = threading.Event()
+
+    class Console(NonBlockingConsoleHandler):
+        def _write(self, payload):
+            entered.set()
+            assert release.wait(5), 'test writer was not released'
+            super()._write(payload)
+
+        def close(self):
+            release.set()
+            super().close()
+
+    with (tmp_path / 'console').open('w+b') as stream:
+        handler = Console(stream)
+        logger = _logger(handler)
+        try:
+            logger.info('first')
+            assert entered.wait(2)
+            for index in range(20):
+                logger.info('queued %d', index)
+            # The standard shutdown holds handler.lock across flush and close.
+            logging.shutdown([weakref.ref(handler)])
+            assert not handler._worker.is_alive()
+            assert handler.dropped_records == 0
+            stream.seek(0)
+            assert stream.read().decode('utf-8').splitlines() == ['first', *(f'queued {i}' for i in range(20))]
+        finally:
+            release.set()
+            handler.close()
+            handler._worker.join(timeout=2)
+
+
+@pytest.mark.asyncio
+async def test_reinitialization_closes_previous_handlers(tmp_path, monkeypatch):
+    monkeypatch.setattr(logmod, 'LOG_DIR', str(tmp_path))
+    with (tmp_path / 'console').open('w+b') as stream:
+        monkeypatch.setattr(logmod.sys, 'stdout', stream)
+        logger = await logmod.init_logging()
+        previous = logger.handlers[:]
+        try:
+            logger.info('before reinitialization')
+            await logmod.init_logging()
+            logger.info('after reinitialization')
+            assert len(logger.handlers) == 2
+            assert all(handler._closed for handler in previous)
+            assert not previous[0]._worker.is_alive()
+            assert previous[1].stream is None
+        finally:
+            for handler in logger.handlers[:]:
+                logger.removeHandler(handler)
+                handler.close()
+
+
+@pytest.mark.parametrize('buffered_startup', [False, True], ids=['direct-init', 'buffered-entrypoint'])
+def test_unread_stdout_keeps_webui_file_cache_and_process_exit_working(tmp_path, buffered_startup):
+    # The parent deliberately never reads stdout until the child exits. This
+    # exercises an actual full OS pipe, including interpreter finalization.
+    script = textwrap.dedent("""\
+        import asyncio
+        import json
+        import logging
+        import os
+        from pathlib import Path
+        import sys
+        import time
+        from types import SimpleNamespace
+        from unittest.mock import AsyncMock
+
+        from quart import Quart
+        from langbot.pkg.core.bootutils import log
+        from langbot.pkg.core.stages.setup_logger import PersistenceHandler
+        from langbot.pkg.utils.logcache import LogCache
+
+        async def main():
+            Path('data/logs').mkdir(parents=True)
+            cache = LogCache()
+            logger = await log.init_logging([PersistenceHandler('cache', SimpleNamespace(log_cache=cache))])
+            console = logger.handlers[0]
+            app = Quart(__name__)
+
+            @app.get('/health')
+            async def health():
+                logger.info('health request completed')
+                return {'ok': True}
+
+            async def produce():
+                for index in range(800):
+                    logger.info('record %d %s', index, 'x' * 2048)
+                    if index % 16 == 0:
+                        await asyncio.sleep(0)
+
+            producer = asyncio.create_task(produce())
+            await producer
+            response = await app.test_client().get('/health')
+            assert response.status_code == 200
+            assert await response.get_json() == {'ok': True}
+            assert console.dropped_records > 0
+            with console._queue.mutex:
+                queued = list(console._queue.queue)
+            assert len(queued) <= log.CONSOLE_QUEUE_CAPACITY
+            assert all(isinstance(item, bytes) and len(item) <= log.CONSOLE_RECORD_MAX_BYTES for item in queued)
+            assert 'health request completed' in cache.get_log_by_pointer(0, 0)[0]
+            assert any('health request completed' in path.read_text(encoding='utf-8')
+                       for path in Path('data/logs').glob('*.log'))
+            started = time.monotonic()
+            logging.shutdown()
+            assert time.monotonic() - started < 1
+            Path('result.json').write_text(json.dumps({'dropped': console.dropped_records}), encoding='utf-8')
+
+        if os.environ['LANGBOT_TEST_BUFFERED_STARTUP'] == '1':
+            from langbot import __main__ as entrypoint
+            from langbot.pkg.core import boot
+            from langbot.pkg.core.bootutils import deps, files
+
+            # Run the actual entrypoint's banner print without starting bots,
+            # installing dependencies, or loading a real configuration.
+            assert not sys.stdout.line_buffering
+            assert not sys.stdout.write_through
+            deps.check_deps = AsyncMock(return_value=[])
+            files.generate_files = AsyncMock(return_value=[])
+
+            async def run_after_startup(loop):
+                await main()
+
+            boot.main = run_after_startup
+            entrypoint.main()
+        else:
+            asyncio.run(main())
+    """)
+    env = dict(
+        os.environ,
+        PYTHONPATH=str(Path(logmod.__file__).resolve().parents[4]),
+        PYTHONIOENCODING='utf-8',
+        LANGBOT_TEST_BUFFERED_STARTUP=str(int(buffered_startup)),
+    )
+    env.pop('PYTHONUNBUFFERED', None)
+    with subprocess.Popen(
+        [sys.executable, '-c', script],
+        cwd=tmp_path,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    ) as process:
+        try:
+            process.wait(timeout=20)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            _, stderr = process.communicate()
+            pytest.fail(f'Logging or process exit blocked on stdout: {stderr.decode(errors="replace")}')
+        _, stderr = process.communicate()
+        assert process.returncode == 0, stderr.decode(errors='replace')
+    assert (tmp_path / 'result.json').exists()