diff --git a/airbyte_cdk/utils/print_buffer.py b/airbyte_cdk/utils/print_buffer.py index 553d412ea6..29bd0e2d11 100644 --- a/airbyte_cdk/utils/print_buffer.py +++ b/airbyte_cdk/utils/print_buffer.py @@ -1,12 +1,68 @@ # Copyright (c) 2024 Airbyte, Inc., all rights reserved. +import atexit +import os import sys +import threading import time +import weakref from io import StringIO from threading import RLock from types import TracebackType from typing import Optional +_REGISTERED: "weakref.WeakSet[PrintBuffer]" = weakref.WeakSet() +_FORK_HELD: list["PrintBuffer"] = [] +_process_hooks_registered = False +_PROCESS_HOOKS_LOCK = threading.Lock() + + +def _register_process_hooks() -> None: + global _process_hooks_registered + if _process_hooks_registered: + return + with _PROCESS_HOOKS_LOCK: + if _process_hooks_registered: + return + atexit.register(_shutdown_all) + if hasattr(os, "register_at_fork"): # not available on Windows + os.register_at_fork( + before=_before_fork, + after_in_parent=_after_fork_in_parent, + after_in_child=_after_fork_in_child, + ) + _process_hooks_registered = True + + +def _shutdown_all() -> None: + for b in list(_REGISTERED): + b._shutdown() + + +def _before_fork() -> None: + buffers = list(_REGISTERED) + for b in buffers: + b.lock.acquire() + _FORK_HELD[:] = buffers + for b in buffers: + try: + b._flush_to_fd() # the child must not inherit (and later re-emit) pending bytes + except (OSError, ValueError): + # best effort: a broken/closed stdout must not block the fork + pass + + +def _after_fork_in_parent() -> None: + for b in _FORK_HELD: + b.lock.release() + _FORK_HELD.clear() + + +def _after_fork_in_child() -> None: + for b in _FORK_HELD: + b._reset_after_fork() + _FORK_HELD.clear() + class PrintBuffer: """ @@ -17,18 +73,27 @@ class PrintBuffer: scenarios where you want to minimize the number of I/O operations by grouping multiple print statements together and flushing them as a single operation. + Buffer content is pushed through to the stdout file descriptor by a daemon + sweeper thread once per flush interval, so callers of `flush()` (such as + logging handlers, which flush after every record) do not each pay for a + file-descriptor flush. + Attributes: buffer (StringIO): A buffer to store the messages before flushing. flush_interval (float): The time interval (in seconds) after which the buffer is flushed. last_flush_time (float): The last time the buffer was flushed. lock (RLock): A reentrant lock to ensure thread-safe operations. + _sweeper (Optional[threading.Thread]): The daemon thread pushing pending output to the file descriptor, or None. + _stop (threading.Event): Signals the sweeper thread to exit. + _pending (threading.Event): Signals the sweeper thread that output is waiting to be pushed. + _sync_flush (bool): When True, `flush()` pushes to the file descriptor synchronously (no sweeper). Methods: write(message: str) -> None: Writes a message to the buffer and flushes if the interval has passed. flush() -> None: - Flushes the buffer content to the standard output. + Drains the buffer to the standard output and marks it pending for the sweeper thread. __enter__() -> "PrintBuffer": Enters the runtime context related to this object, redirecting stdout and stderr. @@ -42,26 +107,119 @@ def __init__(self, flush_interval: float = 0.1): self.flush_interval = flush_interval self.last_flush_time = time.monotonic() self.lock = RLock() + self._sweeper: Optional[threading.Thread] = None + self._stop = threading.Event() + self._pending = threading.Event() + self._sync_flush = False + self._pending_marked = False def write(self, message: str) -> None: with self.lock: self.buffer.write(message) + if not self._pending_marked: + self._mark_pending() current_time = time.monotonic() if (current_time - self.last_flush_time) >= self.flush_interval: self.flush() self.last_flush_time = current_time def flush(self) -> None: - """Writes the buffered content to the real stdout and pushes it through to the file descriptor. + """Drains the buffered content to the real stdout and marks it pending for the sweeper. - `sys.__stdout__` is block-buffered when stdout is a pipe, so without the explicit - `flush()` the messages would stay in the interpreter's ~8 KiB text buffer until it fills up. + `sys.__stdout__` is block-buffered when stdout is a pipe; the sweeper thread + pushes pending output through to the file descriptor within one flush interval. + When `_sync_flush` is set, the push to the file descriptor happens here instead. """ with self.lock: - combined_message = self.buffer.getvalue() - sys.__stdout__.write(combined_message) # type: ignore[union-attr] - sys.__stdout__.flush() # type: ignore[union-attr] - self.buffer = StringIO() + if self._sync_flush: + self._flush_to_fd() + else: + self._drain() + if not self._pending_marked: + self._mark_pending() + + def _drain(self) -> None: + combined_message = self.buffer.getvalue() + sys.__stdout__.write(combined_message) # type: ignore[union-attr] + self.buffer = StringIO() + + def _flush_to_fd(self) -> None: + self._drain() + sys.__stdout__.flush() # type: ignore[union-attr] + + def _mark_pending(self) -> None: + if self._sync_flush: + return + if self._sweeper is None: + self._start_sweeper() + if self._sweeper is None: + return # _start_sweeper fell back to synchronous flush + self._pending.set() + self._pending_marked = True + + def _start_sweeper(self) -> None: + stop, pending = threading.Event(), threading.Event() + thread = threading.Thread( + target=self._sweep, + args=(stop, pending), + name="airbyte-print-buffer-sweeper", + daemon=True, + ) + _register_process_hooks() + _REGISTERED.add(self) + try: + thread.start() + except RuntimeError: # e.g. interpreter shutdown / thread limit + self._sync_flush = True + return + self._sweeper, self._stop, self._pending = thread, stop, pending + + def _sweep(self, stop: threading.Event, pending: threading.Event) -> None: + while not stop.is_set(): + pending.wait() # idle buffers block here without consuming CPU + stop.wait(self.flush_interval) # coalesce for one interval + with self.lock: + pending.clear() + if pending is self._pending: + self._pending_marked = False + try: + self._flush_to_fd() + except Exception: + # broken pipe / closed stdout, or anything unexpected: hand fd flushes back to + # callers so output never waits on a dead sweeper + self._sync_flush = True + return + + def _stop_sweeper(self) -> None: + with self.lock: + thread, stop, pending = self._sweeper, self._stop, self._pending + self._sweeper = None + self._pending_marked = False + if thread is None: + return + stop.set() + pending.set() + if thread is not threading.current_thread(): + thread.join() + + def _shutdown(self) -> None: + with self.lock: + self._sync_flush = True + self._stop_sweeper() + with self.lock: + try: + self._flush_to_fd() + except (OSError, ValueError): + # stdout already broken/closed at shutdown; nothing more can be delivered + pass + + def _reset_after_fork(self) -> None: + # the sweeper thread does not survive fork; locks and events may be inconsistent + self.lock = RLock() + self._sweeper = None + self._stop = threading.Event() + self._pending = threading.Event() + self._pending_marked = False def __enter__(self) -> "PrintBuffer": self.old_stdout, self.old_stderr = sys.stdout, sys.stderr @@ -77,5 +235,7 @@ def __exit__( exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: - self.flush() + self._stop_sweeper() + with self.lock: + self._flush_to_fd() sys.stdout, sys.stderr = self.old_stdout, self.old_stderr diff --git a/unit_tests/utils/test_print_buffer.py b/unit_tests/utils/test_print_buffer.py index b1bf88d3d4..ce47cbcfbe 100644 --- a/unit_tests/utils/test_print_buffer.py +++ b/unit_tests/utils/test_print_buffer.py @@ -2,26 +2,182 @@ # Copyright (c) 2025 Airbyte, Inc., all rights reserved. # +import io +import logging import os +import re import select +import signal import subprocess import sys +import threading import time from unittest.mock import MagicMock +import pytest + from airbyte_cdk.utils.print_buffer import PrintBuffer -def test_flush_pushes_through_underlying_stdout(monkeypatch): +def test_flush_drains_without_fd_flush_and_sweeper_flushes(monkeypatch): + flushed = threading.Event() + mock_stdout = MagicMock() + mock_stdout.flush.side_effect = lambda *args, **kwargs: flushed.set() + monkeypatch.setattr(sys, "__stdout__", mock_stdout) + + with PrintBuffer(flush_interval=1.0) as print_buffer: + print_buffer.write("hello") + print_buffer.flush() + + mock_stdout.write.assert_called_once_with("hello") + # ~1 s of margin before the sweeper would fire + assert not mock_stdout.flush.called + assert flushed.wait(timeout=10) + + +def test_unexpected_sweeper_error_falls_back_to_synchronous_flush(monkeypatch): + sweeper_failed = threading.Event() mock_stdout = MagicMock() + + def flush_fails_once(): + if not sweeper_failed.is_set(): + sweeper_failed.set() + raise RuntimeError("unexpected") + + mock_stdout.flush.side_effect = flush_fails_once monkeypatch.setattr(sys, "__stdout__", mock_stdout) - print_buffer = PrintBuffer() - print_buffer.write("hello") - print_buffer.flush() + with PrintBuffer(flush_interval=0.01) as print_buffer: + print_buffer.write("a") + assert sweeper_failed.wait(timeout=10) + print_buffer._sweeper.join(timeout=10) + assert not print_buffer._sweeper.is_alive() + assert print_buffer._sync_flush + + mock_stdout.flush.reset_mock() + print_buffer.flush() + + mock_stdout.flush.assert_called_once() + + +class _CountingRaw(io.RawIOBase): + def __init__(self): + self.write_calls = 0 + self.data = bytearray() + + def writable(self): + return True + + def write(self, b): + self.write_calls += 1 + self.data.extend(b) + return len(b) + + +def _buffered_counting_stdout(monkeypatch): + raw = _CountingRaw() + monkeypatch.setattr( + sys, "__stdout__", io.TextIOWrapper(io.BufferedWriter(raw, buffer_size=8192)) + ) + return raw + + +def test_exit_flushes_pending_and_stops_sweeper(monkeypatch): + raw = _buffered_counting_stdout(monkeypatch) + + with PrintBuffer(flush_interval=60) as print_buffer: + print_buffer.write("a\n") + print_buffer.flush() + print_buffer.write("b\n") + sweeper = print_buffer._sweeper + + assert bytes(raw.data) == b"a\nb\n" + assert sweeper is not None + assert not sweeper.is_alive() + + +def test_syscall_count_stays_batched_at_one_log_line_per_record(monkeypatch): + raw = _buffered_counting_stdout(monkeypatch) + print_buffer = PrintBuffer(flush_interval=0.1) + logger = logging.Logger("print_buffer_syscalls") + logger.setLevel(logging.INFO) + logger.propagate = False + logger.addHandler(logging.StreamHandler(print_buffer)) + + n = 2000 + record_line = ( + '{"type":"RECORD","record":{"stream":"s","data":{"id":0,' + '"payload":"' + "x" * 200 + '"},"emitted_at":0}}\n' + ) + expected = "" + started = time.monotonic() + with print_buffer: + for i in range(n): + print_buffer.write(record_line) + expected += record_line + logger.info("Read %s records", i) + expected += f"Read {i} records\n" + elapsed = time.monotonic() - started + + total_bytes = len(expected.encode()) + assert bytes(raw.data) == expected.encode() + max_write_calls = total_bytes // 8192 + 1 + int(elapsed / 0.1) + 3 + assert raw.write_calls <= max_write_calls, ( + f"expected batched fd writes, got {raw.write_calls} raw writes " + f"for {n} log lines in {elapsed:.2f}s (bound {max_write_calls})" + ) + + +def _spawn_child(code): + env = os.environ.copy() + env.pop("PYTHONUNBUFFERED", None) + return subprocess.Popen( + [sys.executable, "-W", "ignore::DeprecationWarning", "-c", code], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + start_new_session=True, + ) + + +def _read_until(process, markers, timeout): + """Read child stdout until every marker is seen or the deadline passes. + + Returns (output bytes, {marker: time.monotonic() at first sighting}). + """ + assert process.stdout is not None + stdout_fd = process.stdout.fileno() + os.set_blocking(stdout_fd, False) + output = bytearray() + seen = {} + deadline = time.monotonic() + timeout + while time.monotonic() < deadline and len(seen) < len(markers): + readable, _, _ = select.select( + [process.stdout], [], [], max(0.0, deadline - time.monotonic()) + ) + if not readable: + break + try: + chunk = os.read(stdout_fd, 65536) + except BlockingIOError: + continue + if not chunk: + break + output.extend(chunk) + now = time.monotonic() + for marker in markers: + if marker not in seen and marker in output: + seen[marker] = now + return bytes(output), seen - mock_stdout.write.assert_called_once_with("hello") - mock_stdout.flush.assert_called_once() + +def _kill_process_group(process): + if process.poll() is None: + try: + os.killpg(process.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + process.kill() + process.wait() def test_output_reaches_pipe_before_process_exit(): @@ -84,3 +240,98 @@ def test_output_reaches_pipe_before_process_exit(): "Child exited before output was observed. " f"Received stdout: {bytes(output)!r}; stderr: {stderr!r}" ) + + +@pytest.mark.skipif( + sys.platform != "linux", reason="relies on CLOCK_MONOTONIC being shared across processes" +) +def test_quiet_line_reaches_pipe_within_two_flush_intervals(): + child = r""" +import time +from airbyte_cdk.logger import PRINT_BUFFER, init_logger + +PRINT_BUFFER.flush_interval = 0.5 +logger = init_logger("airbyte") +with PRINT_BUFFER: + print('{"type":"RECORD","record":{"stream":"s","data":{"id":1},"emitted_at":0}}\n', end="") + logger.info(f"quiet-line sent_at={time.monotonic()!r}") + time.sleep(30) +""" + process = _spawn_child(child) + try: + output, seen = _read_until(process, [b"sent_at="], timeout=10) + match = re.search(rb"sent_at=([0-9.]+)", output) + assert match is not None, f"quiet-line not observed. stdout: {output!r}" + sent_at = float(match.group(1)) + latency = seen[b"sent_at="] - sent_at + print(f"\nquiet-line latency: {latency:.3f}s") + assert latency < 2 * 0.5, f"quiet-line took {latency:.3f}s to reach the pipe" + assert process.poll() is None, "child exited before output was observed" + finally: + _kill_process_group(process) + stderr = process.stderr.read() + if stderr: + print(f"\nchild stderr: {stderr!r}") + + +def test_pending_output_flushed_on_exit(): + child = r""" +from airbyte_cdk.logger import PRINT_BUFFER, init_logger + +logger = init_logger("airbyte") +with PRINT_BUFFER: + for i in range(1000): + print('{"type":"RECORD","record":{"stream":"s","data":{"id":%d,"e":0},"emitted_at":0}}\n' % i, end="") + logger.info("log-line-%d-done", i) +logger.info("after-with") +PRINT_BUFFER.write("raw-write-no-flush\n") +""" + process = _spawn_child(child) + try: + stdout, stderr = process.communicate(timeout=60) + finally: + _kill_process_group(process) + + for i in range(1000): + assert stdout.count(b'"id":%d,' % i) == 1, f"RECORD {i} missing or duplicated" + assert stdout.count(b"log-line-%d-done" % i) == 1, f"log line {i} missing or duplicated" + assert b"after-with" in stdout + assert b"raw-write-no-flush" in stdout + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="requires os.fork") +def test_forked_child_gets_a_working_sweeper(): + child = r""" +import os, time +from airbyte_cdk.logger import PRINT_BUFFER, init_logger + +logger = init_logger("airbyte") +with PRINT_BUFFER: + logger.info("before-fork") + pid = os.fork() + if pid == 0: + logger.info("in-fork-child") + time.sleep(30) + os._exit(0) + else: + logger.info("in-fork-parent") + time.sleep(30) +""" + process = _spawn_child(child) + try: + output, seen = _read_until(process, [b"in-fork-child", b"in-fork-parent"], timeout=10) + assert b"in-fork-child" in seen, ( + f"forked child's log line never arrived. stdout: {output!r}" + ) + assert b"in-fork-parent" in seen, ( + f"fork parent's log line never arrived. stdout: {output!r}" + ) + assert process.poll() is None, "child exited before output was observed" + assert output.count(b"before-fork") == 1, ( + f"pre-fork output emitted != 1 time. stdout: {output!r}" + ) + finally: + _kill_process_group(process) + stderr = process.stderr.read() + if stderr: + print(f"\nchild stderr: {stderr!r}")