Fix ConcurrentModificationException in CommandScheduler (#1938)

wpilibc/src/main/native/cpp/frc2/command/CommandScheduler.cpp

@@ -39,6 +39,11 @@ void CommandScheduler::AddButton(wpi::unique_function<void()> button) {
 void CommandScheduler::ClearButtons() { m_buttons.clear(); }
 
 void CommandScheduler::Schedule(bool interruptible, Command* command) {
+  if (m_inRunLoop) {
+    m_toSchedule.try_emplace(command, interruptible);
+    return;
+  }
+
   if (command->IsGrouped()) {
     wpi_setWPIErrorWithContext(CommandIllegalUse,
                                "A command that is part of a command group "
@@ -125,6 +130,7 @@ void CommandScheduler::Run() {
     button();
   }
 
+  m_inRunLoop = true;
   // Run scheduled commands, remove finished commands.
   for (auto iterator = m_scheduledCommands.begin();
        iterator != m_scheduledCommands.end(); iterator++) {
@@ -153,6 +159,18 @@ void CommandScheduler::Run() {
       m_scheduledCommands.erase(iterator);
     }
   }
+  m_inRunLoop = false;
+
+  for (auto&& commandInterruptible : m_toSchedule) {
+    Schedule(commandInterruptible.second, commandInterruptible.first);
+  }
+
+  for (auto&& command : m_toCancel) {
+    Cancel(command);
+  }
+
+  m_toSchedule.clear();
+  m_toCancel.clear();
 
   // Add default commands for un-required registered subsystems.
   for (auto&& subsystem : m_subsystems) {
@@ -198,6 +216,11 @@ Command* CommandScheduler::GetDefaultCommand(const Subsystem* subsystem) const {
 }
 
 void CommandScheduler::Cancel(Command* command) {
+  if (m_inRunLoop) {
+    m_toCancel.emplace_back(command);
+    return;
+  }
+
   auto find = m_scheduledCommands.find(command);
   if (find == m_scheduledCommands.end()) return;
   command->End(true);

wpilibc/src/main/native/include/frc2/command/CommandScheduler.h

@@ -360,6 +360,13 @@ class CommandScheduler final : public frc::Sendable,
   wpi::SmallVector<Action, 4> m_interruptActions;
   wpi::SmallVector<Action, 4> m_finishActions;
 
+  // Flag and queues for avoiding concurrent modification if commands are
+  // scheduled/canceled during run
+
+  bool m_inRunLoop = false;
+  wpi::DenseMap<Command*, bool> m_toSchedule;
+  wpi::SmallVector<Command*, 4> m_toCancel;
+
   friend class CommandTestBase;
 };
 }  // namespace frc2

wpilibc/src/test/native/cpp/frc2/command/ScheduleCommandTest.cpp

@@ -10,6 +10,7 @@
 #include "CommandTestBase.h"
 #include "frc2/command/InstantCommand.h"
 #include "frc2/command/ScheduleCommand.h"
+#include "frc2/command/SequentialCommandGroup.h"
 
 using namespace frc2;
 class ScheduleCommandTest : public CommandTestBase {};

wpilibj/src/main/java/edu/wpi/first/wpilibj2/command/CommandScheduler.java

@@ -81,6 +81,13 @@ public final class CommandScheduler implements Sendable {
   private final List<Consumer<Command>> m_interruptActions = new ArrayList<>();
   private final List<Consumer<Command>> m_finishActions = new ArrayList<>();
 
+  // Flag and queues for avoiding ConcurrentModificationException if commands are
+  // scheduled/canceled during run
+  private boolean m_inRunLoop;
+  private final Map<Command, Boolean> m_toSchedule = new LinkedHashMap<>();
+  private final List<Command> m_toCancel = new ArrayList<>();
+
+
   CommandScheduler() {
     HAL.report(tResourceType.kResourceType_Command, tInstances.kCommand_Scheduler);
     SendableRegistry.addLW(this, "Scheduler");
@@ -132,6 +139,11 @@ public final class CommandScheduler implements Sendable {
    */
   @SuppressWarnings({"PMD.CyclomaticComplexity", "PMD.NPathComplexity"})
   private void schedule(boolean interruptible, Command command) {
+    if (m_inRunLoop) {
+      m_toSchedule.put(command, interruptible);
+      return;
+    }
+
     if (CommandGroupBase.getGroupedCommands().contains(command)) {
       throw new IllegalArgumentException(
           "A command that is part of a command group cannot be independently scheduled");
@@ -222,6 +234,7 @@ public final class CommandScheduler implements Sendable {
       button.run();
     }
 
+    m_inRunLoop = true;
     //Run scheduled commands, remove finished commands.
     for (Iterator<Command> iterator = m_scheduledCommands.keySet().iterator();
          iterator.hasNext(); ) {
@@ -251,6 +264,19 @@ public final class CommandScheduler implements Sendable {
         m_requirements.keySet().removeAll(command.getRequirements());
       }
     }
+    m_inRunLoop = false;
+
+    //Schedule/cancel commands from queues populated during loop
+    for (Map.Entry<Command, Boolean> commandInterruptible : m_toSchedule.entrySet()) {
+      schedule(commandInterruptible.getValue(), commandInterruptible.getKey());
+    }
+
+    for (Command command : m_toCancel) {
+      cancel(command);
+    }
+
+    m_toSchedule.clear();
+    m_toCancel.clear();
 
     //Add default commands for un-required registered subsystems.
     for (Map.Entry<Subsystem, Command> subsystemCommand : m_subsystems.entrySet()) {
@@ -326,6 +352,11 @@ public final class CommandScheduler implements Sendable {
    * @param commands the commands to cancel
    */
   public void cancel(Command... commands) {
+    if (m_inRunLoop) {
+      m_toCancel.addAll(List.of(commands));
+      return;
+    }
+
     for (Command command : commands) {
       if (!m_scheduledCommands.containsKey(command)) {
         continue;

wpilibj/src/test/java/edu/wpi/first/wpilibj2/command/ScheduleCommandTest.java

@@ -9,6 +9,7 @@ package edu.wpi.first.wpilibj2.command;
 
 import org.junit.jupiter.api.Test;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.mockito.Mockito.verify;
 
 class ScheduleCommandTest extends CommandTestBase {
@@ -28,4 +29,19 @@ class ScheduleCommandTest extends CommandTestBase {
     verify(command1).schedule();
     verify(command2).schedule();
   }
+
+  @Test
+  void scheduleCommandDuringRunTest() {
+    CommandScheduler scheduler = CommandScheduler.getInstance();
+
+    InstantCommand toSchedule = new InstantCommand();
+    ScheduleCommand scheduleCommand = new ScheduleCommand(toSchedule);
+    SequentialCommandGroup group =
+        new SequentialCommandGroup(new InstantCommand(), scheduleCommand);
+
+    scheduler.schedule(group);
+    scheduler.schedule(new InstantCommand().perpetually());
+    scheduler.run();
+    assertDoesNotThrow(scheduler::run);
+  }
 }