Skip to content

Commit

Permalink
Add wait_for_shutdown utility (#29)
Browse files Browse the repository at this point in the history
Signed-off-by: Michel Hidalgo <mhidalgo@theaiinstitute.com>
mhidalgo-bdai authored Oct 18, 2023
1 parent 1e70b15 commit a1273c7
Showing 3 changed files with 40 additions and 6 deletions.
24 changes: 24 additions & 0 deletions bdai_ros2_wrappers/bdai_ros2_wrappers/context.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
# Copyright (c) 2023 Boston Dynamics AI Institute Inc. All rights reserved.
import threading
from typing import Optional

from rclpy.context import Context
from rclpy.utilities import get_default_context


def wait_for_shutdown(*, timeout_sec: Optional[float] = None, context: Optional[Context] = None) -> bool:
"""
Wait for context shutdown.
Args:
timeout_sec: optional timeout for wait, wait indefinitely by default.
context: context to wait on, use default context by default.
Returns:
True if shutdown, False on timeout.
"""
if context is None:
context = get_default_context()
event = threading.Event()
context.on_shutdown(event.set)
return event.wait(timeout_sec)
8 changes: 2 additions & 6 deletions bdai_ros2_wrappers/bdai_ros2_wrappers/process.py
Original file line number Diff line number Diff line change
@@ -12,6 +12,7 @@
import rclpy.logging
import rclpy.node

import bdai_ros2_wrappers.context as context
import bdai_ros2_wrappers.scope as scope
from bdai_ros2_wrappers.scope import AnyEntity, AnyEntityFactoryCallable, ROSAwareScope
from bdai_ros2_wrappers.utilities import either_or
@@ -172,12 +173,7 @@ def wait_for_shutdown(self, *, timeout_sec: typing.Optional[float] = None) -> bo
with self._lock:
if self._scope is None:
raise RuntimeError("process is not executing")
context = self._scope.context
if context is None:
context = rclpy.get_default_context()
event = threading.Event()
context.on_shutdown(event.set)
return event.wait(timeout_sec)
return context.wait_for_shutdown(timeout_sec=timeout_sec, context=self._scope.context)

def try_shutdown(self) -> None:
"""Atempts to shutdown the underlying scope context."""
14 changes: 14 additions & 0 deletions bdai_ros2_wrappers/test/test_context.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
# Copyright (c) 2023 Boston Dynamics AI Institute Inc. All rights reserved.

import rclpy

from bdai_ros2_wrappers.context import wait_for_shutdown
from bdai_ros2_wrappers.process import ROSAwareScope


def test_wait_for_shutdown(ros: ROSAwareScope) -> None:
"""Asserts that wait for shutdown works as expected."""
assert not wait_for_shutdown(timeout_sec=1.0)
assert ros.executor is not None
ros.executor.create_task(lambda: rclpy.shutdown())
assert wait_for_shutdown(timeout_sec=10.0)

0 comments on commit a1273c7

Please sign in to comment.