001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package org.wpilib.commands3;
006
007import static edu.wpi.first.units.Units.Microseconds;
008import static edu.wpi.first.units.Units.Milliseconds;
009
010import edu.wpi.first.util.ErrorMessages;
011import edu.wpi.first.util.protobuf.ProtobufSerializable;
012import edu.wpi.first.wpilibj.RobotController;
013import edu.wpi.first.wpilibj.TimedRobot;
014import edu.wpi.first.wpilibj.event.EventLoop;
015import java.util.ArrayList;
016import java.util.Collection;
017import java.util.Collections;
018import java.util.HashSet;
019import java.util.LinkedHashMap;
020import java.util.LinkedHashSet;
021import java.util.List;
022import java.util.Map;
023import java.util.SequencedMap;
024import java.util.SequencedSet;
025import java.util.Set;
026import java.util.Stack;
027import java.util.function.Consumer;
028import java.util.stream.Collectors;
029import org.wpilib.annotation.NoDiscard;
030import org.wpilib.commands3.button.CommandGenericHID;
031import org.wpilib.commands3.proto.SchedulerProto;
032
033/**
034 * Manages the lifecycles of {@link Coroutine}-based {@link Command Commands}. Commands may be
035 * scheduled directly using {@link #schedule(Command)}, or be bound to {@link Trigger Triggers} to
036 * automatically handle scheduling and cancellation based on internal or external events. User code
037 * is responsible for calling {@link #run()} periodically to update trigger conditions and execute
038 * scheduled commands. Most often, this is done by overriding {@link TimedRobot#robotPeriodic()} to
039 * include a call to {@code Scheduler.getDefault().run()}:
040 *
041 * <pre>{@code
042 * public class Robot extends TimedRobot {
043 *   @Override
044 *   public void robotPeriodic() {
045 *     // Update the scheduler on every loop
046 *     Scheduler.getDefault().run();
047 *   }
048 * }
049 * }</pre>
050 *
051 * <h2>Danger</h2>
052 *
053 * <p>The scheduler <i>must</i> be used in a single-threaded program. Commands must be scheduled and
054 * canceled by the same thread that runs the scheduler, and cannot be run in a virtual thread.
055 *
056 * <p><strong>Using the commands framework in a multithreaded environment can cause crashes in the
057 * Java virtual machine at any time, including on an official field during a match.</strong> The
058 * Java JIT compilers make assumptions that rely on coroutines being used on a single thread.
059 * Breaking those assumptions can cause incorrect JIT code to be generated with undefined behavior,
060 * potentially causing control issues or crashes deep in JIT-generated native code.
061 *
062 * <p><strong>Normal concurrency constructs like locks, atomic references, and synchronized blocks
063 * or methods cannot save you.</strong>
064 *
065 * <h2>Lifecycle</h2>
066 *
067 * <p>The {@link #run()} method runs five steps:
068 *
069 * <ol>
070 *   <li>Call {@link #sideload(Consumer) periodic sideload functions}
071 *   <li>Poll all registered triggers to queue and cancel commands
072 *   <li>Queue default commands for any mechanisms without a running command. The queued commands
073 *       can be superseded by any manual scheduling or commands scheduled by triggers in the next
074 *       run.
075 *   <li>Start all queued commands. This happens after all triggers are checked in case multiple
076 *       commands with conflicting requirements are queued in the same update; the last command to
077 *       be queued takes precedence over the rest.
078 *   <li>Loop over all running commands, mounting and calling each in turn until they either exit or
079 *       call {@link Coroutine#yield()}. Commands run in the order in which they were scheduled.
080 * </ol>
081 *
082 * <h2>Telemetry</h2>
083 *
084 * <p>There are two mechanisms for telemetry for a scheduler. A protobuf serializer can be used to
085 * take a snapshot of a scheduler instance, and report what commands are queued (scheduled but have
086 * not yet started to run), commands that are running (along with timing data for each command), and
087 * total time spent in the most recent {@link #run()} call. However, it cannot detect one-shot
088 * commands that are scheduled, run, and complete all in a single {@code run()} invocation -
089 * effectively, commands that never call {@link Coroutine#yield()} are invisible.
090 *
091 * <p>A second telemetry mechanism is provided by {@link #addEventListener(Consumer)}. The scheduler
092 * will issue events to all registered listeners when certain events occur (see {@link
093 * SchedulerEvent} for all event types). These events are emitted immediately and can be used to
094 * detect lifecycle events for all commands, including one-shots that would be invisible to the
095 * protobuf serializer. However, it is up to the user to log those events themselves.
096 */
097public final class Scheduler implements ProtobufSerializable {
098  private final Map<Mechanism, Command> m_defaultCommands = new LinkedHashMap<>();
099
100  /** The set of commands scheduled since the start of the previous run. */
101  private final SequencedSet<CommandState> m_queuedToRun = new LinkedHashSet<>();
102
103  /**
104   * The states of all running commands (does not include on deck commands). We preserve insertion
105   * order to guarantee that child commands run after their parents.
106   */
107  private final SequencedMap<Command, CommandState> m_runningCommands = new LinkedHashMap<>();
108
109  /**
110   * The stack of currently executing commands. Child commands are pushed onto the stack and popped
111   * when they complete. Use {@link #currentState()} and {@link #currentCommand()} to get the
112   * currently executing command or its state.
113   */
114  private final Stack<CommandState> m_currentCommandAncestry = new Stack<>();
115
116  /** The periodic callbacks to run, outside of the command structure. */
117  private final List<Coroutine> m_periodicCallbacks = new ArrayList<>();
118
119  /** Event loop for trigger bindings. */
120  private final EventLoop m_eventLoop = new EventLoop();
121
122  /** The scope for continuations to yield to. */
123  private final ContinuationScope m_scope = new ContinuationScope("coroutine commands");
124
125  // Telemetry
126  /** Protobuf serializer for a scheduler. */
127  public static final SchedulerProto proto = new SchedulerProto();
128
129  private double m_lastRunTimeMs = -1;
130
131  private final Set<Consumer<? super SchedulerEvent>> m_eventListeners = new LinkedHashSet<>();
132
133  /** The default scheduler instance. */
134  private static final Scheduler s_defaultScheduler = new Scheduler();
135
136  /**
137   * Gets the default scheduler instance for use in a robot program. Unless otherwise specified,
138   * triggers and mechanisms will be registered with the default scheduler and require the default
139   * scheduler to run. However, triggers and mechanisms can be constructed to be registered with a
140   * specific scheduler instance, which may be useful for isolation for unit tests.
141   *
142   * @return the default scheduler instance.
143   */
144  @NoDiscard
145  public static Scheduler getDefault() {
146    return s_defaultScheduler;
147  }
148
149  /**
150   * Creates a new scheduler object. Note that most built-in constructs like {@link Trigger} and
151   * {@link CommandGenericHID} will use the {@link #getDefault() default scheduler instance} unless
152   * they were explicitly constructed with a different scheduler instance. Teams should use the
153   * default instance for convenience; however, new scheduler instances can be useful for unit
154   * tests.
155   *
156   * @return a new scheduler instance that is independent of the default scheduler instance.
157   */
158  @NoDiscard
159  public static Scheduler createIndependentScheduler() {
160    return new Scheduler();
161  }
162
163  /** Private constructor. Use static factory methods or the default scheduler instance. */
164  private Scheduler() {}
165
166  /**
167   * Sets the default command for a mechanism. The command must require that mechanism, and cannot
168   * require any other mechanisms.
169   *
170   * @param mechanism the mechanism for which to set the default command
171   * @param defaultCommand the default command to execute on the mechanism
172   * @throws IllegalArgumentException if the command does not meet the requirements for being a
173   *     default command
174   */
175  public void setDefaultCommand(Mechanism mechanism, Command defaultCommand) {
176    if (!defaultCommand.requires(mechanism)) {
177      throw new IllegalArgumentException(
178          "A mechanism's default command must require that mechanism");
179    }
180
181    if (defaultCommand.requirements().size() > 1) {
182      throw new IllegalArgumentException(
183          "A mechanism's default command cannot require other mechanisms");
184    }
185
186    m_defaultCommands.put(mechanism, defaultCommand);
187  }
188
189  /**
190   * Gets the default command set for a mechanism.
191   *
192   * @param mechanism The mechanism
193   * @return The default command, or null if no default command was ever set
194   */
195  public Command getDefaultCommandFor(Mechanism mechanism) {
196    return m_defaultCommands.get(mechanism);
197  }
198
199  /**
200   * Adds a callback to run as part of the scheduler. The callback should not manipulate or control
201   * any mechanisms, but can be used to log information, update data (such as simulations or LED
202   * data buffers), or perform some other helpful task. The callback is responsible for managing its
203   * own control flow and end conditions. If you want to run a single task periodically for the
204   * entire lifespan of the scheduler, use {@link #addPeriodic(Runnable)}.
205   *
206   * <p><strong>Note:</strong> Like commands, any loops in the callback must appropriately yield
207   * control back to the scheduler with {@link Coroutine#yield} or risk stalling your program in an
208   * unrecoverable infinite loop!
209   *
210   * @param callback the callback to sideload
211   */
212  public void sideload(Consumer<Coroutine> callback) {
213    var coroutine = new Coroutine(this, m_scope, callback);
214    m_periodicCallbacks.add(coroutine);
215  }
216
217  /**
218   * Adds a task to run repeatedly for as long as the scheduler runs. This internally handles the
219   * looping and control yielding necessary for proper function. The callback will run at the same
220   * periodic frequency as the scheduler.
221   *
222   * <p>For example:
223   *
224   * <pre>{@code
225   * scheduler.addPeriodic(() -> leds.setData(ledDataBuffer));
226   * scheduler.addPeriodic(() -> {
227   *   SmartDashboard.putNumber("X", getX());
228   *   SmartDashboard.putNumber("Y", getY());
229   * });
230   * }</pre>
231   *
232   * @param callback the periodic function to run
233   */
234  public void addPeriodic(Runnable callback) {
235    sideload(
236        coroutine -> {
237          while (true) {
238            callback.run();
239            coroutine.yield();
240          }
241        });
242  }
243
244  /** Represents possible results of a command scheduling attempt. */
245  public enum ScheduleResult {
246    /** The command was successfully scheduled and added to the queue. */
247    SUCCESS,
248    /** The command is already scheduled or running. */
249    ALREADY_RUNNING,
250    /** The command is a lower priority and conflicts with a command that's already running. */
251    LOWER_PRIORITY_THAN_RUNNING_COMMAND,
252  }
253
254  /**
255   * Schedules a command to run. If one command schedules another (a "parent" and "child"), the
256   * child command will be canceled when the parent command completes. It is not possible to fork a
257   * child command and have it live longer than its parent.
258   *
259   * <p>Does nothing if the command is already scheduled or running, or requires at least one
260   * mechanism already used by a higher priority command.
261   *
262   * @param command the command to schedule
263   * @return the result of the scheduling attempt. See {@link ScheduleResult} for details.
264   * @throws IllegalArgumentException if scheduled by a command composition that has already
265   *     scheduled another command that shares at least one required mechanism
266   */
267  // Implementation detail: a child command will immediately start running when scheduled by a
268  // parent command, skipping the queue entirely. This avoids dead loop cycles where a parent
269  // schedules a child, appending it to the queue, then waits for the next cycle to pick it up and
270  // start it. With deeply nested commands, dead loops could quickly to add up and cause the
271  // innermost commands that actually _do_ something to start running hundreds of milliseconds after
272  // their root ancestor was scheduled.
273  public ScheduleResult schedule(Command command) {
274    // Note: we use a throwable here instead of Thread.currentThread().getStackTrace() for easier
275    //       stack frame filtering and modification.
276    var binding =
277        new Binding(
278            BindingScope.global(), BindingType.IMMEDIATE, command, new Throwable().getStackTrace());
279
280    return schedule(binding);
281  }
282
283  // Package-private for use by Trigger
284  ScheduleResult schedule(Binding binding) {
285    var command = binding.command();
286
287    if (isScheduledOrRunning(command)) {
288      return ScheduleResult.ALREADY_RUNNING;
289    }
290
291    if (lowerPriorityThanConflictingCommands(command)) {
292      return ScheduleResult.LOWER_PRIORITY_THAN_RUNNING_COMMAND;
293    }
294
295    for (var scheduledState : m_queuedToRun) {
296      if (!command.conflictsWith(scheduledState.command())) {
297        // No shared requirements, skip
298        continue;
299      }
300      if (command.isLowerPriorityThan(scheduledState.command())) {
301        // Lower priority than an already-scheduled (but not yet running) command that requires at
302        // one of the same mechanism. Ignore it.
303        return ScheduleResult.LOWER_PRIORITY_THAN_RUNNING_COMMAND;
304      }
305    }
306
307    // Evict conflicting on-deck commands
308    // We check above if the input command is lower priority than any of these,
309    // so at this point we're guaranteed to be >= priority than anything already on deck
310    evictConflictingOnDeckCommands(command);
311
312    // If the binding is scoped to a particular command, that command is the parent. If we're in the
313    // middle of a run cycle and running commands, the parent is whatever command is currently
314    // running. Otherwise, there is no parent command.
315    var parentCommand =
316        binding.scope() instanceof BindingScope.ForCommand scope
317            ? scope.command()
318            : currentCommand();
319    var state = new CommandState(command, parentCommand, buildCoroutine(command), binding);
320
321    emitScheduledEvent(command);
322
323    if (currentState() != null) {
324      // Scheduling a child command while running. Start it immediately instead of waiting a full
325      // cycle. This prevents issues with deeply nested command groups taking many scheduler cycles
326      // to start running the commands that actually _do_ things
327      evictConflictingRunningCommands(state);
328      m_runningCommands.put(command, state);
329      runCommand(state);
330    } else {
331      // Scheduling outside a command, add it to the pending set. If it's not overridden by another
332      // conflicting command being scheduled in the same scheduler loop, it'll be promoted and
333      // start to run when #runCommands() is called
334      m_queuedToRun.add(state);
335    }
336
337    return ScheduleResult.SUCCESS;
338  }
339
340  /**
341   * Checks if a command conflicts with and is a lower priority than any running command. Used when
342   * determining if the command can be scheduled.
343   */
344  private boolean lowerPriorityThanConflictingCommands(Command command) {
345    Set<CommandState> ancestors = new HashSet<>();
346    for (var state = currentState(); state != null; state = m_runningCommands.get(state.parent())) {
347      ancestors.add(state);
348    }
349
350    // Check for conflicts with the commands that are already running
351    for (var state : m_runningCommands.values()) {
352      if (ancestors.contains(state)) {
353        // Can't conflict with an ancestor command
354        continue;
355      }
356
357      var c = state.command();
358      if (c.conflictsWith(command) && command.isLowerPriorityThan(c)) {
359        return true;
360      }
361    }
362
363    return false;
364  }
365
366  private void evictConflictingOnDeckCommands(Command command) {
367    for (var iterator = m_queuedToRun.iterator(); iterator.hasNext(); ) {
368      var scheduledState = iterator.next();
369      var scheduledCommand = scheduledState.command();
370      if (scheduledCommand.conflictsWith(command)) {
371        // Remove the lower priority conflicting command from the on deck commands.
372        // We don't need to call removeOrphanedChildren here because it hasn't started yet,
373        // meaning it hasn't had a chance to schedule any children
374        iterator.remove();
375        emitInterruptedEvent(scheduledCommand, command);
376        emitCanceledEvent(scheduledCommand);
377      }
378    }
379  }
380
381  /**
382   * Cancels all running commands with which an incoming state conflicts. Ancestor commands of the
383   * incoming state will not be canceled.
384   */
385  @SuppressWarnings("PMD.CompareObjectsWithEquals")
386  private void evictConflictingRunningCommands(CommandState incomingState) {
387    // The set of root states with which the incoming state conflicts but does not inherit from
388    Set<CommandState> conflictingRootStates =
389        m_runningCommands.values().stream()
390            .filter(state -> incomingState.command().conflictsWith(state.command()))
391            .filter(state -> !isAncestorOf(state.command(), incomingState))
392            .map(
393                state -> {
394                  // Find the highest level ancestor of the conflicting command from which the
395                  // incoming state does _not_ inherit. If they're totally unrelated, this will
396                  // get the very root ancestor; otherwise, it'll return a direct sibling of the
397                  // incoming command
398                  CommandState root = state;
399                  while (root.parent() != null && root.parent() != incomingState.parent()) {
400                    root = m_runningCommands.get(root.parent());
401                  }
402                  return root;
403                })
404            .collect(Collectors.toSet());
405
406    // Cancel the root commands
407    for (var conflictingState : conflictingRootStates) {
408      emitInterruptedEvent(conflictingState.command(), incomingState.command());
409      cancel(conflictingState.command());
410    }
411  }
412
413  /**
414   * Checks if a particular command is an ancestor of another.
415   *
416   * @param ancestor the potential ancestor for which to search
417   * @param state the state to check
418   * @return true if {@code ancestor} is the direct parent or indirect ancestor, false if not
419   */
420  @SuppressWarnings({"PMD.CompareObjectsWithEquals", "PMD.SimplifyBooleanReturns"})
421  private boolean isAncestorOf(Command ancestor, CommandState state) {
422    if (state.parent() == null) {
423      // No parent, cannot inherit
424      return false;
425    }
426    if (!m_runningCommands.containsKey(ancestor)) {
427      // The given ancestor isn't running
428      return false;
429    }
430    if (state.parent() == ancestor) {
431      // Direct child
432      return true;
433    }
434    // Check if the command's parent inherits from the given ancestor
435    return m_runningCommands.values().stream()
436        .filter(s -> state.parent() == s.command())
437        .anyMatch(s -> isAncestorOf(ancestor, s));
438  }
439
440  /**
441   * Cancels a command and any other command scheduled by it. This occurs immediately and does not
442   * need to wait for a call to {@link #run()}. Any command that it scheduled will also be canceled
443   * to ensure commands within compositions do not continue to run.
444   *
445   * <p>This has no effect if the given command is not currently scheduled or running.
446   *
447   * @param command the command to cancel
448   */
449  @SuppressWarnings("PMD.CompareObjectsWithEquals")
450  public void cancel(Command command) {
451    if (command == currentCommand()) {
452      throw new IllegalArgumentException(
453          "Command `" + command.name() + "` is mounted and cannot be canceled");
454    }
455
456    boolean running = isRunning(command);
457
458    // Evict the command. The next call to run() will schedule the default command for all its
459    // required mechanisms, unless another command requiring those mechanisms is scheduled between
460    // calling cancel() and calling run()
461    m_runningCommands.remove(command);
462    m_queuedToRun.removeIf(state -> state.command() == command);
463
464    if (running) {
465      // Only run the hook if the command was running. If it was on deck or not
466      // even in the scheduler at the time, then there's nothing to do
467      command.onCancel();
468      emitCanceledEvent(command);
469    }
470
471    // Clean up any orphaned child commands; their lifespan must not exceed the parent's
472    removeOrphanedChildren(command);
473  }
474
475  /**
476   * Updates the command scheduler. This will run operations in the following order:
477   *
478   * <ol>
479   *   <li>Run sideloaded functions from {@link #sideload(Consumer)} and {@link
480   *       #addPeriodic(Runnable)}
481   *   <li>Update trigger bindings to queue and cancel bound commands
482   *   <li>Queue default commands for mechanisms that do not have a queued or running command
483   *   <li>Promote queued commands to the running set
484   *   <li>For every command in the running set, mount and run that command until it calls {@link
485   *       Coroutine#yield()} or exits
486   * </ol>
487   *
488   * <p>This method is intended to be called in a periodic loop like {@link
489   * TimedRobot#robotPeriodic()}
490   */
491  public void run() {
492    final long startMicros = RobotController.getTime();
493
494    // Sideloads may change some state that affects triggers. Run them first.
495    runPeriodicSideloads();
496
497    // Poll triggers next to schedule and cancel commands
498    m_eventLoop.poll();
499
500    // Schedule default commands for any mechanisms that don't have a running command and didn't
501    // have a new command scheduled by a sideload function or a trigger
502    scheduleDefaultCommands();
503
504    // Move all scheduled commands to the running set
505    promoteScheduledCommands();
506
507    // Run every command in order until they call Coroutine.yield() or exit
508    runCommands();
509
510    final long endMicros = RobotController.getTime();
511    m_lastRunTimeMs = Milliseconds.convertFrom(endMicros - startMicros, Microseconds);
512  }
513
514  private void promoteScheduledCommands() {
515    // Clear any commands that conflict with the scheduled set
516    for (var queuedState : m_queuedToRun) {
517      evictConflictingRunningCommands(queuedState);
518    }
519
520    // Move any scheduled commands to the running set
521    for (var queuedState : m_queuedToRun) {
522      m_runningCommands.put(queuedState.command(), queuedState);
523    }
524
525    // Clear the set of on-deck commands,
526    // since we just put them all into the set of running commands
527    m_queuedToRun.clear();
528  }
529
530  private void runPeriodicSideloads() {
531    // Update periodic callbacks
532    for (Coroutine coroutine : m_periodicCallbacks) {
533      coroutine.mount();
534      try {
535        coroutine.runToYieldPoint();
536      } finally {
537        Continuation.mountContinuation(null);
538      }
539    }
540
541    // And remove any periodic callbacks that have completed
542    m_periodicCallbacks.removeIf(Coroutine::isDone);
543  }
544
545  private void runCommands() {
546    // Tick every command that hasn't been completed yet
547    // Run in reverse so parent commands can resume in the same loop cycle an awaited child command
548    // completes. Otherwise, parents could only resume on the next loop cycle, introducing a delay
549    // at every layer of nesting.
550    for (var state : List.copyOf(m_runningCommands.values()).reversed()) {
551      runCommand(state);
552    }
553  }
554
555  /**
556   * Mounts and runs a command until it yields or exits.
557   *
558   * @param state The command state to run
559   */
560  @SuppressWarnings("PMD.AvoidCatchingGenericException")
561  private void runCommand(CommandState state) {
562    final var command = state.command();
563    final var coroutine = state.coroutine();
564
565    if (!m_runningCommands.containsKey(command)) {
566      // Probably canceled by an owning composition, do not run
567      return;
568    }
569
570    var previousState = currentState();
571
572    m_currentCommandAncestry.push(state);
573    long startMicros = RobotController.getTime();
574    emitMountedEvent(command);
575    coroutine.mount();
576    try {
577      coroutine.runToYieldPoint();
578    } catch (RuntimeException e) {
579      // Command encountered an uncaught exception.
580      handleCommandException(state, e);
581    } finally {
582      long endMicros = RobotController.getTime();
583      double elapsedMs = Milliseconds.convertFrom(endMicros - startMicros, Microseconds);
584      state.setLastRuntimeMs(elapsedMs);
585
586      if (state.equals(currentState())) {
587        // Remove the command we just ran from the top of the stack
588        m_currentCommandAncestry.pop();
589      }
590
591      if (previousState != null) {
592        // Remount the parent command, if there is one
593        previousState.coroutine().mount();
594      } else {
595        Continuation.mountContinuation(null);
596      }
597    }
598
599    if (coroutine.isDone()) {
600      // Immediately check if the command has completed and remove any children commands.
601      // This prevents child commands from being executed one extra time in the run() loop
602      emitCompletedEvent(command);
603      m_runningCommands.remove(command);
604      removeOrphanedChildren(command);
605    } else {
606      // Yielded
607      emitYieldedEvent(command);
608    }
609  }
610
611  /**
612   * Handles uncaught runtime exceptions from a mounted and running command. The command's ancestor
613   * and child commands will all be canceled and the exception's backtrace will be modified to
614   * include the stack frames of the schedule call site.
615   *
616   * @param state The state of the command that encountered the exception.
617   * @param e The exception that was thrown.
618   * @throws RuntimeException rethrows the exception, with a modified backtrace pointing to the
619   *     schedule site
620   */
621  @SuppressWarnings("PMD.CompareObjectsWithEquals")
622  private void handleCommandException(CommandState state, RuntimeException e) {
623    var command = state.command();
624
625    // Fetch the root command
626    // (needs to be done before removing the failed command from the running set)
627    Command root = command;
628    while (getParentOf(root) != null) {
629      root = getParentOf(root);
630    }
631
632    // Remove it from the running set.
633    m_runningCommands.remove(command);
634
635    // Intercept the exception, inject stack frames from the schedule site, and rethrow it
636    var binding = state.binding();
637    e.setStackTrace(CommandTraceHelper.modifyTrace(e.getStackTrace(), binding.frames()));
638    emitCompletedWithErrorEvent(command, e);
639
640    // Clean up child commands after emitting the event so child Canceled events are emitted
641    // after the parent's CompletedWithError
642    removeOrphanedChildren(command);
643
644    // Bubble up to the root and cancel all commands between the root and this one
645    // Note: Because we remove the command from the running set above, we still need to
646    //       clean up all the failed command's children
647    if (root != null && root != command) {
648      cancel(root);
649    }
650
651    // Then rethrow the exception
652    throw e;
653  }
654
655  /**
656   * Gets the currently executing command state, or null if no command is currently executing.
657   *
658   * @return the currently executing command state
659   */
660  private CommandState currentState() {
661    if (m_currentCommandAncestry.isEmpty()) {
662      // Avoid EmptyStackException
663      return null;
664    }
665
666    return m_currentCommandAncestry.peek();
667  }
668
669  /**
670   * Gets the currently executing command, or null if no command is currently executing.
671   *
672   * @return the currently executing command
673   */
674  public Command currentCommand() {
675    var state = currentState();
676    if (state == null) {
677      return null;
678    }
679
680    return state.command();
681  }
682
683  private void scheduleDefaultCommands() {
684    // Schedule the default commands for every mechanism that doesn't currently have a running or
685    // scheduled command.
686    m_defaultCommands.forEach(
687        (mechanism, defaultCommand) -> {
688          if (m_runningCommands.keySet().stream().noneMatch(c -> c.requires(mechanism))
689              && m_queuedToRun.stream().noneMatch(c -> c.command().requires(mechanism))
690              && defaultCommand != null) {
691            // Nothing currently running or scheduled
692            // Schedule the mechanism's default command, if it has one
693            schedule(defaultCommand);
694          }
695        });
696  }
697
698  /**
699   * Removes all commands descended from a parent command. This is used to ensure that any command
700   * scheduled within a composition or group cannot live longer than any ancestor.
701   *
702   * @param parent the root command whose descendants to remove from the scheduler
703   */
704  @SuppressWarnings("PMD.CompareObjectsWithEquals")
705  private void removeOrphanedChildren(Command parent) {
706    m_runningCommands.entrySet().stream()
707        .filter(e -> e.getValue().parent() == parent)
708        .toList() // copy to an intermediate list to avoid concurrent modification
709        .forEach(e -> cancel(e.getKey()));
710  }
711
712  /**
713   * Builds a coroutine object that the command will be bound to. The coroutine will be scoped to
714   * this scheduler object and cannot be used by another scheduler instance.
715   *
716   * @param command the command for which to build a coroutine
717   * @return the binding coroutine
718   */
719  private Coroutine buildCoroutine(Command command) {
720    return new Coroutine(this, m_scope, command::run);
721  }
722
723  /**
724   * Checks if a command is currently running.
725   *
726   * @param command the command to check
727   * @return true if the command is running, false if not
728   */
729  public boolean isRunning(Command command) {
730    return m_runningCommands.containsKey(command);
731  }
732
733  /**
734   * Checks if a command is currently scheduled to run, but is not yet running.
735   *
736   * @param command the command to check
737   * @return true if the command is scheduled to run, false if not
738   */
739  @SuppressWarnings("PMD.CompareObjectsWithEquals")
740  public boolean isScheduled(Command command) {
741    return m_queuedToRun.stream().anyMatch(state -> state.command() == command);
742  }
743
744  /**
745   * Checks if a command is currently scheduled to run, or is already running.
746   *
747   * @param command the command to check
748   * @return true if the command is scheduled to run or is already running, false if not
749   */
750  public boolean isScheduledOrRunning(Command command) {
751    return isScheduled(command) || isRunning(command);
752  }
753
754  /**
755   * Gets the set of all currently running commands. Commands are returned in the order in which
756   * they were scheduled. The returned set is read-only.
757   *
758   * @return the currently running commands
759   */
760  public Collection<Command> getRunningCommands() {
761    return Collections.unmodifiableSet(m_runningCommands.keySet());
762  }
763
764  /**
765   * Gets all the currently running commands that require a particular mechanism. Commands are
766   * returned in the order in which they were scheduled. The returned list is read-only.
767   *
768   * @param mechanism the mechanism to get the commands for
769   * @return the currently running commands that require the mechanism.
770   */
771  public List<Command> getRunningCommandsFor(Mechanism mechanism) {
772    return m_runningCommands.keySet().stream()
773        .filter(command -> command.requires(mechanism))
774        .toList();
775  }
776
777  /**
778   * Cancels all currently running and scheduled commands. All default commands will be scheduled on
779   * the next call to {@link #run()}, unless a higher priority command is scheduled or triggered
780   * after {@code cancelAll()} is used.
781   */
782  public void cancelAll() {
783    for (var onDeckIter = m_queuedToRun.iterator(); onDeckIter.hasNext(); ) {
784      var state = onDeckIter.next();
785      onDeckIter.remove();
786      emitCanceledEvent(state.command());
787    }
788
789    for (var liveIter = m_runningCommands.entrySet().iterator(); liveIter.hasNext(); ) {
790      var entry = liveIter.next();
791      liveIter.remove();
792      Command canceledCommand = entry.getKey();
793      canceledCommand.onCancel();
794      emitCanceledEvent(canceledCommand);
795    }
796  }
797
798  /**
799   * An event loop used by the scheduler to poll triggers that schedule or cancel commands. This
800   * event loop is always polled on every call to {@link #run()}. Custom event loops need to be
801   * bound to this one for synchronicity with the scheduler; however, they can always be polled
802   * manually before or after the call to {@link #run()}.
803   *
804   * @return the default event loop.
805   */
806  @NoDiscard
807  public EventLoop getDefaultEventLoop() {
808    return m_eventLoop;
809  }
810
811  /**
812   * For internal use.
813   *
814   * @return The commands that have been scheduled but not yet started.
815   */
816  @NoDiscard
817  public Collection<Command> getQueuedCommands() {
818    return m_queuedToRun.stream().map(CommandState::command).toList();
819  }
820
821  /**
822   * For internal use.
823   *
824   * @param command The command to check
825   * @return The command that forked the provided command. Null if the command is not a child of
826   *     another command.
827   */
828  public Command getParentOf(Command command) {
829    var state = m_runningCommands.get(command);
830    if (state == null) {
831      return null;
832    }
833    return state.parent();
834  }
835
836  /**
837   * Gets how long a command took to run in the previous cycle. If the command is not currently
838   * running, returns -1.
839   *
840   * @param command The command to check
841   * @return How long, in milliseconds, the command last took to execute.
842   */
843  public double lastCommandRuntimeMs(Command command) {
844    if (m_runningCommands.containsKey(command)) {
845      return m_runningCommands.get(command).lastRuntimeMs();
846    } else {
847      return -1;
848    }
849  }
850
851  /**
852   * Gets how long a command has taken to run, in aggregate, since it was most recently scheduled.
853   * If the command is not currently running, returns -1.
854   *
855   * @param command The command to check
856   * @return How long, in milliseconds, the command has taken to execute in total
857   */
858  public double totalRuntimeMs(Command command) {
859    if (m_runningCommands.containsKey(command)) {
860      return m_runningCommands.get(command).totalRuntimeMs();
861    } else {
862      // Not running; no data
863      return -1;
864    }
865  }
866
867  /**
868   * Gets the unique run id for a scheduled or running command. If the command is not currently
869   * scheduled or running, this function returns {@code 0}.
870   *
871   * @param command The command to get the run ID for
872   * @return The run of the command
873   */
874  @NoDiscard
875  @SuppressWarnings("PMD.CompareObjectsWithEquals")
876  public int runId(Command command) {
877    if (m_runningCommands.containsKey(command)) {
878      return m_runningCommands.get(command).id();
879    }
880
881    // Check scheduled commands
882    for (var scheduled : m_queuedToRun) {
883      if (scheduled.command() == command) {
884        return scheduled.id();
885      }
886    }
887
888    return 0;
889  }
890
891  /**
892   * Gets how long the scheduler took to process its most recent {@link #run()} invocation, in
893   * milliseconds. Defaults to -1 if the scheduler has not yet run.
894   *
895   * @return How long, in milliseconds, the scheduler last took to execute.
896   */
897  @NoDiscard
898  public double lastRuntimeMs() {
899    return m_lastRunTimeMs;
900  }
901
902  // Event-base telemetry and helpers. The static factories are for convenience to automatically
903  // set the timestamp instead of littering RobotController.getTime() everywhere.
904
905  private void emitScheduledEvent(Command command) {
906    var event = new SchedulerEvent.Scheduled(command, RobotController.getTime());
907    emitEvent(event);
908  }
909
910  private void emitMountedEvent(Command command) {
911    var event = new SchedulerEvent.Mounted(command, RobotController.getTime());
912    emitEvent(event);
913  }
914
915  private void emitYieldedEvent(Command command) {
916    var event = new SchedulerEvent.Yielded(command, RobotController.getTime());
917    emitEvent(event);
918  }
919
920  private void emitCompletedEvent(Command command) {
921    var event = new SchedulerEvent.Completed(command, RobotController.getTime());
922    emitEvent(event);
923  }
924
925  private void emitCompletedWithErrorEvent(Command command, Throwable error) {
926    var event = new SchedulerEvent.CompletedWithError(command, error, RobotController.getTime());
927    emitEvent(event);
928  }
929
930  private void emitCanceledEvent(Command command) {
931    var event = new SchedulerEvent.Canceled(command, RobotController.getTime());
932    emitEvent(event);
933  }
934
935  private void emitInterruptedEvent(Command command, Command interrupter) {
936    var event = new SchedulerEvent.Interrupted(command, interrupter, RobotController.getTime());
937    emitEvent(event);
938  }
939
940  /**
941   * Adds a listener to handle events that are emitted by the scheduler. Events are emitted when
942   * certain actions are taken by user code or by internal processing logic in the scheduler.
943   * Listeners should take care to be quick, simple, and not schedule or cancel commands, as that
944   * may cause inconsistent scheduler behavior or even cause a program crash.
945   *
946   * <p>Listeners are primarily expected to be for data logging and telemetry. In particular, a
947   * one-shot command (one that never calls {@link Coroutine#yield()}) will never appear in the
948   * standard protobuf telemetry because it is scheduled, runs, and finishes all in a single
949   * scheduler cycle. However, {@link SchedulerEvent.Scheduled},{@link SchedulerEvent.Mounted}, and
950   * {@link SchedulerEvent.Completed} events will be emitted corresponding to those actions, and
951   * user code can listen for and log such events.
952   *
953   * @param listener The listener to add. Cannot be null.
954   * @throws NullPointerException if given a null listener
955   */
956  public void addEventListener(Consumer<? super SchedulerEvent> listener) {
957    ErrorMessages.requireNonNullParam(listener, "listener", "addEventListener");
958
959    m_eventListeners.add(listener);
960  }
961
962  private void emitEvent(SchedulerEvent event) {
963    // TODO: Prevent listeners from interacting with the scheduler.
964    //       Scheduling or canceling commands while the scheduler is processing will probably cause
965    //       bugs in user code or even a program crash.
966    for (var listener : m_eventListeners) {
967      listener.accept(event);
968    }
969  }
970}