fix(log): redirect the run log when the first log stream comes back empty

Author: vdusek

src/apify_client/_streamed_log.py

@@ -32,15 +32,24 @@ class StreamedLogBase:
     duration of the run (Impit currently maps it to an effective 24-hour cap) and mirrors the JS client.
     """
 
+    _empty_stream_retry_s: ClassVar[float] = 0.5
+    """Pause before reopening a log stream that ended before the run logged anything.
+
+    The API serves the log of a run that has not logged anything yet as an empty stream that ends at once.
+    """
+
     def __init__(self, to_logger: logging.Logger, *, from_start: bool = True) -> None:
         if self._force_propagate:
             to_logger.propagate = True
         self._to_logger = to_logger
         self._stream_buffer = list[bytes]()
         self._split_marker = re.compile(rb'(?:\n|^)(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z)')
         self._relevancy_time_limit: datetime | None = None if from_start else datetime.now(tz=UTC)
+        self._received_data = False
 
     def _process_new_data(self, data: bytes) -> None:
+        if data:
+            self._received_data = True
         new_chunk = data
         self._stream_buffer.append(new_chunk)
         if re.findall(self._split_marker, new_chunk):
@@ -75,6 +84,12 @@ def _log_buffer_content(self, *, include_last_part: bool = False) -> None:
             message = decoded_marker + decoded_content
             self._to_logger.log(level=self._guess_log_level_from_message(message), msg=message.strip())
 
+    def _process_whole_log(self, log: bytes | None) -> None:
+        """Redirect a log read in one request, for a stop that came before the stream delivered anything."""
+        if log:
+            self._process_new_data(log)
+            self._log_buffer_content(include_last_part=True)
+
     @staticmethod
     def _guess_log_level_from_message(message: str) -> int:
         """Guess the log level from the message."""
@@ -121,6 +136,7 @@ def __init__(self, log_client: LogClient, *, to_logger: logging.Logger, from_sta
         self._streaming_thread: Thread | None = None
         self._log_stream: HttpResponse | None = None
         self._stop_logging = False
+        self._stop_event = threading.Event()
 
     def start(self) -> Thread:
         """Start the streaming thread.
@@ -130,6 +146,7 @@ def start(self) -> Thread:
         if self._streaming_thread and self._streaming_thread.is_alive():
             raise RuntimeError('Streaming thread already active')
         self._stop_logging = False
+        self._stop_event.clear()
         # A daemon thread so a stream still blocked on a read can never hold up interpreter shutdown.
         self._streaming_thread = threading.Thread(target=self._stream_log, daemon=True)
         self._streaming_thread.start()
@@ -140,11 +157,13 @@ def stop(self) -> None:
 
         A thread that outlives the wait is a daemon with `_stop_logging` set, so it exits after at most one more chunk,
         and only then does its buffered tail reach the logger. Its handle is kept while it is alive, so `start` cannot
-        revive it beside a second thread on the same buffer.
+        revive it beside a second thread on the same buffer. If no stream has delivered anything yet, the thread reads
+        the whole log in one request before it ends.
         """
         if not self._streaming_thread:
             raise RuntimeError('Streaming thread is not active')
         self._stop_logging = True
+        self._stop_event.set()
         # Read once; the streaming thread clears the attribute as soon as the stream ends.
         log_stream = self._log_stream
         if log_stream is not None:
@@ -173,39 +192,57 @@ def __exit__(
 
     def _stream_log(self) -> None:
         try:
-            with self._log_client.stream(raw=True, timeout=self._stream_timeout) as log_stream:
-                if not log_stream:
+            # An empty stream means the run has not logged anything yet, so reopen it until the first bytes arrive.
+            while not self._stop_logging:
+                if not self._stream_log_once() or self._received_data:
                     return
-                # Published so `stop` can close the response.
-                self._log_stream = log_stream
-                try:
-                    # `stop` may have run before the response existed for it to close.
-                    if self._stop_logging:
-                        return
-                    for data in log_stream.iter_bytes():
-                        self._process_new_data(data)
-                        if self._stop_logging:
-                            break
-                finally:
-                    self._log_stream = None
-                    try:
-                        # Flush the last buffered part even if the read timed out or was stopped.
-                        self._log_buffer_content(include_last_part=True)
-                    except Exception:
-                        # A truncated stream leaves an undecodable tail, which is worth a traceback even while a stop
-                        # is in progress.
-                        self._to_logger.exception('Log redirection stopped due to unexpected error:')
+                self._stop_event.wait(self._empty_stream_retry_s)
         except Exception as exc:
             if self._stop_logging:
                 # `stop` closed the stream out from under the read, so the failure is expected.
                 self._to_logger.debug('Log streaming stopped while `stop` was in progress: %r', exc)
-                return
-            if self._log_client._http_client.is_timeout_error(exc):  # noqa: SLF001
+            elif self._log_client._http_client.is_timeout_error(exc):  # noqa: SLF001
                 # The stream cannot continue, so warn and let the thread end instead of leaking a traceback.
                 self._to_logger.warning('Log streaming stopped: the log stream request timed out.')
+                return
             else:
                 # Any other failure in log redirection must not escape the background thread; log it instead.
                 self._to_logger.exception('Log redirection stopped due to unexpected error:')
+                return
+        if self._received_data:
+            return
+        # Stopped before any stream delivered a byte, which a run that finishes quickly can cause.
+        try:
+            self._process_whole_log(self._log_client.get_as_bytes(raw=True))
+        except Exception:
+            self._to_logger.exception('Log redirection stopped due to unexpected error:')
+
+    def _stream_log_once(self) -> bool:
+        """Redirect one log stream until it ends or `stop` is called. Return `False` when the log does not exist."""
+        with self._log_client.stream(raw=True, timeout=self._stream_timeout) as log_stream:
+            if not log_stream:
+                return False
+            # Published so `stop` can close the response.
+            self._log_stream = log_stream
+            try:
+                # `stop` may have run before the response existed for it to close. A stream opened this late would
+                # end after its first chunk, so the whole log is read in one request instead.
+                if self._stop_logging:
+                    return True
+                for data in log_stream.iter_bytes():
+                    self._process_new_data(data)
+                    if self._stop_logging:
+                        break
+            finally:
+                self._log_stream = None
+                try:
+                    # Flush the last buffered part even if the read timed out or was stopped.
+                    self._log_buffer_content(include_last_part=True)
+                except Exception:
+                    # A truncated stream leaves an undecodable tail, which is worth a traceback even while a stop is
+                    # in progress.
+                    self._to_logger.exception('Log redirection stopped due to unexpected error:')
+        return True
 
 
 @docs_group('Other')
@@ -244,17 +281,28 @@ def start(self) -> Task:
         return self._streaming_task
 
     async def stop(self) -> None:
-        """Stop the streaming task."""
+        """Stop the streaming task.
+
+        If no stream has delivered anything yet, read the whole log in one request instead.
+        """
         if not self._streaming_task:
             raise RuntimeError('Streaming task is not active')
 
+        was_streaming = not self._streaming_task.done()
         self._streaming_task.cancel()
         try:
             await self._streaming_task
         except asyncio.CancelledError:
             pass
         finally:
             self._streaming_task = None
+        if not was_streaming or self._received_data:
+            return
+        # Stopped before any stream delivered a byte, which a run that finishes quickly can cause.
+        try:
+            self._process_whole_log(await self._log_client.get_as_bytes(raw=True))
+        except Exception:
+            self._to_logger.exception('Log redirection stopped due to unexpected error:')
 
     async def __aenter__(self) -> Self:
         """Start the streaming task within the context. Exiting the context will cancel the streaming task."""
@@ -269,20 +317,25 @@ async def __aexit__(
 
     async def _stream_log(self) -> None:
         try:
-            async with self._log_client.stream(raw=True, timeout=self._stream_timeout) as log_stream:
-                if not log_stream:
-                    return
-                try:
-                    async for data in log_stream.aiter_bytes():
-                        self._process_new_data(data)
-                finally:
+            # An empty stream means the run has not logged anything yet, so reopen it until the first bytes arrive.
+            while True:
+                async with self._log_client.stream(raw=True, timeout=self._stream_timeout) as log_stream:
+                    if not log_stream:
+                        return
                     try:
-                        # Flush the last buffered part even if the task is cancelled by `stop()`.
-                        self._log_buffer_content(include_last_part=True)
-                    except Exception:
-                        # A truncated stream leaves an undecodable tail. Keeping the failure here also keeps the
-                        # cancellation `stop` raised propagating, so the task ends up cancelled as asyncio expects.
-                        self._to_logger.exception('Log redirection stopped due to unexpected error:')
+                        async for data in log_stream.aiter_bytes():
+                            self._process_new_data(data)
+                    finally:
+                        try:
+                            # Flush the last buffered part even if the task is cancelled by `stop()`.
+                            self._log_buffer_content(include_last_part=True)
+                        except Exception:
+                            # A truncated stream leaves an undecodable tail. Keeping the failure here also keeps the
+                            # cancellation `stop` raised propagating, so the task ends up cancelled as asyncio expects.
+                            self._to_logger.exception('Log redirection stopped due to unexpected error:')
+                if self._received_data:
+                    return
+                await asyncio.sleep(self._empty_stream_retry_s)
         except Exception as exc:
             if self._log_client._http_client.is_timeout_error(exc):  # noqa: SLF001
                 # A timeout on the long-lived stream is an expected terminal condition, not an error.

tests/unit/test_logging.py

@@ -4,6 +4,7 @@
 import itertools
 import json
 import logging
+import math
 import threading
 import time
 from datetime import datetime, timedelta
@@ -20,7 +21,7 @@
 from apify_client._streamed_log import StreamedLog, StreamedLogAsync, StreamedLogBase
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from collections.abc import Callable, Iterator
 
     from _pytest.logging import LogCaptureFixture
     from pytest_httpserver import HTTPServer
@@ -1393,6 +1394,103 @@ def test_streamed_log_sync_stop_reports_failing_stream_close(
         streaming_thread.join(timeout=5)
 
 
+def serve_log_after_empty_streams(httpserver: HTTPServer, *, empty_streams: float) -> list[Request]:
+    """Serve the mocked log, but answer the first `empty_streams` stream requests with an empty body.
+
+    That is how the API answers a log stream request before the run has logged anything. Return the list the stream
+    requests are recorded in.
+    """
+    stream_requests: list[Request] = []
+
+    def handler(request: Request) -> Response:
+        if 'stream' in request.args:
+            stream_requests.append(request)
+            if len(stream_requests) <= empty_streams:
+                return Response(b'', status=200, mimetype='application/octet-stream')
+        return Response(b''.join(_MOCKED_ACTOR_LOGS), status=200, mimetype='application/octet-stream')
+
+    httpserver.expect_request(f'/v2/actor-runs/{_MOCKED_RUN_ID}/log', method='GET').respond_with_handler(handler)
+    return stream_requests
+
+
+def wait_until(condition: Callable[[], bool], *, timeout: float = 5) -> None:
+    """Poll `condition` until it holds. Async tests call it through `asyncio.to_thread` to keep the event loop free."""
+    deadline = time.monotonic() + timeout
+    while not condition() and time.monotonic() < deadline:
+        time.sleep(0.01)
+    assert condition(), 'condition not met in time'
+
+
+def test_streamed_log_sync_reopens_empty_stream(caplog: LogCaptureFixture, httpserver: HTTPServer) -> None:
+    """A log stream that ends empty is reopened, and the reopened stream delivers the log."""
+    serve_log_after_empty_streams(httpserver, empty_streams=1)
+    logger = logging.getLogger('apify_client.tests.reopen_empty_stream_sync')
+    api_url = httpserver.url_for('/').removesuffix('/')
+    log_client = ApifyClient(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID).log()
+    streamed_log = StreamedLog(log_client=log_client, to_logger=logger)
+
+    with caplog.at_level(logging.DEBUG, logger=logger.name):
+        streamed_log.start()
+        # Only the reopened stream can deliver the log before `stop` would read it in one request.
+        wait_until(lambda: len(caplog.records) == len(_EXPECTED_MESSAGES_AND_LEVELS))
+        streamed_log.stop()
+
+    assert [(record.message, record.levelno) for record in caplog.records] == list(_EXPECTED_MESSAGES_AND_LEVELS)
+
+
+async def test_streamed_log_async_reopens_empty_stream(caplog: LogCaptureFixture, httpserver: HTTPServer) -> None:
+    """A log stream that ends empty is reopened, and the reopened stream delivers the log."""
+    serve_log_after_empty_streams(httpserver, empty_streams=1)
+    logger = logging.getLogger('apify_client.tests.reopen_empty_stream_async')
+    api_url = httpserver.url_for('/').removesuffix('/')
+    log_client = ApifyClientAsync(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID).log()
+    streamed_log = StreamedLogAsync(log_client=log_client, to_logger=logger)
+
+    with caplog.at_level(logging.DEBUG, logger=logger.name):
+        streamed_log.start()
+        # Only the reopened stream can deliver the log before `stop` would read it in one request.
+        await asyncio.to_thread(wait_until, lambda: len(caplog.records) == len(_EXPECTED_MESSAGES_AND_LEVELS))
+        await streamed_log.stop()
+
+    assert [(record.message, record.levelno) for record in caplog.records] == list(_EXPECTED_MESSAGES_AND_LEVELS)
+
+
+def test_streamed_log_sync_stop_reads_log_when_streams_stay_empty(
+    caplog: LogCaptureFixture, httpserver: HTTPServer
+) -> None:
+    """When every log stream is empty until `stop`, the whole log is read in one request."""
+    stream_requests = serve_log_after_empty_streams(httpserver, empty_streams=math.inf)
+    logger = logging.getLogger('apify_client.tests.empty_streams_sync')
+    api_url = httpserver.url_for('/').removesuffix('/')
+    log_client = ApifyClient(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID).log()
+    streamed_log = StreamedLog(log_client=log_client, to_logger=logger)
+
+    with caplog.at_level(logging.DEBUG, logger=logger.name):
+        streamed_log.start()
+        wait_until(lambda: len(stream_requests) >= 2)
+        streamed_log.stop()
+
+    assert [(record.message, record.levelno) for record in caplog.records] == list(_EXPECTED_MESSAGES_AND_LEVELS)
+
+
+async def test_streamed_log_async_stop_reads_log_when_streams_stay_empty(
+    caplog: LogCaptureFixture, httpserver: HTTPServer
+) -> None:
+    """When every log stream is empty until `stop`, the whole log is read in one request."""
+    stream_requests = serve_log_after_empty_streams(httpserver, empty_streams=math.inf)
+    logger = logging.getLogger('apify_client.tests.empty_streams_async')
+    api_url = httpserver.url_for('/').removesuffix('/')
+    log_client = ApifyClientAsync(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID).log()
+    streamed_log = StreamedLogAsync(log_client=log_client, to_logger=logger)
+
+    with caplog.at_level(logging.DEBUG, logger=logger.name):
+        streamed_log.start()
+        await asyncio.to_thread(wait_until, lambda: len(stream_requests) >= 2)
+        await streamed_log.stop()
+
+    assert [(record.message, record.levelno) for record in caplog.records] == list(_EXPECTED_MESSAGES_AND_LEVELS)
+
+
 def test_logger_once_logs_the_first_call(caplog: LogCaptureFixture) -> None:
     """Test the first call with a given key is logged."""
     logger = logging.getLogger('apify_client.tests.log_once_first')