From df3ae15f8a963f7d78992edefc235bd3f9aac6de Mon Sep 17 00:00:00 2001 From: vikramlc Date: Mon, 26 Oct 2026 14:34:32 +0530 Subject: [PATCH 1/3] fix(odin): Make check-ins resilient to startup failures and transient network errors - Log and treat as fatal any failure building/reporting the startup request instead of letting the ExtractorCheckin daemon thread die silently. - Make the extractor process exit with a dedicated non-zero code on a fatal check-in error and have the runtime exit 1 instead of 0. - Treat raw requests/httpx transport errors and timeouts (including from the OAuth token refresh) as transient: warn, requeue and retry, not fatal. Co-Authored-By: Claude Sonnet 5.5 --- CHANGELOG.md | 6 + cognite/extractorutils/unstable/core/base.py | 12 +- .../unstable/core/checkin_worker.py | 53 ++- .../extractorutils/unstable/core/runtime.py | 24 +- .../test_unstable/test_checkin_resilience.py | 311 ++++++++++++++++++ 5 files changed, 402 insertions(+), 4 deletions(-) create mode 100644 tests/test_unstable/test_checkin_resilience.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b19edb1a..d0e42205 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,12 @@ Changes are grouped as follows - `Fixed` for any bug fixes. - `Security` in case of vulnerabilities. +## Unreleased + +### Fixed +* In the `unstable` package: Fixed the check-in thread dying silently when the startup request could not be built or reported (for example when more than 100 actions are available). The failure is now logged at critical level and treated as fatal, and the extractor process exits with a non-zero exit code when it stops because of a fatal check-in error. +* In the `unstable` package: Fixed transient network errors during check-ins (timeouts, connection errors, including raw `requests`/`httpx` errors from the OAuth token refresh) being treated as fatal and shutting the extractor down. They are now logged as warnings and retried at the next check-in. Only 401/403 errors and bad credentials remain fatal. + ## 7.13.2 ### Fixed diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 61972b05..31b84e52 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -446,7 +446,17 @@ def _get_startup_request(self) -> StartupRequest: ) def _run_checkin(self) -> None: - self._checkin_worker.run_periodic_checkin(self.cancellation_token, self._get_startup_request()) + try: + self._checkin_worker.run_periodic_checkin(self.cancellation_token, self._get_startup_request()) + except Exception as e: + # This runs on a daemon thread: without this, the traceback would only go to stderr and the extractor + # would keep running without ever checking in. Treat it as fatal instead. + self._logger.critical("Check-in failed unexpectedly, shutting down the extractor: %s", e, exc_info=e) + try: + self._checkin_worker.notify_fatal_error(e) + except Exception: + self._logger.exception("Fatal error handler failed") + self.cancellation_token.cancel() def _report_error(self, error: Error) -> None: self._checkin_worker.report_error(error) diff --git a/cognite/extractorutils/unstable/core/checkin_worker.py b/cognite/extractorutils/unstable/core/checkin_worker.py index c9c295d3..89cb8ab3 100644 --- a/cognite/extractorutils/unstable/core/checkin_worker.py +++ b/cognite/extractorutils/unstable/core/checkin_worker.py @@ -16,6 +16,8 @@ from threading import RLock, Thread from time import sleep +import httpx +import requests from cognite.client import CogniteClient from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError from requests import Response @@ -43,6 +45,19 @@ MAX_ACTION_UPDATES_PER_CHECKIN = 100 rng = SystemRandom() +# Errors that mean "could not reach CDF right now" rather than "the request or credentials are wrong". The Cognite SDK +# only wraps some of these (e.g. connection errors from API calls, OAuth2Error from token refresh), so raw +# requests/httpx transport errors and timeouts (notably from the OAuth token refresh) can still surface unwrapped. +# Deliberately narrow: configuration mistakes (invalid URL, missing schema, unsupported protocol) must stay fatal. +_TRANSIENT_ERRORS: tuple[type[BaseException], ...] = ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + httpx.TimeoutException, + httpx.NetworkError, + TimeoutError, + ConnectionError, +) + def _safe_dispatch( dispatcher: Callable[[list[Action]], None], @@ -184,6 +199,18 @@ def set_on_fatal_error_handler(self, on_fatal_error: Callable[[Exception], None] """ self._on_fatal_error = on_fatal_error + def notify_fatal_error(self, error: Exception) -> None: + """ + Invoke the fatal error handler (if any) with the given error. + + Used by callers that fail outside of a check-in request, such as when building the startup request. + + Arguments: + error (Exception): The error that should be treated as fatal. + """ + if self._on_fatal_error is not None: + self._on_fatal_error(error) + def set_retry_startup(self, retry_startup: bool) -> None: """ Set whether to retry reporting startup if it fails. @@ -480,6 +507,15 @@ def _requeue_checkin( self._action_updates = action_updates + self._action_updates def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: + """ + Run a check-in like request and classify any failure. + + Transient network problems are logged and requeued without being fatal. Only authentication/authorization + failures (401/403, bad credentials) and unexpected errors are reported as fatal. + + Returns: + True if the payload of the request should be requeued. + """ exception_to_report: Exception | None = None requeue = False try: @@ -493,9 +529,13 @@ def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: except CogniteAuthError as e: self._logger.error(str(e)) - self._logger.critical("Could not get an access token. Please check your configuration.") - exception_to_report = e requeue = True + if isinstance(e.__cause__, _TRANSIENT_ERRORS): + # The token endpoint could not be reached; this says nothing about the credentials. + self._log_transient_error("Could not reach the token endpoint to get an access token", e.__cause__) + else: + self._logger.critical("Could not get an access token. Please check your configuration.") + exception_to_report = e except CogniteAPIError as e: if e.code == 401: @@ -519,6 +559,10 @@ def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: self._logger.critical(f"Error while connecting to CDF {e!s}") requeue = True + except _TRANSIENT_ERRORS as e: + self._log_transient_error("Could not reach CDF", e) + requeue = True + except Exception as e: self._logger.critical(f"Extractor could not connect to CDF {e!s}") exception_to_report = e @@ -528,3 +572,8 @@ def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: self._on_fatal_error(exception_to_report) return requeue + + def _log_transient_error(self, message: str, error: BaseException) -> None: + self._logger.warning( + "%s (%s: %s). The check-in will be retried at the next check-in.", message, type(error).__name__, error + ) diff --git a/cognite/extractorutils/unstable/core/runtime.py b/cognite/extractorutils/unstable/core/runtime.py index be1f055e..df5aca8b 100644 --- a/cognite/extractorutils/unstable/core/runtime.py +++ b/cognite/extractorutils/unstable/core/runtime.py @@ -34,6 +34,7 @@ def main() -> None: from multiprocessing.synchronize import Event as MpEvent from pathlib import Path from random import randint +from threading import Event as ThreadEvent from threading import Thread from typing import Any, Generic, TypeVar from uuid import uuid4 @@ -68,6 +69,11 @@ def main() -> None: ExtractorType = TypeVar("ExtractorType", bound=Extractor) +# Exit code the extractor process uses to tell the runtime that it stopped because of a fatal error (for example +# bad credentials), as opposed to a requested shutdown. Kept distinct from 1 so that unrelated crashes of the child +# process do not change the runtime's behaviour. +FATAL_ERROR_EXIT_CODE = 3 + @dataclass class _RuntimeControls: @@ -83,7 +89,13 @@ def _extractor_process_entrypoint( ) -> None: logger = logging.getLogger(f"{extractor_class.EXTERNAL_ID}.runtime") checkin_worker.active_revision = config.current_config_revision - checkin_worker.set_on_fatal_error_handler(lambda _: on_fatal_error(controls)) + fatal_error_occurred = ThreadEvent() + + def handle_fatal_error(_: Exception) -> None: + fatal_error_occurred.set() + on_fatal_error(controls) + + checkin_worker.set_on_fatal_error_handler(handle_fatal_error) checkin_worker.set_on_revision_change_handler(lambda _: on_revision_changed(controls)) checkin_worker.set_retry_startup(extractor_class.RETRY_STARTUP) extractor = extractor_class._init_from_runtime(config, checkin_worker) @@ -99,6 +111,11 @@ def _extractor_process_entrypoint( except Exception: logger.exception("Extractor crashed, will attempt restart") controls.message_queue.put(RuntimeMessage.RESTART) + return + + if fatal_error_occurred.is_set(): + logger.critical("Extractor stopped because of a fatal error") + sys.exit(FATAL_ERROR_EXIT_CODE) def on_revision_changed(controls: _RuntimeControls) -> None: @@ -598,6 +615,11 @@ def _main_runtime(self, args: Namespace) -> None: case _: assert_never(message) + elif process.exitcode == FATAL_ERROR_EXIT_CODE: + self.logger.critical("Extractor stopped because of a fatal error, shutting down runtime") + self._cancellation_token.cancel() + sys.exit(1) + else: self.logger.info("Shutting down runtime") self._cancellation_token.cancel() diff --git a/tests/test_unstable/test_checkin_resilience.py b/tests/test_unstable/test_checkin_resilience.py new file mode 100644 index 00000000..31e917ff --- /dev/null +++ b/tests/test_unstable/test_checkin_resilience.py @@ -0,0 +1,311 @@ +""" +Tests for check-in resilience: startup failures are fatal and logged, transient network errors are not fatal. + +These tests do not need a CDF project: the Cognite client is a mock whose ``post`` is controlled by each test. +""" + +import logging +from collections.abc import Callable +from datetime import datetime, timezone +from threading import Event, Thread +from typing import Any +from unittest.mock import MagicMock + +import httpx +import pytest +import requests +from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError +from oauthlib.oauth2 import OAuth2Error + +from cognite.extractorutils.threading import CancellationToken +from cognite.extractorutils.unstable.core._dto import ActionStatus, ActionUpdate +from cognite.extractorutils.unstable.core.base import FullConfig +from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker +from cognite.extractorutils.unstable.core.tasks import ScheduledTask + +from .conftest import TestConfig, TestExtractor + +MAX_AVAILABLE_ACTIONS = 100 + + +def _make_worker(post: Callable[..., Any] | None = None) -> tuple[CheckinWorker, MagicMock]: + client = MagicMock() + client.config.project = "test-project" + if post is not None: + client.post.side_effect = post + worker = CheckinWorker(client, "test-integration", logging.getLogger("test.checkin_resilience")) + worker.set_retry_startup(True) + return worker, client + + +def _make_extractor(worker: CheckinWorker, n_tasks: int = 0) -> TestExtractor: + conn = MagicMock() + conn.integration.external_id = "test-integration" + extractor = TestExtractor( + FullConfig( + connection_config=conn, + application_config=TestConfig(parameter_one=1, parameter_two="a"), + current_config_revision=1, + ), + worker, + ) + extractor._start_time = datetime.now(tz=timezone.utc) + for i in range(n_tasks): + extractor.add_task(ScheduledTask.from_interval(interval="1h", name=f"task-{i}", target=lambda _: None)) + return extractor + + +def _n_builtin_actions() -> int: + worker, _ = _make_worker() + request = _make_extractor(worker)._get_startup_request() + return len(request.available_actions or []) + + +def _ok_response() -> MagicMock: + response = MagicMock() + response.json.return_value = {"externalId": "test-integration"} + return response + + +def _raiser(error: BaseException) -> Callable[..., Any]: + def post(*args: Any, **kwargs: Any) -> Any: + raise error + + return post + + +def _oauth_auth_error() -> CogniteAuthError: + error = CogniteAuthError("Error generating access token: invalid_client") + error.__cause__ = OAuth2Error("invalid_client") + return error + + +def _timeout_auth_error() -> CogniteAuthError: + error = CogniteAuthError("Error generating access token") + error.__cause__ = requests.exceptions.ReadTimeout("token endpoint timed out") + return error + + +# --------------------------------------------------------------------------- +# Startup request failures +# --------------------------------------------------------------------------- + + +def test_too_many_available_actions_is_logged_and_fatal(caplog: pytest.LogCaptureFixture) -> None: + """More than 100 available actions fails StartupRequest validation; this must not kill the thread silently.""" + worker, client = _make_worker(lambda *a, **k: _ok_response()) + fatal_errors: list[Exception] = [] + worker.set_on_fatal_error_handler(fatal_errors.append) + + n_tasks = (MAX_AVAILABLE_ACTIONS - _n_builtin_actions()) // 2 + 1 # one task over the limit + extractor = _make_extractor(worker, n_tasks=n_tasks) + + with caplog.at_level(logging.WARNING): + extractor._run_checkin() # runs synchronously; would raise (thread dies) without the fix + + assert len(fatal_errors) == 1 + assert extractor.cancellation_token.is_cancelled + client.post.assert_not_called() + critical = [r for r in caplog.records if r.levelno >= logging.CRITICAL] + assert critical, "failure must be logged at CRITICAL so it reaches the log file" + assert critical[0].exc_info is not None + + +def test_exactly_100_available_actions_is_fine() -> None: + worker, _ = _make_worker() + n_tasks = (MAX_AVAILABLE_ACTIONS - _n_builtin_actions()) // 2 + extractor = _make_extractor(worker, n_tasks=n_tasks) + request = extractor._get_startup_request() + assert len(request.available_actions or []) <= MAX_AVAILABLE_ACTIONS + + +def test_startup_not_retried_failure_is_fatal(caplog: pytest.LogCaptureFixture) -> None: + """A startup failure when retries are disabled used to kill the check-in thread with no log.""" + worker, _ = _make_worker(_raiser(CogniteConnectionError())) + worker.set_retry_startup(False) + fatal_errors: list[Exception] = [] + worker.set_on_fatal_error_handler(fatal_errors.append) + extractor = _make_extractor(worker) + + with caplog.at_level(logging.WARNING): + extractor._run_checkin() + + assert len(fatal_errors) == 1 + assert any(r.levelno >= logging.CRITICAL for r in caplog.records) + + +# --------------------------------------------------------------------------- +# Transient errors are not fatal +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "error", + [ + pytest.param(requests.exceptions.ReadTimeout("read timed out"), id="requests-read-timeout"), + pytest.param(requests.exceptions.ConnectTimeout("connect timed out"), id="requests-connect-timeout"), + pytest.param(requests.exceptions.ConnectionError("reset"), id="requests-connection-error"), + pytest.param(httpx.ReadTimeout("read timed out"), id="httpx-read-timeout"), + pytest.param(httpx.ConnectTimeout("connect timed out"), id="httpx-connect-timeout"), + pytest.param(httpx.ConnectError("refused"), id="httpx-connect-error"), + pytest.param(TimeoutError("timed out"), id="builtin-timeout"), + pytest.param(CogniteConnectionError(), id="sdk-connection-error"), + pytest.param(CogniteAPIError("Bad gateway", code=502), id="http-502"), + pytest.param(CogniteAPIError("Unavailable", code=503), id="http-503"), + pytest.param(_timeout_auth_error(), id="token-refresh-timeout-wrapped"), + ], +) +def test_transient_error_is_requeued_and_not_fatal(error: BaseException) -> None: + worker, _ = _make_worker(_raiser(error)) + on_fatal = MagicMock() + worker.set_on_fatal_error_handler(on_fatal) + update = ActionUpdate(external_id="act-1", status=ActionStatus.running) + worker.queue_action_update(update) + + worker.try_write_checkin([], []) + + on_fatal.assert_not_called() + assert worker._action_updates == [update], "the update must be requeued for the next check-in" + + +def test_transient_error_during_startup_is_not_fatal_and_is_retried(monkeypatch: pytest.MonkeyPatch) -> None: + calls = 0 + token = CancellationToken() + + def post(*args: Any, **kwargs: Any) -> Any: + nonlocal calls + calls += 1 + if calls == 1: + raise requests.exceptions.ReadTimeout("token refresh timed out") + return _ok_response() + + worker, _ = _make_worker(post) + on_fatal = MagicMock() + worker.set_on_fatal_error_handler(on_fatal) + monkeypatch.setattr("cognite.extractorutils.unstable.core.checkin_worker.sleep", lambda _: None) + extractor = _make_extractor(worker) + + worker._run_startup_report(token, extractor._get_startup_request(), interval=0.01) + + assert calls == 2 + assert worker._has_reported_startup + on_fatal.assert_not_called() + + +# --------------------------------------------------------------------------- +# Auth/config errors stay fatal +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "error", + [ + pytest.param(CogniteAPIError("Unauthorized", code=401), id="401"), + pytest.param(CogniteAPIError("Forbidden", code=403), id="403"), + pytest.param(_oauth_auth_error(), id="bad-credentials"), + pytest.param(CogniteAuthError("Unable to determine device authorization endpoint"), id="auth-error-no-cause"), + pytest.param(ValueError("unexpected"), id="unexpected-error"), + pytest.param(requests.exceptions.InvalidURL("bad url"), id="invalid-url-is-a-config-error"), + pytest.param(requests.exceptions.MissingSchema("no scheme"), id="missing-schema-is-a-config-error"), + pytest.param(httpx.UnsupportedProtocol("bad protocol"), id="unsupported-protocol-is-a-config-error"), + ], +) +def test_auth_and_config_errors_remain_fatal(error: BaseException) -> None: + worker, _ = _make_worker(_raiser(error)) + on_fatal = MagicMock() + worker.set_on_fatal_error_handler(on_fatal) + + worker.try_write_checkin([], []) + + on_fatal.assert_called_once() + assert on_fatal.call_args[0][0] is error + + +def test_periodic_checkin_keeps_running_after_read_timeout() -> None: + """End to end: a timeout in a periodic check-in does not stop the loop, and it recovers on the next cycle.""" + calls = 0 + token = CancellationToken() + recovered = Event() + + def post(*args: Any, **kwargs: Any) -> Any: + nonlocal calls + calls += 1 + if calls == 2: # first periodic check-in after startup + raise httpx.ReadTimeout("timed out") + if calls >= 3: + recovered.set() + return _ok_response() + + worker, _ = _make_worker(post) + on_fatal = MagicMock() + worker.set_on_fatal_error_handler(on_fatal) + extractor = _make_extractor(worker) + + thread = Thread( + target=worker.run_periodic_checkin, args=(token, extractor._get_startup_request(), 0.05), daemon=True + ) + thread.start() + assert recovered.wait(timeout=10) + token.cancel() + thread.join(timeout=10) + + on_fatal.assert_not_called() + assert calls >= 3 + + +# --------------------------------------------------------------------------- +# Exit code when the extractor stops because of a fatal error +# --------------------------------------------------------------------------- + + +class _FakeExtractor: + """Stands in for an Extractor in the child process entrypoint; ``run`` optionally hits a fatal check-in error.""" + + EXTERNAL_ID = "fake" + RETRY_STARTUP = True + raise_fatal = False + + def __init__(self, worker: CheckinWorker) -> None: + self._worker = worker + + @classmethod + def _init_from_runtime(cls, config: Any, checkin_worker: CheckinWorker) -> "_FakeExtractor": + return cls(checkin_worker) + + def _attach_runtime_controls(self, **_: Any) -> None: + pass + + def __enter__(self) -> "_FakeExtractor": + return self + + def __exit__(self, *_: Any) -> bool: + return False + + def run(self) -> None: + if self.raise_fatal: + self._worker.notify_fatal_error(RuntimeError("bad credentials")) + + +@pytest.mark.parametrize("fatal", [True, False]) +def test_extractor_process_exit_code_reflects_fatal_error(fatal: bool) -> None: + from cognite.extractorutils.unstable.core.runtime import ( + FATAL_ERROR_EXIT_CODE, + _extractor_process_entrypoint, + _RuntimeControls, + ) + + class Extractor(_FakeExtractor): + raise_fatal = fatal + + controls = _RuntimeControls(cancel_event=MagicMock(), message_queue=MagicMock()) + config = MagicMock() + config.current_config_revision = 1 + worker, _ = _make_worker() + + if fatal: + with pytest.raises(SystemExit) as exc_info: + _extractor_process_entrypoint(Extractor, controls, config, worker) # type: ignore[arg-type] + assert exc_info.value.code == FATAL_ERROR_EXIT_CODE + controls.cancel_event.set.assert_called_once() + else: + _extractor_process_entrypoint(Extractor, controls, config, worker) # type: ignore[arg-type] From 90affef2e80812d1103a43c9400ad58116874c61 Mon Sep 17 00:00:00 2001 From: vikramlc Date: Mon, 26 Oct 2026 14:34:32 +0530 Subject: [PATCH 2/3] fix(odin): Make check-ins resilient to startup failures and transient network errors - Log and treat as fatal any failure building/reporting the startup request instead of letting the ExtractorCheckin daemon thread die silently. - Make the extractor process exit with a dedicated non-zero code on a fatal check-in error and have the runtime exit 1 instead of 0. - Treat raw requests/httpx transport errors and timeouts (including from the OAuth token refresh) as transient: warn, requeue and retry, not fatal. Co-Authored-By: Claude Sonnet 5.5 --- CHANGELOG.md | 6 + cognite/extractorutils/unstable/core/base.py | 12 +- .../unstable/core/checkin_worker.py | 53 ++- .../extractorutils/unstable/core/runtime.py | 24 +- .../test_unstable/test_checkin_resilience.py | 311 ++++++++++++++++++ 5 files changed, 402 insertions(+), 4 deletions(-) create mode 100644 tests/test_unstable/test_checkin_resilience.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b19edb1a..d0e42205 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,12 @@ Changes are grouped as follows - `Fixed` for any bug fixes. - `Security` in case of vulnerabilities. +## Unreleased + +### Fixed +* In the `unstable` package: Fixed the check-in thread dying silently when the startup request could not be built or reported (for example when more than 100 actions are available). The failure is now logged at critical level and treated as fatal, and the extractor process exits with a non-zero exit code when it stops because of a fatal check-in error. +* In the `unstable` package: Fixed transient network errors during check-ins (timeouts, connection errors, including raw `requests`/`httpx` errors from the OAuth token refresh) being treated as fatal and shutting the extractor down. They are now logged as warnings and retried at the next check-in. Only 401/403 errors and bad credentials remain fatal. + ## 7.13.2 ### Fixed diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 33466ed1..5ddef5e3 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -440,7 +440,17 @@ def _get_startup_request(self) -> StartupRequest: ) def _run_checkin(self) -> None: - self._checkin_worker.run_periodic_checkin(self.cancellation_token, self._get_startup_request()) + try: + self._checkin_worker.run_periodic_checkin(self.cancellation_token, self._get_startup_request()) + except Exception as e: + # This runs on a daemon thread: without this, the traceback would only go to stderr and the extractor + # would keep running without ever checking in. Treat it as fatal instead. + self._logger.critical("Check-in failed unexpectedly, shutting down the extractor: %s", e, exc_info=e) + try: + self._checkin_worker.notify_fatal_error(e) + except Exception: + self._logger.exception("Fatal error handler failed") + self.cancellation_token.cancel() def _report_error(self, error: Error) -> None: self._checkin_worker.report_error(error) diff --git a/cognite/extractorutils/unstable/core/checkin_worker.py b/cognite/extractorutils/unstable/core/checkin_worker.py index c9c295d3..89cb8ab3 100644 --- a/cognite/extractorutils/unstable/core/checkin_worker.py +++ b/cognite/extractorutils/unstable/core/checkin_worker.py @@ -16,6 +16,8 @@ from threading import RLock, Thread from time import sleep +import httpx +import requests from cognite.client import CogniteClient from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError from requests import Response @@ -43,6 +45,19 @@ MAX_ACTION_UPDATES_PER_CHECKIN = 100 rng = SystemRandom() +# Errors that mean "could not reach CDF right now" rather than "the request or credentials are wrong". The Cognite SDK +# only wraps some of these (e.g. connection errors from API calls, OAuth2Error from token refresh), so raw +# requests/httpx transport errors and timeouts (notably from the OAuth token refresh) can still surface unwrapped. +# Deliberately narrow: configuration mistakes (invalid URL, missing schema, unsupported protocol) must stay fatal. +_TRANSIENT_ERRORS: tuple[type[BaseException], ...] = ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + httpx.TimeoutException, + httpx.NetworkError, + TimeoutError, + ConnectionError, +) + def _safe_dispatch( dispatcher: Callable[[list[Action]], None], @@ -184,6 +199,18 @@ def set_on_fatal_error_handler(self, on_fatal_error: Callable[[Exception], None] """ self._on_fatal_error = on_fatal_error + def notify_fatal_error(self, error: Exception) -> None: + """ + Invoke the fatal error handler (if any) with the given error. + + Used by callers that fail outside of a check-in request, such as when building the startup request. + + Arguments: + error (Exception): The error that should be treated as fatal. + """ + if self._on_fatal_error is not None: + self._on_fatal_error(error) + def set_retry_startup(self, retry_startup: bool) -> None: """ Set whether to retry reporting startup if it fails. @@ -480,6 +507,15 @@ def _requeue_checkin( self._action_updates = action_updates + self._action_updates def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: + """ + Run a check-in like request and classify any failure. + + Transient network problems are logged and requeued without being fatal. Only authentication/authorization + failures (401/403, bad credentials) and unexpected errors are reported as fatal. + + Returns: + True if the payload of the request should be requeued. + """ exception_to_report: Exception | None = None requeue = False try: @@ -493,9 +529,13 @@ def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: except CogniteAuthError as e: self._logger.error(str(e)) - self._logger.critical("Could not get an access token. Please check your configuration.") - exception_to_report = e requeue = True + if isinstance(e.__cause__, _TRANSIENT_ERRORS): + # The token endpoint could not be reached; this says nothing about the credentials. + self._log_transient_error("Could not reach the token endpoint to get an access token", e.__cause__) + else: + self._logger.critical("Could not get an access token. Please check your configuration.") + exception_to_report = e except CogniteAPIError as e: if e.code == 401: @@ -519,6 +559,10 @@ def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: self._logger.critical(f"Error while connecting to CDF {e!s}") requeue = True + except _TRANSIENT_ERRORS as e: + self._log_transient_error("Could not reach CDF", e) + requeue = True + except Exception as e: self._logger.critical(f"Extractor could not connect to CDF {e!s}") exception_to_report = e @@ -528,3 +572,8 @@ def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: self._on_fatal_error(exception_to_report) return requeue + + def _log_transient_error(self, message: str, error: BaseException) -> None: + self._logger.warning( + "%s (%s: %s). The check-in will be retried at the next check-in.", message, type(error).__name__, error + ) diff --git a/cognite/extractorutils/unstable/core/runtime.py b/cognite/extractorutils/unstable/core/runtime.py index be1f055e..df5aca8b 100644 --- a/cognite/extractorutils/unstable/core/runtime.py +++ b/cognite/extractorutils/unstable/core/runtime.py @@ -34,6 +34,7 @@ def main() -> None: from multiprocessing.synchronize import Event as MpEvent from pathlib import Path from random import randint +from threading import Event as ThreadEvent from threading import Thread from typing import Any, Generic, TypeVar from uuid import uuid4 @@ -68,6 +69,11 @@ def main() -> None: ExtractorType = TypeVar("ExtractorType", bound=Extractor) +# Exit code the extractor process uses to tell the runtime that it stopped because of a fatal error (for example +# bad credentials), as opposed to a requested shutdown. Kept distinct from 1 so that unrelated crashes of the child +# process do not change the runtime's behaviour. +FATAL_ERROR_EXIT_CODE = 3 + @dataclass class _RuntimeControls: @@ -83,7 +89,13 @@ def _extractor_process_entrypoint( ) -> None: logger = logging.getLogger(f"{extractor_class.EXTERNAL_ID}.runtime") checkin_worker.active_revision = config.current_config_revision - checkin_worker.set_on_fatal_error_handler(lambda _: on_fatal_error(controls)) + fatal_error_occurred = ThreadEvent() + + def handle_fatal_error(_: Exception) -> None: + fatal_error_occurred.set() + on_fatal_error(controls) + + checkin_worker.set_on_fatal_error_handler(handle_fatal_error) checkin_worker.set_on_revision_change_handler(lambda _: on_revision_changed(controls)) checkin_worker.set_retry_startup(extractor_class.RETRY_STARTUP) extractor = extractor_class._init_from_runtime(config, checkin_worker) @@ -99,6 +111,11 @@ def _extractor_process_entrypoint( except Exception: logger.exception("Extractor crashed, will attempt restart") controls.message_queue.put(RuntimeMessage.RESTART) + return + + if fatal_error_occurred.is_set(): + logger.critical("Extractor stopped because of a fatal error") + sys.exit(FATAL_ERROR_EXIT_CODE) def on_revision_changed(controls: _RuntimeControls) -> None: @@ -598,6 +615,11 @@ def _main_runtime(self, args: Namespace) -> None: case _: assert_never(message) + elif process.exitcode == FATAL_ERROR_EXIT_CODE: + self.logger.critical("Extractor stopped because of a fatal error, shutting down runtime") + self._cancellation_token.cancel() + sys.exit(1) + else: self.logger.info("Shutting down runtime") self._cancellation_token.cancel() diff --git a/tests/test_unstable/test_checkin_resilience.py b/tests/test_unstable/test_checkin_resilience.py new file mode 100644 index 00000000..31e917ff --- /dev/null +++ b/tests/test_unstable/test_checkin_resilience.py @@ -0,0 +1,311 @@ +""" +Tests for check-in resilience: startup failures are fatal and logged, transient network errors are not fatal. + +These tests do not need a CDF project: the Cognite client is a mock whose ``post`` is controlled by each test. +""" + +import logging +from collections.abc import Callable +from datetime import datetime, timezone +from threading import Event, Thread +from typing import Any +from unittest.mock import MagicMock + +import httpx +import pytest +import requests +from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError +from oauthlib.oauth2 import OAuth2Error + +from cognite.extractorutils.threading import CancellationToken +from cognite.extractorutils.unstable.core._dto import ActionStatus, ActionUpdate +from cognite.extractorutils.unstable.core.base import FullConfig +from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker +from cognite.extractorutils.unstable.core.tasks import ScheduledTask + +from .conftest import TestConfig, TestExtractor + +MAX_AVAILABLE_ACTIONS = 100 + + +def _make_worker(post: Callable[..., Any] | None = None) -> tuple[CheckinWorker, MagicMock]: + client = MagicMock() + client.config.project = "test-project" + if post is not None: + client.post.side_effect = post + worker = CheckinWorker(client, "test-integration", logging.getLogger("test.checkin_resilience")) + worker.set_retry_startup(True) + return worker, client + + +def _make_extractor(worker: CheckinWorker, n_tasks: int = 0) -> TestExtractor: + conn = MagicMock() + conn.integration.external_id = "test-integration" + extractor = TestExtractor( + FullConfig( + connection_config=conn, + application_config=TestConfig(parameter_one=1, parameter_two="a"), + current_config_revision=1, + ), + worker, + ) + extractor._start_time = datetime.now(tz=timezone.utc) + for i in range(n_tasks): + extractor.add_task(ScheduledTask.from_interval(interval="1h", name=f"task-{i}", target=lambda _: None)) + return extractor + + +def _n_builtin_actions() -> int: + worker, _ = _make_worker() + request = _make_extractor(worker)._get_startup_request() + return len(request.available_actions or []) + + +def _ok_response() -> MagicMock: + response = MagicMock() + response.json.return_value = {"externalId": "test-integration"} + return response + + +def _raiser(error: BaseException) -> Callable[..., Any]: + def post(*args: Any, **kwargs: Any) -> Any: + raise error + + return post + + +def _oauth_auth_error() -> CogniteAuthError: + error = CogniteAuthError("Error generating access token: invalid_client") + error.__cause__ = OAuth2Error("invalid_client") + return error + + +def _timeout_auth_error() -> CogniteAuthError: + error = CogniteAuthError("Error generating access token") + error.__cause__ = requests.exceptions.ReadTimeout("token endpoint timed out") + return error + + +# --------------------------------------------------------------------------- +# Startup request failures +# --------------------------------------------------------------------------- + + +def test_too_many_available_actions_is_logged_and_fatal(caplog: pytest.LogCaptureFixture) -> None: + """More than 100 available actions fails StartupRequest validation; this must not kill the thread silently.""" + worker, client = _make_worker(lambda *a, **k: _ok_response()) + fatal_errors: list[Exception] = [] + worker.set_on_fatal_error_handler(fatal_errors.append) + + n_tasks = (MAX_AVAILABLE_ACTIONS - _n_builtin_actions()) // 2 + 1 # one task over the limit + extractor = _make_extractor(worker, n_tasks=n_tasks) + + with caplog.at_level(logging.WARNING): + extractor._run_checkin() # runs synchronously; would raise (thread dies) without the fix + + assert len(fatal_errors) == 1 + assert extractor.cancellation_token.is_cancelled + client.post.assert_not_called() + critical = [r for r in caplog.records if r.levelno >= logging.CRITICAL] + assert critical, "failure must be logged at CRITICAL so it reaches the log file" + assert critical[0].exc_info is not None + + +def test_exactly_100_available_actions_is_fine() -> None: + worker, _ = _make_worker() + n_tasks = (MAX_AVAILABLE_ACTIONS - _n_builtin_actions()) // 2 + extractor = _make_extractor(worker, n_tasks=n_tasks) + request = extractor._get_startup_request() + assert len(request.available_actions or []) <= MAX_AVAILABLE_ACTIONS + + +def test_startup_not_retried_failure_is_fatal(caplog: pytest.LogCaptureFixture) -> None: + """A startup failure when retries are disabled used to kill the check-in thread with no log.""" + worker, _ = _make_worker(_raiser(CogniteConnectionError())) + worker.set_retry_startup(False) + fatal_errors: list[Exception] = [] + worker.set_on_fatal_error_handler(fatal_errors.append) + extractor = _make_extractor(worker) + + with caplog.at_level(logging.WARNING): + extractor._run_checkin() + + assert len(fatal_errors) == 1 + assert any(r.levelno >= logging.CRITICAL for r in caplog.records) + + +# --------------------------------------------------------------------------- +# Transient errors are not fatal +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "error", + [ + pytest.param(requests.exceptions.ReadTimeout("read timed out"), id="requests-read-timeout"), + pytest.param(requests.exceptions.ConnectTimeout("connect timed out"), id="requests-connect-timeout"), + pytest.param(requests.exceptions.ConnectionError("reset"), id="requests-connection-error"), + pytest.param(httpx.ReadTimeout("read timed out"), id="httpx-read-timeout"), + pytest.param(httpx.ConnectTimeout("connect timed out"), id="httpx-connect-timeout"), + pytest.param(httpx.ConnectError("refused"), id="httpx-connect-error"), + pytest.param(TimeoutError("timed out"), id="builtin-timeout"), + pytest.param(CogniteConnectionError(), id="sdk-connection-error"), + pytest.param(CogniteAPIError("Bad gateway", code=502), id="http-502"), + pytest.param(CogniteAPIError("Unavailable", code=503), id="http-503"), + pytest.param(_timeout_auth_error(), id="token-refresh-timeout-wrapped"), + ], +) +def test_transient_error_is_requeued_and_not_fatal(error: BaseException) -> None: + worker, _ = _make_worker(_raiser(error)) + on_fatal = MagicMock() + worker.set_on_fatal_error_handler(on_fatal) + update = ActionUpdate(external_id="act-1", status=ActionStatus.running) + worker.queue_action_update(update) + + worker.try_write_checkin([], []) + + on_fatal.assert_not_called() + assert worker._action_updates == [update], "the update must be requeued for the next check-in" + + +def test_transient_error_during_startup_is_not_fatal_and_is_retried(monkeypatch: pytest.MonkeyPatch) -> None: + calls = 0 + token = CancellationToken() + + def post(*args: Any, **kwargs: Any) -> Any: + nonlocal calls + calls += 1 + if calls == 1: + raise requests.exceptions.ReadTimeout("token refresh timed out") + return _ok_response() + + worker, _ = _make_worker(post) + on_fatal = MagicMock() + worker.set_on_fatal_error_handler(on_fatal) + monkeypatch.setattr("cognite.extractorutils.unstable.core.checkin_worker.sleep", lambda _: None) + extractor = _make_extractor(worker) + + worker._run_startup_report(token, extractor._get_startup_request(), interval=0.01) + + assert calls == 2 + assert worker._has_reported_startup + on_fatal.assert_not_called() + + +# --------------------------------------------------------------------------- +# Auth/config errors stay fatal +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "error", + [ + pytest.param(CogniteAPIError("Unauthorized", code=401), id="401"), + pytest.param(CogniteAPIError("Forbidden", code=403), id="403"), + pytest.param(_oauth_auth_error(), id="bad-credentials"), + pytest.param(CogniteAuthError("Unable to determine device authorization endpoint"), id="auth-error-no-cause"), + pytest.param(ValueError("unexpected"), id="unexpected-error"), + pytest.param(requests.exceptions.InvalidURL("bad url"), id="invalid-url-is-a-config-error"), + pytest.param(requests.exceptions.MissingSchema("no scheme"), id="missing-schema-is-a-config-error"), + pytest.param(httpx.UnsupportedProtocol("bad protocol"), id="unsupported-protocol-is-a-config-error"), + ], +) +def test_auth_and_config_errors_remain_fatal(error: BaseException) -> None: + worker, _ = _make_worker(_raiser(error)) + on_fatal = MagicMock() + worker.set_on_fatal_error_handler(on_fatal) + + worker.try_write_checkin([], []) + + on_fatal.assert_called_once() + assert on_fatal.call_args[0][0] is error + + +def test_periodic_checkin_keeps_running_after_read_timeout() -> None: + """End to end: a timeout in a periodic check-in does not stop the loop, and it recovers on the next cycle.""" + calls = 0 + token = CancellationToken() + recovered = Event() + + def post(*args: Any, **kwargs: Any) -> Any: + nonlocal calls + calls += 1 + if calls == 2: # first periodic check-in after startup + raise httpx.ReadTimeout("timed out") + if calls >= 3: + recovered.set() + return _ok_response() + + worker, _ = _make_worker(post) + on_fatal = MagicMock() + worker.set_on_fatal_error_handler(on_fatal) + extractor = _make_extractor(worker) + + thread = Thread( + target=worker.run_periodic_checkin, args=(token, extractor._get_startup_request(), 0.05), daemon=True + ) + thread.start() + assert recovered.wait(timeout=10) + token.cancel() + thread.join(timeout=10) + + on_fatal.assert_not_called() + assert calls >= 3 + + +# --------------------------------------------------------------------------- +# Exit code when the extractor stops because of a fatal error +# --------------------------------------------------------------------------- + + +class _FakeExtractor: + """Stands in for an Extractor in the child process entrypoint; ``run`` optionally hits a fatal check-in error.""" + + EXTERNAL_ID = "fake" + RETRY_STARTUP = True + raise_fatal = False + + def __init__(self, worker: CheckinWorker) -> None: + self._worker = worker + + @classmethod + def _init_from_runtime(cls, config: Any, checkin_worker: CheckinWorker) -> "_FakeExtractor": + return cls(checkin_worker) + + def _attach_runtime_controls(self, **_: Any) -> None: + pass + + def __enter__(self) -> "_FakeExtractor": + return self + + def __exit__(self, *_: Any) -> bool: + return False + + def run(self) -> None: + if self.raise_fatal: + self._worker.notify_fatal_error(RuntimeError("bad credentials")) + + +@pytest.mark.parametrize("fatal", [True, False]) +def test_extractor_process_exit_code_reflects_fatal_error(fatal: bool) -> None: + from cognite.extractorutils.unstable.core.runtime import ( + FATAL_ERROR_EXIT_CODE, + _extractor_process_entrypoint, + _RuntimeControls, + ) + + class Extractor(_FakeExtractor): + raise_fatal = fatal + + controls = _RuntimeControls(cancel_event=MagicMock(), message_queue=MagicMock()) + config = MagicMock() + config.current_config_revision = 1 + worker, _ = _make_worker() + + if fatal: + with pytest.raises(SystemExit) as exc_info: + _extractor_process_entrypoint(Extractor, controls, config, worker) # type: ignore[arg-type] + assert exc_info.value.code == FATAL_ERROR_EXIT_CODE + controls.cancel_event.set.assert_called_once() + else: + _extractor_process_entrypoint(Extractor, controls, config, worker) # type: ignore[arg-type] From 1c86a51c37fe5917454b0b37c215be5246cf2458 Mon Sep 17 00:00:00 2001 From: vikramlc Date: Tue, 3 Nov 2026 14:02:19 +0530 Subject: [PATCH 3/3] fix(odin): Addressed review comments --- CHANGELOG.md | 4 +- .../unstable/core/checkin_worker.py | 7 +- .../extractorutils/unstable/core/runtime.py | 5 ++ .../test_unstable/test_checkin_resilience.py | 64 +++++++++++++------ 4 files changed, 56 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d0e42205..f9378143 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,8 +15,8 @@ Changes are grouped as follows ## Unreleased ### Fixed -* In the `unstable` package: Fixed the check-in thread dying silently when the startup request could not be built or reported (for example when more than 100 actions are available). The failure is now logged at critical level and treated as fatal, and the extractor process exits with a non-zero exit code when it stops because of a fatal check-in error. -* In the `unstable` package: Fixed transient network errors during check-ins (timeouts, connection errors, including raw `requests`/`httpx` errors from the OAuth token refresh) being treated as fatal and shutting the extractor down. They are now logged as warnings and retried at the next check-in. Only 401/403 errors and bad credentials remain fatal. +* In the `unstable` package: Fixed the check-in thread dying silently when the startup request could not be built or reported (for example when the startup request is rejected). The failure is now logged at critical level and treated as fatal, and the extractor process exits with a non-zero exit code when it stops because of a fatal check-in error. +* In the `unstable` package: Fixed transient network errors during check-ins (timeouts, connection errors, including raw `requests`/`httpx` errors from the OAuth token refresh) being treated as fatal and shutting the extractor down. They are now logged as warnings and retried at the next check-in. Server-side errors from the token endpoint are retried as well. Only 401/403 errors and bad credentials remain fatal. ## 7.13.2 diff --git a/cognite/extractorutils/unstable/core/checkin_worker.py b/cognite/extractorutils/unstable/core/checkin_worker.py index 89cb8ab3..4afddcb9 100644 --- a/cognite/extractorutils/unstable/core/checkin_worker.py +++ b/cognite/extractorutils/unstable/core/checkin_worker.py @@ -20,6 +20,7 @@ import requests from cognite.client import CogniteClient from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError +from oauthlib.oauth2 import ServerError, TemporarilyUnavailableError # type: ignore[import-untyped] from requests import Response from cognite.extractorutils.threading import CancellationToken @@ -58,6 +59,10 @@ ConnectionError, ) +# The identity provider reporting a problem on its side rather than with the credentials. The SDK wraps these in a +# CogniteAuthError, with the oauthlib error as the cause. +_TRANSIENT_TOKEN_ERRORS: tuple[type[BaseException], ...] = (ServerError, TemporarilyUnavailableError) + def _safe_dispatch( dispatcher: Callable[[list[Action]], None], @@ -530,7 +535,7 @@ def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool: except CogniteAuthError as e: self._logger.error(str(e)) requeue = True - if isinstance(e.__cause__, _TRANSIENT_ERRORS): + if isinstance(e.__cause__, _TRANSIENT_ERRORS + _TRANSIENT_TOKEN_ERRORS): # The token endpoint could not be reached; this says nothing about the credentials. self._log_transient_error("Could not reach the token endpoint to get an access token", e.__cause__) else: diff --git a/cognite/extractorutils/unstable/core/runtime.py b/cognite/extractorutils/unstable/core/runtime.py index df5aca8b..75962c4c 100644 --- a/cognite/extractorutils/unstable/core/runtime.py +++ b/cognite/extractorutils/unstable/core/runtime.py @@ -109,6 +109,11 @@ def handle_fatal_error(_: Exception) -> None: extractor.run() except Exception: + if fatal_error_occurred.is_set(): + # The fatal error cancelled the extractor, which can make running tasks raise while shutting down. Don't + # mistake that for a crash, or the runtime would restart an extractor that can never succeed. + logger.exception("Extractor raised while stopping because of a fatal error") + sys.exit(FATAL_ERROR_EXIT_CODE) logger.exception("Extractor crashed, will attempt restart") controls.message_queue.put(RuntimeMessage.RESTART) return diff --git a/tests/test_unstable/test_checkin_resilience.py b/tests/test_unstable/test_checkin_resilience.py index 31e917ff..dd20f729 100644 --- a/tests/test_unstable/test_checkin_resilience.py +++ b/tests/test_unstable/test_checkin_resilience.py @@ -15,7 +15,7 @@ import pytest import requests from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError -from oauthlib.oauth2 import OAuth2Error +from oauthlib.oauth2 import OAuth2Error, ServerError, TemporarilyUnavailableError from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.unstable.core._dto import ActionStatus, ActionUpdate @@ -25,8 +25,6 @@ from .conftest import TestConfig, TestExtractor -MAX_AVAILABLE_ACTIONS = 100 - def _make_worker(post: Callable[..., Any] | None = None) -> tuple[CheckinWorker, MagicMock]: client = MagicMock() @@ -55,12 +53,6 @@ def _make_extractor(worker: CheckinWorker, n_tasks: int = 0) -> TestExtractor: return extractor -def _n_builtin_actions() -> int: - worker, _ = _make_worker() - request = _make_extractor(worker)._get_startup_request() - return len(request.available_actions or []) - - def _ok_response() -> MagicMock: response = MagicMock() response.json.return_value = {"externalId": "test-integration"} @@ -80,6 +72,12 @@ def _oauth_auth_error() -> CogniteAuthError: return error +def _idp_auth_error(cause: OAuth2Error) -> CogniteAuthError: + error = CogniteAuthError("Error generating access token") + error.__cause__ = cause + return error + + def _timeout_auth_error() -> CogniteAuthError: error = CogniteAuthError("Error generating access token") error.__cause__ = requests.exceptions.ReadTimeout("token endpoint timed out") @@ -91,14 +89,19 @@ def _timeout_auth_error() -> CogniteAuthError: # --------------------------------------------------------------------------- -def test_too_many_available_actions_is_logged_and_fatal(caplog: pytest.LogCaptureFixture) -> None: - """More than 100 available actions fails StartupRequest validation; this must not kill the thread silently.""" +def test_failure_building_startup_request_is_logged_and_fatal( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """Any failure building the startup request must not kill the check-in thread silently.""" worker, client = _make_worker(lambda *a, **k: _ok_response()) fatal_errors: list[Exception] = [] worker.set_on_fatal_error_handler(fatal_errors.append) + extractor = _make_extractor(worker) + + def broken_startup_request() -> None: + raise ValueError("cannot build startup request") - n_tasks = (MAX_AVAILABLE_ACTIONS - _n_builtin_actions()) // 2 + 1 # one task over the limit - extractor = _make_extractor(worker, n_tasks=n_tasks) + monkeypatch.setattr(extractor, "_get_startup_request", broken_startup_request) with caplog.at_level(logging.WARNING): extractor._run_checkin() # runs synchronously; would raise (thread dies) without the fix @@ -111,14 +114,6 @@ def test_too_many_available_actions_is_logged_and_fatal(caplog: pytest.LogCaptur assert critical[0].exc_info is not None -def test_exactly_100_available_actions_is_fine() -> None: - worker, _ = _make_worker() - n_tasks = (MAX_AVAILABLE_ACTIONS - _n_builtin_actions()) // 2 - extractor = _make_extractor(worker, n_tasks=n_tasks) - request = extractor._get_startup_request() - assert len(request.available_actions or []) <= MAX_AVAILABLE_ACTIONS - - def test_startup_not_retried_failure_is_fatal(caplog: pytest.LogCaptureFixture) -> None: """A startup failure when retries are disabled used to kill the check-in thread with no log.""" worker, _ = _make_worker(_raiser(CogniteConnectionError())) @@ -153,6 +148,8 @@ def test_startup_not_retried_failure_is_fatal(caplog: pytest.LogCaptureFixture) pytest.param(CogniteAPIError("Bad gateway", code=502), id="http-502"), pytest.param(CogniteAPIError("Unavailable", code=503), id="http-503"), pytest.param(_timeout_auth_error(), id="token-refresh-timeout-wrapped"), + pytest.param(_idp_auth_error(ServerError("idp failed")), id="token-endpoint-server-error"), + pytest.param(_idp_auth_error(TemporarilyUnavailableError()), id="token-endpoint-temporarily-unavailable"), ], ) def test_transient_error_is_requeued_and_not_fatal(error: BaseException) -> None: @@ -309,3 +306,28 @@ class Extractor(_FakeExtractor): controls.cancel_event.set.assert_called_once() else: _extractor_process_entrypoint(Extractor, controls, config, worker) # type: ignore[arg-type] + + +def test_exception_during_shutdown_after_fatal_error_exits_fatal_not_restart() -> None: + """A fatal error cancels the extractor; tasks raising while stopping must not trigger a runtime restart.""" + from cognite.extractorutils.unstable.core.runtime import ( + FATAL_ERROR_EXIT_CODE, + _extractor_process_entrypoint, + _RuntimeControls, + ) + + class Extractor(_FakeExtractor): + def run(self) -> None: + self._worker.notify_fatal_error(RuntimeError("bad credentials")) + raise RuntimeError("task interrupted by shutdown") + + controls = _RuntimeControls(cancel_event=MagicMock(), message_queue=MagicMock()) + config = MagicMock() + config.current_config_revision = 1 + worker, _ = _make_worker() + + with pytest.raises(SystemExit) as exc_info: + _extractor_process_entrypoint(Extractor, controls, config, worker) # type: ignore[arg-type] + + assert exc_info.value.code == FATAL_ERROR_EXIT_CODE + controls.message_queue.put.assert_not_called()