Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Check if Task(Future) is canceled. #1377

Merged
merged 8 commits into from
Jan 15, 2025
16 changes: 14 additions & 2 deletions rclpy/rclpy/executors.py
Original file line number Diff line number Diff line change
Expand Up @@ -321,14 +321,24 @@ def spin_until_future_complete(
future.add_done_callback(lambda x: self.wake())

if timeout_sec is None or timeout_sec < 0:
while self._context.ok() and not future.done() and not self._is_shutdown:
while (
self._context.ok() and
not future.done() and
not future.cancelled()
and not self._is_shutdown
fujitatomoya marked this conversation as resolved.
Show resolved Hide resolved
):
self._spin_once_until_future_complete(future, timeout_sec)
else:
start = time.monotonic()
end = start + timeout_sec
timeout_left = TimeoutObject(timeout_sec)

while self._context.ok() and not future.done() and not self._is_shutdown:
while (
self._context.ok() and
not future.done() and
fujitatomoya marked this conversation as resolved.
Show resolved Hide resolved
not future.cancelled()
and not self._is_shutdown
):
self._spin_once_until_future_complete(future, timeout_left)
now = time.monotonic()

Expand Down Expand Up @@ -610,6 +620,8 @@ def _wait_for_ready_callbacks(
with self._tasks_lock:
# Get rid of any tasks that are done
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))

# Gather entities that can be waited on
subscriptions: List[Subscription] = []
Expand Down
68 changes: 48 additions & 20 deletions rclpy/rclpy/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from enum import auto, StrEnum
fujitatomoya marked this conversation as resolved.
Show resolved Hide resolved
import inspect
import sys
import threading
Expand All @@ -31,14 +32,19 @@ def _fake_weakref() -> None:
return None


class FutureState(StrEnum):
"""States defining the lifecycle of a future."""

PENDING = auto()
CANCELLED = auto()
FINISHED = auto()


class Future(Generic[T]):
"""Represent the outcome of a task in the future."""

def __init__(self, *, executor: Optional['Executor'] = None) -> None:
# true if the task is done or cancelled
self._done = False
# true if the task is cancelled
self._cancelled = False
self._state = FutureState.PENDING
# the final return value of the handler
self._result: Optional[T] = None
# An exception raised by the handler when called
Expand All @@ -61,15 +67,20 @@ def __del__(self) -> None:

def __await__(self) -> Generator[None, None, Optional[T]]:
# Yield if the task is not finished
while not self._done:
while self._pending():
yield
return self.result()


def _pending(self) -> bool:
return self._state == FutureState.PENDING

def cancel(self) -> None:
"""Request cancellation of the running task if it is not done already."""
with self._lock:
if not self._done:
self._cancelled = True
if not self._pending():
return

self._state = FutureState.CANCELLED
self._schedule_or_invoke_done_callbacks()

def cancelled(self) -> bool:
Expand All @@ -78,15 +89,15 @@ def cancelled(self) -> bool:

:return: True if the task was cancelled
"""
return self._cancelled
return self._state == FutureState.CANCELLED

def done(self) -> bool:
"""
Indicate if the task has finished executing.
Indicate if the task has finished or cancelled executing.

:return: True if the task is finished or raised while it was executing
:return: True if the task is finished, cancelled or raised while it was executing
fujitatomoya marked this conversation as resolved.
Show resolved Hide resolved
"""
return self._done
return self._state == FutureState.FINISHED

def result(self) -> Optional[T]:
"""
Expand Down Expand Up @@ -114,12 +125,17 @@ def set_result(self, result: T) -> None:
"""
Set the result returned by a task.

:raises: RuntimeError if future state is not pending

:param result: The output of a long running task.
"""
with self._lock:
if not self._pending():
raise RuntimeError(f"Cannot set result from state: {self._state}")

fujitatomoya marked this conversation as resolved.
Show resolved Hide resolved
self._result = result
self._done = True
self._cancelled = False
self._state = FutureState.FINISHED

self._schedule_or_invoke_done_callbacks()

def set_exception(self, exception: Exception) -> None:
Expand All @@ -129,10 +145,13 @@ def set_exception(self, exception: Exception) -> None:
:param result: The output of a long running task.
"""
with self._lock:
if not self._pending():
raise RuntimeError(f"Cannot set exception from state: {self._state}")

self._exception = exception
self._exception_fetched = False
self._done = True
self._cancelled = False
self._state = FutureState.FINISHED

self._schedule_or_invoke_done_callbacks()

def _schedule_or_invoke_done_callbacks(self) -> None:
Expand Down Expand Up @@ -181,7 +200,7 @@ def add_done_callback(self, callback: Callable[['Future[T]'], None]) -> None:
"""
invoke = False
with self._lock:
if self._done:
if not self._pending():
assert self._executor is not None
executor = self._executor()
if executor is not None:
Expand Down Expand Up @@ -239,10 +258,14 @@ def __call__(self) -> None:

The return value of the handler is stored as the task result.
"""
if self._done or self._executing or not self._task_lock.acquire(blocking=False):
if (
not self._pending() or
self._executing or
not self._task_lock.acquire(blocking=False)
):
return
try:
if self._done:
if not self._pending():
return
self._executing = True

Expand All @@ -253,7 +276,6 @@ def __call__(self) -> None:
handler.send(None)
except StopIteration as e:
# The coroutine finished; store the result
handler.close()
self.set_result(e.value)
self._complete_task()
except Exception as e:
Expand Down Expand Up @@ -285,3 +307,9 @@ def executing(self) -> bool:
:return: True if the task is currently executing.
"""
return self._executing

def cancel(self) -> None:
if self._pending() and inspect.iscoroutine(self._handler):
self._handler.close()

super().cancel()
20 changes: 20 additions & 0 deletions rclpy/test/test_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -273,6 +273,26 @@ async def coroutine():
self.assertTrue(future.done())
self.assertEqual('Sentinel Result', future.result())

def test_create_task_coroutine_cancel(self) -> None:
fujitatomoya marked this conversation as resolved.
Show resolved Hide resolved
self.assertIsNotNone(self.node.handle)
executor = SingleThreadedExecutor(context=self.context)
executor.add_node(self.node)

async def coroutine():
return 'Sentinel Result'

future = executor.create_task(coroutine)
self.assertFalse(future.done())
self.assertFalse(future.cancelled())

future.cancel()
self.assertTrue(future.cancelled())

executor.spin_until_future_complete(future)
self.assertFalse(future.done())
self.assertTrue(future.cancelled())
self.assertEqual(None, future.result())

def test_create_task_normal_function(self) -> None:
self.assertIsNotNone(self.node.handle)
executor = SingleThreadedExecutor(context=self.context)
Expand Down
23 changes: 23 additions & 0 deletions rclpy/test/test_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -322,6 +322,29 @@ def cb(fut):
f.add_done_callback(cb)
assert called

def test_set_result_on_done_future_raises_exception(self) -> None:
f = Future()
f.set_result(None)
with self.assertRaises(RuntimeError):
f.set_result(None)

def test_set_result_on_cancelled_future_raises_exception(self) -> None:
f = Future()
f.cancel()
with self.assertRaises(RuntimeError):
f.set_result(None)

def test_set_exception_on_done_future_raises_exception(self) -> None:
f = Future()
f.set_result(None)
fujitatomoya marked this conversation as resolved.
Show resolved Hide resolved
with self.assertRaises(RuntimeError):
f.set_exception(Exception())

def test_set_exception_on_cancelled_future_raises_exception(self) -> None:
f = Future()
f.cancel()
with self.assertRaises(RuntimeError):
f.set_exception(Exception())

if __name__ == '__main__':
unittest.main()