Repository: nasa-ammos/aerie
Issue: Stop the ResumableSimulationDriver when it reaches the plan end bound

When using the simulateActivity method, the driver will stop only when the activity finishes which could happen after the plan end.

Changed files:
scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationDriver.java
scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/SimulationFacade.java
scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/constraints/activities/ActivityCreationTemplate.java
scheduler-worker/src/test/java/gov/nasa/jpl/aerie/scheduler/worker/services/SchedulingIntegrationTests.java
scheduler-driver/src/test/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationTest.java

diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/constraints/activities/ActivityCreationTemplate.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/constraints/activities/ActivityCreationTemplate.java
index 7060df2c3..d70203d13 100644
--- a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/constraints/activities/ActivityCreationTemplate.java
+++ b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/constraints/activities/ActivityCreationTemplate.java
@@ -286,41 +286,26 @@ public class ActivityCreationTemplate extends ActivityExpression implements Expr
               true);
           final var lastInsertion = history.getLastEvent();
           Optional<Duration> computedDuration = Optional.empty();
-          if(lastInsertion.isPresent()){
-            try {
-              //remove and insert at the same time to avoid unnecessary potential resimulation.
-              //
-              // Current sim: A -> B1 -> C
-              // Sim at end of next iteration: A -> B2 -> C
-              //If we would do remove() and then insert(), we would simulation A -> C and then again A -> B2 -> C
-              facade.removeAndInsertActivitiesFromSimulation(List.of(lastInsertion.get().activity()), List.of(actToSim));
-              computedDuration = facade.getActivityDuration(actToSim);
-              //record insertion: if successful, it will stay in the simulation, otherwise, it'll get deleted at the next iteration
+          final var toRemove = new ArrayList<SchedulingActivityDirective>();
+          lastInsertion.ifPresent(eventWithActivity -> toRemove.add(eventWithActivity.activity()));
+          try {
+            facade.removeAndInsertActivitiesFromSimulation(toRemove, List.of(actToSim));
+            computedDuration = facade.getActivityDuration(actToSim);
+            if(computedDuration.isPresent()) {
               history.add(new EventWithActivity(start, start.plus(computedDuration.get()), actToSim));
-            } catch (SimulationFacade.SimulationException e) {
-              //still need to record so we can remove the activity at the next iteration
-              history.add(new EventWithActivity(start, null, actToSim));
-            }
-          } else {
-            try {
-              facade.simulateActivity(actToSim);
-              computedDuration = facade.getActivityDuration(actToSim);
-              if(computedDuration.isPresent()) {
-                history.add(new EventWithActivity(start, start.plus(computedDuration.get()), actToSim));
-              } else{
-                logger.debug("No simulation error but activity duration could not be found in simulation, unfinished activity?");
-                history.add(new EventWithActivity(start,  null, actToSim));
-              }
-            } catch (SimulationFacade.SimulationException e) {
-              logger.debug("Simulation error while trying to simulate activities: " + e);
+            } else{
+              logger.debug("No simulation error but activity duration could not be found in simulation, likely caused by unfinished activity.");
               history.add(new EventWithActivity(start,  null, actToSim));
             }
+          } catch (SimulationFacade.SimulationException e) {
+            logger.debug("Simulation error while trying to simulate activities: " + e);
+            history.add(new EventWithActivity(start,  null, actToSim));
           }
           return computedDuration.map(start::plus).orElse(Duration.MAX_VALUE);
         }
 
       };
-      return rootFindingHelper(f, history, solved);
+      return rootFindingHelper(f, history, solved, facade);
       //CASE 2: activity has a controllable duration
     } else if (this.type.getDurationType() instanceof DurationType.Controllable dt) {
       //select earliest start time, STN guarantees satisfiability
@@ -421,7 +406,7 @@ public class ActivityCreationTemplate extends ActivityExpression implements Expr
         }
       };
 
-      return rootFindingHelper(f, history, solved);
+      return rootFindingHelper(f, history, solved, facade);
     } else {
      throw new UnsupportedOperationException("Unsupported duration type found: " + this.type.getDurationType());
     }
@@ -450,7 +435,8 @@ public class ActivityCreationTemplate extends ActivityExpression implements Expr
   private  Optional<SchedulingActivityDirective> rootFindingHelper(
       final EquationSolvingAlgorithms.Function<Duration, HistoryWithActivity> f,
       final HistoryWithActivity history,
-      final TaskNetworkAdapter.TNActData solved
+      final TaskNetworkAdapter.TNActData solved,
+      final SimulationFacade simulationFacade
   ) {
     try {
       var endInterval = solved.end();
@@ -486,6 +472,13 @@ public class ActivityCreationTemplate extends ActivityExpression implements Expr
     } catch (EquationSolvingAlgorithms.NoSolutionException e) {
       logger.debug("No solution");
     }
+    if(!history.events.isEmpty()) {
+      try {
+        simulationFacade.removeActivitiesFromSimulation(List.of(history.getLastEvent().get().activity()));
+      } catch (SimulationFacade.SimulationException e) {
+        throw new RuntimeException("Exception while simulating original plan after activity insertion failure" ,e);
+      }
+    }
     return Optional.empty();
   }
 
diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationDriver.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationDriver.java
index 992c0a2f0..217f8d9ec 100644
--- a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationDriver.java
+++ b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationDriver.java
@@ -251,6 +251,9 @@ public class ResumableSimulationDriver<Model> implements AutoCloseable {
       // Update batch and increment real time, if necessary.
       batch = engine.extractNextJobs(Duration.MAX_VALUE);
       delta = batch.offsetFromStart().minus(curTime);
+      if(batch.offsetFromStart().longerThan(planDuration)){
+        break;
+      }
     }
     lastSimResults = null;
   }
diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/SimulationFacade.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/SimulationFacade.java
index 9f4e69590..6440c5e27 100644
--- a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/SimulationFacade.java
+++ b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/simulation/SimulationFacade.java
@@ -92,9 +92,6 @@ public class SimulationFacade implements AutoCloseable{
     }
     final var duration = driver.getActivityDuration(planActDirectiveIdToSimulationActivityDirectiveId.get(
         schedulingActivityDirective.getId()));
-    if(duration.isEmpty()){
-      logger.error("Simulation is probably outdated, check that no activity is removed between simulation and querying");
-    }
     return duration;
   }
 
@@ -136,10 +133,10 @@ public class SimulationFacade implements AutoCloseable{
       final Collection<SchedulingActivityDirective> activitiesToRemove,
       final Collection<SchedulingActivityDirective> activitiesToAdd) throws SimulationException
   {
-    var atLeastOne = false;
+    var atLeastOneActualRemoval = false;
     for(final var act: activitiesToRemove){
       if(insertedActivities.containsKey(act)){
-        atLeastOne = true;
+        atLeastOneActualRemoval = true;
         insertedActivities.remove(act);
       }
     }
@@ -147,16 +144,16 @@ public class SimulationFacade implements AutoCloseable{
     for(final var act: activitiesToAdd){
       earliestActStartTime = Duration.min(earliestActStartTime, act.startOffset());
     }
+    final var allActivitiesToSimulate = new ArrayList<>(activitiesToAdd);
     //reset resumable simulation
-    if(atLeastOne || earliestActStartTime.shorterThan(this.driver.getCurrentSimulationEndTime())){
-      final var allActivitiesToSimulate = new ArrayList<>(insertedActivities.keySet());
+    if(atLeastOneActualRemoval || earliestActStartTime.shorterThan(this.driver.getCurrentSimulationEndTime())){
+      allActivitiesToSimulate.addAll(insertedActivities.keySet());
       insertedActivities.clear();
       planActDirectiveIdToSimulationActivityDirectiveId.clear();
       if (driver != null) driver.close();
       driver = new ResumableSimulationDriver<>(missionModel, planningHorizon.getAerieHorizonDuration());
-      allActivitiesToSimulate.addAll(activitiesToAdd);
-      simulateActivities(allActivitiesToSimulate);
     }
+    simulateActivities(allActivitiesToSimulate);
   }
 
   public void removeActivitiesFromSimulation(final Collection<SchedulingActivityDirective> activities)
diff --git a/scheduler-driver/src/test/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationTest.java b/scheduler-driver/src/test/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationTest.java
index cd3737aeb..38f141348 100644
--- a/scheduler-driver/src/test/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationTest.java
+++ b/scheduler-driver/src/test/java/gov/nasa/jpl/aerie/scheduler/simulation/ResumableSimulationTest.java
@@ -4,7 +4,7 @@ import gov.nasa.jpl.aerie.merlin.driver.ActivityDirectiveId;
 import gov.nasa.jpl.aerie.merlin.driver.SerializedActivity;
 import gov.nasa.jpl.aerie.merlin.driver.engine.SimulationEngine;
 import gov.nasa.jpl.aerie.merlin.protocol.types.Duration;
-import gov.nasa.jpl.aerie.merlin.protocol.types.InstantiationException;
+import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
 import gov.nasa.jpl.aerie.scheduler.SimulationUtility;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -14,6 +14,7 @@ import java.util.ArrayList;
 import java.util.Map;
 import java.util.concurrent.ThreadPoolExecutor;
 
+import static gov.nasa.jpl.aerie.merlin.protocol.types.Duration.MICROSECOND;
 import static gov.nasa.jpl.aerie.merlin.protocol.types.Duration.SECONDS;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
@@ -24,6 +25,7 @@ public class ResumableSimulationTest {
   Duration endOfLastAct;
 
   private final Duration tenHours = Duration.of(10, Duration.HOURS);
+  private final Duration fiveHours = Duration.of(5, Duration.HOURS);
 
   @BeforeEach
   public void init() {
@@ -76,6 +78,21 @@ public class ResumableSimulationTest {
     assertTrue(act2Dur.get().isEqualTo(Duration.of(2, SECONDS)));
   }
 
+  @Test
+  public void testStopsAtEndOfPlanningHorizon(){
+    final var activity = new TestSimulatedActivity(
+        Duration.of(0, SECONDS),
+        new SerializedActivity("ControllableDurationActivity", Map.of("duration", SerializedValue.of(tenHours.in(MICROSECOND)))),
+        new ActivityDirectiveId(1));
+    final var fooMissionModel = SimulationUtility.getFooMissionModel();
+    resumableSimulationDriver = new ResumableSimulationDriver<>(fooMissionModel, fiveHours);
+    resumableSimulationDriver.initSimulation();
+    resumableSimulationDriver.clearActivitiesInserted();
+    resumableSimulationDriver.simulateActivity(activity.start, activity.activity, null, true, activity.id);
+    assertEquals(fiveHours, resumableSimulationDriver.getCurrentSimulationEndTime());
+    assert(resumableSimulationDriver.getSimulationResults(Instant.now()).unfinishedActivities.size() == 1);
+  }
+
   @Test
   public void testThreadsReleased() {
     final var activity = new TestSimulatedActivity(
diff --git a/scheduler-worker/src/test/java/gov/nasa/jpl/aerie/scheduler/worker/services/SchedulingIntegrationTests.java b/scheduler-worker/src/test/java/gov/nasa/jpl/aerie/scheduler/worker/services/SchedulingIntegrationTests.java
index 3653cb7d5..ccc775b0f 100644
--- a/scheduler-worker/src/test/java/gov/nasa/jpl/aerie/scheduler/worker/services/SchedulingIntegrationTests.java
+++ b/scheduler-worker/src/test/java/gov/nasa/jpl/aerie/scheduler/worker/services/SchedulingIntegrationTests.java
@@ -1774,6 +1774,30 @@ public class SchedulingIntegrationTests {
     assertEquals(96, results.updatedPlan().size());
   }
 
+  @Test
+  void testUnfinishedActivity(){
+    final PlanningHorizon PLANNING_HORIZON = new PlanningHorizon(
+        TimeUtility.fromDOY("2025-090T00:00:00"),
+        TimeUtility.fromDOY("2025-134T00:00:00"));
+    final var results = runScheduler(
+        BANANANATION,
+        List.of(),
+        List.of(new SchedulingGoal(new GoalId(0L), """
+          export default (): Goal => {
+                           return Goal.ActivityRecurrenceGoal({
+                             activityTemplate: ActivityTemplates.parent({
+                               label: "label"
+                           }),
+                             interval: Temporal.Duration.from({ days: 30})
+                           })
+                       }
+            """, true)), PLANNING_HORIZON);
+    //parent takes much more than 134 - 90 = 44 days to finish
+    assertEquals(0, results.updatedPlan.size());
+    final var goalResult = results.scheduleResults.goalResults().get(new GoalId(0L));
+    assertFalse(goalResult.satisfied());
+  }
+
   /**
    * If you passed activities without duration to the scheduler in an initial plan, it would fail
    */