[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/CommandTestBase.h

@@ -18,6 +18,7 @@
 #include "make_vector.h"
 
 namespace frc2 {
+
 class CommandTestBase : public ::testing::Test {
  public:
   CommandTestBase();
@@ -93,4 +94,91 @@ class CommandTestBase : public ::testing::Test {
 
   void SetDSEnabled(bool enabled);
 };
+
+template <typename T>
+class CommandTestBaseWithParam : public ::testing::TestWithParam<T> {
+ public:
+  CommandTestBaseWithParam() {
+    auto& scheduler = CommandScheduler::GetInstance();
+    scheduler.CancelAll();
+    scheduler.Enable();
+    scheduler.GetActiveButtonLoop()->Clear();
+  }
+
+  class TestSubsystem : public SubsystemBase {};
+
+ protected:
+  class MockCommand : public Command {
+   public:
+    MOCK_CONST_METHOD0(GetRequirements, wpi::SmallSet<Subsystem*, 4>());
+    MOCK_METHOD0(IsFinished, bool());
+    MOCK_CONST_METHOD0(RunsWhenDisabled, bool());
+    MOCK_METHOD0(Initialize, void());
+    MOCK_METHOD0(Execute, void());
+    MOCK_METHOD1(End, void(bool interrupted));
+
+    MockCommand() {
+      m_requirements = {};
+      EXPECT_CALL(*this, GetRequirements())
+          .WillRepeatedly(::testing::Return(m_requirements));
+      EXPECT_CALL(*this, IsFinished()).WillRepeatedly(::testing::Return(false));
+      EXPECT_CALL(*this, RunsWhenDisabled())
+          .WillRepeatedly(::testing::Return(true));
+    }
+
+    MockCommand(std::initializer_list<Subsystem*> requirements,
+                bool finished = false, bool runWhenDisabled = true) {
+      m_requirements.insert(requirements.begin(), requirements.end());
+      EXPECT_CALL(*this, GetRequirements())
+          .WillRepeatedly(::testing::Return(m_requirements));
+      EXPECT_CALL(*this, IsFinished())
+          .WillRepeatedly(::testing::Return(finished));
+      EXPECT_CALL(*this, RunsWhenDisabled())
+          .WillRepeatedly(::testing::Return(runWhenDisabled));
+    }
+
+    MockCommand(MockCommand&& other) {
+      EXPECT_CALL(*this, IsFinished())
+          .WillRepeatedly(::testing::Return(other.IsFinished()));
+      EXPECT_CALL(*this, RunsWhenDisabled())
+          .WillRepeatedly(::testing::Return(other.RunsWhenDisabled()));
+      std::swap(m_requirements, other.m_requirements);
+      EXPECT_CALL(*this, GetRequirements())
+          .WillRepeatedly(::testing::Return(m_requirements));
+    }
+
+    MockCommand(const MockCommand& other) : Command{other} {}
+
+    void SetFinished(bool finished) {
+      EXPECT_CALL(*this, IsFinished())
+          .WillRepeatedly(::testing::Return(finished));
+    }
+
+    ~MockCommand() {  // NOLINT
+      auto& scheduler = CommandScheduler::GetInstance();
+      scheduler.Cancel(this);
+    }
+
+   protected:
+    std::unique_ptr<Command> TransferOwnership() && {  // NOLINT
+      return std::make_unique<MockCommand>(std::move(*this));
+    }
+
+   private:
+    wpi::SmallSet<Subsystem*, 4> m_requirements;
+  };
+
+  CommandScheduler GetScheduler() { return CommandScheduler(); }
+
+  void SetUp() override { frc::sim::DriverStationSim::SetEnabled(true); }
+
+  void TearDown() override {
+    CommandScheduler::GetInstance().GetActiveButtonLoop()->Clear();
+  }
+
+  void SetDSEnabled(bool enabled) {
+    frc::sim::DriverStationSim::SetEnabled(enabled);
+  }
+};
+
 }  // namespace frc2

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());