[🛠️] Change Command Scheduler to fix iterator invalidation bugs and get rid of "hacks" in C++ and Java

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

@@ -26,7 +26,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -68,6 +67,9 @@ public static synchronized CommandScheduler getInstance() {
 
   // A set of the currently-running commands.
   private final Set<Command> m_scheduledCommands = new LinkedHashSet<>();
+  // A copy of the currently-running commands, used for iteration stored on class for caching
+  // purposes.
+  private Command[] m_scheduledCommandsCopy = new Command[12]; // 12 is arbitrary, it auto-resizes
 
   // A map from required subsystems to their requiring commands. Also used as a set of the
   // currently-required subsystems.
@@ -89,14 +91,6 @@ public static synchronized CommandScheduler getInstance() {
   private final List<BiConsumer<Command, Optional<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 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, () -> {});
 
   CommandScheduler() {
@@ -187,10 +181,6 @@ private void schedule(Command command) {
       DriverStation.reportWarning("Tried to schedule a null command", true);
       return;
     }
-    if (m_inRunLoop) {
-      m_toSchedule.add(command);
-      return;
-    }
 
     requireNotComposed(command);
 
@@ -274,11 +264,19 @@ public void run() {
     loopCache.poll();
     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();
+    m_scheduledCommandsCopy = m_scheduledCommands.toArray(m_scheduledCommandsCopy);
+    for (Command command : m_scheduledCommandsCopy) {
+      if (command == null) {
+        // at least 1 Command was removed before `run` was called (diff between copy and original).
+        // No more elements to iterate over (see toArray documentation)
+        break;
+      }
+
+      if (!isScheduled(command)) {
+        continue; // Command was canceled in the previous iterations.
+      }
 
       if (isDisabled && !command.runsWhenDisabled()) {
         cancel(command, kNoInterruptor);
@@ -291,32 +289,17 @@ public void run() {
       }
       m_watchdog.addEpoch(command.getName() + ".execute()");
       if (command.isFinished()) {
-        m_endingCommands.add(command);
+        // remove command first so the command calling cancel() doesn't crash...
+        m_scheduledCommands.remove(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());
         m_watchdog.addEpoch(command.getName() + ".end(false)");
       }
     }
-    m_inRunLoop = false;
-
-    // Schedule/cancel commands from queues populated during loop
-    for (Command command : m_toSchedule) {
-      schedule(command);
-    }
-
-    for (int i = 0; i < m_toCancelCommands.size(); i++) {
-      cancel(m_toCancelCommands.get(i), m_toCancelInterruptors.get(i));
-    }
-
-    m_toSchedule.clear();
-    m_toCancelCommands.clear();
-    m_toCancelInterruptors.clear();
 
     // Add default commands for un-required registered subsystems.
     for (Map.Entry<Subsystem, Command> subsystemCommand : m_subsystems.entrySet()) {
@@ -467,25 +450,16 @@ 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);
-      return;
-    }
     if (!isScheduled(command)) {
       return;
     }
 
-    m_endingCommands.add(command);
+    // remove command first so the command calling cancel() on itself doesn't crash...
+    m_scheduledCommands.remove(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)");
   }

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

@@ -54,14 +54,6 @@ class CommandScheduler::Impl {
   wpi::SmallVector<Action, 4> executeActions;
   wpi::SmallVector<InterruptAction, 4> interruptActions;
   wpi::SmallVector<Action, 4> finishActions;
-
-  // 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>
@@ -113,11 +105,6 @@ frc::EventLoop* CommandScheduler::GetDefaultButtonLoop() const {
 }
 
 void CommandScheduler::Schedule(Command* command) {
-  if (m_impl->inRunLoop) {
-    m_impl->toSchedule.emplace_back(command);
-    return;
-  }
-
   RequireUngrouped(command);
 
   if (m_impl->disabled || m_impl->scheduledCommands.contains(command) ||
@@ -197,10 +184,13 @@ void CommandScheduler::Run() {
   loopCache->Poll();
   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) {
+  // create a new set to avoid iterator invalidation.
+  for (Command* command : wpi::SmallSet(m_impl->scheduledCommands)) {
+    if (!IsScheduled(command)) {
+      continue;  // skip as the normal scheduledCommands was modified
+    }
+
     if (isDisabled && !command->RunsWhenDisabled()) {
       Cancel(command, std::nullopt);
       continue;
@@ -213,34 +203,19 @@ void CommandScheduler::Run() {
     m_watchdog.AddEpoch(command->GetName() + ".Execute()");
 
     if (command->IsFinished()) {
-      m_impl->endingCommands.insert(command);
+      m_impl->scheduledCommands.erase(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_watchdog.AddEpoch(command->GetName() + ".End(false)");
     }
   }
-  m_impl->inRunLoop = false;
-
-  for (Command* command : m_impl->toSchedule) {
-    Schedule(command);
-  }
-
-  for (size_t i = 0; i < m_impl->toCancelCommands.size(); i++) {
-    Cancel(m_impl->toCancelCommands[i], m_impl->toCancelInterruptors[i]);
-  }
-
-  m_impl->toSchedule.clear();
-  m_impl->toCancelCommands.clear();
-  m_impl->toCancelInterruptors.clear();
 
   // Add default commands for un-required registered subsystems.
   for (auto&& subsystem : m_impl->subsystems) {
@@ -333,24 +308,14 @@ 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;
-  }
   if (!IsScheduled(command)) {
     return;
   }
-  m_impl->endingCommands.insert(command);
+  m_impl->scheduledCommands.erase(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);

wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandScheduleTest.java

@@ -5,6 +5,7 @@
 package edu.wpi.first.wpilibj2.command;
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.never;
@@ -13,6 +14,7 @@
 
 import edu.wpi.first.networktables.NetworkTableInstance;
 import edu.wpi.first.wpilibj.smartdashboard.SmartDashboard;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.Test;
 
 class CommandScheduleTest extends CommandTestBase {
@@ -109,6 +111,102 @@ void schedulerCancelTest() {
     }
   }
 
+  @Test
+  void cancelNextCommandTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      Command[] commands = new Command[2];
+      var commandRunCounter = new AtomicInteger(0);
+      // due to the ordering of sets being non-deterministic, we cancel the other command in the
+      // first command and check that the second command is not run
+      Command command1 =
+          new RunCommand(
+              () -> {
+                scheduler.cancel(commands[1]);
+                commandRunCounter.incrementAndGet();
+              });
+      Command command2 =
+          new RunCommand(
+              () -> {
+                scheduler.cancel(commands[0]);
+                commandRunCounter.incrementAndGet();
+              });
+
+      commands[0] = command1;
+      commands[1] = command2;
+
+      scheduler.schedule(command1, command2);
+      scheduler.run();
+
+      assertEquals(
+          1, commandRunCounter.get(), "Second command was run when it shouldn't have been");
+
+      // only one of the commands should be canceled.
+      assertFalse(
+          scheduler.isScheduled(command1) && scheduler.isScheduled(command2),
+          "None of the commands were canceled when one should have been");
+      // one of the commands shouldn't be canceled because the other one is canceled first
+      assertTrue(
+          scheduler.isScheduled(command1) || scheduler.isScheduled(command2),
+          "Both commands were canceled when only one should have been");
+
+      scheduler.run();
+      assertEquals(2, commandRunCounter.get());
+    }
+  }
+
+  @Test
+  void commandKnowsWhenEndedTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      Command[] commands = new Command[1];
+      Command command =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              isForced -> {
+                assertFalse(
+                    scheduler.isScheduled(commands[0]),
+                    "Command shouldn't be scheduled when its end is called");
+              },
+              () -> true);
+
+      commands[0] = command;
+      scheduler.schedule(command);
+      scheduler.run();
+      assertFalse(
+          scheduler.isScheduled(command),
+          "Command should be removed from scheduler when its isFinished() returns true");
+    }
+  }
+
+  @Test
+  void scheduleCommandInCommand() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger(0);
+      Command[] commands = new Command[1];
+      Command command1 =
+          new RunCommand(
+              () -> {
+                scheduler.schedule(commands[0]);
+                assertEquals(
+                    1,
+                    counter.get(),
+                    "Command 2's init was not run immediately after getting scheduled");
+              });
+      Command command2 = new InstantCommand(counter::incrementAndGet);
+
+      commands[0] = command2;
+
+      scheduler.schedule(command1);
+      scheduler.run();
+      assertEquals(1, counter.get(), "Command 2 was not run when it should have been");
+      assertTrue(scheduler.isScheduled(command2));
+
+      scheduler.run();
+      assertEquals(1, counter.get(), "Command 2 was run when it shouldn't have been");
+      assertFalse(scheduler.isScheduled(command2), "Command 2 did not end when it should have");
+    }
+  }
+
   @Test
   void notScheduledCancelTest() {
     try (CommandScheduler scheduler = new CommandScheduler()) {