diff --git a/cognite/extractorutils/unstable/core/actions.py b/cognite/extractorutils/unstable/core/actions.py index 477afcdd..f6cb50a6 100644 --- a/cognite/extractorutils/unstable/core/actions.py +++ b/cognite/extractorutils/unstable/core/actions.py @@ -6,6 +6,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Generic +from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.unstable.configuration.models import ConfigType from cognite.extractorutils.unstable.core._dto import ActionStatus, ActionUpdate from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel @@ -27,6 +28,10 @@ class ActionContext(Generic[ConfigType], CogniteLogger): ``external_id`` and ``call_metadata`` come from the pending action payload sent by Odin and are available for use when reporting results back to the server. + + ``cancellation_token`` is cancelled if the user cancels this specific action invocation from Odin + while it is running. Long-running custom actions should check ``cancellation_token.is_cancelled`` + (or use ``cancellation_token.wait(...)`` instead of blocking sleeps) to exit early when cancelled. """ def __init__( @@ -34,6 +39,7 @@ def __init__( action: "CustomAction", extractor: "Extractor[ConfigType]", external_id: str, + cancellation_token: CancellationToken, call_metadata: dict[str, str] | None = None, ) -> None: super().__init__() @@ -41,6 +47,7 @@ def __init__( self._extractor = extractor self.external_id = external_id self.call_metadata = call_metadata + self.cancellation_token = cancellation_token self._result_message: str | None = None self._result_metadata: dict[str, str] | None = None diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 52225dae..d01b1f3b 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -203,6 +203,11 @@ def __init__(self, config: FullConfig[ConfigType], checkin_worker: CheckinWorker self._custom_actions: list[CustomAction] = [] self._running_task_tokens: dict[str, CancellationToken] = {} self._running_task_tokens_lock = RLock() + # Keyed by Action.external_id (not task name) so a `cancel_pending` re-delivery of an + # in-flight start_task/custom action can be cancelled instead of re-dispatched. Populated by + # _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() self._start_time: datetime self.metrics: BaseMetrics = self._load_metrics(config.metrics_class) @@ -538,6 +543,22 @@ def run_task(task_context: TaskContext) -> None: task=lambda: self._run_task_with_token(t), ) + @staticmethod + def _release_if_owned( + tokens: dict[str, CancellationToken], lock: RLock, key: str, token: CancellationToken + ) -> None: + """ + Remove ``key`` from ``tokens`` only if it still maps to ``token``. + + Used when cleaning up a task/action's cancellation-token registration on completion or + error. The identity check guards against clobbering a newer registration that may have + already replaced this one under the same key (e.g. a task restarted with a fresh token + before this cleanup ran). + """ + with lock: + if tokens.get(key) is token: + tokens.pop(key, None) + def _run_task_with_token( self, task: ScheduledTask | ContinuousTask, child_token: CancellationToken | None = None ) -> None: @@ -563,9 +584,7 @@ def _run_task_with_token( try: task.target(TaskContext(task=task, extractor=self, cancellation_token=child_token)) finally: - with self._running_task_tokens_lock: - if self._running_task_tokens.get(task.name) is child_token: - self._running_task_tokens.pop(task.name, None) + self._release_if_owned(self._running_task_tokens, self._running_task_tokens_lock, task.name, child_token) def _launch_continuous_task(self, task: ContinuousTask) -> None: """ @@ -590,9 +609,7 @@ def _launch_continuous_task(self, task: ContinuousTask) -> None: args=(task, child_token), ).start() except Exception as e: - with self._running_task_tokens_lock: - if self._running_task_tokens.get(task.name) is child_token: - self._running_task_tokens.pop(task.name, None) + self._release_if_owned(self._running_task_tokens, self._running_task_tokens_lock, task.name, child_token) message = f"Failed to launch continuous task '{task.name}'" self._logger.log(level=ErrorLevel.fatal.log_level, msg=message, exc_info=e) self._new_error( @@ -612,6 +629,16 @@ def _handle_actions(self, actions: list[Action]) -> None: ).start() def _dispatch_single_action(self, action: Action) -> None: + if action.status == ActionStatus.cancel_pending: + # Odin re-delivers an already-dispatched action once a user cancels it, with its status + # flipped to cancel_pending. This is not a fresh invocation: cancel the token tracking the + # in-flight run (if we have one) and stop, rather than re-running the action's handler. + with self._running_action_tokens_lock: + token = self._running_action_tokens.get(action.external_id) + if token is not None: + token.cancel() + return + actionable_tasks = [t for t in self._tasks if isinstance(t, ACTIONABLE_TASK_TYPES)] scheduled_start_names = {f"Start {t.name}" for t in actionable_tasks} scheduled_stop_names = {f"Stop {t.name}" for t in actionable_tasks} @@ -661,11 +688,14 @@ def _handle_start_task_action(self, action: Action) -> None: return self._running_task_tokens[task_name] = child_token - self._checkin_worker.queue_action_update( - ActionUpdate(external_id=action.external_id, status=ActionStatus.running) - ) - try: + with self._running_action_tokens_lock: + self._running_action_tokens[action.external_id] = child_token + + self._checkin_worker.queue_action_update( + ActionUpdate(external_id=action.external_id, status=ActionStatus.running) + ) + self._run_task_with_token(task, child_token) status = ActionStatus.canceled if child_token.is_cancelled else ActionStatus.succeeded self._checkin_worker.queue_action_update(ActionUpdate(external_id=action.external_id, status=status)) @@ -677,6 +707,14 @@ def _handle_start_task_action(self, action: Action) -> None: result_message=str(e), ) ) + finally: + # Guards against a leaked "running" registration if something between here and + # _run_task_with_token's own cleanup (e.g. the ActionUpdate construction above) ever + # raises before _run_task_with_token gets a chance to run its own finally. + self._release_if_owned(self._running_task_tokens, self._running_task_tokens_lock, task_name, child_token) + self._release_if_owned( + self._running_action_tokens, self._running_action_tokens_lock, action.external_id, child_token + ) def _handle_stop_task_action(self, action: Action) -> None: task_name = action.action_name[len("Stop ") :] @@ -711,32 +749,42 @@ def _handle_custom_action(self, action: Action) -> None: ) return - self._checkin_worker.queue_action_update( - ActionUpdate(external_id=action.external_id, status=ActionStatus.running) - ) - - ctx = ActionContext( - action=custom, - extractor=self, - external_id=action.external_id, - call_metadata=action.call_metadata, - ) + action_token = self.cancellation_token.create_child_token() try: + with self._running_action_tokens_lock: + self._running_action_tokens[action.external_id] = action_token + + self._checkin_worker.queue_action_update( + ActionUpdate(external_id=action.external_id, status=ActionStatus.running) + ) + + ctx = ActionContext( + action=custom, + extractor=self, + external_id=action.external_id, + call_metadata=action.call_metadata, + cancellation_token=action_token, + ) + custom.target(ctx) filtered_metadata, oversized_fields = drop_oversized_metadata_fields(ctx._result_metadata) + completed_status = ActionStatus.canceled if action_token.is_cancelled else ActionStatus.succeeded if oversized_fields: - # The action itself ran to completion — only reporting the full result back to Odin - # failed, because a metadata value is too large to send. Fail the action instead of - # queuing a payload that Odin would reject (which would otherwise poison the checkin - # batch and retry forever, since checkin bundles all pending updates together and - # requeues the whole batch on any rejection). Non-oversized fields are still reported. + # The action itself ran to completion (or was cancelled) — only reporting the full + # result back to Odin is affected, because a metadata value is too large to send. + # Drop the oversized field(s) instead of queuing a payload that Odin would reject + # (which would otherwise poison the checkin batch and retry forever, since checkin + # bundles all pending updates together and requeues the whole batch on any rejection). + # Non-oversized fields are still reported, and the action's real outcome (succeeded or + # canceled) is preserved rather than being overwritten by this reporting issue. + outcome = "was canceled" if completed_status == ActionStatus.canceled else "completed successfully" self._checkin_worker.queue_action_update( ActionUpdate( external_id=action.external_id, - status=ActionStatus.failed, + status=completed_status, result_message=truncate_message( - f"Action '{custom.name}' completed successfully, but metadata field(s) " + f"Action '{custom.name}' {outcome}, but metadata field(s) " f"{', '.join(oversized_fields)} exceeded the {MAX_METADATA_VALUE_BYTES}-byte-per-value " f"limit and were dropped from the reported result" ), @@ -747,17 +795,20 @@ def _handle_custom_action(self, action: Action) -> None: self._checkin_worker.queue_action_update( ActionUpdate( external_id=action.external_id, - status=ActionStatus.succeeded, + status=completed_status, result_message=ctx._result_message, result_metadata=ctx._result_metadata, ) ) except ActionError as e: + # As with start_task actions, a cooperative abort in response to cancellation may surface + # as an ActionError rather than a clean return; report that as canceled, not failed. + status = ActionStatus.canceled if action_token.is_cancelled else ActionStatus.failed filtered_metadata, oversized_fields = drop_oversized_metadata_fields(e.result_metadata) self._checkin_worker.queue_action_update( ActionUpdate( external_id=action.external_id, - status=ActionStatus.failed, + status=status, result_message=( str(e) if not oversized_fields @@ -770,13 +821,18 @@ def _handle_custom_action(self, action: Action) -> None: ) ) except Exception as e: + status = ActionStatus.canceled if action_token.is_cancelled else ActionStatus.failed self._checkin_worker.queue_action_update( ActionUpdate( external_id=action.external_id, - status=ActionStatus.failed, + status=status, result_message=str(e), ) ) + finally: + self._release_if_owned( + self._running_action_tokens, self._running_action_tokens_lock, action.external_id, action_token + ) def start(self) -> None: """ diff --git a/tests/test_unstable/test_action_dispatch.py b/tests/test_unstable/test_action_dispatch.py index bace9c58..1d29dddc 100644 --- a/tests/test_unstable/test_action_dispatch.py +++ b/tests/test_unstable/test_action_dispatch.py @@ -30,8 +30,8 @@ def _queued_updates(extractor: TestExtractor) -> list[ActionUpdate]: return [c[0][0] for c in extractor._checkin_worker.queue_action_update.call_args_list] -def _make_action(external_id: str, action_name: str) -> Action: - return Action(external_id=external_id, action_name=action_name, status=ActionStatus.pending) +def _make_action(external_id: str, action_name: str, status: ActionStatus = ActionStatus.pending) -> Action: + return Action(external_id=external_id, action_name=action_name, status=status) def test_dispatch_unrecognised_action_name_reports_failed() -> None: @@ -317,7 +317,7 @@ def target(ctx: ActionContext) -> None: assert failed.result_message == "bad input" -def test_oversized_result_metadata_fails_action_but_keeps_valid_fields() -> None: +def test_oversized_result_metadata_keeps_completed_status_and_valid_fields() -> None: def target(ctx: ActionContext) -> None: ctx.set_result("done", metadata={"summary": "ok", "blob": "x" * 600}) @@ -327,13 +327,32 @@ def target(ctx: ActionContext) -> None: updates = _queued_updates(extractor) final = updates[-1] - assert final.status == ActionStatus.failed + assert final.status == ActionStatus.succeeded assert final.result_metadata == {"summary": "ok"} assert "big-result" in (final.result_message or "") + assert "completed successfully" in (final.result_message or "") assert "blob" in (final.result_message or "") assert "512" in (final.result_message or "") +def test_oversized_result_metadata_reports_canceled_when_cancelled() -> None: + def target(ctx: ActionContext) -> None: + ctx.cancellation_token.cancel() + ctx.set_result("done", metadata={"summary": "ok", "blob": "x" * 600}) + + extractor = _make_extractor() + extractor.add_action(CustomAction(name="big-result-canceled", target=target)) + extractor._dispatch_single_action(_make_action("act-big-canceled", "big-result-canceled")) + + updates = _queued_updates(extractor) + final = updates[-1] + assert final.status == ActionStatus.canceled + assert final.result_metadata == {"summary": "ok"} + assert "was canceled" in (final.result_message or "") + assert "completed successfully" not in (final.result_message or "") + assert "blob" in (final.result_message or "") + + def test_oversized_action_error_metadata_drops_only_oversized_field() -> None: def target(ctx: ActionContext) -> None: raise ActionError("bad input", error_type="invalid_parameter", details="x" * 600) @@ -383,7 +402,7 @@ def target(ctx: ActionContext) -> None: updates = _queued_updates(extractor) final = updates[-1] - assert final.status == ActionStatus.failed + assert final.status == ActionStatus.succeeded assert final.result_metadata == {"summary": "ok"} assert final.result_message is not None assert len(final.result_message) <= MAX_MESSAGE_LENGTH @@ -531,3 +550,180 @@ def capture(ctx: ActionContext) -> None: extractor._dispatch_single_action(_make_action("act-p", "probe")) assert captured["cdf_client"] is extractor.cognite_client assert captured["integration_external_id"] == "test-integration" + + +def test_cancel_pending_custom_action_cancels_token_without_rerunning_target() -> None: + call_count = {"n": 0} + started = Event() + allow_exit = Event() + + def target(ctx: ActionContext) -> None: + call_count["n"] += 1 + started.set() + allow_exit.wait(timeout=5) + + extractor = _make_extractor() + extractor.add_action(CustomAction(name="long-running", target=target)) + + dispatch_thread = threading.Thread( + target=extractor._dispatch_single_action, + args=(_make_action("act-1", "long-running"),), + daemon=True, + ) + dispatch_thread.start() + assert started.wait(timeout=5) + + with extractor._running_action_tokens_lock: + token = extractor._running_action_tokens.get("act-1") + assert token is not None and not token.is_cancelled + + # Odin re-delivers the same external_id with cancel_pending once the user cancels it. + extractor._dispatch_single_action(_make_action("act-1", "long-running", status=ActionStatus.cancel_pending)) + + assert token.is_cancelled + assert call_count["n"] == 1 # target was not re-invoked by the cancel_pending re-delivery + + allow_exit.set() + dispatch_thread.join(timeout=5) + + +def test_custom_action_reports_canceled_when_cancelled_mid_run() -> None: + def target(ctx: ActionContext) -> None: + ctx.cancellation_token.cancel() + + extractor = _make_extractor() + extractor.add_action(CustomAction(name="cooperative", target=target)) + extractor._dispatch_single_action(_make_action("act-1", "cooperative")) + + updates = _queued_updates(extractor) + assert any(u.status == ActionStatus.canceled and u.external_id == "act-1" for u in updates) + + +def test_cancel_pending_start_task_action_cancels_running_task_instead_of_redispatching() -> None: + extractor = _make_extractor() + task_started = Event() + allow_exit = Event() + + def cancellable(ctx: TaskContext) -> None: + task_started.set() + allow_exit.wait(timeout=5) + + extractor.add_task(ScheduledTask.from_interval(interval="1h", name="worker", target=cancellable)) + + dispatch_thread = threading.Thread( + target=extractor._dispatch_single_action, + args=(_make_action("act-1", "Start worker"),), + daemon=True, + ) + dispatch_thread.start() + assert task_started.wait(timeout=5) + + # Odin re-delivers the same external_id with cancel_pending once the user cancels it. + extractor._dispatch_single_action(_make_action("act-1", "Start worker", status=ActionStatus.cancel_pending)) + + with extractor._running_task_tokens_lock: + token = extractor._running_task_tokens.get("worker") + assert token is not None and token.is_cancelled + + allow_exit.set() + dispatch_thread.join(timeout=5) + + updates = _queued_updates(extractor) + statuses = [u.status for u in updates if u.external_id == "act-1"] + # No spurious "already running" failure from the cancel_pending re-delivery being re-dispatched. + assert ActionStatus.failed not in statuses + assert statuses[-1] == ActionStatus.canceled + + +def test_cancel_pending_unknown_action_is_a_no_op() -> None: + extractor = _make_extractor() + extractor._dispatch_single_action( + _make_action("act-unknown", "does not matter", status=ActionStatus.cancel_pending) + ) + + assert _queued_updates(extractor) == [] + + +def test_custom_action_reports_canceled_when_target_raises_action_error_after_cancellation() -> None: + def target(ctx: ActionContext) -> None: + ctx.cancellation_token.cancel() + raise ActionError("aborted early", error_type="canceled_mid_run") + + extractor = _make_extractor() + extractor.add_action(CustomAction(name="cooperative", target=target)) + extractor._dispatch_single_action(_make_action("act-1", "cooperative")) + + updates = _queued_updates(extractor) + final = updates[-1] + assert final.status == ActionStatus.canceled + assert final.result_message == "aborted early" + + +def test_custom_action_reports_canceled_when_target_raises_generic_exception_after_cancellation() -> None: + def target(ctx: ActionContext) -> None: + ctx.cancellation_token.cancel() + raise RuntimeError("connection reset") + + extractor = _make_extractor() + extractor.add_action(CustomAction(name="cooperative", target=target)) + extractor._dispatch_single_action(_make_action("act-1", "cooperative")) + + updates = _queued_updates(extractor) + final = updates[-1] + assert final.status == ActionStatus.canceled + assert final.result_message == "connection reset" + + +def test_custom_action_reports_failed_when_target_raises_without_cancellation() -> None: + def target(ctx: ActionContext) -> None: + raise ActionError("bad input", error_type="invalid_parameter") + + extractor = _make_extractor() + extractor.add_action(CustomAction(name="strict", target=target)) + extractor._dispatch_single_action(_make_action("act-1", "strict")) + + updates = _queued_updates(extractor) + final = updates[-1] + assert final.status == ActionStatus.failed + assert final.result_message == "bad input" + + +def test_start_task_action_cleans_up_registration_when_running_update_raises() -> None: + # Regression test: registration into _running_task_tokens/_running_action_tokens must happen + # inside the same try/finally that cleans them up, so a failure anywhere before the task + # actually runs (e.g. constructing the "running" ActionUpdate) can't leave the task stuck + # "running" forever. + extractor = _make_extractor() + extractor.add_task(ScheduledTask.from_interval(interval="1h", name="worker", target=lambda _: None)) + + def flaky_queue_update(update: ActionUpdate) -> None: + if update.status == ActionStatus.running: + raise RuntimeError("boom") + + extractor._checkin_worker.queue_action_update.side_effect = flaky_queue_update + + extractor._dispatch_single_action(_make_action("act-1", "Start worker")) + + assert "worker" not in extractor._running_task_tokens + assert "act-1" not in extractor._running_action_tokens + + updates = _queued_updates(extractor) + assert any(u.status == ActionStatus.failed and u.result_message == "boom" for u in updates) + + +def test_custom_action_cleans_up_registration_when_running_update_raises() -> None: + extractor = _make_extractor() + extractor.add_action(CustomAction(name="flaky", target=lambda ctx: None)) + + def flaky_queue_update(update: ActionUpdate) -> None: + if update.status == ActionStatus.running: + raise RuntimeError("boom") + + extractor._checkin_worker.queue_action_update.side_effect = flaky_queue_update + + extractor._dispatch_single_action(_make_action("act-1", "flaky")) + + assert "act-1" not in extractor._running_action_tokens + + updates = _queued_updates(extractor) + assert any(u.status == ActionStatus.failed and u.result_message == "boom" for u in updates) diff --git a/tests/test_unstable/test_actions.py b/tests/test_unstable/test_actions.py index f125f3bf..5648de51 100644 --- a/tests/test_unstable/test_actions.py +++ b/tests/test_unstable/test_actions.py @@ -2,6 +2,7 @@ import pytest +from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.unstable.core.actions import ActionContext, CustomAction from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel @@ -47,6 +48,7 @@ def test_action_context_attributes(mock_extractor: MagicMock, simple_action: Cus action=simple_action, extractor=mock_extractor, external_id="triggered-action-ext-id", + cancellation_token=CancellationToken(), call_metadata={"key": "value"}, ) @@ -55,20 +57,26 @@ def test_action_context_attributes(mock_extractor: MagicMock, simple_action: Cus def test_action_context_call_metadata_none(mock_extractor: MagicMock, simple_action: CustomAction) -> None: - ctx = ActionContext(action=simple_action, extractor=mock_extractor, external_id="ext-id") + ctx = ActionContext( + action=simple_action, extractor=mock_extractor, external_id="ext-id", cancellation_token=CancellationToken() + ) assert ctx.call_metadata is None def test_action_context_logger_name(mock_extractor: MagicMock, simple_action: CustomAction) -> None: - ctx = ActionContext(action=simple_action, extractor=mock_extractor, external_id="ext-id") + ctx = ActionContext( + action=simple_action, extractor=mock_extractor, external_id="ext-id", cancellation_token=CancellationToken() + ) assert ctx._logger.name == "test-extractor.action.myaction" def test_action_context_logger_name_strips_spaces(mock_extractor: MagicMock) -> None: action = CustomAction(name="process data", target=lambda ctx: None) - ctx = ActionContext(action=action, extractor=mock_extractor, external_id="ext-id") + ctx = ActionContext( + action=action, extractor=mock_extractor, external_id="ext-id", cancellation_token=CancellationToken() + ) assert ctx._logger.name == "test-extractor.action.processdata" @@ -89,7 +97,9 @@ def test_action_context_error_task_name( task_name: str | None, expected_task_name: str, ) -> None: - ctx = ActionContext(action=simple_action, extractor=mock_extractor, external_id="ext-id") + ctx = ActionContext( + action=simple_action, extractor=mock_extractor, external_id="ext-id", cancellation_token=CancellationToken() + ) ctx._new_error( level=ErrorLevel.warning, description="Something went wrong", details="some details", task_name=task_name @@ -108,7 +118,9 @@ def my_action(ctx: ActionContext) -> None: called.append(ctx) action = CustomAction(name="test", target=my_action) - ctx = ActionContext(action=action, extractor=mock_extractor, external_id="ext-id") + ctx = ActionContext( + action=action, extractor=mock_extractor, external_id="ext-id", cancellation_token=CancellationToken() + ) assert callable(action.target) action.target(ctx) diff --git a/tests/test_unstable/test_log_upload_action.py b/tests/test_unstable/test_log_upload_action.py index 9f46d008..3b798e38 100644 --- a/tests/test_unstable/test_log_upload_action.py +++ b/tests/test_unstable/test_log_upload_action.py @@ -5,6 +5,7 @@ import pytest +from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.unstable.configuration.models import ( LogFileHandlerConfig, LogLevel, @@ -569,7 +570,9 @@ def test_fetch_logs_action_all_files_missing_still_succeeds(tmp_path: Path) -> N def test_set_result_raises_on_second_call() -> None: extractor = _make_extractor() action = CustomAction(name="test", target=lambda ctx: None) - ctx = ActionContext(action=action, extractor=extractor, external_id="test-123") + ctx = ActionContext( + action=action, extractor=extractor, external_id="test-123", cancellation_token=CancellationToken() + ) ctx.set_result("first result") with pytest.raises(RuntimeError, match="set_result\\(\\) has already been called"): ctx.set_result("second result") @@ -631,6 +634,7 @@ def test_report_progress_queues_running_action_update() -> None: action=CustomAction(name="test-action", target=lambda ctx: None), extractor=extractor, external_id="act-progress", + cancellation_token=CancellationToken(), ) ctx.report_progress("Uploading: 1/3 files complete")