Fix issues with resuming async tasks awaiting a future (backport #1469)

rclpy/rclpy/executors.py

@@ -12,8 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import deque
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
+<<<<<<< HEAD
+=======
+from dataclasses import dataclass
+from functools import partial
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
 import inspect
 import multiprocessing
 from threading import Condition
@@ -23,6 +29,11 @@
 from typing import Any
 from typing import Callable
 from typing import Coroutine
+<<<<<<< HEAD
+=======
+from typing import Deque
+from typing import Dict
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
 from typing import Generator
 from typing import List
 from typing import Optional
@@ -146,7 +157,17 @@ def timeout(self, timeout):
         self._timeout = timeout
 
 
+<<<<<<< HEAD
 class Executor:
+=======
+@dataclass
+class TaskData:
+    source_node: 'Optional[Node]' = None
+    source_entity: 'Optional[Entity]' = None
+
+
+class Executor(ContextManager['Executor']):
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
     """
     The base class for an executor.
 
@@ -165,8 +186,15 @@ def __init__(self, *, context: Context = None) -> None:
         self._context = get_default_context() if context is None else context
         self._nodes: Set[Node] = set()
         self._nodes_lock = RLock()
+<<<<<<< HEAD
         # Tasks to be executed (oldest first) 3-tuple Task, Entity, Node
         self._tasks: List[Tuple[Task, Optional[WaitableEntityType], Optional[Node]]] = []
+=======
+        # all tasks that are not complete or canceled
+        self._pending_tasks: Dict[Task, TaskData] = {}
+        # tasks that are ready to execute
+        self._ready_tasks: Deque[Task[Any]] = deque()
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
         self._tasks_lock = Lock()
         # This is triggered when wait_for_ready_callbacks should rebuild the wait list
         self._guard = GuardCondition(
@@ -200,10 +228,27 @@ def create_task(self, callback: Union[Callable, Coroutine], *args, **kwargs) ->
         """
         task = Task(callback, args, kwargs, executor=self)
         with self._tasks_lock:
+<<<<<<< HEAD
             self._tasks.append((task, None, None))
             self._guard.trigger()
         # Task inherits from Future
         return task
+=======
+            self._pending_tasks[task] = TaskData()
+        self._call_task_in_next_spin(task)
+        return task
+
+    def _call_task_in_next_spin(self, task: Task) -> None:
+        """
+        Add a task to the executor to be executed in the next spin.
+
+        :param task: A task to be run in the executor.
+        """
+        with self._tasks_lock:
+            self._ready_tasks.append(task)
+            if self._guard:
+                self._guard.trigger()
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
 
     def shutdown(self, timeout_sec: float = None) -> bool:
         """
@@ -473,7 +518,10 @@ async def handler(entity, gc, is_shutdown, work_tracker):
             handler, (entity, self._guard, self._is_shutdown, self._work_tracker),
             executor=self)
         with self._tasks_lock:
-            self._tasks.append((task, entity, node))
+            self._pending_tasks[task] = TaskData(
+                source_entity=entity,
+                source_node=node
+            )
         return task
 
     def can_execute(self, entity: WaitableEntityType) -> bool:
@@ -517,8 +565,8 @@ def _wait_for_ready_callbacks(
                 nodes_to_use = self.get_nodes()
 
             # Yield tasks in-progress before waiting for new work
-            tasks = None
             with self._tasks_lock:
+<<<<<<< HEAD
                 tasks = list(self._tasks)
             if tasks:
                 for task, entity, node in reversed(tasks):
@@ -531,7 +579,26 @@ def _wait_for_ready_callbacks(
                     self._tasks = list(filter(lambda t_e_n: not t_e_n[0].done(), self._tasks))
                     # Get rid of any tasks that are cancelled
                     self._tasks = list(filter(lambda t_e_n: not t_e_n[0].cancelled(), self._tasks))
-
+=======
+                # Get rid of any tasks that are done or cancelled
+                for task in list(self._pending_tasks.keys()):
+                    if task.done() or task.cancelled():
+                        del self._pending_tasks[task]
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
+
+                ready_tasks_count = len(self._ready_tasks)
+            for _ in range(ready_tasks_count):
+                task = self._ready_tasks.popleft()
+                task_data = self._pending_tasks[task]
+                node = task_data.source_node
+                if node is None or node in nodes_to_use:
+                    entity = task_data.source_entity
+                    yielded_work = True
+                    yield task, entity, node
+                else:
+                    # Asked not to execute these tasks, so don't do them yet
+                    with self._tasks_lock:
+                        self._ready_tasks.append(task)
             # Gather entities that can be waited on
             subscriptions: List[Subscription] = []
             guards: List[GuardCondition] = []

rclpy/rclpy/task.py

@@ -57,10 +57,17 @@ def __del__(self):
                 'The following exception was never retrieved: ' + str(self._exception),
                 file=sys.stderr)
 
+<<<<<<< HEAD
     def __await__(self):
+=======
+    def __await__(self) -> Generator['Future[T]', None, Optional[T]]:
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
         # Yield if the task is not finished
-        while self._pending():
-            yield
+        if self._pending():
+            # This tells the task to suspend until the future is done
+            yield self
+        if self._pending():
+            raise RuntimeError('Future awaited a second time before it was done')
         return self.result()
 
     def _pending(self):
@@ -249,6 +256,7 @@ def __call__(self):
             self._executing = True
 
             if inspect.iscoroutine(self._handler):
+<<<<<<< HEAD
                 # Execute a coroutine
                 try:
                     self._handler.send(None)
@@ -259,6 +267,9 @@ def __call__(self):
                 except Exception as e:
                     self.set_exception(e)
                     self._complete_task()
+=======
+                self._execute_coroutine_step(self._handler)
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
             else:
                 # Execute a normal function
                 try:
@@ -271,7 +282,52 @@ def __call__(self):
         finally:
             self._task_lock.release()
 
+<<<<<<< HEAD
     def _complete_task(self):
+=======
+    def _execute_coroutine_step(self, coro: Coroutine[Any, Any, T]) -> None:
+        """Execute or resume a coroutine task."""
+        try:
+            result = coro.send(None)
+        except StopIteration as e:
+            # The coroutine finished; store the result
+            self.set_result(e.value)
+            self._complete_task()
+        except Exception as e:
+            # The coroutine raised; store the exception
+            self.set_exception(e)
+            self._complete_task()
+        else:
+            # The coroutine yielded; suspend the task until it is resumed
+            executor = self._executor()
+            if executor is None:
+                raise RuntimeError(
+                    'Task tried to reschedule but no executor was set: '
+                    'tasks should only be initialized through executor.create_task()')
+            elif isinstance(result, Future):
+                # Schedule the task to resume when the future is done
+                self._add_resume_callback(result, executor)
+            elif result is None:
+                # The coroutine yielded None, schedule the task to resume in the next spin
+                executor._call_task_in_next_spin(self)
+            else:
+                raise TypeError(
+                    f'Expected coroutine to yield a Future or None, got: {type(result)}')
+
+    def _add_resume_callback(self, future: Future[T], executor: 'Executor') -> None:
+        future_executor = future._executor()
+        if future_executor is None:
+            # The future is not associated with an executor yet, so associate it with ours
+            future._set_executor(executor)
+        elif future_executor is not executor:
+            raise RuntimeError('A task can only await futures associated with the same executor')
+
+        # The future is associated with the same executor, so we can resume the task directly
+        # in the done callback
+        future.add_done_callback(lambda _: self.__call__())
+
+    def _complete_task(self) -> None:
+>>>>>>> 9695271 (Fix issues with resuming async tasks awaiting a future (#1469))
         """Cleanup after task finished."""
         self._handler = None
         self._args = None