fix(concurrent-source): surface stalled reads instead of hanging silently

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -109,6 +109,19 @@ def __init__(
                     f"Will defer starting this stream if another stream in the same group or its parents are active."
                 )
 
+    def get_in_flight_streams_description(self) -> str:
+        """Return a deterministic description of streams with in-flight work."""
+        generating_streams = sorted(self._streams_currently_generating_partitions)
+        running_partitions = {
+            stream_name: len(self._streams_to_running_partitions[stream_name])
+            for stream_name in sorted(self._streams_to_running_partitions)
+            if self._streams_to_running_partitions[stream_name]
+        }
+        return (
+            f"Streams generating partitions: {generating_streams}; "
+            f"streams with running partitions: {running_partitions}"
+        )
+
     def on_partition_generation_completed(
         self, sentinel: PartitionGenerationCompletedSentinel
     ) -> Iterable[AirbyteMessage]:

airbyte_cdk/sources/concurrent_source/concurrent_source.py

@@ -4,10 +4,12 @@
 
 import concurrent
 import logging
-from queue import Queue
+import os
+import time
+from queue import Empty, Queue
 from typing import Iterable, Iterator, List, Optional
 
-from airbyte_cdk.models import AirbyteMessage
+from airbyte_cdk.models import AirbyteMessage, FailureType
 from airbyte_cdk.sources.concurrent_source.concurrent_read_processor import ConcurrentReadProcessor
 from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import (
     PartitionGenerationCompletedSentinel,
@@ -25,6 +27,7 @@
 )
 from airbyte_cdk.sources.types import Record
 from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger, SliceLogger
+from airbyte_cdk.utils import AirbyteTracedException
 
 
 class ConcurrentSource:
@@ -36,6 +39,7 @@ class ConcurrentSource:
     """
 
     DEFAULT_TIMEOUT_SECONDS = 900
+    QUEUE_POLL_INTERVAL_SECONDS = 60
 
     @staticmethod
     def create(
@@ -46,6 +50,7 @@ def create(
         message_repository: MessageRepository,
         queue: Optional[Queue[QueueItem]] = None,
         timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
+        no_progress_timeout_seconds: Optional[int] = None,
     ) -> "ConcurrentSource":
         if initial_number_of_partitions_to_generate < 1:
             raise ValueError(
@@ -72,6 +77,7 @@ def create(
             message_repository=message_repository,
             initial_number_partitions_to_generate=initial_number_of_partitions_to_generate,
             timeout_seconds=timeout_seconds,
+            no_progress_timeout_seconds=no_progress_timeout_seconds,
         )
 
     def __init__(
@@ -83,21 +89,26 @@ def __init__(
         message_repository: MessageRepository = InMemoryMessageRepository(),
         initial_number_partitions_to_generate: int = 1,
         timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
+        no_progress_timeout_seconds: Optional[int] = None,
     ) -> None:
         """
         :param threadpool: The threadpool to submit tasks to
         :param logger: The logger to log to
         :param slice_logger: The slice logger used to create messages on new slices
         :param message_repository: The repository to emit messages to
         :param initial_number_partitions_to_generate: The initial number of concurrent partition generation tasks. Limiting this number ensures will limit the latency of the first records emitted. While the latency is not critical, emitting the records early allows the platform and the destination to process them as early as possible.
-        :param timeout_seconds: The maximum number of seconds to wait for a record to be read from the queue. If no record is read within this time, the source will stop reading and return.
+        :param timeout_seconds: The interval in seconds between no-progress warning logs.
+        :param no_progress_timeout_seconds: The optional number of seconds without queue progress before stopping the source.
         """
         self._threadpool = threadpool
         self._logger = logger
         self._slice_logger = slice_logger
         self._message_repository = message_repository
         self._initial_number_partitions_to_generate = initial_number_partitions_to_generate
         self._timeout_seconds = timeout_seconds
+        self._no_progress_timeout_seconds = self._get_no_progress_timeout(
+            no_progress_timeout_seconds
+        )
 
         # We set a maxsize to for the main thread to process record items when the queue size grows. This assumes that there are less
         # threads generating partitions that than are max number of workers. If it weren't the case, we could have threads only generating
@@ -135,6 +146,35 @@ def read(
         self._threadpool.check_for_errors_and_shutdown()
         self._logger.info("Finished syncing")
 
+    def _get_no_progress_timeout(self, configured_timeout: Optional[int]) -> Optional[int]:
+        if configured_timeout is not None:
+            if configured_timeout <= 0:
+                self._logger.warning(
+                    "Ignoring non-positive no_progress_timeout_seconds value: %s",
+                    configured_timeout,
+                )
+                return None
+            return configured_timeout
+
+        configured_timeout_from_env = os.getenv("AIRBYTE_NO_PROGRESS_TIMEOUT_SECONDS")
+        if configured_timeout_from_env is None:
+            return None
+        try:
+            timeout = int(configured_timeout_from_env)
+        except ValueError:
+            self._logger.warning(
+                "Ignoring invalid AIRBYTE_NO_PROGRESS_TIMEOUT_SECONDS value: %s",
+                configured_timeout_from_env,
+            )
+            return None
+        if timeout <= 0:
+            self._logger.warning(
+                "Ignoring non-positive AIRBYTE_NO_PROGRESS_TIMEOUT_SECONDS value: %s",
+                configured_timeout_from_env,
+            )
+            return None
+        return timeout
+
     def _submit_initial_partition_generators(
         self, concurrent_stream_processor: ConcurrentReadProcessor
     ) -> Iterable[AirbyteMessage]:
@@ -148,7 +188,49 @@ def _consume_from_queue(
         queue: Queue[QueueItem],
         concurrent_stream_processor: ConcurrentReadProcessor,
     ) -> Iterable[AirbyteMessage]:
-        while airbyte_message_or_record_or_exception := queue.get():
+        last_progress = time.monotonic()
+        last_warning = last_progress
+        poll_interval = min(
+            self.QUEUE_POLL_INTERVAL_SECONDS,
+            self._timeout_seconds,
+            self._no_progress_timeout_seconds or self._timeout_seconds,
+        )
+        while True:
+            try:
+                airbyte_message_or_record_or_exception = queue.get(timeout=poll_interval)
+            except Empty:
+                now = time.monotonic()
+                if now - last_progress >= self._timeout_seconds:
+                    if now - last_warning >= self._timeout_seconds:
+                        in_flight_description = (
+                            concurrent_stream_processor.get_in_flight_streams_description()
+                        )
+                        self._logger.warning(
+                            "No queue progress for %s seconds. %s",
+                            self._timeout_seconds,
+                            in_flight_description,
+                        )
+                        last_warning = now
+                if (
+                    self._no_progress_timeout_seconds is not None
+                    and now - last_progress >= self._no_progress_timeout_seconds
+                ):
+                    in_flight_description = (
+                        concurrent_stream_processor.get_in_flight_streams_description()
+                    )
+                    self._threadpool.shutdown()
+                    raise AirbyteTracedException(
+                        message=f"Source made no progress for {self._no_progress_timeout_seconds} seconds and was stopped.",
+                        internal_message=(
+                            f"No queue progress for {self._no_progress_timeout_seconds} seconds. "
+                            f"{in_flight_description}"
+                        ),
+                        failure_type=FailureType.system_error,
+                    )
+                continue
+            last_progress = time.monotonic()
+            if not airbyte_message_or_record_or_exception:
+                break
             yield from self._handle_item(
                 airbyte_message_or_record_or_exception,
                 concurrent_stream_processor,

airbyte_cdk/sources/concurrent_source/thread_pool_manager.py

@@ -78,6 +78,9 @@ def _shutdown(self) -> None:
         # this imperfect approach because we only do this in case of `self._most_recently_seen_exception` which we don't expect to happen
         self._threadpool.shutdown(wait=False, cancel_futures=True)
 
+    def shutdown(self) -> None:
+        self._shutdown()
+
     def is_done(self) -> bool:
         return all([f.done() for f in self._futures])
 

unit_tests/sources/concurrent_source/test_no_progress_watchdog.py

@@ -0,0 +1,133 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+from queue import Empty
+from unittest.mock import Mock
+
+import pytest
+
+from airbyte_cdk.models import AirbyteMessage, AirbyteRecordMessage, FailureType
+from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource
+from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
+from airbyte_cdk.sources.message import InMemoryMessageRepository
+from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger
+from airbyte_cdk.utils import AirbyteTracedException
+
+
+class _FakeClock:
+    def __init__(self, step: int = 1) -> None:
+        self._now = 0
+        self._step = step
+
+    def __call__(self) -> int:
+        now = self._now
+        self._now += self._step
+        return now
+
+
+def _message() -> AirbyteMessage:
+    return AirbyteMessage(
+        type=MessageType.RECORD,
+        record=AirbyteRecordMessage(stream="stream", data={}, emitted_at=0),
+    )
+
+
+def _source(queue, **kwargs) -> ConcurrentSource:
+    return ConcurrentSource(
+        threadpool=Mock(spec=ThreadPoolManager),
+        logger=Mock(spec=logging.Logger),
+        slice_logger=DebugSliceLogger(),
+        queue=queue,
+        message_repository=InMemoryMessageRepository(),
+        timeout_seconds=1,
+        **kwargs,
+    )
+
+
+def _processor(done: bool = True) -> Mock:
+    processor = Mock()
+    processor.is_done.return_value = done
+    processor.get_in_flight_streams_description.return_value = (
+        "Streams generating partitions: ['stream']; streams with running partitions: {'stream': 1}"
+    )
+    return processor
+
+
+def test_stalled_queue_logs_warning_without_raising(monkeypatch) -> None:
+    queue = Mock()
+    queue.get.side_effect = [Empty, Empty, _message()]
+    queue.empty.return_value = True
+    processor = _processor()
+    source = _source(queue)
+    monkeypatch.setattr(
+        "airbyte_cdk.sources.concurrent_source.concurrent_source.time.monotonic",
+        _FakeClock(),
+    )
+
+    assert list(source._consume_from_queue(queue, processor))
+    assert source._logger.warning.call_count >= 1
+
+
+def test_stalled_queue_raises_with_constructor_timeout(monkeypatch) -> None:
+    queue = Mock()
+    queue.get.side_effect = Empty
+    source = _source(queue, no_progress_timeout_seconds=1)
+    processor = _processor()
+    monkeypatch.setattr(
+        "airbyte_cdk.sources.concurrent_source.concurrent_source.time.monotonic",
+        _FakeClock(),
+    )
+
+    with pytest.raises(AirbyteTracedException) as exc_info:
+        list(source._consume_from_queue(queue, processor))
+
+    assert exc_info.value.failure_type == FailureType.system_error
+    assert exc_info.value.message == "Source made no progress for 1 seconds and was stopped."
+    source._threadpool.shutdown.assert_called_once_with()
+
+
+def test_stalled_queue_raises_with_environment_timeout(monkeypatch) -> None:
+    monkeypatch.setenv("AIRBYTE_NO_PROGRESS_TIMEOUT_SECONDS", "1")
+    queue = Mock()
+    queue.get.side_effect = Empty
+    source = _source(queue)
+    processor = _processor()
+    monkeypatch.setattr(
+        "airbyte_cdk.sources.concurrent_source.concurrent_source.time.monotonic",
+        _FakeClock(),
+    )
+
+    with pytest.raises(AirbyteTracedException) as exc_info:
+        list(source._consume_from_queue(queue, processor))
+
+    assert source._no_progress_timeout_seconds == 1
+    assert exc_info.value.failure_type == FailureType.system_error
+
+
+def test_normal_read_completes_unchanged() -> None:
+    queue = Mock()
+    queue.get.return_value = _message()
+    queue.empty.return_value = True
+    processor = _processor()
+    source = _source(queue, no_progress_timeout_seconds=2)
+
+    assert list(source._consume_from_queue(queue, processor))
+
+
+def test_slow_progressing_read_does_not_trip_watchdog(monkeypatch) -> None:
+    queue = Mock()
+    queue.get.side_effect = [_message(), Empty, _message()]
+    queue.empty.side_effect = [True, True]
+    processor = _processor()
+    processor.is_done.side_effect = [False, True]
+    source = _source(queue, no_progress_timeout_seconds=2)
+
+    monkeypatch.setattr(
+        "airbyte_cdk.sources.concurrent_source.concurrent_source.time.monotonic",
+        _FakeClock(),
+    )
+    assert len(list(source._consume_from_queue(queue, processor))) == 2
+    source._threadpool.shutdown.assert_not_called()