diff --git a/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/Command.java b/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/Command.java index d695006ba54..8a3f801fe30 100644 --- a/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/Command.java +++ b/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/Command.java @@ -90,6 +90,14 @@ public interface Command { */ void run(Coroutine coroutine); + /** + * An optional lifecycle hook that can be implemented to execute some code whenever this command + * is cancelled before it naturally completes. + */ + default void onCancel() { + // NOP by default + } + /** * The name of the command. * diff --git a/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/CommandBuilder.java b/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/CommandBuilder.java index 9822a69f834..5d171a19bb9 100644 --- a/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/CommandBuilder.java +++ b/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/CommandBuilder.java @@ -10,6 +10,7 @@ public class CommandBuilder { private final Set requirements = new HashSet<>(); private Consumer impl; + private Runnable onCancel = () -> {}; private String name; private int priority = Command.DEFAULT_PRIORITY; private Command.RobotDisabledBehavior disabledBehavior = @@ -68,6 +69,11 @@ public CommandBuilder executing(Consumer impl) { return this; } + public CommandBuilder whenCanceled(Runnable onCancel) { + this.onCancel = onCancel; + return this; + } + public Command make() { Objects.requireNonNull(name, "Name was not specified"); Objects.requireNonNull(impl, "Command logic was not specified"); @@ -78,6 +84,11 @@ public void run(Coroutine coroutine) { impl.accept(coroutine); } + @Override + public void onCancel() { + onCancel.run(); + } + @Override public String name() { return name; diff --git a/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/Scheduler.java b/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/Scheduler.java index f969edd2130..c224998ddf6 100644 --- a/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/Scheduler.java +++ b/commandsv3/src/main/java/edu/wpi/first/wpilibj/commandsv3/Scheduler.java @@ -343,6 +343,8 @@ public void cancel(Command command) { onDeck.removeIf(state -> state.command() == command); suspendedStates.removeIf(state -> state.command() == command); + command.onCancel(); + // Clean up any orphaned child commands; their lifespan must not exceed the parent's removeOrphanedChildren(command); } @@ -499,12 +501,7 @@ private void removeOrphanedChildren(Command parent) { .stream() .filter(e -> e.getValue().parent() == parent) .toList() // copy to an intermediate list to avoid concurrent modification - .forEach(e -> { - commandStates.entrySet().remove(e); - suspendedStates.removeIf(s -> s.parent() == e.getKey()); - onDeck.removeIf(s -> s.parent() == e.getKey()); - removeOrphanedChildren(e.getKey()); // recursion - }); + .forEach(e -> cancel(e.getKey())); } /** diff --git a/commandsv3/src/test/java/edu/wpi/first/wpilibj/commandsv3/SchedulerTest.java b/commandsv3/src/test/java/edu/wpi/first/wpilibj/commandsv3/SchedulerTest.java index 1b121ec0a67..0788722a05c 100644 --- a/commandsv3/src/test/java/edu/wpi/first/wpilibj/commandsv3/SchedulerTest.java +++ b/commandsv3/src/test/java/edu/wpi/first/wpilibj/commandsv3/SchedulerTest.java @@ -3,12 +3,11 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import java.util.ArrayList; -import java.util.List; +import java.util.Collections; import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -456,6 +455,75 @@ void compositionsCannotAwaitConflictingCommands() { } } + @Test + void doesNotRunOnCancelWhenInterruptingOnDeck() { + var ran = new AtomicBoolean(false); + + var resource = new RequireableResource("The Resource", scheduler); + var cmd = resource.run(Coroutine::yield).whenCanceled(() -> ran.set(true)).named("cmd"); + var interrupter = resource.run(Coroutine::yield).named("Interrupter"); + scheduler.schedule(cmd); + scheduler.schedule(interrupter); + scheduler.run(); + + assertFalse(ran.get(), "onCancel ran when it shouldn't have!"); + } + + @Test + void doesNotRunOnCancelWhenInterruptingCommand() { + var ran = new AtomicBoolean(false); + + var resource = new RequireableResource("The Resource", scheduler); + var cmd = resource.run(Coroutine::yield).whenCanceled(() -> ran.set(true)).named("cmd"); + var interrupter = resource.run(Coroutine::yield).named("Interrupter"); + scheduler.schedule(cmd); + scheduler.run(); + scheduler.schedule(interrupter); + + assertFalse(ran.get(), "onCancel ran when it shouldn't have!"); + } + + @Test + void doesNotRunOnCancelWhenCompleting() { + var ran = new AtomicBoolean(false); + + var resource = new RequireableResource("The Resource", scheduler); + var cmd = resource.run(Coroutine::yield).whenCanceled(() -> ran.set(true)).named("cmd"); + scheduler.schedule(cmd); + scheduler.run(); + scheduler.run(); + + assertFalse(scheduler.isScheduledOrRunning(cmd)); + assertFalse(ran.get(), "onCancel ran when it shouldn't have!"); + } + + @Test + void runsOnCancelWhenCancelling() { + var ran = new AtomicBoolean(false); + + var resource = new RequireableResource("The Resource", scheduler); + var cmd = resource.run(Coroutine::yield).whenCanceled(() -> ran.set(true)).named("cmd"); + scheduler.schedule(cmd); + scheduler.cancel(cmd); + + assertTrue(ran.get(), "onCancel should have run!"); + } + + @Test + void runsOnCancelWhenCancellingParent() { + var ran = new AtomicBoolean(false); + + var resource = new RequireableResource("The Resource", scheduler); + var cmd = resource.run(Coroutine::yield).whenCanceled(() -> ran.set(true)).named("cmd"); + + var group = new Sequence("Seq", Collections.singletonList(cmd)); + scheduler.schedule(group); + scheduler.run(); + scheduler.cancel(group); + + assertTrue(ran.get(), "onCancel should have run!"); + } + record PriorityCommand(int priority, RequireableResource... subsystems) implements Command { @Override public void run(Coroutine coroutine) {