diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 3b8a6d9d..a8549b09 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -85,7 +85,10 @@ def my_task_function(self, task_context: TaskContext) -> None: create_state_store, ) from cognite.extractorutils.unstable.core._dto import ( + Action, + ActionStatus, ActionType, + ActionUpdate, AvailableActionWrite, CogniteModel, ExtractorInfo, @@ -96,7 +99,7 @@ def my_task_function(self, task_context: TaskContext) -> None: Task as DtoTask, ) from cognite.extractorutils.unstable.core._messaging import RuntimeMessage -from cognite.extractorutils.unstable.core.actions import CustomAction +from cognite.extractorutils.unstable.core.actions import ActionContext, CustomAction from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel from cognite.extractorutils.unstable.core.logger import CogniteLogger, RobustFileHandler @@ -514,17 +517,25 @@ def run_task(task_context: TaskContext) -> None: task=lambda: self._run_task_with_token(t), ) - def _run_task_with_token(self, task: ScheduledTask) -> None: + def _run_task_with_token(self, task: ScheduledTask, child_token: CancellationToken | None = None) -> None: """ Create a fresh child cancellation token for a scheduled task run. Stores the token in ``_running_task_tokens`` for external cancellation (e.g. a stop_task action), then executes the task. Called by both the scheduler and any future start_task action handler. + + If ``child_token`` is supplied (action dispatch path), the caller has already + registered it in ``_running_task_tokens`` atomically; this method skips the + registration step and goes straight to execution. """ - child_token = self.cancellation_token.create_child_token() - with self._running_task_tokens_lock: - self._running_task_tokens[task.name] = child_token + if child_token is None: + with self._running_task_tokens_lock: + if task.name in self._running_task_tokens: + self._logger.warning("Task '%s' is already running, skipping scheduled run.", task.name) + return + child_token = self.cancellation_token.create_child_token() + self._running_task_tokens[task.name] = child_token try: task.target(TaskContext(task=task, extractor=self, cancellation_token=child_token)) finally: @@ -532,6 +543,139 @@ def _run_task_with_token(self, task: ScheduledTask) -> None: if self._running_task_tokens.get(task.name) is child_token: self._running_task_tokens.pop(task.name, None) + def _handle_actions(self, actions: list[Action]) -> None: + for action in actions: + Thread( + target=self._dispatch_single_action, + args=(action,), + name=f"Action-{action.external_id}", + daemon=True, + ).start() + + def _dispatch_single_action(self, action: Action) -> None: + scheduled_start_names = {f"Start {t.name}" for t in self._tasks if isinstance(t, ScheduledTask)} + scheduled_stop_names = {f"Stop {t.name}" for t in self._tasks if isinstance(t, ScheduledTask)} + custom_names = {a.name for a in self._custom_actions} + + if action.action_name in scheduled_start_names: + self._handle_start_task_action(action) + elif action.action_name in scheduled_stop_names: + self._handle_stop_task_action(action) + elif action.action_name in custom_names: + self._handle_custom_action(action) + else: + self._checkin_worker.queue_action_update( + ActionUpdate( + external_id=action.external_id, + status=ActionStatus.failed, + result_message=f"No action named '{action.action_name}' registered", + ) + ) + + def _handle_start_task_action(self, action: Action) -> None: + task_name = action.action_name[len("Start ") :] + task = next( + (t for t in self._tasks if t.name == task_name and isinstance(t, ScheduledTask)), + None, + ) + if task is None: + self._checkin_worker.queue_action_update( + ActionUpdate( + external_id=action.external_id, + status=ActionStatus.failed, + result_message=f"No scheduled task named '{task_name}' found", + ) + ) + return + + child_token = self.cancellation_token.create_child_token() + with self._running_task_tokens_lock: + if task_name in self._running_task_tokens: + self._checkin_worker.queue_action_update( + ActionUpdate( + external_id=action.external_id, + status=ActionStatus.failed, + result_message=f"Task '{task_name}' is already running", + ) + ) + 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: + 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)) + except Exception as e: + self._checkin_worker.queue_action_update( + ActionUpdate( + external_id=action.external_id, + status=ActionStatus.failed, + result_message=str(e), + ) + ) + + def _handle_stop_task_action(self, action: Action) -> None: + task_name = action.action_name[len("Stop ") :] + with self._running_task_tokens_lock: + token = self._running_task_tokens.get(task_name) + if token is not None: + token.cancel() + + if token is None: + self._checkin_worker.queue_action_update( + ActionUpdate( + external_id=action.external_id, + status=ActionStatus.failed, + result_message=f"Task '{task_name}' is not currently running", + ) + ) + return + + self._checkin_worker.queue_action_update( + ActionUpdate(external_id=action.external_id, status=ActionStatus.canceled) + ) + + def _handle_custom_action(self, action: Action) -> None: + custom = next((a for a in self._custom_actions if a.name == action.action_name), None) + if custom is None: + self._checkin_worker.queue_action_update( + ActionUpdate( + external_id=action.external_id, + status=ActionStatus.failed, + result_message=f"No custom action named '{action.action_name}' registered", + ) + ) + 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, + ) + + try: + custom.target(ctx) + self._checkin_worker.queue_action_update( + ActionUpdate(external_id=action.external_id, status=ActionStatus.succeeded) + ) + except Exception as e: + self._checkin_worker.queue_action_update( + ActionUpdate( + external_id=action.external_id, + status=ActionStatus.failed, + result_message=str(e), + ) + ) + def start(self) -> None: """ Start the extractor. @@ -545,6 +689,8 @@ def start(self) -> None: self._load_state_store() self.state_store.start() + self._checkin_worker.set_action_dispatcher(self._handle_actions) + Thread(target=self._run_checkin, name="ExtractorCheckin", daemon=True).start() if self.metrics_push_manager: self.metrics_push_manager.start() diff --git a/tests/test_unstable/test_action_dispatch.py b/tests/test_unstable/test_action_dispatch.py new file mode 100644 index 00000000..e885d617 --- /dev/null +++ b/tests/test_unstable/test_action_dispatch.py @@ -0,0 +1,249 @@ +import threading +from collections.abc import Callable +from threading import Event +from unittest.mock import MagicMock + +import pytest + +from cognite.extractorutils.unstable.core._dto import Action, ActionStatus, ActionUpdate +from cognite.extractorutils.unstable.core.actions import ActionContext, CustomAction +from cognite.extractorutils.unstable.core.base import FullConfig +from cognite.extractorutils.unstable.core.tasks import ScheduledTask, TaskContext + +from .conftest import TestConfig, TestExtractor + + +def _make_extractor(extractor_cls: type[TestExtractor] = TestExtractor) -> TestExtractor: + conn = MagicMock() + conn.integration.external_id = "test-integration" + full_config = FullConfig( + connection_config=conn, + application_config=TestConfig(parameter_one=1, parameter_two="a"), + current_config_revision=1, + ) + return extractor_cls(full_config, MagicMock()) + + +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 test_dispatch_unrecognised_action_name_reports_failed() -> None: + extractor = _make_extractor() + extractor._dispatch_single_action(_make_action("act-1", "DoesNotExist")) + + updates = _queued_updates(extractor) + assert len(updates) == 1 + assert updates[0].status == ActionStatus.failed + assert "DoesNotExist" in (updates[0].result_message or "") + + +@pytest.mark.parametrize( + "register,action_name,expected_status", + [ + pytest.param( + lambda ext: ext.add_task(ScheduledTask.from_interval(interval="1h", name="sync", target=lambda _: None)), + "Start sync", + ActionStatus.running, + id="start_task", + ), + pytest.param( + lambda ext: ext.add_action(CustomAction(name="flush", target=lambda ctx: None)), + "flush", + ActionStatus.succeeded, + id="custom_action", + ), + ], +) +def test_dispatch_routes_to_correct_handler( + register: Callable, action_name: str, expected_status: ActionStatus +) -> None: + extractor = _make_extractor() + register(extractor) + extractor._dispatch_single_action(_make_action("act-1", action_name)) + assert any(u.status == expected_status for u in _queued_updates(extractor)) + + +def test_start_task_action_already_running_reports_failed() -> None: + extractor = _make_extractor() + task_started = Event() + allow_finish = Event() + + def blocking(ctx: TaskContext) -> None: + task_started.set() + allow_finish.wait(timeout=5) + + extractor.add_task(ScheduledTask.from_interval(interval="1h", name="worker", target=blocking)) + extractor._scheduler.trigger("worker") + task_started.wait(timeout=5) + + extractor._dispatch_single_action(_make_action("act-dup", "Start worker")) + + allow_finish.set() + + updates = _queued_updates(extractor) + assert any(u.status == ActionStatus.failed and "already running" in (u.result_message or "") for u in updates) + + +def test_start_task_action_valid_idle_task_reports_running_then_succeeded() -> None: + extractor = _make_extractor() + ran = [] + + def quick(ctx: TaskContext) -> None: + ran.append(True) + + extractor.add_task(ScheduledTask.from_interval(interval="1h", name="quick", target=quick)) + extractor._dispatch_single_action(_make_action("act-1", "Start quick")) + + assert ran == [True] + updates = _queued_updates(extractor) + statuses = [u.status for u in updates if u.external_id == "act-1"] + assert ActionStatus.running in statuses + assert ActionStatus.succeeded in statuses + assert statuses.index(ActionStatus.running) < statuses.index(ActionStatus.succeeded) + + +def test_stop_task_action_not_running_reports_failed() -> None: + extractor = _make_extractor() + extractor.add_task(ScheduledTask.from_interval(interval="1h", name="worker", target=lambda _: None)) + + extractor._dispatch_single_action(_make_action("act-stop", "Stop worker")) + + updates = _queued_updates(extractor) + assert len(updates) == 1 + assert updates[0].status == ActionStatus.failed + assert "not currently running" in (updates[0].result_message or "") + + +def test_stop_task_action_cancels_child_token_and_reports_canceled() -> 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)) + extractor._scheduler.trigger("worker") + task_started.wait(timeout=5) + + with extractor._running_task_tokens_lock: + token = extractor._running_task_tokens.get("worker") + + extractor._dispatch_single_action(_make_action("act-stop", "Stop worker")) + + updates = _queued_updates(extractor) + assert any(u.status == ActionStatus.canceled and u.external_id == "act-stop" for u in updates) + assert token is not None and token.is_cancelled + + allow_exit.set() + + +@pytest.mark.parametrize( + "raises,expected_final,expected_message", + [ + (False, ActionStatus.succeeded, None), + (True, ActionStatus.failed, "something went wrong"), + ], +) +def test_custom_action_status_lifecycle( + raises: bool, expected_final: ActionStatus, expected_message: str | None +) -> None: + def target(ctx: ActionContext) -> None: + if raises: + raise ValueError("something went wrong") + + extractor = _make_extractor() + extractor.add_action(CustomAction(name="act", target=target)) + extractor._dispatch_single_action(_make_action("act-1", "act")) + + updates = _queued_updates(extractor) + statuses = [u.status for u in updates if u.external_id == "act-1"] + assert statuses[0] == ActionStatus.running + assert statuses[-1] == expected_final + if expected_message: + assert expected_message in (updates[-1].result_message or "") + + +def test_custom_action_receives_call_metadata_in_context() -> None: + received_metadata: list[dict | None] = [] + + def capture(ctx: ActionContext) -> None: + received_metadata.append(ctx.call_metadata) + + extractor = _make_extractor() + extractor.add_action(CustomAction(name="greet", target=capture)) + + action = Action( + external_id="act-meta", + action_name="greet", + status=ActionStatus.pending, + call_metadata={"key": "value"}, + ) + extractor._dispatch_single_action(action) + + assert received_metadata == [{"key": "value"}] + + +def test_handle_actions_spawns_daemon_thread_named_after_external_id() -> None: + extractor = _make_extractor() + captured: dict = {} + done = Event() + + def target(ctx: ActionContext) -> None: + t = threading.current_thread() + captured["name"] = t.name + captured["daemon"] = t.daemon + done.set() + + extractor.add_action(CustomAction(name="work", target=target)) + extractor._handle_actions([_make_action("xid-42", "work")]) + + done.wait(timeout=5) + assert captured["name"] == "Action-xid-42" + assert captured["daemon"] is True + + +def test_handle_actions_multiple_actions_run_concurrently() -> None: + extractor = _make_extractor() + a1_started = Event() + a2_started = Event() + release = Event() + + def slow(ctx: ActionContext) -> None: + if ctx.external_id == "act-1": + a1_started.set() + else: + a2_started.set() + release.wait(timeout=5) + + extractor.add_action(CustomAction(name="op-1", target=slow)) + extractor.add_action(CustomAction(name="op-2", target=slow)) + + extractor._handle_actions( + [ + _make_action("act-1", "op-1"), + _make_action("act-2", "op-2"), + ] + ) + + assert a1_started.wait(timeout=5), "op-1 never started" + assert a2_started.wait(timeout=5), "op-2 never started" + + release.set() + + +def test_start_registers_handle_actions_as_dispatcher() -> None: + extractor = _make_extractor() + with extractor: + pass + + extractor._checkin_worker.set_action_dispatcher.assert_called_once() + registered = extractor._checkin_worker.set_action_dispatcher.call_args[0][0] + assert registered.__func__.__name__ == "_handle_actions" + assert registered.__self__ is extractor diff --git a/tests/test_unstable/test_action_registration.py b/tests/test_unstable/test_action_registration.py index da86456c..6d3142d3 100644 --- a/tests/test_unstable/test_action_registration.py +++ b/tests/test_unstable/test_action_registration.py @@ -29,12 +29,20 @@ def _startup_request(extractor: Extractor) -> StartupRequest: return extractor._get_startup_request() -# -- available_actions population -- - - -def test_no_scheduled_tasks_no_custom_actions_sends_available_actions_none() -> None: - # TestExtractor has one StartupTask; StartupTasks do not produce available_actions entries. +@pytest.mark.parametrize( + "extra_tasks", + [ + pytest.param([], id="no_tasks"), + pytest.param( + [ContinuousTask(name="cont", target=lambda _: None), StartupTask(name="init", target=lambda _: None)], + id="non_scheduled_tasks", + ), + ], +) +def test_available_actions_none_without_scheduled_tasks(extra_tasks: list) -> None: extractor = _make_extractor() + for task in extra_tasks: + extractor.add_task(task) assert _startup_request(extractor).available_actions is None @@ -65,13 +73,6 @@ def test_scheduled_task_action_entry_has_correct_type_and_task_ref( assert by_name[action_name].task == expected_task -def test_continuous_and_startup_tasks_do_not_produce_available_actions() -> None: - extractor = _make_extractor() - extractor.add_task(ContinuousTask(name="cont", target=lambda _: None)) - extractor.add_task(StartupTask(name="init", target=lambda _: None)) - assert _startup_request(extractor).available_actions is None - - def test_scheduled_and_custom_actions_combined_ordering() -> None: extractor = _make_extractor() extractor.add_task(ScheduledTask.from_interval(interval="1h", name="sync", target=lambda _: None)) @@ -79,7 +80,6 @@ def test_scheduled_and_custom_actions_combined_ordering() -> None: req = _startup_request(extractor) assert req.available_actions is not None assert len(req.available_actions) == 3 - # Scheduled task entries appear before custom actions names = [a.name for a in req.available_actions] assert names == ["Start sync", "Stop sync", "flush"] @@ -94,10 +94,7 @@ def test_custom_action_appears_with_correct_type_and_description() -> None: assert actions[0].description == "Clears state" -# -- __init_actions__ hook and add_action -- - - -def test_init_actions_hook_called_after_init_tasks() -> None: +def test_init_actions_hook_called_after_init_tasks_and_can_register_actions() -> None: call_order: list[str] = [] class _Ext(TestExtractor): @@ -106,20 +103,10 @@ def __init_tasks__(self) -> None: def __init_actions__(self) -> None: call_order.append("actions") - - _make_extractor(_Ext) - assert call_order == ["tasks", "actions"] - - -def test_add_action_from_init_actions_subclass_hook() -> None: - class _Ext(TestExtractor): - def __init_tasks__(self) -> None: - pass - - def __init_actions__(self) -> None: self.add_action(CustomAction(name="ping", target=lambda _: None)) extractor = _make_extractor(_Ext) + assert call_order == ["tasks", "actions"] assert len(extractor._custom_actions) == 1 assert extractor._custom_actions[0].name == "ping" @@ -139,36 +126,36 @@ def test_add_action_raises_on_duplicate_name() -> None: @pytest.mark.parametrize("conflicting_name", ["Start sync", "Stop sync"]) -def test_add_action_raises_on_conflict_with_scheduled_task_action_name(conflicting_name: str) -> None: - extractor = _make_extractor() - extractor.add_task(ScheduledTask.from_interval(interval="1h", name="sync", target=lambda _: None)) +def test_action_name_conflict_is_bidirectional(conflicting_name: str) -> None: + ext = _make_extractor() + ext.add_task(ScheduledTask.from_interval(interval="1h", name="sync", target=lambda _: None)) with pytest.raises(ValueError, match=conflicting_name): - extractor.add_action(CustomAction(name=conflicting_name, target=lambda _: None)) + ext.add_action(CustomAction(name=conflicting_name, target=lambda _: None)) + ext2 = _make_extractor() + ext2.add_action(CustomAction(name=conflicting_name, target=lambda _: None)) + with pytest.raises(ValueError, match="sync"): + ext2.add_task(ScheduledTask.from_interval(interval="1h", name="sync", target=lambda _: None)) -# -- _running_task_tokens lifecycle -- - -def test_token_present_in_running_task_tokens_during_execution() -> None: +def test_token_in_running_tokens_during_execution_is_child() -> None: extractor = _make_extractor() - token_present: list[bool] = [] - task_running = Event() - appended = Event() + captured: list = [] + done = Event() allow_finish = Event() def target(_: TaskContext) -> None: - task_running.set() - token_present.append("the-task" in extractor._running_task_tokens) - appended.set() + captured.append(extractor._running_task_tokens.get("the-task")) + done.set() allow_finish.wait(timeout=5) extractor.add_task(ScheduledTask.from_interval(interval="1h", name="the-task", target=target)) extractor._scheduler.trigger("the-task") - task_running.wait(timeout=5) + done.wait(timeout=5) allow_finish.set() - appended.wait(timeout=5) - assert token_present == [True] + assert captured[0] is not None + assert captured[0]._parent is extractor.cancellation_token @pytest.mark.parametrize("raises", [False, True]) @@ -196,7 +183,6 @@ def test_token_cleanup_does_not_clobber_replacement_token() -> None: replacement = extractor.cancellation_token.create_child_token() def target(_: TaskContext) -> None: - # Simulate a subsequent invocation overwriting the entry before this run finishes extractor._running_task_tokens["the-task"] = replacement done.set() @@ -207,22 +193,4 @@ def target(_: TaskContext) -> None: while any(j.name == "the-task" for j in extractor._scheduler._running) and time.monotonic() < deadline: time.sleep(0.005) - # Identity check must preserve the replacement — the old blind pop would remove it assert extractor._running_task_tokens.get("the-task") is replacement - - -def test_scheduled_task_token_is_child_of_extractor_cancellation_token() -> None: - extractor = _make_extractor() - captured: list = [] - done = Event() - - def target(_: TaskContext) -> None: - captured.append(extractor._running_task_tokens.get("the-task")) - done.set() - - extractor.add_task(ScheduledTask.from_interval(interval="1h", name="the-task", target=target)) - extractor._scheduler.trigger("the-task") - done.wait(timeout=5) - - assert len(captured) == 1 and captured[0] is not None - assert captured[0]._parent is extractor.cancellation_token