Respect simulateAfter when serializing scheduling sim results

procedural/scheduling/src/main/kotlin/gov/nasa/ammos/aerie/procedural/scheduling/plan/EditablePlan.kt

@@ -25,8 +25,7 @@ interface EditablePlan: Plan {
       type: String,
       start: DirectiveStart,
       arguments: Map<String, SerializedValue>
-  ) = create(
-      NewDirective(
+  ) = create(NewDirective(
       AnyDirective(arguments),
       "Unnamed Activity",
       type,

scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/Procedure.java

@@ -27,8 +27,8 @@ public class Procedure extends Goal {
   private final Path jarPath;
   private final SerializedValue args;
 
-  public Procedure(final PlanningHorizon planningHorizon, Path jarPath, SerializedValue args) {
-    this.simulateAfter = true;
+  public Procedure(final PlanningHorizon planningHorizon, Path jarPath, SerializedValue args, boolean simulateAfter) {
+    this.simulateAfter = simulateAfter;
     this.planHorizon = planningHorizon;
     this.jarPath = jarPath;
     this.args = args;

scheduler-server/src/main/java/gov/nasa/jpl/aerie/scheduler/server/remotes/postgres/GoalBuilder.java

@@ -161,7 +161,7 @@ public static Goal goalOfGoalSpecifier(
       }
 
       case SchedulingDSL.GoalSpecifier.Procedure g -> {
-        return new Procedure(planningHorizon, g.jarPath(), g.arguments());
+        return new Procedure(planningHorizon, g.jarPath(), g.arguments(), simulateAfter);
       }
     }
   }

scheduler-worker/src/main/java/gov/nasa/jpl/aerie/scheduler/worker/services/SynchronousSchedulerAgent.java

@@ -11,10 +11,8 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -59,7 +57,6 @@
 import gov.nasa.jpl.aerie.scheduler.server.models.ExternalProfiles;
 import gov.nasa.jpl.aerie.scheduler.server.models.GoalId;
 import gov.nasa.jpl.aerie.scheduler.server.models.GoalRecord;
-import gov.nasa.jpl.aerie.scheduler.server.models.GoalSource;
 import gov.nasa.jpl.aerie.scheduler.server.models.GoalType;
 import gov.nasa.jpl.aerie.scheduler.server.models.MerlinPlan;
 import gov.nasa.jpl.aerie.scheduler.server.models.PlanId;
@@ -77,6 +74,7 @@
 import gov.nasa.jpl.aerie.scheduler.server.services.SpecificationService;
 import gov.nasa.jpl.aerie.scheduler.simulation.CheckpointSimulationFacade;
 import gov.nasa.jpl.aerie.scheduler.simulation.InMemoryCachedEngineStore;
+import gov.nasa.jpl.aerie.scheduler.simulation.SimulationData;
 import gov.nasa.jpl.aerie.scheduler.simulation.SimulationFacade;
 import gov.nasa.jpl.aerie.scheduler.solver.PrioritySolver;
 import org.apache.commons.collections4.BidiMap;
@@ -280,13 +278,21 @@ public void schedule(
 
       final var planMetadataAfterChanges = merlinService.getPlanMetadata(specification.planId());
       Optional<DatasetId> datasetId = initialSimulationResultsAndDatasetId.map(Pair::getRight);
-      if(planMetadataAfterChanges.planRev() != specification.planRevision()) {
+
+      final var lastGoalSimulateAfter = !problem.getGoals().isEmpty() && problem.getGoals().getLast().simulateAfter;
+      if(lastGoalSimulateAfter && planMetadataAfterChanges.planRev() != specification.planRevision()) {
         datasetId = storeSimulationResults(
-            solutionPlan,
-            planningHorizon,
-            simulationFacade,
+            simulationFacade.simulateWithResults(solutionPlan, planningHorizon.getEndAerie()),
+            planMetadataAfterChanges,
+            instancesToIds
+        );
+      } else if (simulationFacade.getLatestSimulationData().isPresent() && simulationFacade.getLatestSimulationData() != problem.getInitialSimulationResults()) {
+        final var latest = simulationFacade.getLatestSimulationData().get();
+        datasetId = storeSimulationResults(
+            latest,
             planMetadataAfterChanges,
-            instancesToIds);
+            instancesToIds
+        );
       }
       //collect results and notify subscribers of success
       final var results = collectResults(solutionPlan, instancesToIds, goals);
@@ -368,36 +374,28 @@ private ExternalProfiles loadExternalProfiles(final PlanId planId)
   }
 
   private Optional<DatasetId> storeSimulationResults(
-      final Plan plan,
-      PlanningHorizon planningHorizon,
-      SimulationFacade simulationFacade,
+      SimulationData simulationData,
       PlanMetadata planMetadata,
       final Map<SchedulingActivityDirective, ActivityDirectiveId> schedDirectiveToMerlinId)
   throws MerlinServiceException, IOException, SchedulingInterruptedException
   {
-    //finish simulation until end of horizon before posting results
-    try {
-      final var simulationData = simulationFacade.simulateWithResults(plan, planningHorizon.getEndAerie());
-      final var schedID_to_MerlinID =
-          schedDirectiveToMerlinId.entrySet().stream()
-                                  .collect(Collectors.toMap(
-                                      (a) -> new SchedulingActivityDirectiveId(a.getKey().id().id()), Map.Entry::getValue));
-      final var schedID_to_simID =
-          simulationData.mapSchedulingIdsToActivityIds().get();
-      final var simID_to_MerlinID =
-          schedID_to_simID.entrySet().stream().collect(Collectors.toMap(
-              Map.Entry::getValue,
-              (a) -> schedID_to_MerlinID.get(a.getKey())));
-      if(simID_to_MerlinID.values().containsAll(schedDirectiveToMerlinId.values()) && schedDirectiveToMerlinId.values().containsAll(simID_to_MerlinID.values())){
-        return Optional.of(merlinService.storeSimulationResults(planMetadata,
-                                                                simulationData.driverResults(),
-                                                                simID_to_MerlinID));
-      } else{
-        //schedule in simulation is inconsistent with current state of the plan (user probably disabled simulation for some of the goals)
-        return Optional.empty();
-      }
-    } catch (SimulationFacade.SimulationException e) {
-      throw new RuntimeException("Error while running simulation before storing simulation results after scheduling", e);
+    final var schedID_to_MerlinID =
+        schedDirectiveToMerlinId.entrySet().stream()
+                                .collect(Collectors.toMap(
+                                    (a) -> new SchedulingActivityDirectiveId(a.getKey().id().id()), Map.Entry::getValue));
+    final var schedID_to_simID =
+        simulationData.mapSchedulingIdsToActivityIds().get();
+    final var simID_to_MerlinID =
+        schedID_to_simID.entrySet().stream().collect(Collectors.toMap(
+            Map.Entry::getValue,
+            (a) -> schedID_to_MerlinID.get(a.getKey())));
+    if(simID_to_MerlinID.values().containsAll(schedDirectiveToMerlinId.values()) && schedDirectiveToMerlinId.values().containsAll(simID_to_MerlinID.values())){
+      return Optional.of(merlinService.storeSimulationResults(planMetadata,
+                                                              simulationData.driverResults(),
+                                                              simID_to_MerlinID));
+    } else{
+      //schedule in simulation is inconsistent with current state of the plan (user probably disabled simulation for some of the goals)
+      return Optional.empty();
     }
   }