Merge pull request #1846 from NASA-AMMOS/fix/scheduler-illegalstate-exception

Fix `IllegalStateException` when Simulating during Scheduling

Author: dandelany

e2e-tests/src/test/java/gov/nasa/jpl/aerie/e2e/procedural/scheduling/BasicSchedulingTests.java

@@ -10,6 +10,7 @@
 
 import javax.json.Json;
 import java.io.IOException;
+import java.util.Comparator;
 import java.util.Objects;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -34,7 +35,8 @@ void localBeforeEach() throws IOException {
         "Test Scheduling Procedure 1",
         dumbRecurrenceGoalJarId,
         specId,
-        0
+        0,
+        true
     );
   }
 
@@ -184,7 +186,7 @@ void saveActivityName() throws IOException {
    * Run a spec that includes unfinished activities
    */
   @Test
-  void IncludesUnfinishedActivities() throws IOException {
+  void includesUnfinishedActivities() throws IOException {
     // Disable the recurrence goal to simplify scheduling results
     hasura.updateSchedulingSpecEnabled(dumbRecurrenceGoalId.invocationId(), false);
 
@@ -226,4 +228,47 @@ void IncludesUnfinishedActivities() throws IOException {
       assertNull(act.duration());
     }
   }
+
+  /**
+   * Scheduling can handle two activities with the same parameters located that start at the same time without
+   * throwing an IllegalStateException.
+   *
+   * Uses DumbRecurrenceGoal for testing, as it will always place duplicate directives when the spec is rescheduled.
+   */
+  @Test
+  void duplicateActivitiesSupportedOnRerun() throws IOException {
+    final var args = Json.createObjectBuilder().add("quantity", 2).add("biteSize", 1).build();
+    hasura.updateSchedulingSpecGoalArguments(dumbRecurrenceGoalId.invocationId(), args);
+
+    hasura.awaitScheduling(planId);
+
+    final var initialActivities = hasura.getPlan(planId).activityDirectives();
+
+    // Rerun scheduling
+    hasura.updatePlanRevisionSchedulingSpec(planId);
+    hasura.awaitScheduling(planId);
+
+    final var updatedActivities = hasura.getPlan(planId).activityDirectives();
+
+    // Two new directives should have been placed on the plan
+    assertEquals(2, initialActivities.size());
+    assertEquals(4, updatedActivities.size());
+
+    // Sort the activities by start time
+    initialActivities.sort(Comparator.comparing(Plan.ActivityDirective::startOffset));
+    updatedActivities.sort(Comparator.comparing(Plan.ActivityDirective::startOffset));
+
+    // Check that the four activities are grouped into two sets of two.
+    final var firstStartTime = initialActivities.getFirst().startOffset();
+    assertEquals(2, updatedActivities.stream()
+                                     .filter(dir -> dir.startOffset().equals(firstStartTime))
+                                     .toList()
+                                     .size());
+
+    final var secondStartTime = initialActivities.getLast().startOffset();
+    assertEquals(2, updatedActivities.stream()
+                                     .filter(dir -> dir.startOffset().equals(secondStartTime))
+                                     .toList()
+                                     .size());
+  }
 }

scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/SimulationFacade.java

@@ -9,14 +9,15 @@
 import gov.nasa.jpl.aerie.types.ActivityDirective;
 import gov.nasa.jpl.aerie.types.ActivityDirectiveId;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
-import java.util.stream.Collectors;
 
 public interface SimulationFacade {
   void setInitialSimResults(SimulationData simulationData);
@@ -82,28 +83,68 @@ public boolean equals(Object other){
      * `map` contains a mapping between *only the IDs that are different* between the two plans.
      */
     public Optional<Map<ActivityDirectiveId, ActivityDirectiveId>> equalsWithIdMap(PlanSimCorrespondence other) {
+      // If the simulations have a different amount of directives, they must not be equal
+      if(directiveIdActivityDirectiveMap.size() != other.directiveIdActivityDirectiveMap.size()) {
+        return Optional.empty();
+      }
+
+      // Build the maps of directives -> ids from the ids -> directives map
+      // Because multiple activities on a plan can be identical sans id, the inverted maps to a list of ids
+      final HashMap<ActivityDirective, List<ActivityDirectiveId>> thisInverted = HashMap.newHashMap(directiveIdActivityDirectiveMap.size());
+      directiveIdActivityDirectiveMap.forEach(
+          (key, value) -> {
+            thisInverted.putIfAbsent(value, new ArrayList<>());
+            thisInverted.get(value).add(key);
+          });
+
+      final HashMap<ActivityDirective, List<ActivityDirectiveId>> otherInverted = HashMap.newHashMap(other.directiveIdActivityDirectiveMap.size());
+      other.directiveIdActivityDirectiveMap.forEach(
+          (key, value) -> {
+            otherInverted.putIfAbsent(value, new ArrayList<>());
+            otherInverted.get(value).add(key);
+          });
+
+      // If these maps have different sizes, then the simulations must not be equal
+      if(thisInverted.size() != otherInverted.size()) {
+        return Optional.empty();
+      }
+
+      // Check every entry for equality while building up the return mapping
       final var result = new HashMap<ActivityDirectiveId, ActivityDirectiveId>();
-      final var thisInverse = directiveIdActivityDirectiveMap.entrySet()
-                                 .stream()
-                                 .collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
-      final var otherInverse = other.directiveIdActivityDirectiveMap.entrySet()
-                                                             .stream()
-                                                             .collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
-      for (var entry : thisInverse.entrySet()) {
-        final var otherId = otherInverse.remove(entry.getKey());
-        if (otherId == null) {
+      for (final var entry : thisInverted.entrySet()) {
+        // Check that at least one instance of this directive is on the other plan
+        if(!otherInverted.containsKey(entry.getKey())){
+          return Optional.empty();
+        }
+
+        final var thisIds = entry.getValue();
+        final var otherIds = otherInverted.get(entry.getKey());
+
+        // Check that there is the same number of instances of this directive on the other plan
+        if (thisIds.size() != otherIds.size()) {
           return Optional.empty();
         }
-        if (!otherId.equals(entry.getValue())) {
-          result.put(entry.getValue(), otherId);
+
+        // If the directive is unique (the most common case), then the two ids can be directly compared
+        if(thisIds.size() == 1) {
+          if(!thisIds.getFirst().equals(otherIds.getFirst())) {
+            result.put(thisIds.getFirst(), otherIds.getFirst());
+          }
+          continue;
         }
-      }
 
-      if (otherInverse.isEmpty()) {
-        return Optional.of(result);
-      } else {
-        return Optional.empty();
+        // Filter out the ids that are on both lists, as this method does not provide a mapping for ids that have not changed
+        final List<ActivityDirectiveId> toRemove = thisIds.stream().filter(otherIds::contains).toList();
+        thisIds.removeAll(toRemove);
+        otherIds.removeAll(toRemove);
+
+        // Add the remaining mappings to the result
+        for(int i = 0; i < thisIds.size(); ++i) {
+          result.put(thisIds.get(i), otherIds.get(i));
+        }
       }
+
+      return Optional.of(result);
     }
   }
 }