add tail call optimization for tasks

NASA-AMMOS · Mar 2, 2024 · 43959cf · 43959cf
1 parent 7029aeb
commit 43959cf
Show file tree

Hide file tree

Showing 11 changed files with 72 additions and 10 deletions.
diff --git a/merlin-driver/src/main/java/gov/nasa/jpl/aerie/merlin/driver/engine/SimulationEngine.java b/merlin-driver/src/main/java/gov/nasa/jpl/aerie/merlin/driver/engine/SimulationEngine.java
@@ -219,14 +219,19 @@ private <Return> void stepEffectModel(
       this.tasks.put(task, progress.continueWith(s.continuation()));
       this.scheduledJobs.schedule(JobId.forTask(task), SubInstant.Tasks.at(currentTime.plus(s.delay())));
     } else if (status instanceof TaskStatus.CallingTask<Return> s) {
-      final var target = TaskId.generate();
-      SimulationEngine.this.tasks.put(target, new ExecutionState.InProgress<>(currentTime, s.child().create(this.executor)));
-      SimulationEngine.this.taskParent.put(target, task);
-      SimulationEngine.this.taskChildren.computeIfAbsent(task, $ -> new HashSet<>()).add(target);
-      frame.signal(JobId.forTask(target));
-
-      this.tasks.put(task, progress.continueWith(s.continuation()));
-      this.waitingTasks.subscribeQuery(task, Set.of(SignalId.forTask(target)));
+      if (s.tailCall()) {
+        this.tasks.put(task, new ExecutionState.InProgress<>(progress.startOffset, s.child().create(this.executor)));
+        this.scheduledJobs.schedule(JobId.forTask(task), SubInstant.Tasks.at(currentTime));
+      } else {
+        final var target = TaskId.generate();
+        this.tasks.put(target, new ExecutionState.InProgress<>(currentTime, s.child().create(this.executor)));
+        this.taskParent.put(target, task);
+        this.taskChildren.computeIfAbsent(task, $ -> new HashSet<>()).add(target);
+        frame.signal(JobId.forTask(target));
+
+        this.tasks.put(task, progress.continueWith(s.continuation()));
+        this.waitingTasks.subscribeQuery(task, Set.of(SignalId.forTask(target)));
+      }
     } else if (status instanceof TaskStatus.AwaitingCondition<Return> s) {
       final var condition = ConditionId.generate();
       this.conditions.put(condition, s.condition());

diff --git a/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/Context.java b/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/Context.java
@@ -30,8 +30,12 @@ enum ContextType { Initializing, Reacting, Querying }
   <Event> void emit(Event event, Topic<Event> topic);
 
   void spawn(TaskFactory<?> task);
+
   <Return> void call(TaskFactory<Return> task);
 
+  <Return> void tailCall(TaskFactory<Return> task);
+
   void delay(Duration duration);
+
   void waitUntil(Condition condition);
 }
diff --git a/...in-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/InitializationContext.java b/...in-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/InitializationContext.java
@@ -61,6 +61,11 @@ public <Return> void call(final TaskFactory<Return> task) {
     throw new IllegalStateException("Cannot yield during initialization");
   }
 
+  @Override
+  public <Return> void tailCall(final TaskFactory<Return> task) {
+    throw new IllegalStateException("Cannot yield during initialization");
+  }
+
   @Override
   public void delay(final Duration duration) {
     throw new IllegalStateException("Cannot yield during initialization");

diff --git a/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ModelActions.java b/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ModelActions.java
@@ -69,6 +69,18 @@ public static <T> void call(final TaskFactory<T> task) {
     context.get().call(task);
   }
 
+  public static void tailCall(final Runnable task) {
+    tailCall(threaded(task));
+  }
+
+  public static <T> void tailCall(final Supplier<T> task) {
+    tailCall(threaded(task));
+  }
+
+  public static <T> void tailCall(final TaskFactory<T> task) {
+    context.get().tailCall(task);
+  }
+
   public static void defer(final Duration duration, final Runnable task) {
     spawn(replaying(() -> { delay(duration); spawn(task); }));
   }

diff --git a/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/QueryContext.java b/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/QueryContext.java
@@ -52,6 +52,11 @@ public <Return> void call(final TaskFactory<Return> task) {
     throw new IllegalStateException("Cannot schedule tasks in a query-only context");
   }
 
+  @Override
+  public <Return> void tailCall(final TaskFactory<Return> task) {
+    throw new IllegalStateException("Cannot schedule tasks in a query-only context");
+  }
+
   @Override
   public void delay(final Duration duration) {
     throw new IllegalStateException("Cannot yield in a query-only context");

diff --git a/...framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ReplayingReactionContext.java b/...framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ReplayingReactionContext.java
@@ -78,6 +78,14 @@ public <T> void call(final TaskFactory<T> task) {
     });
   }
 
+  @Override
+  public <T> void tailCall(final TaskFactory<T> task) {
+    this.memory.doOnce(() -> {
+      this.scheduler = null;  // Relinquish the current scheduler before yielding, in case an exception is thrown.
+      this.scheduler = this.handle.tailCall(task);
+    });
+  }
+
   @Override
   public void delay(final Duration duration) {
     this.memory.doOnce(() -> {

diff --git a/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ReplayingTask.java b/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ReplayingTask.java
@@ -56,6 +56,11 @@ public Scheduler call(final TaskFactory<?> child) {
       return this.yield(TaskStatus.calling(child, ReplayingTask.this));
     }
 
+    @Override
+    public Scheduler tailCall(final TaskFactory<?> child) {
+      return this.yield(TaskStatus.tailCalling(child, ReplayingTask.this));
+    }
+
     @Override
     public Scheduler await(final gov.nasa.jpl.aerie.merlin.protocol.model.Condition condition) {
       return this.yield(TaskStatus.awaiting(condition, ReplayingTask.this));

diff --git a/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/TaskHandle.java b/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/TaskHandle.java
@@ -9,5 +9,7 @@ public interface TaskHandle {
 
   Scheduler call(TaskFactory<?> child);
 
+  Scheduler tailCall(TaskFactory<?> child);
+
   Scheduler await(gov.nasa.jpl.aerie.merlin.protocol.model.Condition condition);
 }
diff --git a/...-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ThreadedReactionContext.java b/...-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ThreadedReactionContext.java
@@ -63,6 +63,12 @@ public <T> void call(final TaskFactory<T> task) {
     this.scheduler = this.handle.call(task);
   }
 
+  @Override
+  public <T> void tailCall(final TaskFactory<T> task) {
+    this.scheduler = null;  // Relinquish the current scheduler before yielding, in case an exception is thrown.
+    this.scheduler = this.handle.tailCall(task);
+  }
+
   @Override
   public void delay(final Duration duration) {
     this.scheduler = null;  // Relinquish the current scheduler before yielding, in case an exception is thrown.

diff --git a/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ThreadedTask.java b/merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/ThreadedTask.java
@@ -207,6 +207,11 @@ public Scheduler call(final TaskFactory<?> child) {
       return this.yield(TaskStatus.calling(child, ThreadedTask.this));
     }
 
+    @Override
+    public Scheduler tailCall(final TaskFactory<?> child) {
+      return this.yield(TaskStatus.tailCalling(child, ThreadedTask.this));
+    }
+
     @Override
     public Scheduler await(final gov.nasa.jpl.aerie.merlin.protocol.model.Condition condition) {
       return this.yield(TaskStatus.awaiting(condition, ThreadedTask.this));

diff --git a/merlin-sdk/src/main/java/gov/nasa/jpl/aerie/merlin/protocol/types/TaskStatus.java b/merlin-sdk/src/main/java/gov/nasa/jpl/aerie/merlin/protocol/types/TaskStatus.java
@@ -9,7 +9,8 @@ record Completed<Return>(Return returnValue) implements TaskStatus<Return> {}
 
   record Delayed<Return>(Duration delay, Task<Return> continuation) implements TaskStatus<Return> {}
 
-  record CallingTask<Return>(TaskFactory<?> child, Task<Return> continuation) implements TaskStatus<Return> {}
+  record CallingTask<Return>(TaskFactory<?> child, Task<Return> continuation, boolean tailCall)
+    implements TaskStatus<Return> {}
 
   record AwaitingCondition<Return>(Condition condition, Task<Return> continuation) implements TaskStatus<Return> {}
 
@@ -23,7 +24,11 @@ static <Return> Delayed<Return> delayed(final Duration delay, final Task<Return>
   }
 
   static <Return> CallingTask<Return> calling(final TaskFactory<?> child, final Task<Return> continuation) {
-    return new CallingTask<>(child, continuation);
+    return new CallingTask<>(child, continuation, false);
+  }
+
+  static <Return> CallingTask<Return> tailCalling(final TaskFactory<?> child, final Task<Return> continuation) {
+    return new CallingTask<>(child, continuation, true);
   }
 
   static <Return> AwaitingCondition<Return> awaiting(final Condition condition, final Task<Return> continuation) {