You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2026/05/10 18:45:01 UTC
(airflow) branch v3-2-test updated: Fix /tmp file leak when API server streams large task logs (#66450) (#66667)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch v3-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-2-test by this push:
new ad08d07aea5 Fix /tmp file leak when API server streams large task logs (#66450) (#66667)
ad08d07aea5 is described below
commit ad08d07aea5546d510c7062b9cc9e78952d71b14
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Sun May 10 20:44:53 2026 +0200
Fix /tmp file leak when API server streams large task logs (#66450) (#66667)
(cherry picked from commit 31d4837c5dc1c40f8386f0ba7093d5167fe7dcda)
Co-authored-by: Dheeraj Turaga <dh...@gmail.com>
---
.../airflow/utils/log/log_stream_accumulator.py | 52 ++++++++++--
.../unit/utils/log/test_stream_accumulator.py | 96 ++++++++++++++++++++--
2 files changed, 136 insertions(+), 12 deletions(-)
diff --git a/airflow-core/src/airflow/utils/log/log_stream_accumulator.py b/airflow-core/src/airflow/utils/log/log_stream_accumulator.py
index 953b47dd971..343a9f35139 100644
--- a/airflow-core/src/airflow/utils/log/log_stream_accumulator.py
+++ b/airflow-core/src/airflow/utils/log/log_stream_accumulator.py
@@ -18,8 +18,10 @@
from __future__ import annotations
+import logging
import os
import tempfile
+import weakref
from itertools import islice
from typing import IO, TYPE_CHECKING
@@ -31,6 +33,24 @@ if TYPE_CHECKING:
StructuredLogStream,
)
+logger = logging.getLogger(__name__)
+
+
+def _safe_remove(path: str) -> None:
+ """
+ Remove path, swallowing missing-file errors and logging others.
+
+ Module-level rather than a method so weakref.finalize can hold it
+ without retaining a strong reference to the owning accumulator (which
+ would defeat the backstop).
+ """
+ try:
+ os.remove(path)
+ except FileNotFoundError:
+ pass
+ except OSError as exc:
+ logger.debug("LogStreamAccumulator: could not remove temp file %s: %s", path, exc)
+
class LogStreamAccumulator:
"""
@@ -69,11 +89,17 @@ class LogStreamAccumulator:
self._buffer: list[StructuredLogMessage] = []
self._disk_lines: int = 0
self._tmpfile: IO[str] | None = None
+ self._stream_accessed: bool = False
+ self._finalizer: weakref.finalize | None = None
def _flush_buffer_to_disk(self) -> None:
"""Flush the buffer contents to a temporary file on disk."""
if self._tmpfile is None:
self._tmpfile = tempfile.NamedTemporaryFile(delete=False, mode="w+", encoding="utf-8")
+ # Backstop: guarantee the spill file is removed when this accumulator
+ # is garbage-collected, even if the caller abandons the stream
+ # generator (e.g. client disconnect on a StreamingResponse).
+ self._finalizer = weakref.finalize(self, _safe_remove, self._tmpfile.name)
self._disk_lines += len(self._buffer)
self._tmpfile.writelines(f"{log.model_dump_json()}\n" for log in self._buffer)
@@ -91,12 +117,16 @@ class LogStreamAccumulator:
self._flush_buffer_to_disk()
def _cleanup(self) -> None:
- """Clean up the temporary file if it exists."""
+ """Release the temp file and detach the finalizer; safe to call repeatedly."""
self._buffer.clear()
- if self._tmpfile:
+ if self._tmpfile is not None:
+ path = self._tmpfile.name
self._tmpfile.close()
- os.remove(self._tmpfile.name)
+ _safe_remove(path)
self._tmpfile = None
+ if self._finalizer is not None:
+ self._finalizer.detach()
+ self._finalizer = None
@property
def total_lines(self) -> int:
@@ -119,6 +149,10 @@ class LogStreamAccumulator:
Returns:
A stream of the captured log messages.
"""
+ self._stream_accessed = True
+ return self._iter_stream()
+
+ def _iter_stream(self) -> StructuredLogStream:
try:
if not self._tmpfile:
# if no temporary file was created, return from the buffer
@@ -147,9 +181,13 @@ class LogStreamAccumulator:
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
"""
- Context manager exit that doesn't perform resource cleanup.
+ Context manager exit.
- Note: Resources are not cleaned up here. Cleanup is deferred until
- get_stream() is called and fully consumed, ensuring all logs are properly
- yielded before cleanup occurs.
+ Cleans up immediately when an exception is propagating, or when the
+ caller never accessed .stream (so no generator exists to run the
+ deferred finally block). When the caller did access .stream,
+ cleanup is left to the generator's finally (happy path) with
+ weakref.finalize as a backstop if the generator is abandoned.
"""
+ if exc_type is not None or not self._stream_accessed:
+ self._cleanup()
diff --git a/airflow-core/tests/unit/utils/log/test_stream_accumulator.py b/airflow-core/tests/unit/utils/log/test_stream_accumulator.py
index b1d2df09c88..e813c5f19cd 100644
--- a/airflow-core/tests/unit/utils/log/test_stream_accumulator.py
+++ b/airflow-core/tests/unit/utils/log/test_stream_accumulator.py
@@ -17,7 +17,9 @@
from __future__ import annotations
+import gc
import os
+import weakref
from typing import TYPE_CHECKING
from unittest import mock
@@ -74,15 +76,18 @@ class TestLogStreamAccumulator:
with accumulator:
mock_setup.assert_called_once()
- def test__flush_buffer_to_disk(self, structured_logs):
+ def test__flush_buffer_to_disk(self, structured_logs, tmp_path):
"""Test flush-to-disk behavior with a small threshold."""
threshold = 6
- # Mock the temporary file to verify it's being written to
- with (
- mock.patch("tempfile.NamedTemporaryFile") as mock_tmpfile,
- ):
+ # Use a real path on tmp_path so __exit__ cleanup (which now runs when
+ # the stream is never accessed) can remove the file without raising.
+ fake_spill = tmp_path / "fake_spill.json"
+ fake_spill.touch()
+
+ with mock.patch("tempfile.NamedTemporaryFile") as mock_tmpfile:
mock_file = mock.MagicMock()
+ mock_file.name = str(fake_spill)
mock_tmpfile.return_value = mock_file
with LogStreamAccumulator(structured_logs, threshold) as accumulator:
@@ -163,3 +168,84 @@ class TestLogStreamAccumulator:
# After fully consuming the stream, cleanup should be called
mock_cleanup.assert_called_once()
+
+ def test_cleanup_when_stream_never_accessed(self, structured_logs):
+ """Temp file is removed on __exit__ when caller never reads .stream (e.g. uses only total_lines)."""
+
+ with LogStreamAccumulator(structured_logs, 5) as accumulator:
+ assert accumulator._tmpfile is not None
+ tmpfile_name = accumulator._tmpfile.name
+ assert os.path.exists(tmpfile_name)
+ assert accumulator.total_lines == LOG_COUNT
+
+ assert accumulator._tmpfile is None
+ assert not os.path.exists(tmpfile_name)
+
+ def test_cleanup_on_exception_in_with_block(self, structured_logs):
+ """An exception propagating through __exit__ overrides the deferred-cleanup path."""
+
+ accumulator = LogStreamAccumulator(structured_logs, 5)
+ accumulator.__enter__()
+ _ = accumulator.stream
+ assert accumulator._tmpfile is not None
+ tmpfile_name = accumulator._tmpfile.name
+ assert os.path.exists(tmpfile_name)
+
+ accumulator.__exit__(RuntimeError, RuntimeError("boom"), None)
+
+ assert not os.path.exists(tmpfile_name)
+ assert accumulator._tmpfile is None
+
+ def test_cleanup_deferred_when_stream_returned_then_iterated(self, structured_logs):
+ """Stream returned from inside the with block must remain readable after __exit__."""
+
+ with LogStreamAccumulator(structured_logs, 5) as accumulator:
+ returned_stream = accumulator.stream
+ assert accumulator._tmpfile is not None
+ tmpfile_name = accumulator._tmpfile.name
+ assert os.path.exists(tmpfile_name)
+
+ assert os.path.exists(tmpfile_name)
+
+ self.validate_log_stream(returned_stream)
+
+ assert not os.path.exists(tmpfile_name)
+
+ def test_cleanup_via_finalizer_on_abandoned_generator(self, structured_logs):
+ """When _cleanup is suppressed, the weakref backstop still removes the spill on GC."""
+
+ def make_and_abandon() -> tuple[str, weakref.ref]:
+ acc = LogStreamAccumulator(structured_logs, 5)
+ # Disable explicit cleanup so only the weakref backstop can remove the file.
+ acc._cleanup = lambda: None # type: ignore[method-assign]
+ acc.__enter__()
+ _gen = acc.stream
+ assert acc._tmpfile is not None
+ path = acc._tmpfile.name
+ acc.__exit__(None, None, None)
+ # _gen, acc go out of scope on return — finalizer should fire.
+ return path, weakref.ref(acc)
+
+ tmpfile_name, acc_ref = make_and_abandon()
+
+ gc.collect()
+
+ assert acc_ref() is None, "accumulator should have been garbage collected"
+ assert not os.path.exists(tmpfile_name), "weakref.finalize backstop should have removed the file"
+
+ def test_cleanup_idempotent(self, structured_logs):
+ """Repeated _cleanup calls do not raise."""
+
+ accumulator = LogStreamAccumulator(structured_logs, 5)
+ accumulator._capture()
+ assert accumulator._tmpfile is not None
+ tmpfile_name = accumulator._tmpfile.name
+ assert os.path.exists(tmpfile_name)
+
+ accumulator._cleanup()
+ assert not os.path.exists(tmpfile_name)
+ assert accumulator._tmpfile is None
+ assert accumulator._finalizer is None
+
+ accumulator._cleanup()
+ accumulator._cleanup()