SAR-302 Proof of concept for a enqueue and suspend task

src/ska_tango_base/base/task_queue_manager.py

@@ -295,6 +295,7 @@ class QueueManager:
             logger: logging.Logger,
             stopping_event: Event,
             aborting_event: Event,
+            suspend_event: Event,
             result_callback: Callable,
             update_command_state_callback: Callable,
             update_progress_callback: Callable,
@@ -312,6 +313,8 @@ class QueueManager:
             :type stopping_event: Event
             :param aborting_event: Indicates whether the queue is being aborted
             :type aborting_event: Event
+            :param suspend_event: Indicates whether to suspend task retrieval
+            :type suspend_event: Event
             :param update_command_state_callback: Callback to update command state
             :type update_command_state_callback: Callable
             """
@@ -320,6 +323,7 @@ class QueueManager:
             self._logger = logger
             self.stopping_event = stopping_event
             self.aborting_event = aborting_event
+            self.suspend_event = suspend_event
             self._result_callback = result_callback
             self._update_command_state_callback = update_command_state_callback
             self._update_progress_callback = update_progress_callback
@@ -341,6 +345,10 @@ class QueueManager:
                     self.current_task_id = None
                     self.current_task_progress = ""
 
+                    # Don't pull new tasks off of the queue until unsuspended.
+                    while self.suspend_event.is_set():
+                        time.sleep(self._queue_fetch_timeout)
+
                     if self.aborting_event.is_set():
                         # Drain the Queue since self.aborting_event is set
                         while not self._work_queue.empty():
@@ -459,6 +467,7 @@ class QueueManager:
         self._push_change_event = push_change_event
         self.stopping_event = threading.Event()
         self.aborting_event = threading.Event()
+        self.suspend_event = threading.Event()
         self._property_update_lock = threading.Lock()
         self._logger = logger if logger else logging.getLogger(__name__)
 
@@ -467,6 +476,17 @@ class QueueManager:
         self._task_status: Dict[str, str] = {}  # unique_id, status
         self._threads = []
 
+        self._long_running_properties = [
+            "longRunningCommandsInQueue",
+            "longRunningCommandStatus",
+            "longRunningCommandProgress",
+            "longRunningCommandIDsInQueue",
+            "longRunningCommandResult",
+        ]
+        self._property_change_callbacks = {}
+        for prop in self._long_running_properties:
+            self._property_change_callbacks[prop] = []
+
         # If there's no queue, don't start threads
         if not self._max_queue_size:
             return
@@ -477,6 +497,7 @@ class QueueManager:
                 self._logger,
                 self.stopping_event,
                 self.aborting_event,
+                self.suspend_event,
                 self.result_callback,
                 self.update_task_state_callback,
                 self.update_progress_callback,
@@ -496,7 +517,7 @@ class QueueManager:
         return self._work_queue.full()
 
     @property
-    def task_result(self) -> Union[Tuple[str, str, str], Tuple[()]]:
+    def task_result(self) -> Tuple[str, str, str]:
         """Return the last task result.
 
         :return: Last task result
@@ -649,6 +670,8 @@ class QueueManager:
         :param property_name: The property value
         :type property_name: Any
         """
+        for callback in self._property_change_callbacks[property_name]:
+            callback(property_value)
         if self._push_change_event:
             self._push_change_event(property_name, property_value)
 
@@ -664,6 +687,25 @@ class QueueManager:
         """Set stopping_event on each thread so it exists out. Killing the thread."""
         self.stopping_event.set()
 
+    def suspend_task_dequeue(self):
+        """Stop pulling new tasks off the queue to execute.
+
+        - Tasks enqueued after this method will not be dequeued.
+        - Existing tasks may be dequeued during this method.
+        """
+        self.suspend_event.set()
+        # Wait a little longer than fetch timeout otherwise one of the worker
+        # threads that is waiting in `queue.get`
+        # will pick up a task (enqueued immediately after this method) before
+        # the suspend takes effect.
+        time.sleep(self._queue_fetch_timeout + 0.1)
+        self._logger.info("Queue task execution suspended")
+
+    def unsuspend_task_dequeue(self):
+        """Undo suspend_task_dequeue."""
+        self.suspend_event.clear()
+        self._logger.info("Queue task execution unsuspended")
+
     @property
     def is_aborting(self) -> bool:
         """Return whether we are in aborting state."""
@@ -701,6 +743,29 @@ class QueueManager:
 
         return TaskState.NOT_FOUND
 
+    def add_property_change_callback(self, attribute: str, update_callback: Callable):
+        """Add a callback that will be executed when the attribute changes.
+
+        :param attribute: The attribute name
+        :type attribute: str
+        :param update_callback: The function to execute
+        :type update_callback: Callable
+        """
+        assert (
+            attribute in self._long_running_properties
+        ), f"[{attribute}] is not supported, should be one of [{self._long_running_properties}]"
+        self._property_change_callbacks[attribute].append(update_callback)
+
+    def remove_property_change_callback(self, update_callback: Callable):
+        """Remove the callback.
+
+        :param update_callback: The function to execute
+        :type update_callback: Callable
+        """
+        for callbacks in self._property_change_callbacks.values():
+            if update_callback in callbacks:
+                callbacks.remove(update_callback)
+
     def __len__(self) -> int:
         """Approximate length of the queue.
 

src/ska_tango_base/utils.py

@@ -6,6 +6,7 @@ import inspect
 import json
 import logging
 import pydoc
+from queue import Empty, Queue
 import traceback
 import sys
 import uuid
@@ -29,8 +30,9 @@ from tango import (
 )
 from tango import DevState
 from contextlib import contextmanager
+from ska_tango_base.commands import BaseCommand
 from ska_tango_base.faults import GroupDefinitionsError, SKABaseError
-from ska_tango_base.base.task_queue_manager import TaskResult
+from ska_tango_base.base.task_queue_manager import QueueManager, TaskResult
 
 int_types = {
     tango._tango.CmdArgType.DevUShort,
@@ -631,7 +633,7 @@ class LongRunningDeviceInterface:
             - Clean up
         """
         if ev.err:
-            self._logger.error("Event system DevError(s) occured: %s", str(ev.errors))
+            self._logger.error("Event system DevError(s) occurred: %s", str(ev.errors))
             return
 
         if ev.attr_value and ev.attr_value.name == "longrunningcommandresult":
@@ -666,7 +668,7 @@ class LongRunningDeviceInterface:
                     command_name = stored_command_group[0].command_name
 
                     # Trigger the callback, send command_name and command_ids
-                    # as paramater
+                    # as parameter
                     self._stored_callbacks[key](command_name, command_ids)
                     # Remove callback as the group completed
 
@@ -684,10 +686,10 @@ class LongRunningDeviceInterface:
     ):
         """Execute the long running command with an argument if any.
 
-        Once the commmand completes, then the `on_completion_callback`
+        Once the command completes, then the `on_completion_callback`
         will be executed with the EventData as parameter.
         This class keeps track of the command ID and events
-        used to determine when this commmand has completed.
+        used to determine when this command has completed.
 
         :param command_name: A long running command that exists on the
             target Tango device.
@@ -714,3 +716,70 @@ class LongRunningDeviceInterface:
                     False,
                 )
             )
+
+
+class EnqueueSuspend:
+    """Context manager that will enqueue a command and then suspend new tasks from being taken off the queue."""
+
+    def __init__(
+        self,
+        queue_manager: QueueManager,
+        command: BaseCommand,
+        args: Any = None,
+        retries: int = 5,
+        timeout: float = 0.5,
+    ) -> None:
+        """Context manager to enqueue a task and suspend new tasks.
+
+        :param queue_manager: The queue manager
+        :type queue_manager: QueueManager
+        :param command: Command to execute
+        :type command: BaseCommand
+        :param args: Argument for the command, defaults to None
+        :type args: Any, optional
+        :param retries: Number of times to retry waiting for the command to start, defaults to 5
+        :type retries: int, optional
+        :param timeout: Time to wait for a status change, defaults to 0.5
+        :type timeout: float, optional
+        """
+        self.queue_manager = queue_manager
+        self.command = command
+        self.args = args
+        self.retries = retries
+        self.timeout = timeout
+        self.unique_id = ""
+        self.event_queue = Queue()
+
+    def _on_status_change_callback(self, new_value):
+        """Add status change to queue."""
+        self.event_queue.put(new_value)
+
+    def wait_for_id(self, unique_id) -> bool:
+        """Wait for the enqueued task to start."""
+        while True:
+            try:
+                while self.retries > 0:
+                    value = self.event_queue.get(timeout=self.timeout)
+                    if unique_id in value:
+                        return True
+                    self.retries -= 1
+            except Empty:
+                return False
+
+    def __enter__(self):
+        """Add callback enqueue task and suspend."""
+        self.queue_manager.add_property_change_callback(
+            "longRunningCommandStatus", self._on_status_change_callback
+        )
+        self.unique_id, _ = self.queue_manager.enqueue_task(self.command, self.args)
+        if not self.wait_for_id(self.unique_id):
+            raise Exception("Command not started")
+        self.queue_manager.suspend_task_dequeue()
+        return self.unique_id
+
+    def __exit__(self, _exc_type, _exc_value, _exc_traceback):
+        """Clear callback and unsuspend."""
+        self.queue_manager.remove_property_change_callback(
+            self._on_status_change_callback
+        )
+        self.queue_manager.unsuspend_task_dequeue()

tests/long_running_tasks/test_task_queue_manager.py

@@ -12,6 +12,7 @@ from ska_tango_base.base.task_queue_manager import (
 )
 from ska_tango_base.base.reference_component_manager import QueueWorkerComponentManager
 from ska_tango_base.commands import BaseCommand
+from ska_tango_base.utils import EnqueueSuspend
 
 logger = logging.getLogger(__name__)
 
@@ -108,6 +109,20 @@ def simple_task():
     return get_task
 
 
+@pytest.fixture
+def noop_task():
+    """Fixture for a very simple task."""
+
+    def get_task():
+        class NoopTask(BaseCommand):
+            def do(self):
+                return True
+
+        return NoopTask(target=None)
+
+    return get_task
+
+
 @pytest.fixture
 def abort_task():
     """Fixture for a task that aborts."""
@@ -524,6 +539,37 @@ class TestComponentManager:
         )
         assert cm.task_ids_in_queue == ()
 
+    def test_enqueue_suspend_util(self, simple_task, noop_task, caplog):
+        """Test that EnqueueSuspend enqueues and suspends."""
+        caplog.set_level(logging.INFO)
+        cm = QueueWorkerComponentManager(
+            op_state_model=None,
+            logger=logger,
+            max_queue_size=3,
+            num_workers=3,
+            push_change_event=None,
+            child_devices=[],
+        )
+        second_task_id = ""
+        with EnqueueSuspend(cm._queue_manager, simple_task(), args=1) as unique_id:
+            # Make sure the command exists
+            assert unique_id
+            # Make sure pulling new tasks off the queue is paused
+            assert cm._queue_manager.suspend_event.is_set()
+            # Make sure any new work is not pulled off queue
+            second_task_id, _ = cm.enqueue(noop_task())
+            for wait_time in [0.1, 0.2, 0.3]:
+                assert (
+                    second_task_id in cm.task_ids_in_queue
+                ), "Task should not be taken off the queue while suspend is in effect"
+                time.sleep(wait_time)
+
+        # Make sure work is unsuspended
+        assert not cm._queue_manager.suspend_event.is_set()
+        # Make sure the enqueued task during suspension is taken off the queue
+        time.sleep(0.5)
+        assert second_task_id not in cm.task_ids_in_queue
+
 
 @pytest.mark.forked
 class TestStress: