Improved clock.sleep()

Signed-off-by: Nadav Elkabets <elnadav12@gmail.com>

Author: nadavelkabets

rclpy/rclpy/experimental/asyncio/async_clock.py

@@ -13,27 +13,59 @@
 # limitations under the License.
 
 import asyncio
-from typing import Optional, Set
+from typing import Dict, Optional
 
-from rclpy.clock import BaseClock, ClockChange, JumpThreshold, TimeJump
+from rclpy.clock import BaseClock, ClockChange, JumpHandle, JumpThreshold, TimeJump
 from rclpy.clock_type import ClockType
 from rclpy.duration import Duration
 from rclpy.exceptions import TimeSourceChangedError
+from rclpy.time import Time
 
 
 class AsyncClock(BaseClock):
     """Clock with asyncio-compatible sleep support."""
 
     def __init__(self, *, clock_type: ClockType = ClockType.SYSTEM_TIME) -> None:
         super().__init__(clock_type=clock_type)
-        self._pending_sleeps: Set[asyncio.Future] = set()
+        self._pending_sleeps: Dict[asyncio.Future, Optional[Time]] = {}
         self._destroyed = False
 
+        threshold = JumpThreshold(
+            min_forward=Duration(nanoseconds=1),
+            min_backward=None,
+            on_clock_change=True,
+        )
+        self._jump_handle: JumpHandle = self.create_jump_callback(
+            threshold, post_callback=self._on_jump)
+
+    @staticmethod
+    def _resolve_future(future: asyncio.Future) -> None:
+        if not future.done():
+            future.set_result(None)
+
+    def _on_jump(self, time_jump: TimeJump) -> None:
+        if time_jump.clock_change in (
+            ClockChange.ROS_TIME_ACTIVATED,
+            ClockChange.ROS_TIME_DEACTIVATED,
+        ):
+            for future in self._pending_sleeps:
+                if not future.done():
+                    future.set_exception(TimeSourceChangedError())
+        elif time_jump.clock_change == ClockChange.ROS_TIME_NO_CHANGE:
+            now = self.now()
+            for future, target in self._pending_sleeps.items():
+                if target is not None and now >= target:
+                    self._resolve_future(future)
+
     def _destroy(self) -> None:
         """Cancel all pending sleeps. Called by AsyncNode.destroy_node()."""
+        if self._destroyed:
+            return
         self._destroyed = True
-        for future in list(self._pending_sleeps):
+        self._jump_handle.unregister()
+        for future in self._pending_sleeps:
             future.cancel()
+        self.handle.destroy_when_not_in_use()
 
     async def sleep(self, duration_sec: float) -> None:
         """
@@ -51,40 +83,18 @@ async def sleep(self, duration_sec: float) -> None:
         loop = asyncio.get_running_loop()
         future: asyncio.Future[None] = loop.create_future()
         timer_handle: Optional[asyncio.TimerHandle] = None
-        target = None
-
-        def _resolve() -> None:
-            if not future.done():
-                future.set_result(None)
-
-        def _reject() -> None:
-            if not future.done():
-                future.set_exception(TimeSourceChangedError())
+        target: Optional[Time] = None
 
         if self.ros_time_is_active:
             target = self.now() + Duration(nanoseconds=int(duration_sec * 1e9))
         else:
-            timer_handle = loop.call_later(duration_sec, _resolve)
+            timer_handle = loop.call_later(
+                duration_sec, AsyncClock._resolve_future, future)
 
-        def _on_jump(time_jump: TimeJump) -> None:
-            if time_jump.clock_change in (
-                ClockChange.ROS_TIME_ACTIVATED,
-                ClockChange.ROS_TIME_DEACTIVATED,
-            ):
-                _reject()
-            elif target is not None and self.now() >= target:
-                _resolve()
-
-        threshold = JumpThreshold(
-            min_forward=Duration(nanoseconds=1),
-            min_backward=None,
-            on_clock_change=True,
-        )
-        with self.create_jump_callback(threshold, post_callback=_on_jump):
-            self._pending_sleeps.add(future)
-            try:
-                await future
-            finally:
-                self._pending_sleeps.discard(future)
-                if timer_handle is not None:
-                    timer_handle.cancel()
+        self._pending_sleeps[future] = target
+        try:
+            await future
+        finally:
+            self._pending_sleeps.pop(future)
+            if timer_handle is not None:
+                timer_handle.cancel()