Commit fork, rewind, delay and attribute changes with their checkpoint

transact/src/main/java/dev/dbos/transact/DBOSClient.java

@@ -1610,7 +1610,7 @@ public void cancelWorkflows(@NonNull List<String> workflowIds, boolean cancelChi
    */
   public void updateWorkflowAttributes(
       @NonNull String workflowId, @Nullable Map<String, Object> attributes) {
-    systemDatabase.updateWorkflowAttributes(workflowId, attributes);
+    systemDatabase.updateWorkflowAttributes(workflowId, attributes, null);
   }
 
   /**
@@ -1735,7 +1735,8 @@ public void deleteWorkflows(@NonNull List<String> workflowIds, boolean deleteChi
    */
   public <T, E extends Exception> @NonNull WorkflowHandle<T, E> forkWorkflow(
       @NonNull String originalWorkflowId, int startStep, @NonNull ForkOptions options) {
-    var forkedWorkflowId = systemDatabase.forkWorkflow(originalWorkflowId, startStep, options);
+    var forkedWorkflowId =
+        systemDatabase.forkWorkflow(originalWorkflowId, startStep, options, null, null);
     return retrieveWorkflow(forkedWorkflowId);
   }
 
@@ -1763,7 +1764,7 @@ public void deleteWorkflows(@NonNull List<String> workflowIds, boolean deleteChi
    */
   public <T, E extends Exception> @NonNull WorkflowHandle<T, E> rewindWorkflow(
       @NonNull String workflowId, int startStep, @NonNull RewindOptions options) {
-    systemDatabase.rewindWorkflow(workflowId, startStep, options);
+    systemDatabase.rewindWorkflow(workflowId, startStep, options, null, null);
     return retrieveWorkflow(workflowId);
   }
 
@@ -2039,7 +2040,7 @@ public void applySchedules(@NonNull WorkflowSchedule... schedules) {
    */
   public void setWorkflowDelay(@NonNull String workflowId, @NonNull Duration delay) {
     var wfDelay = new WorkflowDelay.Delay(Objects.requireNonNull(delay, "delay must not be null"));
-    systemDatabase.setWorkflowDelay(workflowId, wfDelay);
+    systemDatabase.setWorkflowDelay(workflowId, wfDelay, null);
   }
 
   /**
@@ -2055,7 +2056,7 @@ public void setWorkflowDelay(@NonNull String workflowId, @NonNull Instant delayU
     var wfDelay =
         new WorkflowDelay.DelayUntil(
             Objects.requireNonNull(delayUntil, "delayUntil must not be null"));
-    systemDatabase.setWorkflowDelay(workflowId, wfDelay);
+    systemDatabase.setWorkflowDelay(workflowId, wfDelay, null);
   }
 
   /**

transact/src/main/java/dev/dbos/transact/database/SystemDatabase.java

@@ -37,6 +37,7 @@
 import dev.dbos.transact.workflow.WorkflowDelay;
 import dev.dbos.transact.workflow.WorkflowIdReusePolicy;
 import dev.dbos.transact.workflow.WorkflowSchedule;
+import dev.dbos.transact.workflow.WorkflowState;
 import dev.dbos.transact.workflow.WorkflowStatus;
 import dev.dbos.transact.workflow.internal.StepResult;
 import dev.dbos.transact.workflow.internal.WorkflowStatusInternal;
@@ -1694,9 +1695,21 @@ public void resumeWorkflows(
     }
   }
 
-  public void updateWorkflowAttributes(String workflowId, Map<String, Object> attributes) {
-    var validated = Validation.validateAttributes(attributes);
-    dbRetry(() -> WorkflowDAO.updateWorkflowAttributes(ctx, workflowId, validated));
+  /**
+   * Replaces a workflow's attributes as the step {@code caller} makes; see {@link
+   * #runTransactionalStep}.
+   */
+  public void updateWorkflowAttributes(
+      String workflowId, Map<String, Object> attributes, @Nullable StepCaller caller) {
+    runTransactionalStep(
+        caller,
+        "DBOS.updateWorkflowAttributes",
+        conn -> {
+          var validated = Validation.validateAttributes(attributes);
+          WorkflowDAO.updateWorkflowAttributes(conn, ctx.schema(), workflowId, validated);
+          logger.info("Updating attributes of workflow {}", workflowId);
+          return null;
+        });
   }
 
   /**
@@ -1903,33 +1916,120 @@ private void recordStepError(
     }
   }
 
-  public String forkWorkflow(String originalWorkflowId, int startStep, ForkOptions options) {
-    return dbRetryIncludingSerializationError(
-        "forkWorkflow",
-        () -> WorkflowDAO.forkWorkflow(ctx, originalWorkflowId, startStep, options));
+  /**
+   * Forks a workflow as the step {@code caller} makes, and returns the fork's ID; see {@link
+   * #runTransactionalStep}. {@code validate}, if given, runs first inside the step, so a replay
+   * returns the recorded fork without repeating it.
+   *
+   * <p>The ID is chosen before the transaction, so a retry that re-runs the work, after a commit
+   * whose acknowledgement was lost, fails on the fork it already wrote rather than forking again
+   * under a second ID.
+   */
+  public String forkWorkflow(
+      String originalWorkflowId,
+      int startStep,
+      @Nullable ForkOptions options,
+      @Nullable Runnable validate,
+      @Nullable StepCaller caller) {
+    var resolved = Objects.requireNonNullElseGet(options, ForkOptions::new);
+    var forkedWorkflowId =
+        Objects.requireNonNullElseGet(
+            resolved.forkedWorkflowId(), () -> UUID.randomUUID().toString());
+    return runTransactionalStep(
+        caller,
+        "DBOS.forkWorkflow",
+        conn -> {
+          if (validate != null) {
+            validate.run();
+          }
+          WorkflowDAO.forkWorkflows(
+              ctx,
+              conn,
+              List.of(new WorkflowFork(originalWorkflowId, forkedWorkflowId, startStep)),
+              resolved.applicationVersion(),
+              resolved.queueName(),
+              resolved.queuePartitionKey(),
+              resolved.timeout());
+          logger.info(
+              "Forking workflow {} from step {} as {}",
+              originalWorkflowId,
+              startStep,
+              forkedWorkflowId);
+          return forkedWorkflowId;
+        });
   }
 
   public void forkWorkflows(
       List<WorkflowFork> forks,
       @Nullable String applicationVersion,
       @Nullable String queueName,
       @Nullable String queuePartitionKey) {
-    dbRetryIncludingSerializationError(
+    runTransactionalStep(
+        null,
         "forkWorkflows",
-        () -> {
+        conn -> {
           WorkflowDAO.forkWorkflows(
-              ctx, forks, applicationVersion, queueName, queuePartitionKey, null);
+              ctx, conn, forks, applicationVersion, queueName, queuePartitionKey, null);
           return null;
         });
   }
 
-  public void rewindWorkflow(String workflowId, int startStep, RewindOptions options) {
-    dbRetryIncludingSerializationError(
-        "rewindWorkflow", () -> WorkflowDAO.rewindWorkflow(ctx, workflowId, startStep, options));
+  /**
+   * Rewinds a terminal workflow to {@code startStep} as the step {@code caller} makes; see {@link
+   * #runTransactionalStep}.
+   *
+   * <p>{@code refusal}, if given, is a refusal the caller reached before the transaction, while
+   * checking the call or clearing what lives outside the system database. It is thrown from inside
+   * the step, so it is recorded as the step's error, unless the step is already recorded, which
+   * then replays.
+   */
+  public void rewindWorkflow(
+      String workflowId,
+      int startStep,
+      RewindOptions options,
+      @Nullable RuntimeException refusal,
+      @Nullable StepCaller caller) {
+    runTransactionalStep(
+        caller,
+        "DBOS.rewindWorkflow",
+        conn -> {
+          if (refusal != null) {
+            throw refusal;
+          }
+          WorkflowDAO.rewindWorkflow(conn, ctx.schema(), workflowId, startStep, options);
+          logger.info("Rewinding workflow {} to step {}", workflowId, startStep);
+          return null;
+        });
   }
 
-  public void checkRewindable(String workflowId, int startStep) {
-    dbRetry(() -> WorkflowDAO.checkRewindable(ctx, workflowId, startStep));
+  /**
+   * What a rewind of {@code workflowId} by {@code caller} must know before its step's transaction,
+   * in one round trip: whether the step would run its work, rather than replay its checkpoint or be
+   * refused because the caller has lost its workflow, and the status of the workflow to rewind.
+   */
+  public record RewindPrecheck(boolean runsWork, @Nullable WorkflowState targetState) {
+    /** Refuses the rewind unless the workflow exists and is in a terminal state. */
+    public void checkRewindable(String workflowId, int startStep) {
+      WorkflowDAO.checkRewindable(workflowId, startStep, targetState);
+    }
+  }
+
+  public RewindPrecheck rewindPrecheck(@Nullable StepCaller caller, String workflowId) {
+    var read =
+        dbRetry(
+            () -> {
+              try (var conn = ctx.getConnection()) {
+                return WorkflowDAO.rewindPrecheck(
+                    conn,
+                    ctx.schema(),
+                    caller == null ? null : caller.workflowId(),
+                    caller instanceof StepCaller.FromWorkflow step ? step.stepId() : -1,
+                    caller == null ? null : caller.ownerXid(),
+                    workflowId);
+              }
+            });
+    boolean runsWork = caller == null || (!read.recorded() && read.callerOwns());
+    return new RewindPrecheck(runsWork, read.targetState());
   }
 
   public void createApplicationVersion(String versionName, @Nullable String applicationName) {
@@ -1983,18 +2083,51 @@ public void garbageCollect(Instant cutoff, Long rowsThreshold, int batchSize) {
     dbRetry(() -> WorkflowDAO.runRetentionRound(ctx, cutoff, rowsThreshold, batchSize));
   }
 
-  public void setWorkflowDelay(String workflowId, WorkflowDelay delay) {
+  /**
+   * Moves a DELAYED workflow's start as the step {@code caller} makes; see {@link
+   * #runTransactionalStep}. A relative delay is resolved once, before the transaction and its
+   * retries, as {@link #resolveDelayUntil(Duration)} requires. A delay that cannot be resolved is
+   * refused inside the step, so the refusal is recorded as the step's error.
+   */
+  public void setWorkflowDelay(
+      String workflowId, WorkflowDelay delay, @Nullable StepCaller caller) {
+    Long resolved;
+    RuntimeException refused;
+    try {
+      resolved = resolveDelayUntil(delay);
+      refused = null;
+    } catch (RuntimeException e) {
+      resolved = null;
+      refused = e;
+    }
+    final Long delayUntilEpochMs = resolved;
+    final RuntimeException refusal = refused;
+    runTransactionalStep(
+        caller,
+        "DBOS.setWorkflowDelay",
+        conn -> {
+          if (refusal != null) {
+            throw refusal;
+          }
+          WorkflowDAO.setWorkflowDelay(conn, ctx.schema(), workflowId, delayUntilEpochMs);
+          logger.info("Delaying workflow {} until {}", workflowId, delayUntilEpochMs);
+          return null;
+        });
+  }
+
+  /**
+   * The absolute end of {@code delay}: a relative one from this JVM's clock, an absolute as given.
+   */
+  private static long resolveDelayUntil(WorkflowDelay delay) {
     Objects.requireNonNull(delay, "delay must not be null");
-    long delayUntilEpochMs;
     if (delay instanceof WorkflowDelay.Delay d) {
-      delayUntilEpochMs = resolveDelayUntil(d.delay());
-    } else if (delay instanceof WorkflowDelay.DelayUntil du) {
+      return resolveDelayUntil(d.delay());
+    }
+    if (delay instanceof WorkflowDelay.DelayUntil du) {
       // An absolute delay is the caller's instant, taken as given.
-      delayUntilEpochMs = du.delayUntil().toEpochMilli();
-    } else {
-      throw new IllegalArgumentException("Unexpected WorkflowDelay value");
+      return du.delayUntil().toEpochMilli();
     }
-    dbRetry(() -> WorkflowDAO.setWorkflowDelay(ctx, workflowId, delayUntilEpochMs));
+    throw new IllegalArgumentException("Unexpected WorkflowDelay value");
   }
 
   /** See {@link WorkflowDAO#cancelTimedOutWorkflows}. */