[wpilib] Support scheduling functions more often than robot loop (#2766)

Currently, teams have to make a Notifier to run feedback controllers
more often than the TimedRobot loop period of 20ms (running TimedRobot
more often than this is not advised). This lets users add callbacks to
the main robot loop that run at a user-defined period. This allows
running feedback controllers more often, but does so synchronously with
TimedRobot so there aren't any thread safety issues.

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

@@ -31,21 +31,37 @@ void TimedRobot::StartCompetition() {
   // Tell the DS that the robot is ready to be enabled
   HAL_ObserveUserProgramStarting();
 
-  m_expirationTime = units::second_t{Timer::GetFPGATimestamp()} + m_period;
-  UpdateAlarm();
-
   // Loop forever, calling the appropriate mode-dependent function
   while (true) {
+    // We don't have to check there's an element in the queue first because
+    // there's always at least one (the constructor adds one). It's reenqueued
+    // at the end of the loop.
+    auto callback = m_callbacks.pop();
+
     int32_t status = 0;
+    HAL_UpdateNotifierAlarm(
+        m_notifier, static_cast<uint64_t>(callback.expirationTime * 1e6),
+        &status);
+    wpi_setHALError(status);
+
     uint64_t curTime = HAL_WaitForNotifierAlarm(m_notifier, &status);
     if (curTime == 0 || status != 0) break;
 
-    m_expirationTime += m_period;
+    callback.func();
 
-    UpdateAlarm();
+    callback.expirationTime += callback.period;
+    m_callbacks.push(std::move(callback));
 
-    // Call callback
-    LoopFunc();
+    // Process all other callbacks that are ready to run
+    while (static_cast<uint64_t>(m_callbacks.top().expirationTime * 1e6) <=
+           curTime) {
+      callback = m_callbacks.pop();
+
+      callback.func();
+
+      callback.expirationTime += callback.period;
+      m_callbacks.push(std::move(callback));
+    }
   }
 }
 
@@ -61,6 +77,9 @@ units::second_t TimedRobot::GetPeriod() const {
 TimedRobot::TimedRobot(double period) : TimedRobot(units::second_t(period)) {}
 
 TimedRobot::TimedRobot(units::second_t period) : IterativeRobotBase(period) {
+  m_startTime = frc2::Timer::GetFPGATimestamp();
+  AddPeriodic([=] { LoopFunc(); }, period);
+
   int32_t status = 0;
   m_notifier = HAL_InitializeNotifier(&status);
   wpi_setHALError(status);
@@ -79,9 +98,7 @@ TimedRobot::~TimedRobot() {
   HAL_CleanNotifier(m_notifier, &status);
 }
 
-void TimedRobot::UpdateAlarm() {
-  int32_t status = 0;
-  HAL_UpdateNotifierAlarm(
-      m_notifier, static_cast<uint64_t>(m_expirationTime * 1e6), &status);
-  wpi_setHALError(status);
+void TimedRobot::AddPeriodic(std::function<void()> callback,
+                             units::second_t period, units::second_t offset) {
+  m_callbacks.emplace(callback, m_startTime, period, offset);
 }

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

@@ -7,12 +7,18 @@
 
 #pragma once
 
+#include <functional>
+#include <vector>
+
 #include <hal/Types.h>
+#include <units/math.h>
 #include <units/time.h>
 #include <wpi/deprecated.h>
+#include <wpi/priority_queue.h>
 
 #include "frc/ErrorBase.h"
 #include "frc/IterativeRobotBase.h"
+#include "frc2/Timer.h"
 
 namespace frc {
 
@@ -67,16 +73,57 @@ class TimedRobot : public IterativeRobotBase, public ErrorBase {
   TimedRobot(TimedRobot&&) = default;
   TimedRobot& operator=(TimedRobot&&) = default;
 
- private:
-  hal::Handle<HAL_NotifierHandle> m_notifier;
-
-  // The absolute expiration time
-  units::second_t m_expirationTime{0};
-
   /**
-   * Update the HAL alarm time.
+   * Add a callback to run at a specific period with a starting time offset.
+   *
+   * This is scheduled on TimedRobot's Notifier, so TimedRobot and the callback
+   * run synchronously. Interactions between them are thread-safe.
+   *
+   * @param callback The callback to run.
+   * @param period   The period at which to run the callback.
+   * @param offset   The offset from the common starting time. This is useful
+   *                 for scheduling a callback in a different timeslot relative
+   *                 to TimedRobot.
    */
-  void UpdateAlarm();
+  void AddPeriodic(std::function<void()> callback, units::second_t period,
+                   units::second_t offset = 0_s);
+
+ private:
+  class Callback {
+   public:
+    std::function<void()> func;
+    units::second_t period;
+    units::second_t expirationTime;
+
+    /**
+     * Construct a callback container.
+     *
+     * @param func      The callback to run.
+     * @param startTime The common starting point for all callback scheduling.
+     * @param period    The period at which to run the callback.
+     * @param offset    The offset from the common starting time.
+     */
+    Callback(std::function<void()> func, units::second_t startTime,
+             units::second_t period, units::second_t offset)
+        : func{func},
+          period{period},
+          expirationTime{
+              startTime + offset +
+              units::math::floor((frc2::Timer::GetFPGATimestamp() - startTime) /
+                                 period) *
+                  period +
+              period} {}
+
+    bool operator>(const Callback& rhs) const {
+      return expirationTime > rhs.expirationTime;
+    }
+  };
+
+  hal::Handle<HAL_NotifierHandle> m_notifier;
+  units::second_t m_startTime;
+
+  wpi::priority_queue<Callback, std::vector<Callback>, std::greater<Callback>>
+      m_callbacks;
 };
 
 }  // namespace frc