diff --git a/agent/thread_scoped_output.py b/agent/thread_scoped_output.py index e46608e492..3c4a7be891 100644 --- a/agent/thread_scoped_output.py +++ b/agent/thread_scoped_output.py @@ -30,6 +30,20 @@ _install_lock = threading.Lock() # Maps the proxy we installed for a given attribute ("stdout"/"stderr") so we # never double-wrap and so we can recover the original stream. _installed: dict[str, "_ThreadRoutingStream"] = {} +# One process-lifetime sink per stream. Temporary process-global redirects can +# displace and later restore a routing proxy; they must not allocate another +# permanent /dev/null descriptor every time that happens. +_sinks: dict[str, TextIO] = {} +_routing_states: dict[str, "_RoutingState"] = {} + + +class _RoutingState: + """Silencing registry shared by every proxy generation for one stream.""" + + def __init__(self, sink: TextIO) -> None: + self.sink = sink + self.silenced: dict[int, int] = {} + self.lock = threading.Lock() class _ThreadRoutingStream: @@ -42,32 +56,27 @@ class _ThreadRoutingStream: ``.fileno()`` behave like the underlying stream for the calling thread. """ - def __init__(self, passthrough: TextIO, sink: TextIO) -> None: + def __init__(self, passthrough: TextIO, state: _RoutingState) -> None: self._passthrough = passthrough - self._sink = sink - # ident -> nesting depth. A thread is silenced while depth > 0, so - # nested ``thread_scoped_silence()`` on the same thread composes - # correctly (the inner exit decrements rather than fully clearing). - self._silenced: dict[int, int] = {} - self._lock = threading.Lock() + self._state = state def _target(self) -> TextIO: - if self._silenced.get(threading.get_ident(), 0) > 0: - return self._sink + if self._state.silenced.get(threading.get_ident(), 0) > 0: + return self._state.sink return self._passthrough # --- registration ----------------------------------------------------- def silence(self, ident: int) -> None: - with self._lock: - self._silenced[ident] = self._silenced.get(ident, 0) + 1 + with self._state.lock: + self._state.silenced[ident] = self._state.silenced.get(ident, 0) + 1 def unsilence(self, ident: int) -> None: - with self._lock: - depth = self._silenced.get(ident, 0) - 1 + with self._state.lock: + depth = self._state.silenced.get(ident, 0) - 1 if depth > 0: - self._silenced[ident] = depth + self._state.silenced[ident] = depth else: - self._silenced.pop(ident, None) + self._state.silenced.pop(ident, None) # --- file-like surface ------------------------------------------------ def write(self, data): # type: ignore[no-untyped-def] @@ -109,14 +118,28 @@ def _ensure_installed(attr: str, passthrough: TextIO) -> "_ThreadRoutingStream": with _install_lock: proxy = _installed.get(attr) current = getattr(sys, attr, None) + if isinstance(current, _ThreadRoutingStream): + # A redirect context can restore an older routing proxy after a + # temporary replacement. Adopt it instead of wrapping it and + # growing an unbounded proxy chain. + _installed[attr] = current + _routing_states[attr] = current._state + return current if proxy is not None and current is proxy: return proxy # Capture whatever is currently bound as the passthrough. If a prior # global redirect_stdout is active, route non-silenced threads to that # stream to preserve the old behavior. passthrough = current if current is not None else passthrough - sink = open(os.devnull, "w", encoding="utf-8") - proxy = _ThreadRoutingStream(passthrough, sink) + sink = _sinks.get(attr) + if sink is None or sink.closed: + sink = open(os.devnull, "w", encoding="utf-8") + _sinks[attr] = sink + state = _routing_states.get(attr) + if state is None or state.sink is not sink: + state = _RoutingState(sink) + _routing_states[attr] = state + proxy = _ThreadRoutingStream(passthrough, state) setattr(sys, attr, proxy) _installed[attr] = proxy return proxy diff --git a/tests/agent/test_thread_scoped_output.py b/tests/agent/test_thread_scoped_output.py index 7f85e7d5f8..fa22921dd9 100644 --- a/tests/agent/test_thread_scoped_output.py +++ b/tests/agent/test_thread_scoped_output.py @@ -7,11 +7,13 @@ context. This is the property the old process-global ``contextlib.redirect_stdout(devnull)`` violated (issue #55769 / #55925). """ +import contextlib import io import sys import threading import time +import agent.thread_scoped_output as thread_output from agent.thread_scoped_output import thread_scoped_silence @@ -94,3 +96,73 @@ def test_repeated_contexts_never_write_to_a_closed_sink(): sys.stdout.fileno() finally: sys.stdout = original + + +def test_temporary_global_redirects_do_not_allocate_new_sinks(monkeypatch): + """A displaced proxy is temporary, not a reason to leak another FD pair.""" + opened_sinks = [] + + def fake_open(*_args, **_kwargs): + sink = io.StringIO() + opened_sinks.append(sink) + return sink + + monkeypatch.setattr(thread_output, "_installed", {}) + monkeypatch.setattr(thread_output, "_sinks", {}, raising=False) + monkeypatch.setattr(thread_output, "open", fake_open, raising=False) + original_stdout, original_stderr = sys.stdout, sys.stderr + sys.stdout, sys.stderr = io.StringIO(), io.StringIO() + try: + with thread_scoped_silence(): + pass + assert len(opened_sinks) == 2 + original_proxies = dict(thread_output._installed) + + for _ in range(20): + with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + with thread_scoped_silence(): + print("hidden") + + with thread_scoped_silence(): + pass + assert len(opened_sinks) == 2 + assert thread_output._installed == original_proxies + finally: + sys.stdout, sys.stderr = original_stdout, original_stderr + + +def test_silence_survives_redirect_restoring_an_older_proxy(monkeypatch): + """Silencing is stream-wide, even when a redirect swaps proxy generations.""" + monkeypatch.setattr(thread_output, "_installed", {}) + monkeypatch.setattr(thread_output, "_sinks", {}, raising=False) + original_stdout, original_stderr = sys.stdout, sys.stderr + passthrough = io.StringIO() + sys.stdout = passthrough + entered = threading.Event() + release = threading.Event() + + try: + with thread_scoped_silence(): + pass + + def worker(): + with thread_scoped_silence(): + entered.set() + assert release.wait(timeout=10) + print("must-stay-silenced") + + redirected = io.StringIO() + with contextlib.redirect_stdout(redirected): + thread = threading.Thread(target=worker) + thread.start() + assert entered.wait(timeout=10) + + release.set() + thread.join(timeout=10) + + assert not thread.is_alive() + assert "must-stay-silenced" not in passthrough.getvalue() + assert "must-stay-silenced" not in redirected.getvalue() + finally: + release.set() + sys.stdout, sys.stderr = original_stdout, original_stderr