Modularize mechanism reservation tracking

src/main/java/com/team766/framework3/ContextImpl.java

@@ -263,11 +263,6 @@ private void transferControl(final ControlOwner thisOwner, final ControlOwner de
             }
             // Pass the baton.
             m_controlOwner = desiredOwner;
-            if (m_controlOwner == ControlOwner.SUBROUTINE) {
-                SchedulerMonitor.currentCommand = this;
-            } else {
-                SchedulerMonitor.currentCommand = null;
-            }
             m_threadSync.notifyAll();
             // Wait for the baton to be passed back.
             waitForControl(thisOwner);
@@ -298,7 +293,6 @@ private void threadFunction() {
         } finally {
             synchronized (m_threadSync) {
                 m_state = State.DONE;
-                SchedulerMonitor.currentCommand = null;
                 m_threadSync.notifyAll();
             }
         }
@@ -390,31 +384,44 @@ public boolean isFinished() {
 
     @Override
     public void end(boolean interrupted) {
-        synchronized (m_threadSync) {
-            if (m_state == State.DONE) {
-                return;
-            }
-            Logger.get(Category.FRAMEWORK)
-                    .logRaw(Severity.DEBUG, "Stopping requested of " + getContextName());
-            m_state = State.CANCELED;
-            if (m_controlOwner == ControlOwner.SUBROUTINE) {
-                throw new IllegalStateException("A Procedure should not cancel() its own Context");
-            }
-            transferControl(ControlOwner.MAIN_THREAD, ControlOwner.SUBROUTINE);
-            if (m_state != State.DONE) {
+        ReservingCommand.enterCommand(this);
+        try {
+            synchronized (m_threadSync) {
+                if (m_state == State.DONE) {
+                    return;
+                }
                 Logger.get(Category.FRAMEWORK)
-                        .logRaw(Severity.ERROR, getContextName() + " did not end when requested");
+                        .logRaw(Severity.DEBUG, "Stopping requested of " + getContextName());
+                m_state = State.CANCELED;
+                if (m_controlOwner == ControlOwner.SUBROUTINE) {
+                    throw new IllegalStateException(
+                            "A Procedure should not cancel() its own Context");
+                }
+                transferControl(ControlOwner.MAIN_THREAD, ControlOwner.SUBROUTINE);
+                if (m_state != State.DONE) {
+                    Logger.get(Category.FRAMEWORK)
+                            .logRaw(
+                                    Severity.ERROR,
+                                    getContextName() + " did not end when requested");
+                }
             }
+        } finally {
+            ReservingCommand.exitCommand(this);
         }
     }
 
     @Override
     public void execute() {
-        if (m_state == State.DONE) {
-            return;
-        }
-        if (m_blockingPredicate == null || m_blockingPredicate.getAsBoolean()) {
-            transferControl(ControlOwner.MAIN_THREAD, ControlOwner.SUBROUTINE);
+        ReservingCommand.enterCommand(this);
+        try {
+            if (m_state == State.DONE) {
+                return;
+            }
+            if (m_blockingPredicate == null || m_blockingPredicate.getAsBoolean()) {
+                transferControl(ControlOwner.MAIN_THREAD, ControlOwner.SUBROUTINE);
+            }
+        } finally {
+            ReservingCommand.exitCommand(this);
         }
     }
 

src/main/java/com/team766/framework3/InstantCommand.java

@@ -13,11 +13,11 @@ public InstantCommand(InstantProcedure procedure) {
 
     @Override
     public void execute() {
+        ReservingCommand.enterCommand(this);
         try {
-            SchedulerMonitor.currentCommand = this;
             procedure.run();
         } finally {
-            SchedulerMonitor.currentCommand = null;
+            ReservingCommand.exitCommand(this);
         }
     }
 

src/main/java/com/team766/framework3/Mechanism.java

@@ -1,11 +1,8 @@
 package com.team766.framework3;
 
 import com.team766.logging.Category;
-import com.team766.logging.Logger;
 import com.team766.logging.LoggerExceptionUtils;
-import com.team766.logging.Severity;
 import edu.wpi.first.wpilibj2.command.Command;
-import edu.wpi.first.wpilibj2.command.CommandScheduler;
 import edu.wpi.first.wpilibj2.command.SubsystemBase;
 
 public abstract class Mechanism<R extends Request<?>> extends SubsystemBase implements LoggingBase {
@@ -28,14 +25,16 @@ public void initialize() {
             try {
                 final var r = getIdleRequest();
                 if (r != null) {
-                    SchedulerMonitor.currentCommand = this;
-                    setRequest(r);
+                    ReservingCommand.enterCommand(this);
+                    try {
+                        setRequest(r);
+                    } finally {
+                        ReservingCommand.exitCommand(this);
+                    }
                 }
             } catch (Exception ex) {
                 ex.printStackTrace();
                 LoggerExceptionUtils.logException(ex);
-            } finally {
-                SchedulerMonitor.currentCommand = null;
             }
         }
 
@@ -88,22 +87,10 @@ protected R getIdleRequest() {
     }
 
     protected void checkContextReservation() {
-        var owningCommand = CommandScheduler.getInstance().requiring(this);
-        if ((owningCommand == null || SchedulerMonitor.currentCommand != owningCommand)
-                && m_runningPeriodic == null) {
-            final String commandName =
-                    SchedulerMonitor.currentCommand != null
-                            ? SchedulerMonitor.currentCommand.getName()
-                            : "non-Procedure code";
-            String message = getName() + " tried to be used by " + commandName;
-            if (owningCommand != null) {
-                message += " while reserved by " + owningCommand.getName();
-            } else {
-                message += " without reserving it";
-            }
-            Logger.get(Category.FRAMEWORK).logRaw(Severity.ERROR, message);
-            throw new IllegalStateException(message);
+        if (m_runningPeriodic != null) {
+            return;
         }
+        ReservingCommand.checkCurrentCommandHasReservation(this);
     }
 
     @Override

src/main/java/com/team766/framework3/ReservingCommand.java

@@ -0,0 +1,102 @@
+package com.team766.framework3;
+
+import com.team766.framework.StackTraceUtils;
+import com.team766.logging.Category;
+import com.team766.logging.Logger;
+import com.team766.logging.Severity;
+import edu.wpi.first.wpilibj2.command.Command;
+import edu.wpi.first.wpilibj2.command.CommandScheduler;
+import edu.wpi.first.wpilibj2.command.Subsystem;
+import edu.wpi.first.wpilibj2.command.WrapperCommand;
+import java.util.ArrayDeque;
+
+public final class ReservingCommand extends WrapperCommand {
+    /*
+     * The currently-executing Command (Context/InstantCommand).
+     */
+    private static ArrayDeque<Command> currentCommands = new ArrayDeque<>();
+
+    public static void checkCurrentCommandHasReservation(Subsystem subsystem) {
+        if (!currentCommandHasReservation(subsystem)) {
+            var exception =
+                    new IllegalStateException(
+                            subsystem.getName() + " tried to be used without reserving it");
+            Logger.get(Category.FRAMEWORK)
+                    .logRaw(
+                            Severity.ERROR,
+                            exception.getMessage()
+                                    + "\n"
+                                    + StackTraceUtils.getStackTrace(exception.getStackTrace()));
+            throw exception;
+        }
+    }
+
+    public static boolean currentCommandHasReservation(Subsystem subsystem) {
+        return currentCommands.size() > 0 && currentCommands.peekLast().hasRequirement(subsystem);
+    }
+
+    public static void enterCommand(Command command) {
+        if (CommandScheduler.getInstance().isScheduled(command)) {
+            if (!currentCommands.isEmpty()) {
+                throw new IllegalStateException();
+            }
+        } else {
+            checkProcedureReservationsSubset(currentCommands.getLast(), command);
+        }
+        currentCommands.addLast(command);
+    }
+
+    public static void exitCommand(Command command) {
+        if (currentCommands.removeLast() != command) {
+            throw new IllegalStateException();
+        }
+    }
+
+    private static void checkProcedureReservationsSubset(Command parent, Command child) {
+        final var thisReservations = parent.getRequirements();
+        for (var req : child.getRequirements()) {
+            if (!thisReservations.contains(req)) {
+                throw new IllegalArgumentException(
+                        parent.getName()
+                                + " tried to run "
+                                + child.getName()
+                                + " but is missing the reservation on "
+                                + req.getName());
+            }
+        }
+    }
+
+    public ReservingCommand(Command command) {
+        super(command);
+    }
+
+    @Override
+    public void initialize() {
+        try {
+            enterCommand(m_command);
+            super.initialize();
+        } finally {
+            exitCommand(m_command);
+        }
+    }
+
+    @Override
+    public void execute() {
+        try {
+            enterCommand(m_command);
+            super.execute();
+        } finally {
+            exitCommand(m_command);
+        }
+    }
+
+    @Override
+    public void end(boolean interrupted) {
+        try {
+            enterCommand(m_command);
+            super.end(interrupted);
+        } finally {
+            exitCommand(m_command);
+        }
+    }
+}

src/main/java/com/team766/framework3/SchedulerMonitor.java

@@ -4,21 +4,10 @@
 import com.team766.logging.Category;
 import com.team766.logging.Logger;
 import com.team766.logging.Severity;
-import edu.wpi.first.wpilibj2.command.Command;
 import edu.wpi.first.wpilibj2.command.CommandScheduler;
 import java.util.stream.Collectors;
 
 public class SchedulerMonitor {
-    /*
-     * The currently-executing Command (Context/InstantCommand).
-     *
-     * This is maintained for things like checking Mechanism ownership, but
-     * intentionally only has package-private visibility - code outside of the
-     * framework should pass around references to the current context object
-     * rather than cheating with this static accessor.
-     */
-    static Command currentCommand = null;
-
     private static Thread c_thread = null;
     private static int c_iterationCount = 0;
 
@@ -44,25 +33,21 @@ private static void monitor() {
             }
 
             if (c_iterationCount == lastIterationCount) {
-                final String commandName =
-                        currentCommand != null ? currentCommand.getName() : "non-Procedure code";
                 Logger.get(Category.FRAMEWORK)
                         .logRaw(
                                 Severity.ERROR,
-                                "The code has gotten stuck in "
-                                        + commandName
-                                        + ". You probably have an unintended infinite loop or need to add a call to context.yield()");
-                Logger.get(Category.FRAMEWORK)
-                        .logRaw(
-                                Severity.INFO,
-                                Thread.getAllStackTraces().entrySet().stream()
-                                        .map(
-                                                e ->
-                                                        e.getKey().getName()
-                                                                + ":\n"
-                                                                + StackTraceUtils.getStackTrace(
-                                                                        e.getValue()))
-                                        .collect(Collectors.joining("\n")));
+                                "The code has gotten stuck. You probably have an unintended infinite "
+                                        + "loop or need to add a call to context.yield() in a Procedure.\n"
+                                        + Thread.getAllStackTraces().entrySet().stream()
+                                                .map(
+                                                        e ->
+                                                                e.getKey().getName()
+                                                                        + ":\n"
+                                                                        + StackTraceUtils
+                                                                                .getStackTrace(
+                                                                                        e
+                                                                                                .getValue()))
+                                                .collect(Collectors.joining("\n")));
             }
 
             lastIterationCount = c_iterationCount;

src/test/java/com/team766/framework3/MechanismTest.java

@@ -97,16 +97,14 @@ protected void run(FakeRequest request, boolean isRequestNew) {}
         cmd.schedule();
         step();
         assertThat(thrownException.get())
-                .matches(
-                        "DummyMechanism tried to be used by .*MechanismTest\\$\\$Lambda.* without reserving it");
+                .matches("DummyMechanism tried to be used without reserving it");
 
         var cmd2 = new ContextImpl(new FakeProcedure(1, Set.of(mech)));
         cmd2.schedule();
         cmd.schedule();
         step();
         assertThat(thrownException.get())
-                .matches(
-                        "DummyMechanism tried to be used by .*MechanismTest\\$\\$Lambda.* while reserved by .*FakeProcedure.*");
+                .matches("DummyMechanism tried to be used without reserving it");
     }
 
     /// Test that checkContextReservation succeeds when called from within the Mechanism's own run()