diff --git a/docs/reference/rest_api.json b/docs/reference/rest_api.json index c20738a..48a1ce6 100644 --- a/docs/reference/rest_api.json +++ b/docs/reference/rest_api.json @@ -1175,6 +1175,15 @@ ], "title": "OpaConfig" }, + "PauseReason": { + "type": "string", + "enum": [ + "Pause requested by user", + "Paused as queue completed", + "Paused as last task errored" + ], + "title": "PauseReason" + }, "PlanSource": { "properties": { "module": { @@ -1201,29 +1210,29 @@ "paused": { "type": "boolean", "title": "Paused" + }, + "last_pause_reason": { + "$ref": "#/components/schemas/PauseReason" } }, "type": "object", "required": [ - "paused" + "paused", + "last_pause_reason" ], "title": "QueueState" }, "QueueStateUpdate": { "properties": { "paused": { - "anyOf": [ - { - "type": "boolean" - }, - { - "type": "null" - } - ], + "type": "boolean", "title": "Paused" } }, "type": "object", + "required": [ + "paused" + ], "title": "QueueStateUpdate" }, "RestConfig": { diff --git a/src/daq_queuing_service/api/api.py b/src/daq_queuing_service/api/api.py index da3aec3..0348946 100644 --- a/src/daq_queuing_service/api/api.py +++ b/src/daq_queuing_service/api/api.py @@ -13,6 +13,7 @@ from daq_queuing_service.task import Experiment, Status, Task from daq_queuing_service.task_queue.queue import ( QUEUE_EVENTS, + PauseReason, QueueState, TaskQueue, TaskWithPosition, @@ -22,7 +23,7 @@ class QueueStateUpdate(BaseModel): - paused: bool | None = None + paused: bool class TaskCancelRequest(BaseModel): @@ -61,7 +62,10 @@ def get_config() -> AppConfig: @router.patch("/queue/state") async def update_queue_state(payload: QueueStateUpdate) -> QueueState: - return await queue.update_state(**payload.model_dump(exclude_none=True)) + if payload.paused: + return await queue.pause_queue(PauseReason.USER_REQUESTED) + else: + return await queue.resume_queue() @router.get("/queue/state") def get_queue_state() -> QueueState: diff --git a/src/daq_queuing_service/task_queue/queue.py b/src/daq_queuing_service/task_queue/queue.py index c7ee549..3d1b278 100644 --- a/src/daq_queuing_service/task_queue/queue.py +++ b/src/daq_queuing_service/task_queue/queue.py @@ -1,5 +1,6 @@ import asyncio from collections.abc import Callable, Sequence +from enum import StrEnum from types import TracebackType from typing import Any, Literal @@ -30,8 +31,15 @@ def __missing__(self, task_id: str) -> Task: raise TaskNotFoundError(f"No task found matching id: {task_id}") +class PauseReason(StrEnum): + USER_REQUESTED = "Pause requested by user" + EMPTY_QUEUE = "Paused as queue completed" + ERROR = "Paused as last task errored" + + class QueueState(BaseModel): paused: bool + last_pause_reason: PauseReason class Modifying(asyncio.Condition): @@ -67,7 +75,9 @@ def __init__(self, convert: Converter, broadcaster: Broadcaster[QUEUE_EVENTS]): self._call_queue: list[BlueapiCall] = [] self._call_history: list[BlueapiCall] = [] self._queue_history: list[BlueapiCall] = [] - self._state: QueueState = QueueState(paused=True) + self._state: QueueState = QueueState( + paused=True, last_pause_reason=PauseReason.EMPTY_QUEUE + ) self._convert = convert self._modifying = Modifying(on_exit=self._sync) self._broadcaster = broadcaster @@ -118,6 +128,9 @@ def _sync(self): if call.parent_task_id: self._tasks[call.parent_task_id].blueapi_calls.append(call) + if not self._call_queue: + self._pause_queue(PauseReason.EMPTY_QUEUE) + self._broadcast_changes() self._modifying.notify_all() @@ -360,21 +373,36 @@ async def clear_history(self): self._history.clear() LOGGER.info("Successfully cleared history") - async def update_state(self, paused: bool | None = None) -> QueueState: - """Update the state of the queue. + def _pause_queue(self, reason: PauseReason): + self._state = QueueState(paused=True, last_pause_reason=reason) + self._broadcaster.broadcast(Event(type="state_update", data=self._state)) + + async def pause_queue(self, reason: PauseReason) -> QueueState: + """Pause the queue if running. Args: - paused (bool | None, optional): Whether or not the queue should be paused. + paused (PauseReason): The reason for pausing the queue. + + Returns: + QueueState: The new state of the queue. + """ + async with self._modifying: + self._pause_queue(reason) + LOGGER.info(f"Successfully paused queue due to {reason}") + return self._state + + async def resume_queue(self) -> QueueState: + """Resume the queue if paused. Returns: QueueState: The new state of the queue. """ async with self._modifying: self._state = QueueState( - paused=self._state.paused if paused is None else paused + paused=False, last_pause_reason=self._state.last_pause_reason ) self._broadcaster.broadcast(Event(type="state_update", data=self._state)) - LOGGER.info(f"Successfully updated queue state to {self._state}") + LOGGER.info("Successfully resumed queue") return self._state @property diff --git a/src/daq_queuing_service/worker/worker.py b/src/daq_queuing_service/worker/worker.py index 11ed2bf..9897708 100644 --- a/src/daq_queuing_service/worker/worker.py +++ b/src/daq_queuing_service/worker/worker.py @@ -16,7 +16,7 @@ from daq_queuing_service.blueapi_interaction.blueapi_adapter import BlueapiClientAdapter from daq_queuing_service.blueapi_interaction.blueapi_call import BlueapiCall, CallStatus from daq_queuing_service.log import HANDLER -from daq_queuing_service.task_queue.queue import TaskQueue +from daq_queuing_service.task_queue.queue import PauseReason, TaskQueue LOGGER = logging.getLogger("Queue Worker") LOGGER.addHandler(HANDLER) @@ -85,7 +85,7 @@ async def _process_call(self, call: BlueapiCall): await self._queue.complete_call(call, task_status.result) case TaskError(): LOGGER.debug(f"Call {call} failed: {task_status.result}. Pausing queue") - await self._queue.update_state(paused=True) + await self._queue.pause_queue(PauseReason.ERROR) await self._queue.fail_call(call, [task_status.result]) @staticmethod diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index bf045e7..68b97e8 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -47,8 +47,8 @@ def tasks() -> list[Task]: @pytest.fixture async def task_queue(tasks: list[Task]): queue = TaskQueue(convert=construct_blueapi_call_list, broadcaster=Broadcaster()) - await queue.update_state(paused=False) await queue.add_tasks(tasks) + await queue.resume_queue() return queue @@ -76,8 +76,8 @@ def construct_blueapi_call_list( return call_list queue = TaskQueue(convert=construct_blueapi_call_list, broadcaster=Broadcaster()) - await queue.update_state(paused=False) await queue.add_tasks(tasks) + await queue.resume_queue() return queue diff --git a/tests/unit_tests/test_api.py b/tests/unit_tests/test_api.py index e7bce36..70153c7 100644 --- a/tests/unit_tests/test_api.py +++ b/tests/unit_tests/test_api.py @@ -33,7 +33,7 @@ TaskKind, TaskWithPosition, ) -from daq_queuing_service.task_queue.queue import QUEUE_EVENTS, TaskQueue +from daq_queuing_service.task_queue.queue import QUEUE_EVENTS, PauseReason, TaskQueue from daq_queuing_service.task_queue.queue_utils import QueueError @@ -95,15 +95,32 @@ def test_healthz_returns_healthy(test_client: TestClient): def test_get_queue_state_returns_queue_state(test_client: TestClient): response = test_client.get("/queue/state") assert response.status_code == 200 - assert response.json() == {"paused": False} + assert response.json() == { + "paused": False, + "last_pause_reason": "Paused as queue completed", + } -def test_update_queue_state_changes_queue_state_and_returns_new_state( +def test_resume_changes_queue_state_and_returns_new_state( test_client: TestClient, ): response = test_client.patch("/queue/state", json={"paused": False}) assert response.status_code == 200 - assert response.json() == {"paused": False} + assert response.json() == { + "paused": False, + "last_pause_reason": PauseReason.EMPTY_QUEUE, + } + + +def test_pause_changes_queue_state_and_returns_new_state( + test_client: TestClient, +): + response = test_client.patch("/queue/state", json={"paused": True}) + assert response.status_code == 200 + assert response.json() == { + "paused": True, + "last_pause_reason": PauseReason.USER_REQUESTED, + } def test_get_queued_tasks_returns_queued_task(test_client: TestClient): diff --git a/tests/unit_tests/test_queue.py b/tests/unit_tests/test_queue.py index 64c912a..c0e30c9 100644 --- a/tests/unit_tests/test_queue.py +++ b/tests/unit_tests/test_queue.py @@ -24,6 +24,7 @@ TaskKind, ) from daq_queuing_service.task_queue.queue import ( + PauseReason, TaskQueue, TaskWithPosition, ) @@ -633,14 +634,14 @@ async def test_clear_history_removes_history_and_removes_completed_tasks_from_re async def test_pausing_queue_prevents_task_from_being_claimed(task_queue: TaskQueue): - await task_queue.update_state(paused=True) + await task_queue.pause_queue(PauseReason.EMPTY_QUEUE) assert task_queue.state.paused with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(task_queue.get_next_call_once_available(), timeout=0.05) async def test_unpausing_queue_allows_tasks_to_being_claimed(task_queue: TaskQueue): - await task_queue.update_state(paused=False) + await task_queue.resume_queue() assert not task_queue.state.paused await task_queue.get_next_call_once_available() @@ -699,7 +700,7 @@ async def test_wait_until_call_available_waits_if_next_task_is_in_progress( async def test_wait_until_call_available_waits_if_queue_paused( task_queue: TaskQueue, ): - await task_queue.update_state(paused=True) + await task_queue.pause_queue(PauseReason.EMPTY_QUEUE) with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(task_queue.wait_until_call_available(), timeout=0.05) @@ -962,3 +963,17 @@ async def test__sync_skips_subsequent_calls_if_previous_one_failed_within_a_task task_queue_one_to_many._sync() assert a_task.blueapi_calls[0].status == Status.ERROR assert all(call.status == CallStatus.SKIPPED for call in a_task.blueapi_calls[1:]) + + +async def test__sync_pauses_queue_if_no_more_items( + task_queue: TaskQueue, +): + task_queue._modifying = MagicMock() + task_queue.state.paused = False + task_queue._queue = [] + task_queue._call_queue = [] + + task_queue._sync() + + assert task_queue.state.paused + assert task_queue.state.last_pause_reason == PauseReason.EMPTY_QUEUE diff --git a/tests/unit_tests/test_worker.py b/tests/unit_tests/test_worker.py index a3d62d2..65f1b46 100644 --- a/tests/unit_tests/test_worker.py +++ b/tests/unit_tests/test_worker.py @@ -22,7 +22,12 @@ ) from daq_queuing_service.blueapi_interaction.blueapi_call import CallStatus from daq_queuing_service.task import Status -from daq_queuing_service.task_queue.queue import TaskError, TaskQueue, TaskResult +from daq_queuing_service.task_queue.queue import ( + PauseReason, + TaskError, + TaskQueue, + TaskResult, +) from daq_queuing_service.worker.worker import LOGGER, QueueWorker @@ -251,6 +256,7 @@ async def test_when_plan_error_then_queue_paused_and_task_failed_and_errors_adde await worker_with_plan_error.run_loop() assert worker_with_plan_error._queue.state.paused is True + assert worker_with_plan_error._queue.state.last_pause_reason == PauseReason.ERROR first_call = first_task.blueapi_calls[0] assert first_call.status == CallStatus.ERROR assert first_call.errors == [