refactor(cdk): switch memory monitor to logging-only trial mode

Per pnilan's review feedback:
- Remove AirbyteTracedException raise at critical threshold — log only
- Change check_interval default from 1000 to 5000 messages
- Log WARNING on every check above 90% (remove one-shot flags)
- Simplify constructor (remove threshold params and state flags)
- Remove try/finally from entrypoint read() (no longer raising exceptions)
- Add TODO comment for cgroup v1 removal
- Update tests: remove exception/one-shot tests, add repeated logging test

Co-Authored-By: bot_apk <apk@cognition.ai>

Author: devin-ai-integration[bot]

airbyte_cdk/entrypoint.py

@@ -279,13 +279,11 @@ def read(
 
         # The Airbyte protocol dictates that counts be expressed as float/double to better protect against integer overflows
         stream_message_counter: DefaultDict[HashableStreamDescriptor, float] = defaultdict(float)
-        try:
-            for message in self.source.read(self.logger, config, catalog, state):
-                yield self.handle_record_counts(message, stream_message_counter)
-                self._memory_monitor.check_memory_usage()
-        finally:
-            for message in self._emit_queued_messages(self.source):
-                yield self.handle_record_counts(message, stream_message_counter)
+        for message in self.source.read(self.logger, config, catalog, state):
+            yield self.handle_record_counts(message, stream_message_counter)
+            self._memory_monitor.check_memory_usage()
+        for message in self._emit_queued_messages(self.source):
+            yield self.handle_record_counts(message, stream_message_counter)
 
     @staticmethod
     def handle_record_counts(

airbyte_cdk/utils/memory_monitor.py

@@ -2,55 +2,49 @@
 # Copyright (c) 2026 Airbyte, Inc., all rights reserved.
 #
 
-"""Source-side memory introspection to emit controlled error messages before OOM kills."""
+"""Source-side memory introspection to log memory usage approaching container limits."""
 
 import logging
 from pathlib import Path
 from typing import Optional
 
-from airbyte_cdk.models import FailureType
-from airbyte_cdk.utils.traced_exception import AirbyteTracedException
-
 logger = logging.getLogger("airbyte")
 
 # cgroup v2 paths
 _CGROUP_V2_CURRENT = Path("/sys/fs/cgroup/memory.current")
 _CGROUP_V2_MAX = Path("/sys/fs/cgroup/memory.max")
 
-# cgroup v1 paths
+# cgroup v1 paths — TODO: remove if all deployments are confirmed cgroup v2
 _CGROUP_V1_USAGE = Path("/sys/fs/cgroup/memory/memory.usage_in_bytes")
 _CGROUP_V1_LIMIT = Path("/sys/fs/cgroup/memory/memory.limit_in_bytes")
 
-# Default thresholds
-_DEFAULT_WARNING_THRESHOLD = 0.85
-_DEFAULT_CRITICAL_THRESHOLD = 0.95
+# Log when usage is at or above 90%
+_MEMORY_THRESHOLD = 0.90
 
 # Check interval (every N messages)
-_DEFAULT_CHECK_INTERVAL = 1000
+_DEFAULT_CHECK_INTERVAL = 5000
 
 
 class MemoryMonitor:
-    """Monitors container memory usage via cgroup files and emits warnings before OOM kills.
+    """Monitors container memory usage via cgroup files and logs warnings when usage is high.
 
     Lazily probes cgroup v2 then v1 files on the first call to
     ``check_memory_usage()``.  Caches which version exists.
     If neither is found (local dev / CI), all subsequent calls are instant no-ops.
+
+    Logs a WARNING on every check interval (default 5000 messages) when memory
+    usage is at or above 90% of the container limit.  This gives breadcrumb
+    trails showing whether memory is climbing, plateauing, or sawtoothing.
     """
 
     def __init__(
         self,
-        warning_threshold: float = _DEFAULT_WARNING_THRESHOLD,
-        critical_threshold: float = _DEFAULT_CRITICAL_THRESHOLD,
         check_interval: int = _DEFAULT_CHECK_INTERVAL,
     ) -> None:
         if check_interval < 1:
             raise ValueError(f"check_interval must be >= 1, got {check_interval}")
-        self._warning_threshold = warning_threshold
-        self._critical_threshold = critical_threshold
         self._check_interval = check_interval
         self._message_count = 0
-        self._warning_emitted = False
-        self._critical_raised = False
         self._cgroup_version: Optional[int] = None
         self._probed = False
 
@@ -108,17 +102,15 @@ def _read_memory(self) -> Optional[tuple[int, int]]:
             return None
 
     def check_memory_usage(self) -> None:
-        """Check memory usage against thresholds.
+        """Check memory usage and log when above 90%.
 
         Intended to be called on every message. The monitor internally tracks
         a message counter and only reads cgroup files every ``check_interval``
-        messages (default 1000) to minimise I/O overhead.
+        messages (default 5000) to minimise I/O overhead.
 
-        At the warning threshold (default 85%), logs a warning message.
-        At the critical threshold (default 95%), raises ``AirbyteTracedException``
-        to trigger a graceful shutdown with an actionable error message.
+        Logs a WARNING on every check above 90% to provide breadcrumb trails
+        showing memory trends over the sync lifetime.
 
-        Each threshold triggers at most once per sync to avoid log spam.
         This method is a no-op if cgroup files are unavailable.
         """
         self._probe_cgroup()
@@ -139,19 +131,9 @@ def check_memory_usage(self) -> None:
         usage_gb = usage_bytes / (1024**3)
         limit_gb = limit_bytes / (1024**3)
 
-        if usage_ratio >= self._critical_threshold and not self._critical_raised:
-            self._critical_raised = True
-            raise AirbyteTracedException(
-                internal_message=f"Memory usage is {usage_percent}% ({usage_gb:.2f} / {limit_gb:.2f} GB). "
-                f"Critical threshold is {int(self._critical_threshold * 100)}%.",
-                message=f"Source exceeded memory limit ({usage_percent}% used) and must shut down to avoid an out-of-memory crash.",
-                failure_type=FailureType.system_error,
-            )
-
-        if usage_ratio >= self._warning_threshold and not self._warning_emitted:
-            self._warning_emitted = True
+        if usage_ratio >= _MEMORY_THRESHOLD:
             logger.warning(
-                "Source memory usage reached %d%% of container limit (%.2f / %.2f GB).",
+                "Source memory usage at %d%% of container limit (%.2f / %.2f GB).",
                 usage_percent,
                 usage_gb,
                 limit_gb,

unit_tests/test_entrypoint.py

@@ -835,97 +835,6 @@ def test_handle_record_counts(
         )
 
 
-def test_memory_limit_exceeded_flushes_queued_messages(mocker, spec_mock, config_mock):
-    """When AirbyteTracedException is raised mid-read, queued messages should still be flushed.
-
-    The read() try/finally ensures _emit_queued_messages runs even when
-    AirbyteTracedException propagates.  The exception still surfaces to the
-    caller, but all messages yielded before (records) and during (finally-
-    block state messages) the exception are available to the consumer.
-    """
-    queued_state = AirbyteMessage(
-        type=Type.STATE,
-        state=AirbyteStateMessage(
-            type=AirbyteStateType.STREAM,
-            stream=AirbyteStreamState(
-                stream_descriptor=StreamDescriptor(name="stream"),
-                stream_state=AirbyteStateBlob(updated_at="2026-01-01"),
-            ),
-        ),
-    )
-    message_repository = MagicMock()
-    # consume_queue calls:
-    #   1. run() preamble → initial queued control message
-    #   2. read() finally block → queued state (the key assertion)
-    #   3. run() outer finally → nothing
-    message_repository.consume_queue.side_effect = [
-        [MESSAGE_FROM_REPOSITORY],
-        [queued_state],
-        [],
-    ]
-    mocker.patch.object(
-        MockSource,
-        "message_repository",
-        new_callable=mocker.PropertyMock,
-        return_value=message_repository,
-    )
-    entrypoint = AirbyteEntrypoint(MockSource())
-
-    record = AirbyteMessage(
-        type=Type.RECORD,
-        record=AirbyteRecordMessage(stream="stream", data={"id": "1"}, emitted_at=1),
-    )
-    mocker.patch.object(MockSource, "read_state", return_value={})
-    mocker.patch.object(MockSource, "read_catalog", return_value={})
-    mocker.patch.object(MockSource, "read", return_value=[record, record])
-
-    from airbyte_cdk.utils.traced_exception import AirbyteTracedException
-
-    call_count = 0
-
-    def _raise_on_second_call() -> None:
-        nonlocal call_count
-        call_count += 1
-        if call_count >= 2:
-            raise AirbyteTracedException(
-                internal_message="Memory at 96%",
-                message="Source exceeded memory limit (96% used) and must shut down to avoid an out-of-memory crash.",
-                failure_type=FailureType.system_error,
-            )
-
-    mocker.patch.object(
-        entrypoint._memory_monitor, "check_memory_usage", side_effect=_raise_on_second_call
-    )
-
-    parsed_args = Namespace(
-        command="read", config="config_path", state="statepath", catalog="catalogpath"
-    )
-
-    # The generator yields messages until AirbyteTracedException propagates.
-    # Collect everything yielded before the exception surfaces.
-    messages: list[str] = []
-    with pytest.raises(AirbyteTracedException):
-        for msg in entrypoint.run(parsed_args):
-            messages.append(msg)
-
-    # 1. Both records were yielded before the exception — the memory check
-    #    runs after yield so every message pulled from the source is emitted.
-    record_messages = [m for m in messages if "RECORD" in m]
-    assert len(record_messages) == 2, "Both records should be yielded before AirbyteTracedException"
-
-    # 2. The queued state message was flushed by the finally block
-    state_messages = [m for m in messages if "STATE" in m]
-    assert len(state_messages) >= 1, (
-        "Queued state message should be flushed even after AirbyteTracedException"
-    )
-
-    # 3. The flushed state has sourceStats.recordCount set by handle_record_counts.
-    #    Both records are yielded (and counted) before the second check_memory_usage
-    #    raises, so the counter is 2.0 at flush time.
-    state_json = orjson.loads(state_messages[0])
-    assert state_json["state"]["sourceStats"]["recordCount"] == 2.0
-
-
 def test_given_serialization_error_using_orjson_then_fallback_on_json(
     entrypoint: AirbyteEntrypoint, mocker, spec_mock, config_mock
 ):