[command] Reorder Scheduler operations (#4261)

Fixes several cases where calling scheduler operations from a command callback could result in NPEs or other issues.

Co-authored-by: Tyler Veness <calcmogul@gmail.com>

wpilibNewCommands/src/test/native/cpp/frc2/command/SchedulingRecursionTest.cpp

@@ -0,0 +1,97 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+#include "CommandTestBase.h"
+#include "frc2/command/CommandHelper.h"
+#include "frc2/command/RunCommand.h"
+#include "gtest/gtest.h"
+
+using namespace frc2;
+
+class SchedulingRecursionTest : public CommandTestBaseWithParam<bool> {};
+
+class SelfCancellingCommand
+    : public CommandHelper<CommandBase, SelfCancellingCommand> {
+ public:
+  SelfCancellingCommand(CommandScheduler* scheduler, Subsystem* requirement)
+      : m_scheduler(scheduler) {
+    AddRequirements(requirement);
+  }
+
+  void Initialize() override { m_scheduler->Cancel(this); }
+
+ private:
+  CommandScheduler* m_scheduler;
+};
+
+/**
+ * Checks <a
+ * href="https://github.com/wpilibsuite/allwpilib/issues/4259">wpilibsuite/allwpilib#4259</a>.
+ */
+TEST_P(SchedulingRecursionTest, CancelFromInitialize) {
+  CommandScheduler scheduler = GetScheduler();
+  bool hasOtherRun = false;
+  TestSubsystem requirement;
+  SelfCancellingCommand selfCancels{&scheduler, &requirement};
+  RunCommand other =
+      RunCommand([&hasOtherRun] { hasOtherRun = true; }, {&requirement});
+
+  scheduler.Schedule(GetParam(), &selfCancels);
+  scheduler.Run();
+  // interruptibility of new arrival isn't checked
+  scheduler.Schedule(&other);
+
+  EXPECT_FALSE(scheduler.IsScheduled(&selfCancels));
+  EXPECT_TRUE(scheduler.IsScheduled(&other));
+  scheduler.Run();
+  EXPECT_TRUE(hasOtherRun);
+}
+
+TEST_P(SchedulingRecursionTest, DefaultCommand) {
+  CommandScheduler scheduler = GetScheduler();
+  bool hasOtherRun = false;
+  TestSubsystem requirement;
+  SelfCancellingCommand selfCancels{&scheduler, &requirement};
+  RunCommand other =
+      RunCommand([&hasOtherRun] { hasOtherRun = true; }, {&requirement});
+  scheduler.SetDefaultCommand(&requirement, std::move(other));
+
+  scheduler.Schedule(GetParam(), &selfCancels);
+  scheduler.Run();
+  scheduler.Run();
+  EXPECT_FALSE(scheduler.IsScheduled(&selfCancels));
+  EXPECT_TRUE(scheduler.IsScheduled(scheduler.GetDefaultCommand(&requirement)));
+  scheduler.Run();
+  EXPECT_TRUE(hasOtherRun);
+}
+
+class CancelEndCommand : public CommandHelper<CommandBase, CancelEndCommand> {
+ public:
+  CancelEndCommand(CommandScheduler* scheduler, int& counter)
+      : m_scheduler(scheduler), m_counter(counter) {}
+
+  void End(bool interrupted) override {
+    m_counter++;
+    m_scheduler->Cancel(this);
+  }
+
+ private:
+  CommandScheduler* m_scheduler;
+  int& m_counter;
+};
+
+TEST_F(SchedulingRecursionTest, CancelFromEnd) {
+  CommandScheduler scheduler = GetScheduler();
+  int counter = 0;
+  CancelEndCommand selfCancels{&scheduler, counter};
+
+  scheduler.Schedule(&selfCancels);
+
+  EXPECT_NO_THROW({ scheduler.Cancel(&selfCancels); });
+  EXPECT_EQ(1, counter);
+  EXPECT_FALSE(scheduler.IsScheduled(&selfCancels));
+}
+
+INSTANTIATE_TEST_SUITE_P(SchedulingRecursionTests, SchedulingRecursionTest,
+                         testing::Bool());