Fix CloudWatch trigger stream lifecycle

providers/amazon/src/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py

@@ -24,6 +24,7 @@
 import logging
 import os
 import shutil
+import threading
 from collections.abc import Generator
 from datetime import date, datetime, timedelta, timezone
 from functools import cached_property
@@ -97,6 +98,13 @@ class CloudWatchRemoteLogIO(LoggingMixin):  # noqa: D101
     _cached_handler: watchtower.CloudWatchLogHandler | None = attrs.field(
         init=False, default=None, repr=False
     )
+    _stream_handlers: dict[str, watchtower.CloudWatchLogHandler] = attrs.field(
+        init=False, factory=dict, repr=False
+    )
+    _stream_lock: threading.RLock = attrs.field(init=False, factory=threading.RLock, repr=False)
+    _closing_streams: set[str] = attrs.field(init=False, factory=set, repr=False)
+    _building_stream_handler: bool = attrs.field(init=False, default=False, repr=False)
+    _streaming_by_path: bool = attrs.field(init=False, default=False, repr=False)
     _closed: bool = attrs.field(init=False, default=False, repr=False)
 
     @log_group.default
@@ -148,11 +156,14 @@ def hook(self):
             aws_conn_id=conf.get("logging", "remote_log_conn_id"), region_name=self.region_name
         )
 
-    def _build_handler(self) -> watchtower.CloudWatchLogHandler:
+    def _build_handler(self, stream_name: str | None = None) -> watchtower.CloudWatchLogHandler:
+        if stream_name is None:
+            stream_name = self.log_stream_name
         _json_serialize = conf.getimport("aws", "cloudwatch_task_handler_json_serializer", fallback=None)
         return watchtower.CloudWatchLogHandler(
             log_group_name=self.log_group,
-            log_stream_name=self.log_stream_name,
+            log_stream_name=stream_name,
+            create_log_group=stream_name is None,
             use_queues=True,
             boto3_client=self.hook.get_conn(),
             json_serialize_default=_json_serialize or json_serialize_legacy,
@@ -168,32 +179,73 @@ def handler(self) -> watchtower.CloudWatchLogHandler:
         Rebuild only while the IO is live: once :meth:`close` has run, keep the closed handler
         so a late record is dropped instead of spawning an orphan handler and background thread.
         """
-        if self._cached_handler is None or (not self._closed and self._cached_handler.shutting_down):
-            self._cached_handler = self._build_handler()
-        return self._cached_handler
+        with self._stream_lock:
+            if self._cached_handler is None or (not self._closed and self._cached_handler.shutting_down):
+                self._cached_handler = self._build_handler()
+            return self._cached_handler
+
+    def _get_stream_handler(self, stream_name: str) -> watchtower.CloudWatchLogHandler | None:
+        """Return the live handler for ``stream_name`` while holding ``_stream_lock``."""
+        if self._closed or self._building_stream_handler or stream_name in self._closing_streams:
+            return None
+
+        handler = self._stream_handlers.get(stream_name)
+        if handler is not None and not handler.shutting_down:
+            return handler
+
+        self._stream_handlers.pop(stream_name, None)
+        if (
+            not self._stream_handlers
+            and self._cached_handler is not None
+            and not self._cached_handler.shutting_down
+        ):
+            handler = self._cached_handler
+            handler.log_stream_name = stream_name
+        else:
+            self._building_stream_handler = True
+            try:
+                handler = self._build_handler(stream_name)
+            finally:
+                self._building_stream_handler = False
+        self._stream_handlers[stream_name] = handler
+        self._cached_handler = handler
+        return handler
+
+    def _close_stream(self, stream_name: str) -> None:
+        with self._stream_lock:
+            handler = self._stream_handlers.pop(stream_name, None)
+            if handler is not None:
+                self._closing_streams.add(stream_name)
+                if self._cached_handler is handler:
+                    self._cached_handler = next(reversed(self._stream_handlers.values()), None)
+
+        if handler is None:
+            self.log.debug("No active CloudWatch handler for completed stream %s", stream_name)
+            return
+
+        try:
+            handler.close()
+        finally:
+            with self._stream_lock:
+                self._closing_streams.discard(stream_name)
 
     @cached_property
     def processors(self) -> tuple[structlog.typing.Processor, ...]:
         from logging import getLogRecordFactory
 
         import structlog.stdlib
 
+        self._streaming_by_path = True
         logRecordFactory = getLogRecordFactory()
         # The handler MUST be initted here, before the processor is actually used to log anything.
         # Otherwise, logging that occurs during the creation of the handler can create infinite loops.
         _ = self.handler
         from airflow.sdk.log import relative_path_from_logger
 
         def proc(logger: structlog.typing.WrappedLogger, method_name: str, event: structlog.typing.EventDict):
-            if not logger or not (stream_name := relative_path_from_logger(logger)):
+            if not logger or not (stream_path := relative_path_from_logger(logger)):
                 return event
-            # Resolve the handler on every record: configure_logging() may have
-            # closed the one built above, in which case ``handler`` rebuilds it.
-            handler = self.handler
-            # We can't set the log stream name in the above init handler because
-            # the log path isn't known at that stage.
-            # Instead, we should always rely on the path (log stream name) provided by the logger.
-            handler.log_stream_name = stream_name.as_posix().replace(":", "_")
+            stream_name = stream_path.as_posix().replace(":", "_")
             name = event.get("logger_name") or event.get("logger", "")
             level = structlog.stdlib.NAME_TO_LEVEL.get(method_name.lower(), logging.INFO)
             msg = copy.copy(event)
@@ -208,7 +260,9 @@ def proc(logger: structlog.typing.WrappedLogger, method_name: str, event: struct
                 ct = created.timestamp()
                 record.created = ct
                 record.msecs = int((ct - int(ct)) * 1000) + 0.0  # Copied from stdlib logging
-            handler.handle(record)
+            with self._stream_lock:
+                if handler := self._get_stream_handler(stream_name):
+                    handler.handle(record)
             return event
 
         return (proc,)
@@ -217,21 +271,30 @@ def close(self):
         """
         Flush pending events one last time and mark the IO closed.
 
-        Only ever called from :meth:`upload`. Mark the IO closed first so ``handler`` stops
-        rebuilding: a record arriving after teardown must be dropped, not revive a fresh
-        handler. Read the cached handler directly so we never build one just to flush it.
+        Called from :meth:`upload` when logs are not streamed by path. Mark the IO closed
+        first so ``handler`` stops rebuilding: a record arriving after teardown must be
+        dropped, not revive a fresh handler. Read the cached handlers directly so we never
+        build one just to flush it.
         """
-        self._closed = True
-        handler = self._cached_handler
-        if handler is None or handler.shutting_down:
-            return
-
-        handler.flush()
+        with self._stream_lock:
+            self._closed = True
+            handlers = {
+                id(handler): handler
+                for handler in (*self._stream_handlers.values(), self._cached_handler)
+                if handler is not None
+            }.values()
+        for handler in handlers:
+            if not handler.shutting_down:
+                handler.flush()
 
     def upload(self, path: os.PathLike | str, ti: RuntimeTI | None = None) -> None:
         """Upload the given log path to the remote storage."""
-        # No batch upload — logs stream in real-time. Flush pending events and clean up.
-        self.close()
+        # No batch upload — logs stream in real-time. Close the completed stream and clean up.
+        if self._streaming_by_path:
+            stream_name = Path(path).as_posix().replace(":", "_")
+            self._close_stream(stream_name)
+        else:
+            self.close()
         if self.delete_local_copy:
             base = self.base_log_folder.resolve()
             raw = Path(path)

providers/amazon/tests/unit/amazon/aws/log/test_cloudwatch_task_handler.py

@@ -21,6 +21,7 @@
 import logging
 import os
 import textwrap
+import threading
 import time
 from datetime import datetime as dt, timedelta, timezone as std_timezone
 from pathlib import Path
@@ -30,6 +31,7 @@
 import boto3
 import pendulum
 import pytest
+import structlog
 import time_machine
 from botocore.exceptions import ClientError
 from moto import mock_aws
@@ -386,6 +388,164 @@ def test_handler_not_rebuilt_after_close(self):
             assert self.subject.handler is original
             assert self.subject.handler.shutting_down is True
 
+    def test_upload_releases_only_the_completed_stream_handler(self):
+        completed_path = "dag_id=a/completed.log"
+        active_path = "dag_id=a/active.log"
+        completed = mock.create_autospec(CloudWatchLogHandler, instance=True)
+        active = mock.create_autospec(CloudWatchLogHandler, instance=True)
+        completed.shutting_down = False
+        active.shutting_down = False
+        self.subject._streaming_by_path = True
+        self.subject._stream_handlers = {completed_path: completed, active_path: active}
+        self.subject._cached_handler = active
+        self.subject.delete_local_copy = False
+
+        with mock.patch.object(self.subject.log, "debug", autospec=True) as log_debug:
+            self.subject.upload(completed_path, self.ti)
+            self.subject.upload(completed_path, self.ti)
+
+        completed.close.assert_called_once_with()
+        active.close.assert_not_called()
+        assert self.subject._stream_handlers == {active_path: active}
+        assert self.subject._cached_handler is active
+        with self.subject._stream_lock:
+            assert self.subject._get_stream_handler(active_path) is active
+        log_debug.assert_called_once_with(
+            "No active CloudWatch handler for completed stream %s", completed_path
+        )
+
+    def test_completed_stream_path_can_be_reused(self):
+        stream_name = "dag_id=a/reused.log"
+        first = mock.create_autospec(CloudWatchLogHandler, instance=True)
+        second = mock.create_autospec(CloudWatchLogHandler, instance=True)
+        first.shutting_down = False
+        second.shutting_down = False
+        self.subject._cached_handler = None
+
+        with mock.patch.object(self.subject, "_build_handler", side_effect=[first, second]):
+            with self.subject._stream_lock:
+                assert self.subject._get_stream_handler(stream_name) is first
+            self.subject._close_stream(stream_name)
+            with self.subject._stream_lock:
+                assert self.subject._get_stream_handler(stream_name) is second
+
+        first.close.assert_called_once_with()
+        second.close.assert_not_called()
+
+    def test_many_completed_streams_do_not_accumulate_handlers(self):
+        handlers = []
+        processor = self.subject.processors[0]
+        if initial_handler := self.subject._cached_handler:
+            initial_handler.close()
+        self.subject._cached_handler = None
+        self.subject.delete_local_copy = False
+
+        def build_handler(_stream_name):
+            handler = mock.create_autospec(CloudWatchLogHandler, instance=True)
+            handler.shutting_down = False
+            handlers.append(handler)
+            return handler
+
+        with (
+            conf_vars({("logging", "base_log_folder"): self.local_log_location.as_posix()}),
+            mock.patch.object(self.subject, "_build_handler", side_effect=build_handler),
+        ):
+            for index in range(100):
+                stream_name = f"dag_id=a/trigger-{index}.log"
+                local_path = self.local_log_location / stream_name
+                local_path.parent.mkdir(parents=True, exist_ok=True)
+                with local_path.open("w") as log_file:
+                    processor(structlog.PrintLogger(log_file), "info", {"event": f"message-{index}"})
+                self.subject.upload(stream_name, self.ti)
+
+        assert self.subject._stream_handlers == {}
+        assert self.subject._cached_handler is None
+        assert all(handler.close.call_count == 1 for handler in handlers)
+        assert all(handler.handle.call_count == 1 for handler in handlers)
+
+    def test_record_arriving_while_stream_closes_is_dropped(self):
+        stream_name = "dag_id=a/closing.log"
+        local_path = self.local_log_location / stream_name
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+        local_path.touch()
+        processor = self.subject.processors[0]
+        handler = mock.create_autospec(CloudWatchLogHandler, instance=True)
+        handler.shutting_down = False
+        close_started = threading.Event()
+        finish_close = threading.Event()
+
+        def close_handler():
+            close_started.set()
+            assert finish_close.wait(timeout=5)
+
+        handler.close.side_effect = close_handler
+        self.subject._stream_handlers = {stream_name: handler}
+        self.subject._cached_handler = handler
+        self.subject.delete_local_copy = False
+
+        upload_thread = threading.Thread(target=self.subject.upload, args=(stream_name, self.ti))
+        upload_thread.start()
+        assert close_started.wait(timeout=5)
+        try:
+            with (
+                conf_vars({("logging", "base_log_folder"): self.local_log_location.as_posix()}),
+                mock.patch.object(self.subject, "_build_handler") as build_handler,
+                local_path.open("w") as log_file,
+            ):
+                processor(structlog.PrintLogger(log_file), "info", {"event": "late message"})
+                build_handler.assert_not_called()
+        finally:
+            finish_close.set()
+            upload_thread.join(timeout=5)
+
+        assert not upload_thread.is_alive()
+        handler.handle.assert_not_called()
+        assert self.subject._closing_streams == set()
+
+    def test_record_during_stream_handler_construction_is_dropped(self):
+        stream_name = "dag_id=a/building.log"
+        local_path = self.local_log_location / stream_name
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+        processor = self.subject.processors[0]
+        if initial_handler := self.subject._cached_handler:
+            initial_handler.close()
+        self.subject._cached_handler = None
+        handler = mock.create_autospec(CloudWatchLogHandler, instance=True)
+        handler.shutting_down = False
+
+        with (
+            conf_vars({("logging", "base_log_folder"): self.local_log_location.as_posix()}),
+            local_path.open("w") as log_file,
+        ):
+            logger = structlog.PrintLogger(log_file)
+
+            def build_handler(_stream_name):
+                processor(logger, "info", {"event": "record from handler construction"})
+                return handler
+
+            with mock.patch.object(self.subject, "_build_handler", side_effect=build_handler) as build:
+                processor(logger, "info", {"event": "task record"})
+
+        build.assert_called_once_with(stream_name)
+        handler.handle.assert_called_once_with(ANY)
+        assert self.subject._stream_handlers == {stream_name: handler}
+
+    def test_close_flushes_each_active_stream_handler_once(self):
+        first = mock.create_autospec(CloudWatchLogHandler, instance=True)
+        second = mock.create_autospec(CloudWatchLogHandler, instance=True)
+        first.shutting_down = False
+        second.shutting_down = False
+        self.subject._stream_handlers = {"first.log": first, "second.log": second}
+        self.subject._cached_handler = second
+
+        self.subject.close()
+
+        first.flush.assert_called_once_with()
+        second.flush.assert_called_once_with()
+        first.close.assert_not_called()
+        second.close.assert_not_called()
+        assert self.subject._closed is True
+
 
 @pytest.mark.db_test
 class TestCloudwatchTaskHandler: