diff --git a/packages/reflex-base/news/+json-output-slow-consumer.bugfix.md b/packages/reflex-base/news/+json-output-slow-consumer.bugfix.md new file mode 100644 index 00000000000..a65cae0e572 --- /dev/null +++ b/packages/reflex-base/news/+json-output-slow-consumer.bugfix.md @@ -0,0 +1 @@ +Avoid premature truncation of JSON logs when consumers read slowly during shutdown, while limiting output draining to 30 seconds after the command exits. diff --git a/packages/reflex-base/src/reflex_base/utils/log.py b/packages/reflex-base/src/reflex_base/utils/log.py index abcc0a51bb0..742cb3f79ba 100644 --- a/packages/reflex-base/src/reflex_base/utils/log.py +++ b/packages/reflex-base/src/reflex_base/utils/log.py @@ -272,9 +272,12 @@ def _write_json(payload: dict, *, stderr: bool): # stops waiting for descendants that still hold the pipe. _DRAIN_IDLE_SECONDS = 0.5 -# How long the supervisor drains the pipes after the child exits, at most. +# Maximum drain time after the child exits, excluding time forwarding output. _DRAIN_MAX_SECONDS = 5 +# Stop shutdown after 30 seconds even when a consumer blocks every write. +_DRAIN_WALL_SECONDS = 30 + # Line ends in child output; a lone ``\r`` ends a progress-bar update. _LINE_END = re.compile(rb"\r\n|\r|\n") @@ -390,6 +393,9 @@ def __init__(self, read_fd: int, write_fd: int, level: str, name: str): self.logger_name = name # When the reader started waiting for data; None while it works. self.idle_since: float | None = None + # Cumulative write time and the start of any ongoing write. + self._write_state: tuple[float, float | None] = (0.0, None) + self._write_lock = threading.Lock() def run(self): """Forward the pipe until every writer closed it.""" @@ -420,27 +426,53 @@ def _write(self, lines: list[bytes], traceback: list[str], final: bool): final: Whether the output ended. """ data = _to_records(lines, traceback, self.level, self.logger_name, final) + with self._write_lock: + elapsed, _ = self._write_state + started = time.monotonic() + self._write_state = (elapsed, started) # Keep reading when the consumer is gone: a full pipe blocks the child. with contextlib.suppress(OSError): _write_all(self.write_fd, data) + with self._write_lock: + self._write_state = (elapsed + time.monotonic() - started, None) + + def writing_snapshot(self) -> tuple[float, float]: + """Capture the current time and cumulative write duration together. - def drain(self, exited_at: float): + Returns: + The monotonic time and seconds spent writing up to that time. + """ + with self._write_lock: + elapsed, started = self._write_state + now = time.monotonic() + return now, elapsed if started is None else elapsed + now - started + + def drain(self, exited_at: float, writing_at_exit: float, wall_deadline: float): """Wait until the pipe is drained after the child exited. Returns early when the reader sat idle for a while after the exit, or - when the drain took too long: only a descendant that outlived the - child still holds the pipe. + when the drain took too long. Time spent forwarding output does not + count toward the active limit. A separate wall deadline bounds shutdown + even when a consumer stops reading. Args: - exited_at: The monotonic time the child exited. + exited_at: The monotonic time of the pump's snapshot after child exit. + writing_at_exit: The cumulative write time at that snapshot. + wall_deadline: The shared monotonic deadline for shutdown. """ while self.is_alive(): self.join(0.05) - now = time.monotonic() + now, writing_elapsed = self.writing_snapshot() idle_since = self.idle_since - if now - exited_at > _DRAIN_MAX_SECONDS or ( - idle_since is not None - and now - max(idle_since, exited_at) > _DRAIN_IDLE_SECONDS + # Seconds charged against the five-second active drain budget. + drain_elapsed = now - exited_at - (writing_elapsed - writing_at_exit) + if ( + now > wall_deadline + or drain_elapsed > _DRAIN_MAX_SECONDS + or ( + idle_since is not None + and now - max(idle_since, exited_at) > _DRAIN_IDLE_SECONDS + ) ): return @@ -492,9 +524,10 @@ def supervise_output(args: list[str]) -> int: except KeyboardInterrupt: # The child gets the same interrupt and shuts down on its own. continue - exited_at = time.monotonic() - for pump in pumps: - pump.drain(exited_at) + wall_deadline = time.monotonic() + _DRAIN_WALL_SECONDS + snapshots = [pump.writing_snapshot() for pump in pumps] + for pump, (exited_at, writing_at_exit) in zip(pumps, snapshots, strict=True): + pump.drain(exited_at, writing_at_exit, wall_deadline) # A child killed by a signal reports -signum; shells report 128 + signum. return returncode if returncode >= 0 else 128 - returncode diff --git a/tests/units/reflex_base/utils/test_log.py b/tests/units/reflex_base/utils/test_log.py index 5748ccd9510..0c062ffab36 100644 --- a/tests/units/reflex_base/utils/test_log.py +++ b/tests/units/reflex_base/utils/test_log.py @@ -1002,32 +1002,117 @@ def test_supervise_output_turns_every_line_into_json(tmp_path): from reflex_base.utils import log if not log.is_output_supervised(): - sys.exit(log.supervise_output([sys.executable, __file__])) + sys.exit(log.supervise_output([sys.executable, __file__, *sys.argv[1:]])) for i in range(200): - print(f"line {i:03d} " + "x" * 1000) + print(f"line {i:03d} " + "x" * 1000, file=getattr(sys, sys.argv[1])) """ -def test_supervise_output_drains_everything_for_a_slow_consumer(tmp_path): +@pytest.mark.parametrize("drain_max_seconds", [log._DRAIN_MAX_SECONDS, 0.05]) +@pytest.mark.parametrize("stream", ["stdout", "stderr"]) +def test_supervise_output_drains_everything_for_a_slow_consumer( + tmp_path, drain_max_seconds, stream +): """Output still in the pipes at exit reaches a consumer that reads late.""" script = tmp_path / "burst.py" - script.write_text(_BURST_SCRIPT) + script.write_text( + f"from reflex_base.utils import log\n" + f"log._DRAIN_MAX_SECONDS = {drain_max_seconds}\n" + _BURST_SCRIPT + ) with subprocess.Popen( - [sys.executable, str(script)], - stdout=subprocess.PIPE, + [sys.executable, str(script), stream], + stdout=subprocess.PIPE if stream == "stdout" else None, + stderr=subprocess.PIPE if stream == "stderr" else None, env={**os.environ, "REFLEX_LOG_JSON": "true"}, ) as proc: - assert proc.stdout is not None + output = proc.stdout if stream == "stdout" else proc.stderr + assert output is not None lines = [] # A consumer slower than the writer keeps the pipes full at exit. - for line in proc.stdout: + for line in output: lines.append(line) time.sleep(0.015) proc.wait(timeout=30) + assert proc.returncode == 0 assert len(lines) == 200 assert json.loads(lines[-1])["message"].startswith("line 199 ") +def test_output_pump_write_snapshots_use_a_consistent_cutoff(monkeypatch): + """Snapshots measure a blocked or completed write at their own cutoff.""" + pump = log._OutputPump(0, 1, "info", "stdout") + snapshots = [] + + def write(fd, data): + """Capture a snapshot while the simulated write is in progress. + + Args: + fd: The output file descriptor. + data: The bytes to write. + + Returns: + The number of bytes written. + """ + snapshots.append(pump.writing_snapshot()) + return len(data) + + monkeypatch.setattr( + time, "monotonic", mock.Mock(side_effect=[10.0, 15.0, 20.0, 25.0]) + ) + monkeypatch.setattr(os, "write", write) + pump._write([b"line"], [], final=False) + assert snapshots == [(15.0, 5.0)] + assert pump.writing_snapshot() == (25.0, 10.0) + + +_STALLED_CONSUMER_SCRIPT = """ +import os +import subprocess +import sys +from pathlib import Path + +from reflex_base.utils import log + +log._DRAIN_WALL_SECONDS = 0.5 +if not log.is_output_supervised(): + sys.exit(log.supervise_output([sys.executable, __file__])) +writer = subprocess.Popen([ + sys.executable, + "-c", + "import os\\nwhile True:\\n os.write(int(os.environ['TEST_OUTPUT_FD']), (b'x' * 1023 + b'\\\\n') * 64)", +]) +Path("writer.pid").write_text(str(writer.pid)) +""" + + +@pytest.mark.parametrize("stream_fd", [1, 2]) +def test_supervise_output_bounds_shutdown_for_a_stalled_consumer(tmp_path, stream_fd): + """A lingering writer and an unread output pipe cannot prevent shutdown.""" + script = tmp_path / "stalled.py" + script.write_text(_STALLED_CONSUMER_SCRIPT) + with subprocess.Popen( + [sys.executable, str(script)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + cwd=tmp_path, + env={ + **os.environ, + "REFLEX_LOG_JSON": "true", + "TEST_OUTPUT_FD": str(stream_fd), + }, + ) as proc: + try: + assert proc.wait(timeout=5) == 0 + finally: + if proc.poll() is None: + proc.kill() + proc.wait(timeout=5) + pid_file = tmp_path / "writer.pid" + if pid_file.exists(): + with contextlib.suppress(OSError): + os.kill(int(pid_file.read_text()), signal.SIGTERM) + + _LINGERING_SCRIPT = """ import subprocess import sys