From 75d6bc1bc39d114719b52a5f829b11e6ddfec68d Mon Sep 17 00:00:00 2001
From: Binit Shah <bshah@hello-robot.com>
Date: Mon, 29 Jan 2024 20:54:52 -0800
Subject: [PATCH] Fix wait_command() timeout bug #255

---
 body/stretch_body/base.py              | 24 +++++++++++----
 body/stretch_body/dynamixel_X_chain.py |  7 ++++-
 body/stretch_body/robot.py             | 41 +++++++++++++++++++++-----
 3 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/body/stretch_body/base.py b/body/stretch_body/base.py
index c0545ac7..44cec755 100644
--- a/body/stretch_body/base.py
+++ b/body/stretch_body/base.py
@@ -102,15 +102,27 @@ def wait_for_contact(self, timeout=5.0):
         return False
 
     def wait_while_is_moving(self,timeout=15.0):
-        left_done_moving = self.left_wheel.wait_while_is_moving(timeout=timeout)
-        right_done_moving = self.right_wheel.wait_while_is_moving(timeout=timeout)
-        return left_done_moving and right_done_moving
+        done = []
+        def check_wait(wait_method):
+            done.append(wait_method(timeout))
+        threads = []
+        threads.append(threading.Thread(target=check_wait, args=(self.left_wheel.wait_while_is_moving,)))
+        threads.append(threading.Thread(target=check_wait, args=(self.right_wheel.wait_while_is_moving,)))
+        [thread.start() for thread in threads]
+        [thread.join() for thread in threads]
+        return all(done)
 
     def wait_until_at_setpoint(self, timeout=15.0):
         #Assume both are in motion. This will exit once both are at setpoints
-        left_at_setpoint = self.left_wheel.wait_until_at_setpoint(timeout)
-        right_at_setpoint = self.right_wheel.wait_until_at_setpoint(timeout)
-        return left_at_setpoint and right_at_setpoint
+        at_setpoint = []
+        def check_wait(wait_method):
+            at_setpoint.append(wait_method(timeout))
+        threads = []
+        threads.append(threading.Thread(target=check_wait, args=(self.left_wheel.wait_until_at_setpoint,)))
+        threads.append(threading.Thread(target=check_wait, args=(self.right_wheel.wait_until_at_setpoint,)))
+        [done_thread.start() for done_thread in threads]
+        [done_thread.join() for done_thread in threads]
+        return all(at_setpoint)
 
     def contact_thresh_to_motor_current(self,is_translate,contact_thresh):
         """
diff --git a/body/stretch_body/dynamixel_X_chain.py b/body/stretch_body/dynamixel_X_chain.py
index cdf3d2fd..15821481 100644
--- a/body/stretch_body/dynamixel_X_chain.py
+++ b/body/stretch_body/dynamixel_X_chain.py
@@ -114,8 +114,13 @@ def stop(self):
 
     def wait_until_at_setpoint(self, timeout=15.0):
         at_setpoint = []
+        def check_wait(wait_method):
+            at_setpoint.append(wait_method(timeout))
+        threads = []
         for motor in self.motors:
-            at_setpoint.append(self.motors[motor].wait_until_at_setpoint(timeout=timeout))
+            threads.append(threading.Thread(target=check_wait, args=(self.motors[motor].wait_until_at_setpoint,)))
+        [done_thread.start() for done_thread in threads]
+        [done_thread.join() for done_thread in threads]
         return all(at_setpoint)
 
     def is_trajectory_active(self):
diff --git a/body/stretch_body/robot.py b/body/stretch_body/robot.py
index 17222c09..72b1691d 100644
--- a/body/stretch_body/robot.py
+++ b/body/stretch_body/robot.py
@@ -418,13 +418,40 @@ def disable_collision_mgmt(self):
         self.collision.disable()
 
     def wait_command(self, timeout=15.0):
-        time.sleep(0.5)
-        base_done = self.base.wait_while_is_moving(timeout=timeout)
-        arm_done = self.arm.wait_while_is_moving(timeout=timeout)
-        lift_done = self.lift.wait_while_is_moving(timeout=timeout)
-        head_done = self.head.wait_until_at_setpoint(timeout=timeout)
-        eoa_done = self.end_of_arm.wait_until_at_setpoint(timeout=timeout)
-        return base_done and arm_done and lift_done and head_done and eoa_done
+        """Pause program execution until all motion complete.
+
+        Queuing up motion and pushing it to the hardware with
+        push_command() is designed to be asynchronous, enabling
+        reactive control of the robot. However, you might want
+        sychronous control, where each command's motion is completed
+        entirely before the program moves on to the next command.
+        This is where you would use wait_command()
+
+        Parameters
+        ----------
+        timeout : float
+            How long to wait for motion to complete. Must be > 0.1 sec.
+
+        Returns
+        -------
+        bool
+            True if motion completed, False if timed out before motion completed
+        """
+        time.sleep(0.1)
+        timeout = max(0.0, timeout - 0.1)
+        done = []
+        def check_wait(wait_method):
+            done.append(wait_method(timeout))
+        start = time.time()
+        threads = []
+        threads.append(threading.Thread(target=check_wait, args=(self.base.wait_while_is_moving,)))
+        threads.append(threading.Thread(target=check_wait, args=(self.arm.wait_while_is_moving,)))
+        threads.append(threading.Thread(target=check_wait, args=(self.lift.wait_while_is_moving,)))
+        threads.append(threading.Thread(target=check_wait, args=(self.head.wait_until_at_setpoint,)))
+        threads.append(threading.Thread(target=check_wait, args=(self.end_of_arm.wait_until_at_setpoint,)))
+        [thread.start() for thread in threads]
+        [thread.join() for thread in threads]
+        return all(done)
 
     # ######### Waypoint Trajectory Interface ##############################