fix(odin): Shut down cleanly with running actions and orphaned extractor processes (EDG-1033)

CHANGELOG.md

@@ -15,6 +15,8 @@ Changes are grouped as follows
 ## Unreleased
 
 ### Fixed
+* In the `unstable` package: Fixed actions left in `running` at Integrations API after the extractor shut down (SIGINT, service stop). On shutdown running actions are now cancelled, given up to 10 seconds to finish, and any that do not stop are reported as `canceled`; the final check-in is then sent with a fresh timeout instead of the already cancelled token.
+* In the `unstable` package: Fixed the extractor process continuing to run and check in as an orphan when the runtime process is killed (for example with SIGKILL). The extractor process now watches its parent and shuts down gracefully within a few seconds.
 * In the `unstable` package: Action status updates (`running`, final statuses) are now reported to Integrations API shortly after they are queued (about one second, debounced) instead of waiting for the next periodic check-in.
 * 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.

cognite/extractorutils/unstable/core/base.py

@@ -55,6 +55,8 @@ def my_task_function(self, task_context: TaskContext) -> None:
 """
 
 import logging
+import os
+import sys
 import time
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
@@ -65,6 +67,7 @@ def my_task_function(self, task_context: TaskContext) -> None:
 from types import TracebackType
 from typing import Generic, TypeVar
 
+import psutil
 from humps import pascalize
 from typing_extensions import Self, assert_never
 
@@ -127,6 +130,30 @@ def my_task_function(self, task_context: TaskContext) -> None:
 """Task types that can support Start/Stop actions.
 Whether a given task supports actions is set by its ``actionable`` flag."""
 
+SHUTDOWN_ACTION_JOIN_TIMEOUT_SECONDS = 10.0
+"""How long ``__exit__`` waits for running action threads to finish after cancelling them. Actions still running
+after this are abandoned (they run on daemon threads) and reported as ``canceled``."""
+
+SHUTDOWN_FLUSH_TIMEOUT_SECONDS = 10.0
+"""How long ``__exit__`` waits for the final check-in (which carries the final action statuses) to be sent."""
+
+PARENT_WATCHDOG_INTERVAL_SECONDS = 1.0
+"""How often the extractor process checks that the runtime process that spawned it is still alive."""
+
+SHUTDOWN_ACTION_MESSAGE = "Extractor shut down while the action was running"
+
+
+def _parent_is_alive(parent_pid: int) -> bool:
+    """
+    Check whether the process with the given PID, which spawned this process, is still running.
+
+    On POSIX an orphaned process is re-parented, so a changed parent PID means the original parent is gone. Windows
+    does not re-parent, so there the parent is looked up by PID (psutil also guards against PID reuse).
+    """
+    if sys.platform != "win32":
+        return os.getppid() == parent_pid
+    return psutil.pid_exists(parent_pid)  # pragma: no cover - Windows only
+
 
 class FullConfig(Generic[_T]):
     """
@@ -207,6 +234,9 @@ def __init__(self, config: FullConfig[ConfigType], checkin_worker: CheckinWorker
         # _handle_start_task_action/_handle_custom_action, consulted by _dispatch_single_action.
         self._running_action_tokens: dict[str, CancellationToken] = {}
         self._running_action_tokens_lock = RLock()
+        # Threads running dispatched actions, joined (bounded) on shutdown so they can report a final status.
+        self._action_threads: list[Thread] = []
+        self._action_threads_lock = RLock()
         self._start_time: datetime
 
         self.metrics: BaseMetrics = self._load_metrics(config.metrics_class)
@@ -399,9 +429,32 @@ def add_action(self, action: CustomAction) -> None:
     def _set_runtime_message_queue(self, queue: Queue) -> None:
         self._runtime_messages = queue
 
-    def _attach_runtime_controls(self, *, cancel_event: MpEvent, message_queue: Queue) -> None:
+    def _attach_runtime_controls(
+        self, *, cancel_event: MpEvent, message_queue: Queue, parent_pid: int | None = None
+    ) -> None:
         self._set_runtime_message_queue(message_queue)
         self._setup_cancellation_watcher(cancel_event)
+        if parent_pid is not None:
+            self._setup_parent_watchdog(parent_pid)
+
+    def _setup_parent_watchdog(self, parent_pid: int) -> None:
+        """
+        Shut down gracefully if the runtime process that spawned this process dies.
+
+        The runtime stops the extractor through a multiprocessing event, which is never set if the runtime is killed
+        (for example with SIGKILL), so the extractor would otherwise keep running and checking in as an orphan.
+        """
+
+        def watchdog() -> None:
+            while not self.cancellation_token.wait(PARENT_WATCHDOG_INTERVAL_SECONDS):
+                if not _parent_is_alive(parent_pid):
+                    self._logger.critical(
+                        "Runtime process (PID %d) is gone, shutting down the extractor gracefully.", parent_pid
+                    )
+                    self.cancellation_token.cancel()
+                    return
+
+        Thread(target=watchdog, name="ParentProcessWatchdog", daemon=True).start()
 
     def _get_startup_request(self) -> StartupRequest:
         available_actions: list[AvailableActionWrite] = []
@@ -626,12 +679,64 @@ def _launch_continuous_task(self, task: ContinuousTask) -> None:
 
     def _handle_actions(self, actions: list[Action]) -> None:
         for action in actions:
-            Thread(
+            thread = Thread(
                 target=self._dispatch_single_action,
                 args=(action,),
                 name=f"Action-{action.external_id}",
                 daemon=True,
-            ).start()
+            )
+            with self._action_threads_lock:
+                self._action_threads = [t for t in self._action_threads if t.is_alive()]
+                self._action_threads.append(thread)
+            thread.start()
+
+    def _finish_running_actions(self) -> None:
+        """
+        Cancel running actions, wait (bounded) for them, and queue ``canceled`` for any that did not finish.
+
+        Action handlers report their own final status once they return, so a well-behaved action needs no help. An
+        action that is still registered after the timeout is stuck in a blocking call: its daemon thread is abandoned
+        and the shutdown is reported on its behalf so Odin does not see it as running until the next startup.
+        """
+        with self._running_action_tokens_lock:
+            tokens = list(self._running_action_tokens.values())
+        for token in tokens:
+            token.cancel()
+
+        with self._action_threads_lock:
+            threads = list(self._action_threads)
+        deadline = time.monotonic() + SHUTDOWN_ACTION_JOIN_TIMEOUT_SECONDS
+        for thread in threads:
+            thread.join(max(0.0, deadline - time.monotonic()))
+
+        with self._running_action_tokens_lock:
+            abandoned = list(self._running_action_tokens)
+        for external_id in abandoned:
+            self._logger.warning("Action %s did not stop in time, reporting it as canceled", external_id)
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(
+                    external_id=external_id, status=ActionStatus.canceled, result_message=SHUTDOWN_ACTION_MESSAGE
+                )
+            )
+
+    def _final_flush(self) -> None:
+        """
+        Send everything still queued, using a fresh token because ``self.cancellation_token`` is already cancelled.
+
+        Bounded by ``SHUTDOWN_FLUSH_TIMEOUT_SECONDS`` so an unreachable CDF cannot hang the shutdown.
+        """
+
+        def flush() -> None:
+            try:
+                self._checkin_worker.flush(CancellationToken())
+            except Exception:
+                self._logger.exception("Final check-in on shutdown failed")
+
+        thread = Thread(target=flush, name="ExtractorFinalFlush", daemon=True)
+        thread.start()
+        thread.join(SHUTDOWN_FLUSH_TIMEOUT_SECONDS)
+        if thread.is_alive():
+            self._logger.warning("Final check-in did not complete within %.0f seconds", SHUTDOWN_FLUSH_TIMEOUT_SECONDS)
 
     def _dispatch_single_action(self, action: Action) -> None:
         if action.status == ActionStatus.cancel_pending:
@@ -897,11 +1002,12 @@ def __exit__(
         Stop the extractor when exiting the context manager.
         """
         self.stop()
+        self._finish_running_actions()
 
         if self.state_store:
             self.state_store.synchronize()
 
-        self._checkin_worker.flush(self.cancellation_token)
+        self._final_flush()
         return exc_val is None
 
     def run(self) -> None:

cognite/extractorutils/unstable/core/runtime.py

@@ -79,6 +79,8 @@ def main() -> None:
 class _RuntimeControls:
     cancel_event: MpEvent
     message_queue: Queue
+    parent_pid: int | None = None
+    """PID of the runtime process; the extractor process shuts down if it disappears."""
 
 
 def _extractor_process_entrypoint(
@@ -102,6 +104,7 @@ def handle_fatal_error(_: Exception) -> None:
     extractor._attach_runtime_controls(
         cancel_event=controls.cancel_event,
         message_queue=controls.message_queue,
+        parent_pid=controls.parent_pid,
     )
 
     try:
@@ -337,6 +340,7 @@ def _spawn_extractor(
         controls = _RuntimeControls(
             cancel_event=self._cancel_event,
             message_queue=self._message_queue,
+            parent_pid=os.getpid(),
         )
 
         process = Process(

tests/test_unstable/test_clean_shutdown.py

@@ -0,0 +1,163 @@
+"""Tests for clean shutdown: final action statuses are sent on exit, and the parent-process watchdog."""
+
+import logging
+import os
+import sys
+import threading
+from datetime import datetime, timezone
+from threading import Event
+from time import monotonic, sleep
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from cognite.extractorutils.statestore import NoStateStore
+from cognite.extractorutils.unstable.core import base
+from cognite.extractorutils.unstable.core._dto import Action, ActionStatus
+from cognite.extractorutils.unstable.core.actions import ActionContext, CustomAction
+from cognite.extractorutils.unstable.core.base import FullConfig
+from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker
+
+from .conftest import TestConfig, TestExtractor
+
+
+def _make_extractor() -> tuple[TestExtractor, list[dict[str, Any]]]:
+    """Extractor with a real CheckinWorker whose HTTP posts are captured in the returned list."""
+    sent: list[dict[str, Any]] = []
+    response = MagicMock()
+    response.json.return_value = {"externalId": "test-integration"}
+
+    def post(url: str, json: dict[str, Any], **_: Any) -> MagicMock:
+        sent.append(json)
+        return response
+
+    client = MagicMock()
+    client.config.project = "test-project"
+    client.post.side_effect = post
+    worker = CheckinWorker(client, "test-integration", logging.getLogger("test.shutdown"))
+    worker._has_reported_startup = True
+
+    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)
+    extractor.state_store = NoStateStore()
+    return extractor, sent
+
+
+def _action_updates(sent: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
+    return {u["externalId"]: u for req in sent for u in req.get("actionUpdates", [])}
+
+
+def _start_action(extractor: TestExtractor, name: str, external_id: str) -> None:
+    extractor._handle_actions([Action(external_id=external_id, action_name=name, status=ActionStatus.pending)])
+
+
+def _wait_for(event: Event) -> None:
+    assert event.wait(5)
+
+
+def test_exit_sends_canceled_for_cooperative_running_action() -> None:
+    extractor, sent = _make_extractor()
+    started = Event()
+
+    def target(ctx: ActionContext) -> None:
+        started.set()
+        ctx.cancellation_token.wait()  # returns once shutdown cancels the token
+
+    extractor.add_action(CustomAction(name="long", target=target))
+    _start_action(extractor, "long", "act-1")
+    _wait_for(started)
+
+    extractor.__exit__(None, None, None)
+
+    updates = _action_updates(sent)
+    assert updates["act-1"]["status"] == "canceled"
+    assert extractor._checkin_worker._action_updates == [], "update must actually be sent, not just requeued"
+
+
+def test_exit_reports_canceled_for_stuck_action_after_bounded_wait(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(base, "SHUTDOWN_ACTION_JOIN_TIMEOUT_SECONDS", 0.3)
+    extractor, sent = _make_extractor()
+    started = Event()
+    release = Event()
+
+    def target(ctx: ActionContext) -> None:
+        started.set()
+        release.wait(30)  # ignores cancellation, like a blocking call
+
+    extractor.add_action(CustomAction(name="stuck", target=target))
+    _start_action(extractor, "stuck", "act-2")
+    _wait_for(started)
+
+    begin = monotonic()
+    try:
+        extractor.__exit__(None, None, None)
+        elapsed = monotonic() - begin
+    finally:
+        release.set()
+
+    assert elapsed < 5, "shutdown must stay bounded"
+    update = _action_updates(sent)["act-2"]
+    assert update["status"] == "canceled"
+    assert update["resultMessage"] == base.SHUTDOWN_ACTION_MESSAGE
+
+
+def test_exit_without_actions_still_flushes_and_does_not_hang() -> None:
+    extractor, sent = _make_extractor()
+    extractor.__exit__(None, None, None)
+    assert _action_updates(sent) == {}
+
+
+def test_exit_flush_is_bounded_when_cdf_hangs(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(base, "SHUTDOWN_FLUSH_TIMEOUT_SECONDS", 0.3)
+    extractor, _ = _make_extractor()
+    release = Event()
+    extractor._checkin_worker._cognite_client.post.side_effect = lambda *a, **k: release.wait(30)  # type: ignore[attr-defined]
+    begin = monotonic()
+    try:
+        extractor.__exit__(None, None, None)
+    finally:
+        release.set()
+    assert monotonic() - begin < 5
+
+
+# ---------------------------------------------------------------------------
+# Parent process watchdog
+# ---------------------------------------------------------------------------
+
+
+def test_watchdog_cancels_token_when_parent_disappears(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(base, "PARENT_WATCHDOG_INTERVAL_SECONDS", 0.05)
+    extractor, _ = _make_extractor()
+    alive = Event()
+    alive.set()
+    monkeypatch.setattr(base, "_parent_is_alive", lambda pid: alive.is_set())
+
+    extractor._setup_parent_watchdog(4242)
+    sleep(0.3)
+    assert not extractor.cancellation_token.is_cancelled
+
+    alive.clear()  # simulate the parent dying (the real check sees a changed os.getppid())
+    assert extractor.cancellation_token.wait(5)
+    assert extractor.cancellation_token.is_cancelled
+
+
+def test_attach_runtime_controls_without_parent_pid_starts_no_watchdog() -> None:
+    extractor, _ = _make_extractor()
+    extractor._attach_runtime_controls(cancel_event=MagicMock(), message_queue=MagicMock())
+    assert not [t for t in threading.enumerate() if t.name == "ParentProcessWatchdog"]
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="POSIX re-parenting semantics")
+def test_parent_is_alive_posix() -> None:
+    assert base._parent_is_alive(os.getppid())
+    assert not base._parent_is_alive(os.getppid() + 1)