[🛠️] 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);
+ // for parity with C++ where the ordering of sets is 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 commandToGetScheduled = new InstantCommand(counter::incrementAndGet);
+ Command command =
+ new RunCommand(
+ () -> {
+ scheduler.schedule(commandToGetScheduled);
+ assertEquals(
+ 1,
+ counter.get(),
+ "Scheduled command's init was not run immediately after getting scheduled");
+ });
+
+ scheduler.schedule(command);
+ scheduler.run();
+ assertEquals(1, counter.get(), "Command 2 was not run when it should have been");
+ assertTrue(scheduler.isScheduled(commandToGetScheduled));
+
+ scheduler.run();
+ assertEquals(1, counter.get(), "Command 2 was run when it shouldn't have been");
+ assertFalse(
+ scheduler.isScheduled(commandToGetScheduled),
+ "Command 2 did not end when it should have");
+ }
+ }
+
  @Test
  void notScheduledCancelTest() {
  try (CommandScheduler scheduler = new CommandScheduler()) {