Simulation result staleness checking

scheduler-driver/src/main/kotlin/gov/nasa/jpl/aerie/scheduler/plan/InMemoryEditablePlan.kt

@@ -11,14 +11,41 @@ import gov.nasa.ammos.aerie.procedural.timeline.payloads.activities.AnyDirective
 import gov.nasa.ammos.aerie.procedural.timeline.payloads.activities.Directive
 import gov.nasa.ammos.aerie.procedural.timeline.payloads.activities.DirectiveStart
 import gov.nasa.ammos.aerie.procedural.timeline.plan.Plan
+import gov.nasa.ammos.aerie.procedural.timeline.plan.SimulationResults
 import gov.nasa.jpl.aerie.merlin.protocol.types.DurationType
 import gov.nasa.jpl.aerie.scheduler.DirectiveIdGenerator
 import gov.nasa.jpl.aerie.scheduler.model.*
 import gov.nasa.jpl.aerie.types.ActivityDirectiveId
+import java.lang.ref.WeakReference
 import java.time.Instant
 import kotlin.jvm.optionals.getOrNull
 import gov.nasa.ammos.aerie.procedural.timeline.plan.SimulationResults as TimelineSimResults
 
+/*
+ * An implementation of [EditablePlan] that stores the plan in memory for use in the internal scheduler.
+ *
+ * ## Staleness checking
+ *
+ * The editable plan instance keeps track of sim results that it has produced using weak references, and can dynamically
+ * update their staleness if the plan is changed after it was simulated. The process is this:
+ *
+ * 1. [InMemoryEditablePlan] has a set of weak references to simulation results objects that are currently up-to-date.
+ *    I used weak references because if the user can't access it anymore, staleness doesn't matter and we might as well
+ *    let it get gc'ed.
+ * 2. When the user gets simulation results, either through simulation or by getting the latest, it always checks for
+ *    plan equality between the returned results and the current plan, even if we just simulated. If it is up-to-date, a
+ *    weak ref is added to the set.
+ * 3. When an edit is made, the sim results in the current set are marked stale; then the set is reset to new reference
+ *    to an empty set.
+ * 4. When a commit is made, the commit object takes *shared ownership* of the set. If a new simulation is run (step 2)
+ *    the plan can still add to the set while it is still jointly owned by the commit. Then when an edit is made (step 3)
+ *    the commit will become the sole owner of the set.
+ * 5. When changes are rolled back, any sim results currently in the plan's set are marked stale, the previous commit's
+ *    sim results are marked not stale, then the plan will resume joint ownership of the previous commit's set.
+ *
+ * The joint ownership freaks me out a wee bit, but I think it's safe because the commits are only used to keep the
+ * previous sets from getting gc'ed in the event of a rollback. Only the plan object actually mutates the set.
+ */
 data class InMemoryEditablePlan(
     private val missionModel: MissionModel<*>,
     private var idGenerator: DirectiveIdGenerator,
@@ -27,16 +54,39 @@ data class InMemoryEditablePlan(
     private val lookupActivityType: (String) -> ActivityType
 ) : EditablePlan, Plan by plan {
 
-  private val commits = mutableListOf<Commit>()
+  private data class Commit(
+    val diff: List<Edit>,
+
+    /**
+     * A record of the simulation results objects that were up-to-date when the commit
+     * was created.
+     *
+     * This has SHARED OWNERSHIP with [InMemoryEditablePlan]; the editable plan may add more to
+     * this list AFTER the commit is created.
+     */
+    val upToDateSimResultsSet: MutableSet<WeakReference<MerlinToProcedureSimulationResultsAdapter>>
+  )
+
+  private var committedChanges = Commit(listOf(), mutableSetOf())
   var uncommittedChanges = mutableListOf<Edit>()
     private set
 
   val totalDiff: List<Edit>
-    get() = commits.flatMap { it.diff }
+    get() = committedChanges.diff
+
+  // Jointly owned set of up-to-date simulation results. See class-level comment for algorithm explanation.
+  private var upToDateSimResultsSet: MutableSet<WeakReference<MerlinToProcedureSimulationResultsAdapter>> = mutableSetOf()
+
+  override fun latestResults(): SimulationResults? {
+    val merlinResults = simulationFacade.latestSimulationData.getOrNull() ?: return null
 
-  override fun latestResults() =
-    simulationFacade.latestSimulationData.getOrNull()
-      ?.let { MerlinToProcedureSimulationResultsAdapter(it.driverResults, false, plan) }
+    // kotlin checks structural equality by default, not referential equality.
+    val isStale = merlinResults.plan.activities != plan.activities
+
+    val results = MerlinToProcedureSimulationResultsAdapter(merlinResults.driverResults, isStale, plan)
+    if (!isStale) upToDateSimResultsSet.add(WeakReference(results))
+    return results
+  }
 
   override fun create(directive: NewDirective): ActivityDirectiveId {
     class ParentSearchException(id: ActivityDirectiveId, size: Int): Exception("Expected one parent activity with id $id, found $size")
@@ -55,16 +105,33 @@ data class InMemoryEditablePlan(
     uncommittedChanges.add(Edit.Create(resolved))
     resolved.validateArguments(lookupActivityType)
     plan.add(resolved.toSchedulingActivity(lookupActivityType, true))
+
+    for (simResults in upToDateSimResultsSet) {
+      simResults.get()?.stale = true
+    }
+    // create a new list instead of `.clear` because commit objects have the same reference
+    upToDateSimResultsSet = mutableSetOf()
+
     return id
   }
 
   override fun commit() {
-    val committedEdits = uncommittedChanges
+    // Early return if there are no changes. This prevents multiple commits from sharing ownership of the set,
+    // because new sets are only created when edits are made.
+    // Probably unnecessary, but shared ownership freaks me out enough already.
+    if (uncommittedChanges.isEmpty()) return
+
+    val newCommittedChanges = uncommittedChanges
     uncommittedChanges = mutableListOf()
-    commits.add(Commit(committedEdits))
+
+    // Create a commit that shares ownership of the simResults set.
+    committedChanges = Commit(committedChanges.diff + newCommittedChanges, upToDateSimResultsSet)
   }
 
   override fun rollback(): List<Edit> {
+    // Early return if there are no changes, to keep staleness accuracy
+    if (uncommittedChanges.isEmpty()) return emptyList()
+
     val result = uncommittedChanges
     uncommittedChanges = mutableListOf()
     for (edit in result) {
@@ -74,6 +141,13 @@ data class InMemoryEditablePlan(
         }
       }
     }
+    for (simResult in upToDateSimResultsSet) {
+      simResult.get()?.stale = true
+    }
+    for (simResult in committedChanges.upToDateSimResultsSet) {
+      simResult.get()?.stale = false
+    }
+    upToDateSimResultsSet = committedChanges.upToDateSimResultsSet
     return result
   }
 

scheduler-driver/src/main/kotlin/gov/nasa/jpl/aerie/scheduler/plan/MerlinToProcedureSimulationResultsAdapter.kt

@@ -18,7 +18,7 @@ import kotlin.jvm.optionals.getOrNull
 
 class MerlinToProcedureSimulationResultsAdapter(
     private val results: gov.nasa.jpl.aerie.merlin.driver.SimulationResults,
-    private val stale: Boolean,
+    var stale: Boolean,
     private val plan: Plan
 ): SimulationResults {
 

scheduler-driver/src/test/java/gov/nasa/jpl/aerie/scheduler/EditablePlanStalenessTest.java

@@ -0,0 +1,147 @@
+package gov.nasa.jpl.aerie.scheduler;
+
+import gov.nasa.ammos.aerie.procedural.scheduling.plan.EditablePlan;
+import gov.nasa.ammos.aerie.procedural.timeline.payloads.activities.DirectiveStart;
+import gov.nasa.jpl.aerie.merlin.driver.MissionModel;
+import gov.nasa.jpl.aerie.merlin.protocol.types.Duration;
+import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
+import gov.nasa.jpl.aerie.scheduler.model.PlanInMemory;
+import gov.nasa.jpl.aerie.scheduler.model.PlanningHorizon;
+import gov.nasa.jpl.aerie.scheduler.model.Problem;
+import gov.nasa.jpl.aerie.scheduler.plan.InMemoryEditablePlan;
+import gov.nasa.jpl.aerie.scheduler.plan.SchedulerToProcedurePlanAdapter;
+import gov.nasa.jpl.aerie.scheduler.simulation.CheckpointSimulationFacade;
+import gov.nasa.jpl.aerie.scheduler.simulation.SimulationFacade;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.time.Instant;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class EditablePlanStalenessTest {
+
+  MissionModel<?> missionModel;
+  Problem problem;
+  SimulationFacade facade;
+  EditablePlan plan;
+
+  private static final Instant start = TestUtility.timeFromEpochMillis(0);
+  private static final Instant end = TestUtility.timeFromEpochDays(1);
+
+  private static final PlanningHorizon horizon = new PlanningHorizon(start, end);
+
+  @BeforeEach
+  public void setUp() {
+    missionModel = SimulationUtility.getBananaMissionModel();
+    final var schedulerModel = SimulationUtility.getBananaSchedulerModel();
+    facade = new CheckpointSimulationFacade(horizon, missionModel, schedulerModel);
+    problem = new Problem(missionModel, horizon, facade, schedulerModel);
+    plan = new InMemoryEditablePlan(
+        missionModel,
+        new DirectiveIdGenerator(0),
+        new SchedulerToProcedurePlanAdapter(
+            new PlanInMemory(
+
+            ),
+            horizon,
+            Map.of(), Map.of(), Map.of()
+        ),
+        facade,
+        problem::getActivityType
+    );
+  }
+
+  @AfterEach
+  public void tearDown() {
+    missionModel = null;
+    problem = null;
+    facade = null;
+    plan = null;
+  }
+
+  @Test
+  public void simResultMarkedStale() {
+    plan.create(
+        "BiteBanana",
+        new DirectiveStart.Absolute(Duration.MINUTE),
+        Map.of("biteSize", SerializedValue.of(1))
+    );
+
+    final var simResults = plan.simulate();
+
+    assertFalse(simResults.isStale());
+
+    plan.create(
+        "GrowBanana",
+        new DirectiveStart.Absolute(Duration.HOUR),
+        Map.of(
+            "growingDuration", SerializedValue.of(10),
+            "quantity", SerializedValue.of(1)
+        )
+    );
+
+    assertTrue(simResults.isStale());
+  }
+
+  @Test
+  public void simResultMarkedNotStaleAfterRollback_CommitThenSimulate() {
+    plan.create(
+        "BiteBanana",
+        new DirectiveStart.Absolute(Duration.MINUTE),
+        Map.of("biteSize", SerializedValue.of(1))
+    );
+
+    plan.commit();
+    final var simResults = plan.simulate();
+
+    assertFalse(simResults.isStale());
+
+    plan.create(
+        "GrowBanana",
+        new DirectiveStart.Absolute(Duration.HOUR),
+        Map.of(
+            "growingDuration", SerializedValue.of(10),
+            "quantity", SerializedValue.of(1)
+        )
+    );
+
+    assertTrue(simResults.isStale());
+
+    plan.rollback();
+
+    assertFalse(simResults.isStale());
+  }
+
+  @Test
+  public void simResultMarkedNotStaleAfterRollback_SimulateThenCommit() {
+    plan.create(
+        "BiteBanana",
+        new DirectiveStart.Absolute(Duration.MINUTE),
+        Map.of("biteSize", SerializedValue.of(1))
+    );
+
+    final var simResults = plan.simulate();
+    plan.commit();
+
+    assertFalse(simResults.isStale());
+
+    plan.create(
+        "GrowBanana",
+        new DirectiveStart.Absolute(Duration.HOUR),
+        Map.of(
+            "growingDuration", SerializedValue.of(10),
+            "quantity", SerializedValue.of(1)
+        )
+    );
+
+    assertTrue(simResults.isStale());
+
+    plan.rollback();
+
+    assertFalse(simResults.isStale());
+  }
+}