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

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

@@ -25,6 +25,7 @@ class SchedulingRecursionTest extends CommandTestBase {
   void cancelFromInitialize(InterruptionBehavior interruptionBehavior) {
     try (CommandScheduler scheduler = new CommandScheduler()) {
       AtomicBoolean hasOtherRun = new AtomicBoolean();
+      AtomicInteger counter = new AtomicInteger();
       Subsystem requirement = new SubsystemBase() {};
       Command selfCancels =
           new Command() {
@@ -37,6 +38,11 @@ class SchedulingRecursionTest extends CommandTestBase {
               scheduler.cancel(this);
             }
 
+            @Override
+            public void end(boolean interrupted) {
+              counter.incrementAndGet();
+            }
+
             @Override
             public InterruptionBehavior getInterruptionBehavior() {
               return interruptionBehavior;
@@ -52,6 +58,47 @@ class SchedulingRecursionTest extends CommandTestBase {
           });
       assertFalse(scheduler.isScheduled(selfCancels));
       assertTrue(scheduler.isScheduled(other));
+      assertEquals(1, counter.get());
+      scheduler.run();
+      assertTrue(hasOtherRun.get());
+    }
+  }
+
+  @EnumSource(InterruptionBehavior.class)
+  @ParameterizedTest
+  void cancelFromInitializeAction(InterruptionBehavior interruptionBehavior) {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicBoolean hasOtherRun = new AtomicBoolean();
+      AtomicInteger counter = new AtomicInteger();
+      Subsystem requirement = new Subsystem() {};
+      Command selfCancels =
+          new Command() {
+            {
+              addRequirements(requirement);
+            }
+
+            @Override
+            public void end(boolean interrupted) {
+              counter.incrementAndGet();
+            }
+
+            @Override
+            public InterruptionBehavior getInterruptionBehavior() {
+              return interruptionBehavior;
+            }
+          };
+      Command other = new RunCommand(() -> hasOtherRun.set(true), requirement);
+
+      assertDoesNotThrow(
+          () -> {
+            scheduler.onCommandInitialize(cmd -> scheduler.cancel(selfCancels));
+            scheduler.schedule(selfCancels);
+            scheduler.run();
+            scheduler.schedule(other);
+          });
+      assertFalse(scheduler.isScheduled(selfCancels));
+      assertTrue(scheduler.isScheduled(other));
+      assertEquals(1, counter.get());
       scheduler.run();
       assertTrue(hasOtherRun.get());
     }
@@ -62,6 +109,7 @@ class SchedulingRecursionTest extends CommandTestBase {
   void defaultCommandGetsRescheduledAfterSelfCanceling(InterruptionBehavior interruptionBehavior) {
     try (CommandScheduler scheduler = new CommandScheduler()) {
       AtomicBoolean hasOtherRun = new AtomicBoolean();
+      AtomicInteger counter = new AtomicInteger();
       Subsystem requirement = new SubsystemBase() {};
       Command selfCancels =
           new Command() {
@@ -74,6 +122,11 @@ class SchedulingRecursionTest extends CommandTestBase {
               scheduler.cancel(this);
             }
 
+            @Override
+            public void end(boolean interrupted) {
+              counter.incrementAndGet();
+            }
+
             @Override
             public InterruptionBehavior getInterruptionBehavior() {
               return interruptionBehavior;
@@ -90,6 +143,7 @@ class SchedulingRecursionTest extends CommandTestBase {
       scheduler.run();
       assertFalse(scheduler.isScheduled(selfCancels));
       assertTrue(scheduler.isScheduled(other));
+      assertEquals(1, counter.get());
       scheduler.run();
       assertTrue(hasOtherRun.get());
     }
@@ -115,24 +169,177 @@ class SchedulingRecursionTest extends CommandTestBase {
     }
   }
 
+  @Test
+  void cancelFromInterruptAction() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger();
+      Command selfCancels = new RunCommand(() -> {});
+      scheduler.onCommandInterrupt(
+          cmd -> {
+            counter.incrementAndGet();
+            scheduler.cancel(selfCancels);
+          });
+      scheduler.schedule(selfCancels);
+
+      assertDoesNotThrow(() -> scheduler.cancel(selfCancels));
+      assertEquals(1, counter.get());
+      assertFalse(scheduler.isScheduled(selfCancels));
+    }
+  }
+
+  @Test
+  void cancelFromEndLoop() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger();
+      FunctionalCommand dCancelsAll =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.cancelAll();
+              },
+              () -> true);
+      FunctionalCommand cCancelsD =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.cancel(dCancelsAll);
+              },
+              () -> true);
+      FunctionalCommand bCancelsC =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.cancel(cCancelsD);
+              },
+              () -> true);
+      FunctionalCommand aCancelsB =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.cancel(bCancelsC);
+              },
+              () -> true);
+
+      scheduler.schedule(aCancelsB);
+      scheduler.schedule(bCancelsC);
+      scheduler.schedule(cCancelsD);
+      scheduler.schedule(dCancelsAll);
+
+      assertDoesNotThrow(() -> scheduler.cancel(aCancelsB));
+      assertEquals(4, counter.get());
+      assertFalse(scheduler.isScheduled(aCancelsB));
+      assertFalse(scheduler.isScheduled(bCancelsC));
+      assertFalse(scheduler.isScheduled(cCancelsD));
+      assertFalse(scheduler.isScheduled(dCancelsAll));
+    }
+  }
+
+  @Test
+  void cancelFromEndLoopWhileInRunLoop() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger();
+      FunctionalCommand dCancelsAll =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.cancelAll();
+              },
+              () -> true);
+      FunctionalCommand cCancelsD =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.cancel(dCancelsAll);
+              },
+              () -> true);
+      FunctionalCommand bCancelsC =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.cancel(cCancelsD);
+              },
+              () -> true);
+      FunctionalCommand aCancelsB =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.cancel(bCancelsC);
+              },
+              () -> true);
+
+      scheduler.schedule(aCancelsB);
+      scheduler.schedule(bCancelsC);
+      scheduler.schedule(cCancelsD);
+      scheduler.schedule(dCancelsAll);
+
+      assertDoesNotThrow(() -> scheduler.run());
+      assertEquals(4, counter.get());
+      assertFalse(scheduler.isScheduled(aCancelsB));
+      assertFalse(scheduler.isScheduled(bCancelsC));
+      assertFalse(scheduler.isScheduled(cCancelsD));
+      assertFalse(scheduler.isScheduled(dCancelsAll));
+    }
+  }
+
+  @Test
+  void multiCancelFromEnd() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger();
+      FunctionalCommand bIncrementsCounter =
+          new FunctionalCommand(
+              () -> {}, () -> {}, interrupted -> counter.incrementAndGet(), () -> true);
+      Command aCancelsB =
+          new Command() {
+            @Override
+            public void end(boolean interrupted) {
+              counter.incrementAndGet();
+              scheduler.cancel(bIncrementsCounter);
+              scheduler.cancel(this);
+            }
+          };
+
+      scheduler.schedule(aCancelsB);
+      scheduler.schedule(bIncrementsCounter);
+
+      assertDoesNotThrow(() -> scheduler.cancel(aCancelsB));
+      assertEquals(2, counter.get());
+      assertFalse(scheduler.isScheduled(aCancelsB));
+      assertFalse(scheduler.isScheduled(bIncrementsCounter));
+    }
+  }
+
   @Test
   void scheduleFromEndCancel() {
     try (CommandScheduler scheduler = new CommandScheduler()) {
       AtomicInteger counter = new AtomicInteger();
       Subsystem requirement = new SubsystemBase() {};
       InstantCommand other = new InstantCommand(() -> {}, requirement);
-      Command selfCancels =
-          new Command() {
-            {
-              addRequirements(requirement);
-            }
-
-            @Override
-            public void end(boolean interrupted) {
-              counter.incrementAndGet();
-              scheduler.schedule(other);
-            }
-          };
+      FunctionalCommand selfCancels =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.schedule(other);
+              },
+              () -> false,
+              requirement);
 
       scheduler.schedule(selfCancels);
 
@@ -148,19 +355,38 @@ class SchedulingRecursionTest extends CommandTestBase {
       AtomicInteger counter = new AtomicInteger();
       Subsystem requirement = new SubsystemBase() {};
       InstantCommand other = new InstantCommand(() -> {}, requirement);
-      Command selfCancels =
-          new Command() {
-            {
-              addRequirements(requirement);
-            }
+      FunctionalCommand selfCancels =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.schedule(other);
+              },
+              () -> false,
+              requirement);
 
-            @Override
-            public void end(boolean interrupted) {
-              counter.incrementAndGet();
-              scheduler.schedule(other);
-            }
-          };
+      scheduler.schedule(selfCancels);
 
+      assertDoesNotThrow(() -> scheduler.schedule(other));
+      assertEquals(1, counter.get());
+      assertFalse(scheduler.isScheduled(selfCancels));
+      assertTrue(scheduler.isScheduled(other));
+    }
+  }
+
+  @Test
+  void scheduleFromEndInterruptAction() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger();
+      Subsystem requirement = new Subsystem() {};
+      InstantCommand other = new InstantCommand(() -> {}, requirement);
+      InstantCommand selfCancels = new InstantCommand(() -> {}, requirement);
+      scheduler.onCommandInterrupt(
+          cmd -> {
+            counter.incrementAndGet();
+            scheduler.schedule(other);
+          });
       scheduler.schedule(selfCancels);
 
       assertDoesNotThrow(() -> scheduler.schedule(other));
@@ -178,18 +404,16 @@ class SchedulingRecursionTest extends CommandTestBase {
       Subsystem requirement = new SubsystemBase() {};
       Command other =
           new InstantCommand(() -> {}, requirement).withInterruptBehavior(interruptionBehavior);
-      Command defaultCommand =
-          new Command() {
-            {
-              addRequirements(requirement);
-            }
-
-            @Override
-            public void initialize() {
-              counter.incrementAndGet();
-              scheduler.schedule(other);
-            }
-          };
+      FunctionalCommand defaultCommand =
+          new FunctionalCommand(
+              () -> {
+                counter.incrementAndGet();
+                scheduler.schedule(other);
+              },
+              () -> {},
+              interrupted -> {},
+              () -> false,
+              requirement);
 
       scheduler.setDefaultCommand(requirement, defaultCommand);
 
@@ -201,4 +425,41 @@ class SchedulingRecursionTest extends CommandTestBase {
       assertTrue(scheduler.isScheduled(other));
     }
   }
+
+  @Test
+  void cancelDefaultCommandFromEnd() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger();
+      Subsystem requirement = new Subsystem() {};
+      Command defaultCommand =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> counter.incrementAndGet(),
+              () -> false,
+              requirement);
+      Command other = new InstantCommand(() -> {}, requirement);
+      Command cancelDefaultCommand =
+          new FunctionalCommand(
+              () -> {},
+              () -> {},
+              interrupted -> {
+                counter.incrementAndGet();
+                scheduler.schedule(other);
+              },
+              () -> false);
+
+      assertDoesNotThrow(
+          () -> {
+            scheduler.schedule(cancelDefaultCommand);
+            scheduler.setDefaultCommand(requirement, defaultCommand);
+
+            scheduler.run();
+            scheduler.cancel(cancelDefaultCommand);
+          });
+      assertEquals(2, counter.get());
+      assertFalse(scheduler.isScheduled(defaultCommand));
+      assertTrue(scheduler.isScheduled(other));
+    }
+  }
 }