[commands] Make command scheduling order consistent (#5470)

wpilibsuite · Oct 10, 2023 · ff18490 · ff18490
1 parent 58e8474
commit ff18490
Show file tree

Hide file tree

Showing 4 changed files with 566 additions and 64 deletions.
diff --git a/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/CommandScheduler.java b/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/CommandScheduler.java
@@ -92,6 +92,7 @@ public static synchronized CommandScheduler getInstance() {
   private final Set<Command> m_toSchedule = new LinkedHashSet<>();
   private final List<Command> m_toCancelCommands = new ArrayList<>();
   private final List<Optional<Command>> m_toCancelInterruptors = new ArrayList<>();
+  private final Set<Command> m_endingCommands = new LinkedHashSet<>();
 
   private final Watchdog m_watchdog = new Watchdog(TimedRobot.kDefaultPeriod, () -> {});
 
@@ -271,18 +272,13 @@ public void run() {
     m_watchdog.addEpoch("buttons.run()");
 
     m_inRunLoop = true;
+    boolean isDisabled = RobotState.isDisabled();
     // Run scheduled commands, remove finished commands.
     for (Iterator<Command> iterator = m_scheduledCommands.iterator(); iterator.hasNext(); ) {
       Command command = iterator.next();
 
-      if (!command.runsWhenDisabled() && RobotState.isDisabled()) {
-        command.end(true);
-        for (BiConsumer<Command, Optional<Command>> action : m_interruptActions) {
-          action.accept(command, kNoInterruptor);
-        }
-        m_requirements.keySet().removeAll(command.getRequirements());
-        iterator.remove();
-        m_watchdog.addEpoch(command.getName() + ".end(true)");
+      if (isDisabled && !command.runsWhenDisabled()) {
+        cancel(command, kNoInterruptor);
         continue;
       }
 
@@ -292,10 +288,12 @@ public void run() {
       }
       m_watchdog.addEpoch(command.getName() + ".execute()");
       if (command.isFinished()) {
+        m_endingCommands.add(command);
         command.end(false);
         for (Consumer<Command> action : m_finishActions) {
           action.accept(command);
         }
+        m_endingCommands.remove(command);
         iterator.remove();
 
         m_requirements.keySet().removeAll(command.getRequirements());
@@ -466,6 +464,9 @@ private void cancel(Command command, Optional<Command> interruptor) {
       DriverStation.reportWarning("Tried to cancel a null command", true);
       return;
     }
+    if (m_endingCommands.contains(command)) {
+      return;
+    }
     if (m_inRunLoop) {
       m_toCancelCommands.add(command);
       m_toCancelInterruptors.add(interruptor);
@@ -475,12 +476,14 @@ private void cancel(Command command, Optional<Command> interruptor) {
       return;
     }
 
-    m_scheduledCommands.remove(command);
-    m_requirements.keySet().removeAll(command.getRequirements());
+    m_endingCommands.add(command);
     command.end(true);
     for (BiConsumer<Command, Optional<Command>> action : m_interruptActions) {
       action.accept(command, interruptor);
     }
+    m_endingCommands.remove(command);
+    m_scheduledCommands.remove(command);
+    m_requirements.keySet().removeAll(command.getRequirements());
     m_watchdog.addEpoch(command.getName() + ".end(true)");
   }
 

diff --git a/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp b/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp
@@ -53,11 +53,11 @@ class CommandScheduler::Impl {
 
   // Flag and queues for avoiding concurrent modification if commands are
   // scheduled/canceled during run
-
   bool inRunLoop = false;
   wpi::SmallVector<Command*, 4> toSchedule;
   wpi::SmallVector<Command*, 4> toCancelCommands;
   wpi::SmallVector<std::optional<Command*>, 4> toCancelInterruptors;
+  wpi::SmallSet<Command*, 4> endingCommands;
 };
 
 template <typename TMap, typename TKey>
@@ -194,9 +194,10 @@ void CommandScheduler::Run() {
   m_watchdog.AddEpoch("buttons.Run()");
 
   m_impl->inRunLoop = true;
+  bool isDisabled = frc::RobotState::IsDisabled();
   // Run scheduled commands, remove finished commands.
   for (Command* command : m_impl->scheduledCommands) {
-    if (!command->RunsWhenDisabled() && frc::RobotState::IsDisabled()) {
+    if (isDisabled && !command->RunsWhenDisabled()) {
       Cancel(command, std::nullopt);
       continue;
     }
@@ -208,16 +209,18 @@ void CommandScheduler::Run() {
     m_watchdog.AddEpoch(command->GetName() + ".Execute()");
 
     if (command->IsFinished()) {
+      m_impl->endingCommands.insert(command);
       command->End(false);
       for (auto&& action : m_impl->finishActions) {
         action(*command);
       }
+      m_impl->endingCommands.erase(command);
 
+      m_impl->scheduledCommands.erase(command);
       for (auto&& requirement : command->GetRequirements()) {
         m_impl->requirements.erase(requirement);
       }
 
-      m_impl->scheduledCommands.erase(command);
       m_watchdog.AddEpoch(command->GetName() + ".End(false)");
     }
   }
@@ -326,27 +329,29 @@ void CommandScheduler::Cancel(Command* command,
   if (!m_impl) {
     return;
   }
-
+  if (m_impl->endingCommands.contains(command)) {
+    return;
+  }
   if (m_impl->inRunLoop) {
     m_impl->toCancelCommands.emplace_back(command);
     m_impl->toCancelInterruptors.emplace_back(interruptor);
     return;
   }
-
-  auto find = m_impl->scheduledCommands.find(command);
-  if (find == m_impl->scheduledCommands.end()) {
+  if (!IsScheduled(command)) {
     return;
   }
-  m_impl->scheduledCommands.erase(*find);
+  m_impl->endingCommands.insert(command);
+  command->End(true);
+  for (auto&& action : m_impl->interruptActions) {
+    action(*command, interruptor);
+  }
+  m_impl->endingCommands.erase(command);
+  m_impl->scheduledCommands.erase(command);
   for (auto&& requirement : m_impl->requirements) {
     if (requirement.second == command) {
       m_impl->requirements.erase(requirement.first);
     }
   }
-  command->End(true);
-  for (auto&& action : m_impl->interruptActions) {
-    action(*command, interruptor);
-  }
   m_watchdog.AddEpoch(command->GetName() + ".End(true)");
 }