Add ability to end startCompetition() main loop (#2032)

This is useful for both cleanly exiting from simulation and for unit testing
at a framework level.

This change required removing move constructor/assignment from IterativeRobot.

wpilibc/src/main/native/cpp/DriverStation.cpp

@@ -473,6 +473,13 @@ double DriverStation::GetBatteryVoltage() const {
   return voltage;
 }
 
+void DriverStation::WakeupWaitForData() {
+  std::scoped_lock waitLock(m_waitForDataMutex);
+  // Nofify all threads
+  m_waitForDataCounter++;
+  m_waitForDataCond.notify_all();
+}
+
 void DriverStation::GetData() {
   {
     // Compute the pressed and released buttons
@@ -494,13 +501,7 @@ void DriverStation::GetData() {
     }
   }
 
-  {
+  WakeupWaitForData();
-    std::scoped_lock waitLock(m_waitForDataMutex);
-    // Nofify all threads
-    m_waitForDataCounter++;
-    m_waitForDataCond.notify_all();
-  }
-
   SendMatchData();
 }
 

wpilibc/src/main/native/cpp/IterativeRobot.cpp

@@ -30,7 +30,13 @@ void IterativeRobot::StartCompetition() {
   while (true) {
     // Wait for driver station data so the loop doesn't hog the CPU
     DriverStation::GetInstance().WaitForData();
+    if (m_exit) break;
 
     LoopFunc();
   }
 }
+
+void IterativeRobot::EndCompetition() {
+  m_exit = true;
+  DriverStation::GetInstance().WakeupWaitForData();
+}

wpilibc/src/main/native/cpp/TimedRobot.cpp

@@ -43,6 +43,11 @@ void TimedRobot::StartCompetition() {
   }
 }
 
+void TimedRobot::EndCompetition() {
+  int32_t status = 0;
+  HAL_StopNotifier(m_notifier, &status);
+}
+
 units::second_t TimedRobot::GetPeriod() const {
   return units::second_t(m_period);
 }

wpilibc/src/main/native/include/frc/DriverStation.h

@@ -394,6 +394,11 @@ class DriverStation : public ErrorBase {
    */
   void InTest(bool entering) { m_userInTest = entering; }
 
+  /**
+   * Forces WaitForData() to return immediately.
+   */
+  void WakeupWaitForData();
+
  protected:
   /**
    * Copy data from the DS task for the user.

wpilibc/src/main/native/include/frc/IterativeRobot.h

@@ -1,5 +1,5 @@
 /*----------------------------------------------------------------------------*/
-/* Copyright (c) 2008-2018 FIRST. All Rights Reserved.                        */
+/* Copyright (c) 2008-2019 FIRST. All Rights Reserved.                        */
 /* Open Source Software - may be modified and shared by FRC teams. The code   */
 /* must be accompanied by the FIRST BSD license file in the root directory of */
 /* the project.                                                               */
@@ -7,6 +7,8 @@
 
 #pragma once
 
+#include <atomic>
+
 #include "frc/IterativeRobotBase.h"
 
 namespace frc {
@@ -28,9 +30,6 @@ class IterativeRobot : public IterativeRobotBase {
   IterativeRobot();
   virtual ~IterativeRobot() = default;
 
-  IterativeRobot(IterativeRobot&&) = default;
-  IterativeRobot& operator=(IterativeRobot&&) = default;
-
   /**
    * Provide an alternate "main loop" via StartCompetition().
    *
@@ -38,6 +37,14 @@ class IterativeRobot : public IterativeRobotBase {
    * with the DS packets.
    */
   void StartCompetition() override;
+
+  /**
+   * Ends the main loop in StartCompetition().
+   */
+  void EndCompetition() override;
+
+ private:
+  std::atomic<bool> m_exit{false};
 };
 
 }  // namespace frc

wpilibc/src/main/native/include/frc/RobotBase.h

@@ -7,9 +7,12 @@
 
 #pragma once
 
+#include <chrono>
 #include <thread>
 
 #include <hal/Main.h>
+#include <wpi/condition_variable.h>
+#include <wpi/mutex.h>
 #include <wpi/raw_ostream.h>
 
 #include "frc/Base.h"
@@ -23,9 +26,13 @@ int RunHALInitialization();
 namespace impl {
 
 template <class Robot>
-void RunRobot() {
+void RunRobot(wpi::mutex& m, Robot** robot) {
-  static Robot robot;
+  static Robot theRobot;
-  robot.StartCompetition();
+  {
+    std::scoped_lock lock{m};
+    *robot = &theRobot;
+  }
+  theRobot.StartCompetition();
 }
 
 }  // namespace impl
@@ -36,20 +43,50 @@ int StartRobot() {
   if (halInit != 0) {
     return halInit;
   }
+
+  static wpi::mutex m;
+  static wpi::condition_variable cv;
+  static Robot* robot = nullptr;
+  static bool exited = false;
+
   if (HAL_HasMain()) {
-    std::thread([] {
+    std::thread thr([] {
       try {
-        impl::RunRobot<Robot>();
+        impl::RunRobot<Robot>(m, &robot);
       } catch (...) {
         HAL_ExitMain();
+        {
+          std::scoped_lock lock{m};
+          robot = nullptr;
+          exited = true;
+        }
+        cv.notify_all();
         throw;
       }
+
       HAL_ExitMain();
-    })
+      {
-        .detach();
+        std::scoped_lock lock{m};
+        robot = nullptr;
+        exited = true;
+      }
+      cv.notify_all();
+    });
+
     HAL_RunMain();
+
+    // signal loop to exit
+    if (robot) robot->EndCompetition();
+
+    // prefer to join, but detach to exit if it doesn't exit in a timely manner
+    using namespace std::chrono_literals;
+    std::unique_lock lock{m};
+    if (cv.wait_for(lock, 1s, [] { return exited; }))
+      thr.join();
+    else
+      thr.detach();
   } else {
-    impl::RunRobot<Robot>();
+    impl::RunRobot<Robot>(m, &robot);
   }
 
   return 0;
@@ -127,6 +164,8 @@ class RobotBase {
 
   virtual void StartCompetition() = 0;
 
+  virtual void EndCompetition() = 0;
+
   static constexpr bool IsReal() {
 #ifdef __FRC_ROBORIO__
     return true;

wpilibc/src/main/native/include/frc/TimedRobot.h

@@ -34,6 +34,11 @@ class TimedRobot : public IterativeRobotBase, public ErrorBase {
    */
   void StartCompetition() override;
 
+  /**
+   * Ends the main loop in StartCompetition().
+   */
+  void EndCompetition() override;
+
   /**
    * Get the time period between calls to Periodic() functions.
    */

wpilibcExamples/src/main/cpp/templates/robotbaseskeleton/cpp/Robot.cpp

@@ -32,7 +32,7 @@ void Robot::StartCompetition() {
   // Tell the DS that the robot is ready to be enabled
   HAL_ObserveUserProgramStarting();
 
-  while (true) {
+  while (!m_exit) {
     if (IsDisabled()) {
       m_ds.InDisabled(true);
       Disabled();
@@ -61,6 +61,8 @@ void Robot::StartCompetition() {
   }
 }
 
+void Robot::EndCompetition() { m_exit = true; }
+
 #ifndef RUNNING_FRC_TESTS
 int main() { return frc::StartRobot<Robot>(); }
 #endif

wpilibcExamples/src/main/cpp/templates/robotbaseskeleton/include/Robot.h

@@ -7,6 +7,8 @@
 
 #pragma once
 
+#include <atomic>
+
 #include <frc/RobotBase.h>
 
 class Robot : public frc::RobotBase {
@@ -18,4 +20,8 @@ class Robot : public frc::RobotBase {
   void Test();
 
   void StartCompetition() override;
+  void EndCompetition() override;
+
+ private:
+  std::atomic<bool> m_exit{false};
 };

wpilibj/src/main/java/edu/wpi/first/wpilibj/DriverStation.java

@@ -1008,6 +1008,16 @@ public class DriverStation {
     m_matchDataSender.controlWord.setDouble(HAL.nativeGetControlWord());
   }
 
+  /**
+   * Forces waitForData() to return immediately.
+   */
+  public void wakeupWaitForData() {
+    m_waitForDataMutex.lock();
+    m_waitForDataCount++;
+    m_waitForDataCond.signalAll();
+    m_waitForDataMutex.unlock();
+  }
+
   /**
    * Copy data from the DS task for the user. If no new data exists, it will just be returned,
    * otherwise the data will be copied from the DS polling loop.
@@ -1061,11 +1071,7 @@ public class DriverStation {
       m_cacheDataMutex.unlock();
     }
 
-    m_waitForDataMutex.lock();
+    wakeupWaitForData();
-    m_waitForDataCount++;
-    m_waitForDataCond.signalAll();
-    m_waitForDataMutex.unlock();
-
     sendMatchData();
   }
 

wpilibj/src/main/java/edu/wpi/first/wpilibj/IterativeRobot.java

@@ -25,6 +25,7 @@ import edu.wpi.first.hal.HAL;
 @Deprecated
 public class IterativeRobot extends IterativeRobotBase {
   private static final double kPacketPeriod = 0.02;
+  private volatile boolean m_exit;
 
   /**
    * Create a new IterativeRobot.
@@ -49,8 +50,20 @@ public class IterativeRobot extends IterativeRobotBase {
     while (!Thread.currentThread().isInterrupted()) {
       // Wait for new data to arrive
       m_ds.waitForData();
+      if (m_exit) {
+        break;
+      }
 
       loopFunc();
     }
   }
+
+  /**
+   * Ends the main loop in startCompetition().
+   */
+  @Override
+  public void endCompetition() {
+    m_exit = true;
+    m_ds.wakeupWaitForData();
+  }
 }

wpilibj/src/main/java/edu/wpi/first/wpilibj/RobotBase.java

@@ -12,6 +12,7 @@ import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Supplier;
 
 import edu.wpi.cscore.CameraServerJNI;
@@ -189,6 +190,11 @@ public abstract class RobotBase implements AutoCloseable {
    */
   public abstract void startCompetition();
 
+  /**
+   * Ends the main loop in startCompetition().
+   */
+  public abstract void endCompetition();
+
   @SuppressWarnings("JavadocMethod")
   public static boolean getBooleanProperty(String name, boolean defaultValue) {
     String propVal = System.getProperty(name);
@@ -204,6 +210,10 @@ public abstract class RobotBase implements AutoCloseable {
     }
   }
 
+  private static final ReentrantLock m_runMutex = new ReentrantLock();
+  private static RobotBase m_robotCopy;
+  private static boolean m_suppressExitWarning;
+
   /**
    * Run the robot main loop.
    */
@@ -232,6 +242,10 @@ public abstract class RobotBase implements AutoCloseable {
       return;
     }
 
+    m_runMutex.lock();
+    m_robotCopy = robot;
+    m_runMutex.unlock();
+
     if (isReal()) {
       try {
         final File file = new File("/tmp/frc_versions/FRC_Lib_Version.ini");
@@ -265,18 +279,32 @@ public abstract class RobotBase implements AutoCloseable {
           throwable.getStackTrace());
       errorOnExit = true;
     } finally {
-      // startCompetition never returns unless exception occurs....
+      m_runMutex.lock();
-      DriverStation.reportWarning("Robots should not quit, but yours did!", false);
+      boolean suppressExitWarning = m_suppressExitWarning;
-      if (errorOnExit) {
+      m_runMutex.unlock();
-        DriverStation.reportError(
+      if (!suppressExitWarning) {
-            "The startCompetition() method (or methods called by it) should have "
+        // startCompetition never returns unless exception occurs....
-                + "handled the exception above.", false);
+        DriverStation.reportWarning("Robots should not quit, but yours did!", false);
-      } else {
+        if (errorOnExit) {
-        DriverStation.reportError("Unexpected return from startCompetition() method.", false);
+          DriverStation.reportError(
+              "The startCompetition() method (or methods called by it) should have "
+                  + "handled the exception above.", false);
+        } else {
+          DriverStation.reportError("Unexpected return from startCompetition() method.", false);
+        }
       }
     }
   }
 
+  /**
+   * Suppress the "Robots should not quit" message.
+   */
+  public static void suppressExitWarning(boolean value) {
+    m_runMutex.lock();
+    m_suppressExitWarning = value;
+    m_runMutex.unlock();
+  }
+
   /**
    * Starting point for the applications.
    */
@@ -299,6 +327,18 @@ public abstract class RobotBase implements AutoCloseable {
       thread.setDaemon(true);
       thread.start();
       HAL.runMain();
+      suppressExitWarning(true);
+      m_runMutex.lock();
+      RobotBase robot = m_robotCopy;
+      m_runMutex.unlock();
+      if (robot != null) {
+        robot.endCompetition();
+      }
+      try {
+        thread.join(1000);
+      } catch (InterruptedException ex) {
+        Thread.currentThread().interrupt();
+      }
     } else {
       runRobot(robotSupplier);
     }

wpilibj/src/main/java/edu/wpi/first/wpilibj/TimedRobot.java

@@ -1,5 +1,5 @@
 /*----------------------------------------------------------------------------*/
-/* Copyright (c) 2017-2018 FIRST. All Rights Reserved.                        */
+/* Copyright (c) 2017-2019 FIRST. All Rights Reserved.                        */
 /* Open Source Software - may be modified and shared by FRC teams. The code   */
 /* must be accompanied by the FIRST BSD license file in the root directory of */
 /* the project.                                                               */
@@ -82,6 +82,14 @@ public class TimedRobot extends IterativeRobotBase {
     }
   }
 
+  /**
+   * Ends the main loop in startCompetition().
+   */
+  @Override
+  public void endCompetition() {
+    NotifierJNI.stopNotifier(m_notifier);
+  }
+
   /**
    * Get time period between calls to Periodic() functions.
    */

wpilibjExamples/src/main/java/edu/wpi/first/wpilibj/templates/robotbaseskeleton/Robot.java

@@ -33,6 +33,8 @@ public class Robot extends RobotBase {
   public void test() {
   }
 
+  private volatile boolean m_exit;
+
   @SuppressWarnings("PMD.CyclomaticComplexity")
   @Override
   public void startCompetition() {
@@ -41,7 +43,7 @@ public class Robot extends RobotBase {
     // Tell the DS that the robot is ready to be enabled
     HAL.observeUserProgramStarting();
 
-    while (!Thread.currentThread().isInterrupted()) {
+    while (!Thread.currentThread().isInterrupted() && !m_exit) {
       if (isDisabled()) {
         m_ds.InDisabled(true);
         disabled();
@@ -77,4 +79,9 @@ public class Robot extends RobotBase {
       }
     }
   }
+
+  @Override
+  public void endCompetition() {
+    m_exit = true;
+  }
 }