fix: preserve JSON output for slow consumers at shutdown (#7428)

* fix: preserve JSON output for slow consumers at shutdown

* fix: bound output draining and synchronize shutdown snapshots

* ci: rerun PR checks

Author: masenf

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.

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
 

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