Resume/Fork on specified queue (#343)

fixes #305

Author: devhawk

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

@@ -614,33 +614,73 @@ public <E extends Exception> void runStep(
   }
 
   /**
-   * Resume a workflow starting from the step after the last complete step
+   * Resume a workflow starting from the step after the last complete step. This method allows
+   * resuming workflows that were previously interrupted, failed, or canceled. The workflow will
+   * continue execution from where it left off, replaying any completed steps deterministically.
    *
    * @param <T> Return type of the workflow function
-   * @param <E> Checked exception thrown by the workflow function, if any
-   * @param workflowId id of the workflow
-   * @return A handle to the workflow
+   * @param <E> Type of checked exception thrown by the workflow function, if any
+   * @param workflowId ID of the workflow to resume; must not be null
+   * @param queueName optional queue name to enqueue the resumed workflow to; if null, the workflow
+   *     will be resumed in the default execution context
+   * @return A handle to the resumed workflow
+   * @throws IllegalStateException if called before DBOS is launched
    */
   @SuppressWarnings("unchecked")
   public <T, E extends Exception> @NonNull WorkflowHandle<T, E> resumeWorkflow(
-      @NonNull String workflowId) {
-    var handles = resumeWorkflows(List.of(workflowId));
+      @NonNull String workflowId, @Nullable String queueName) {
+    var handles = resumeWorkflows(List.of(workflowId), queueName);
     assert (handles.size() == 1);
     return (WorkflowHandle<T, E>) handles.get(0);
   }
 
+  /**
+   * Resume a workflow starting from the step after the last complete step using the default queue.
+   * This method is equivalent to calling {@code resumeWorkflow(workflowId, null)}. The workflow
+   * will continue execution from where it left off, replaying any completed steps
+   * deterministically.
+   *
+   * @param <T> Return type of the workflow function
+   * @param <E> Type of checked exception thrown by the workflow function, if any
+   * @param workflowId ID of the workflow to resume; must not be null
+   * @return A handle to the resumed workflow
+   * @throws IllegalStateException if called before DBOS is launched
+   */
+  public <T, E extends Exception> @NonNull WorkflowHandle<T, E> resumeWorkflow(
+      @NonNull String workflowId) {
+    return resumeWorkflow(workflowId, null);
+  }
+
   /**
    * Resume multiple workflows starting from the step after the last complete step for each
-   * workflow. This method allows bulk resumption of workflows that were previously interrupted or
-   * failed.
+   * workflow. This method allows bulk resumption of workflows that were previously interrupted,
+   * failed, or canceled. Each workflow will continue execution from where it left off, replaying
+   * any completed steps deterministically.
+   *
+   * @param workflowIds a list of workflow IDs to resume; must not be null
+   * @param queueName optional queue name to enqueue the resumed workflows to; if null, the
+   *     workflows will be resumed in the default execution context
+   * @return A list of handles to the resumed workflows
+   * @throws IllegalStateException if called before DBOS is launched
+   */
+  public @NonNull List<WorkflowHandle<Object, Exception>> resumeWorkflows(
+      @NonNull List<String> workflowIds, @Nullable String queueName) {
+    return ensureLaunched("resumeWorkflow").resumeWorkflows(workflowIds, queueName);
+  }
+
+  /**
+   * Resume multiple workflows starting from the step after the last complete step for each workflow
+   * using the default queue. This method is equivalent to calling {@code
+   * resumeWorkflows(workflowIds, null)}. Each workflow will continue execution from where it left
+   * off, replaying any completed steps deterministically.
    *
    * @param workflowIds a list of workflow IDs to resume; must not be null
    * @return A list of handles to the resumed workflows
    * @throws IllegalStateException if called before DBOS is launched
    */
   public @NonNull List<WorkflowHandle<Object, Exception>> resumeWorkflows(
       @NonNull List<String> workflowIds) {
-    return ensureLaunched("resumeWorkflow").resumeWorkflows(workflowIds);
+    return resumeWorkflows(workflowIds, null);
   }
 
   /***

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

@@ -676,30 +676,67 @@ public void cancelWorkflows(@NonNull List<String> workflowIds) {
   }
 
   /**
-   * Resume a canceled workflow, providing a handle to the workflow
+   * Resume a workflow starting from the step after the last complete step. This method allows
+   * resuming workflows that were previously interrupted, failed, or canceled. The workflow will
+   * continue execution from where it left off, replaying any completed steps deterministically.
    *
    * @param <T> Type of the workflow's return value
    * @param <E> Type of any checked exception thrown by the workflow
-   * @param workflowId ID of the workflow to resume
-   * @return `WorkflowHandle` for the resumed workflow
+   * @param workflowId ID of the workflow to resume; must not be null
+   * @param queueName optional queue name to enqueue the resumed workflow to; if null, the workflow
+   *     will be resumed in the default execution context
+   * @return WorkflowHandle for the resumed workflow
    */
   public <T, E extends Exception> @NonNull WorkflowHandle<T, E> resumeWorkflow(
-      @NonNull String workflowId) {
-    systemDatabase.resumeWorkflows(List.of(workflowId));
+      @NonNull String workflowId, @Nullable String queueName) {
+    systemDatabase.resumeWorkflows(List.of(workflowId), queueName);
     return retrieveWorkflow(workflowId);
   }
 
   /**
-   * Resume multiple workflows starting from the step after the last complete step for each
-   * workflow. This method allows bulk resumption of workflows that were previously interrupted or
-   * failed.
+   * Resume a workflow starting from the step after the last complete step using the default queue.
+   * This method is equivalent to calling {@code resumeWorkflow(workflowId, null)}. The workflow
+   * will continue execution from where it left off, replaying any completed steps
+   * deterministically.
+   *
+   * @param <T> Type of the workflow's return value
+   * @param <E> Type of any checked exception thrown by the workflow
+   * @param workflowId ID of the workflow to resume; must not be null
+   * @return WorkflowHandle for the resumed workflow
+   */
+  public <T, E extends Exception> @NonNull WorkflowHandle<T, E> resumeWorkflow(
+      @NonNull String workflowId) {
+    return resumeWorkflow(workflowId, null);
+  }
+
+  /**
+   * Resume multiple workflows starting from the step after the last complete step for each workflow
+   * using the default queue. This method is equivalent to calling {@code
+   * resumeWorkflows(workflowIds, null)}. Each workflow will continue execution from where it left
+   * off, replaying any completed steps deterministically.
    *
    * @param workflowIds a list of workflow IDs to resume; must not be null
    * @return A list of handles to the resumed workflows
    */
   public @NonNull List<WorkflowHandle<Object, Exception>> resumeWorkflows(
       @NonNull List<String> workflowIds) {
-    systemDatabase.resumeWorkflows(workflowIds);
+    return resumeWorkflows(workflowIds, null);
+  }
+
+  /**
+   * Resume multiple workflows starting from the step after the last complete step for each
+   * workflow. This method allows bulk resumption of workflows that were previously interrupted,
+   * failed, or canceled. Each workflow will continue execution from where it left off, replaying
+   * any completed steps deterministically.
+   *
+   * @param workflowIds a list of workflow IDs to resume; must not be null
+   * @param queueName optional queue name to enqueue the resumed workflows to; if null, the
+   *     workflows will be resumed in the default execution context
+   * @return A list of handles to the resumed workflows
+   */
+  public @NonNull List<WorkflowHandle<Object, Exception>> resumeWorkflows(
+      @NonNull List<String> workflowIds, @Nullable String queueName) {
+    systemDatabase.resumeWorkflows(workflowIds, queueName);
     return workflowIds.stream().map(this::retrieveWorkflow).toList();
   }
 

transact/src/main/java/dev/dbos/transact/admin/AdminServer.java

@@ -239,7 +239,7 @@ private void resume(HttpExchange exchange, String wfid) throws IOException {
 
     logger.info("resume workflow {}", wfid);
 
-    dbosExecutor.resumeWorkflows(List.of(wfid));
+    dbosExecutor.resumeWorkflows(List.of(wfid), null);
     exchange.sendResponseHeaders(204, 0);
   }
 
@@ -248,7 +248,9 @@ private void fork(HttpExchange exchange, String wfid) throws IOException {
 
     var request = mapper.readValue(exchange.getRequestBody(), ForkRequest.class);
     int startStep = request.start_step == null ? 0 : request.start_step;
-    var options = new ForkOptions(request.new_workflow_id, request.application_version, null);
+    var options =
+        new ForkOptions(request.new_workflow_id)
+            .withApplicationVersion(request.application_version);
 
     logger.info("fork workflow {} step {}", wfid, startStep);
     var handle = dbosExecutor.forkWorkflow(wfid, startStep, options);

transact/src/main/java/dev/dbos/transact/conductor/Conductor.java

@@ -778,7 +778,7 @@ static CompletableFuture<BaseResponse> handleResume(Conductor conductor, BaseMes
                   ? request.workflow_ids
                   : List.of(request.workflow_id);
           try {
-            conductor.dbosExecutor.resumeWorkflows(ids);
+            conductor.dbosExecutor.resumeWorkflows(ids, request.queue_name);
             return new SuccessResponse(request, true);
           } catch (Exception e) {
             logger.error("Exception encountered when resuming workflow(s) {}", ids, e);
@@ -811,12 +811,9 @@ static CompletableFuture<BaseResponse> handleFork(Conductor conductor, BaseMessa
             return new ForkWorkflowResponse(request, null, "Invalid Fork Workflow Request");
           }
           try {
-            var options =
-                new ForkOptions(
-                    request.body.new_workflow_id, request.body.application_version, null);
             WorkflowHandle<?, ?> handle =
                 conductor.dbosExecutor.forkWorkflow(
-                    request.body.workflow_id, request.body.start_step, options);
+                    request.body.workflow_id, request.body.start_step, request.toOptions());
             return new ForkWorkflowResponse(request, handle.workflowId());
           } catch (Exception e) {
             logger.error("Exception encountered when forking workflow {}", request, e);

transact/src/main/java/dev/dbos/transact/conductor/protocol/ForkWorkflowRequest.java

@@ -1,5 +1,7 @@
 package dev.dbos.transact.conductor.protocol;
 
+import dev.dbos.transact.workflow.ForkOptions;
+
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 
 public class ForkWorkflowRequest extends BaseMessage {
@@ -28,5 +30,16 @@ public static class ForkWorkflowBody {
     public Integer start_step;
     public String application_version; // optional
     public String new_workflow_id; // optional
+    public String queue_name; // optional
+    public String queue_partition_key; // optional
+  }
+
+  public ForkOptions toOptions() {
+    return new ForkOptions(
+        body.new_workflow_id,
+        body.application_version,
+        null,
+        body.queue_name,
+        body.queue_partition_key);
   }
 }

transact/src/main/java/dev/dbos/transact/conductor/protocol/ResumeRequest.java

@@ -5,6 +5,7 @@
 public class ResumeRequest extends BaseMessage {
   public String workflow_id;
   public List<String> workflow_ids;
+  public String queue_name;
 
   public ResumeRequest() {}
 

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

@@ -396,8 +396,8 @@ public void cancelWorkflows(List<String> workflowIds) {
     dbRetry(() -> workflowDAO.cancelWorkflows(workflowIds));
   }
 
-  public void resumeWorkflows(List<String> workflowIds) {
-    dbRetry(() -> workflowDAO.resumeWorkflows(workflowIds));
+  public void resumeWorkflows(List<String> workflowIds, String queueName) {
+    dbRetry(() -> workflowDAO.resumeWorkflows(workflowIds, queueName));
   }
 
   public void deleteWorkflows(List<String> workflowIds, boolean deleteChildren) {

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

@@ -803,11 +803,12 @@ AND status NOT IN (?, ?)
     }
   }
 
-  void resumeWorkflows(List<String> workflowIds) throws SQLException {
+  void resumeWorkflows(List<String> workflowIds, String queueName) throws SQLException {
     List<String> filtered = filterNullsAndBlanks(workflowIds);
     if (filtered.isEmpty()) {
       return;
     }
+
     String sql =
         """
           UPDATE "%s".workflow_status
@@ -827,7 +828,7 @@ AND status NOT IN (?, ?)
       Array array = conn.createArrayOf("text", filtered.toArray(String[]::new));
       try {
         stmt.setString(1, WorkflowState.ENQUEUED.name());
-        stmt.setString(2, Constants.DBOS_INTERNAL_QUEUE);
+        stmt.setString(2, Objects.requireNonNullElse(queueName, Constants.DBOS_INTERNAL_QUEUE));
         stmt.setArray(3, array);
         stmt.setString(4, WorkflowState.SUCCESS.name());
         stmt.setString(5, WorkflowState.ERROR.name());
@@ -915,14 +916,11 @@ String forkWorkflow(String originalWorkflowId, int startStep, ForkOptions option
     }
 
     String forkedWorkflowId =
-        options.forkedWorkflowId() == null
-            ? UUID.randomUUID().toString()
-            : options.forkedWorkflowId();
+        Objects.requireNonNullElseGet(
+            options.forkedWorkflowId(), () -> UUID.randomUUID().toString());
 
     logger.debug("forkWorkflow Original id {} forked id {}", originalWorkflowId, forkedWorkflowId);
 
-    String applicationVersion = options.applicationVersion();
-
     var timeout = Objects.requireNonNullElseGet(options.timeout(), Timeout::inherit);
     Long timeoutMS = null;
     if (timeout instanceof Timeout.Inherit) {
@@ -941,8 +939,10 @@ String forkWorkflow(String originalWorkflowId, int startStep, ForkOptions option
             originalWorkflowId,
             forkedWorkflowId,
             status,
-            applicationVersion,
+            options.applicationVersion(),
             timeoutMS,
+            options.queueName(),
+            options.queuePartitionKey(),
             this.schema,
             this.serializer);
 
@@ -969,6 +969,8 @@ private static void insertForkedWorkflowStatus(
       WorkflowStatus originalStatus,
       String applicationVersion,
       Long timeoutMS,
+      String queueName,
+      String queuePartitionKey,
       String schema,
       DBOSSerializer serializer)
       throws SQLException {
@@ -978,8 +980,9 @@ private static void insertForkedWorkflowStatus(
         """
           INSERT INTO "%s".workflow_status (
             workflow_uuid, status, name, class_name, config_name, application_version, application_id,
-            authenticated_user, authenticated_roles, assumed_role, queue_name, inputs, workflow_timeout_ms, forked_from, serialization
-          ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            authenticated_user, authenticated_roles, assumed_role, queue_name, queue_partition_key, inputs,
+            workflow_timeout_ms, forked_from, serialization
+          ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
         """
             .formatted(schema);
 
@@ -998,15 +1001,16 @@ private static void insertForkedWorkflowStatus(
               ? null
               : JSONUtil.toJson(originalStatus.authenticatedRoles()));
       stmt.setString(10, originalStatus.assumedRole());
-      stmt.setString(11, Constants.DBOS_INTERNAL_QUEUE);
+      stmt.setString(11, Objects.requireNonNullElse(queueName, Constants.DBOS_INTERNAL_QUEUE));
+      stmt.setString(12, queuePartitionKey);
       stmt.setString(
-          12,
+          13,
           SerializationUtil.serializeArgs(
                   originalStatus.input(), null, originalStatus.serialization(), serializer)
               .serializedValue());
-      stmt.setObject(13, timeoutMS);
-      stmt.setString(14, originalWorkflowId);
-      stmt.setString(15, originalStatus.serialization());
+      stmt.setObject(14, timeoutMS);
+      stmt.setString(15, originalWorkflowId);
+      stmt.setString(16, originalStatus.serialization());
 
       stmt.executeUpdate();
     }