partition queues (#245)

fixes #206

Author: devhawk

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

@@ -114,7 +114,8 @@ public record EnqueueOptions(
       Duration timeout,
       Instant deadline,
       String deduplicationId,
-      Integer priority) {
+      Integer priority,
+      String queuePartitionKey) {
 
     public EnqueueOptions {
       if (Objects.requireNonNull(workflowName, "EnqueueOptions workflowName must not be null")
@@ -132,6 +133,16 @@ public record EnqueueOptions(
         throw new IllegalArgumentException("EnqueueOptions className must not be empty");
       }
 
+      if (queuePartitionKey != null && queuePartitionKey.isEmpty()) {
+        throw new IllegalArgumentException(
+            "EnqueueOptions queuePartitionKey must not be empty if not null");
+      }
+
+      if (deduplicationId != null && deduplicationId.isEmpty()) {
+        throw new IllegalArgumentException(
+            "EnqueueOptions deduplicationId must not be empty if not null");
+      }
+
       if (instanceName == null) instanceName = "";
 
       if (timeout != null) {
@@ -149,7 +160,7 @@ public record EnqueueOptions(
 
     /** Construct `EnqueueOptions` with a minimum set of required options */
     public EnqueueOptions(String className, String workflowName, String queueName) {
-      this(workflowName, queueName, className, "", null, null, null, null, null, null);
+      this(workflowName, queueName, className, "", null, null, null, null, null, null, null);
     }
 
     /**
@@ -169,7 +180,8 @@ public EnqueueOptions withClassName(String className) {
           this.timeout,
           this.deadline,
           this.deduplicationId,
-          this.priority);
+          this.priority,
+          this.queuePartitionKey);
     }
 
     /**
@@ -190,7 +202,8 @@ public EnqueueOptions withWorkflowId(String workflowId) {
           this.timeout,
           this.deadline,
           this.deduplicationId,
-          this.priority);
+          this.priority,
+          this.queuePartitionKey);
     }
 
     /**
@@ -211,7 +224,8 @@ public EnqueueOptions withAppVersion(String appVersion) {
           this.timeout,
           this.deadline,
           this.deduplicationId,
-          this.priority);
+          this.priority,
+          this.queuePartitionKey);
     }
 
     /**
@@ -232,7 +246,8 @@ public EnqueueOptions withTimeout(Duration timeout) {
           timeout,
           this.deadline,
           this.deduplicationId,
-          this.priority);
+          this.priority,
+          this.queuePartitionKey);
     }
 
     /**
@@ -253,7 +268,8 @@ public EnqueueOptions withDeadline(Instant deadline) {
           this.timeout,
           deadline,
           this.deduplicationId,
-          this.priority);
+          this.priority,
+          this.queuePartitionKey);
     }
 
     /**
@@ -274,7 +290,8 @@ public EnqueueOptions withDeduplicationId(String deduplicationId) {
           this.timeout,
           this.deadline,
           deduplicationId,
-          this.priority);
+          this.priority,
+          this.queuePartitionKey);
     }
 
     /**
@@ -295,7 +312,8 @@ public EnqueueOptions withInstanceName(String instName) {
           this.timeout,
           this.deadline,
           this.deduplicationId,
-          this.priority);
+          this.priority,
+          this.queuePartitionKey);
     }
 
     /**
@@ -315,7 +333,31 @@ public EnqueueOptions withPriority(Integer priority) {
           this.timeout,
           this.deadline,
           this.deduplicationId,
-          priority);
+          priority,
+          this.queuePartitionKey);
+    }
+
+    /**
+     * Creates a new EnqueueOptions instance with the specified queue partition key. The partition
+     * key is used to determine which partition of the queue the workflow should be enqueued to,
+     * allowing for better load distribution and ordering guarantees.
+     *
+     * @param partitionKey the partition key to use for queue partitioning, can be null
+     * @return a new EnqueueOptions instance with the specified partition key
+     */
+    public EnqueueOptions withQueuePartitionKey(String partitionKey) {
+      return new EnqueueOptions(
+          this.workflowName,
+          this.queueName,
+          this.className,
+          this.instanceName,
+          this.workflowId,
+          this.appVersion,
+          this.timeout,
+          this.deadline,
+          this.deduplicationId,
+          this.priority,
+          partitionKey);
     }
 
     /**
@@ -356,6 +398,7 @@ public <T, E extends Exception> WorkflowHandle<T, E> enqueueWorkflow(
                 options.queueName(), "EnqueueOptions queueName must not be null"),
             options.deduplicationId,
             options.priority,
+            options.queuePartitionKey,
             false,
             false),
         null,
@@ -379,8 +422,9 @@ public void send(String destinationId, Object message, String topic, String idem
     var workflowId = "%s-%s".formatted(destinationId, idempotencyKey);
 
     var status =
-        new WorkflowStatusInternal(workflowId, WorkflowState.SUCCESS)
-            .withName("temp_workflow-send-client");
+        WorkflowStatusInternal.builder(workflowId, WorkflowState.SUCCESS)
+            .name("temp_workflow-send-client")
+            .build();
     systemDatabase.initWorkflowStatus(status, null, false, false);
     systemDatabase.send(status.workflowId(), 0, destinationId, message, topic);
   }

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

@@ -23,14 +23,17 @@
  *     enqueued workflows.
  * @param priority If `queueName` is specified and refers to a queue with priority enabled, the
  *     priority to assign.
+ * @param queuePartitionKey If `queueName` is specified, an optional partition key used to
+ *     distribute workflows across queue partitions for load balancing and ordered processing.
  */
 public record StartWorkflowOptions(
     String workflowId,
     Timeout timeout,
     Instant deadline,
     String queueName,
     String deduplicationId,
-    Integer priority) {
+    Integer priority,
+    String queuePartitionKey) {
 
   public StartWorkflowOptions {
     if (timeout instanceof Timeout.Explicit explicit) {
@@ -44,16 +47,26 @@ public record StartWorkflowOptions(
             "StartWorkflowOptions explicit timeout and deadline cannot both be set");
       }
     }
+
+    if (queuePartitionKey != null && queuePartitionKey.isEmpty()) {
+      throw new IllegalArgumentException(
+          "EnqueueOptions queuePartitionKey must not be empty if not null");
+    }
+
+    if (deduplicationId != null && deduplicationId.isEmpty()) {
+      throw new IllegalArgumentException(
+          "EnqueueOptions deduplicationId must not be empty if not null");
+    }
   }
 
   /** Construct with default options */
   public StartWorkflowOptions() {
-    this(null, null, null, null, null, null);
+    this(null, null, null, null, null, null, null);
   }
 
   /** Construct with a specified workflow ID */
   public StartWorkflowOptions(String workflowId) {
-    this(workflowId, null, null, null, null, null);
+    this(workflowId, null, null, null, null, null, null);
   }
 
   /** Produces a new StartWorkflowOptions that overrides the ID assigned to the started workflow */
@@ -64,7 +77,8 @@ public StartWorkflowOptions withWorkflowId(String workflowId) {
         this.deadline,
         this.queueName,
         this.deduplicationId,
-        this.priority);
+        this.priority,
+        this.queuePartitionKey);
   }
 
   /** Produces a new StartWorkflowOptions that overrides timeout value for the started workflow */
@@ -75,7 +89,8 @@ public StartWorkflowOptions withTimeout(Timeout timeout) {
         this.deadline,
         this.queueName,
         this.deduplicationId,
-        this.priority);
+        this.priority,
+        this.queuePartitionKey);
   }
 
   /** Produces a new StartWorkflowOptions that overrides timeout value for the started workflow */
@@ -96,7 +111,8 @@ public StartWorkflowOptions withNoTimeout() {
         this.deadline,
         this.queueName,
         this.deduplicationId,
-        this.priority);
+        this.priority,
+        this.queuePartitionKey);
   }
 
   /** Produces a new StartWorkflowOptions that overrides deadline value for the started workflow */
@@ -107,13 +123,20 @@ public StartWorkflowOptions withDeadline(Instant deadline) {
         deadline,
         this.queueName,
         this.deduplicationId,
-        this.priority);
+        this.priority,
+        this.queuePartitionKey);
   }
 
   /** Produces a new StartWorkflowOptions that assigns the started workflow to a queue */
   public StartWorkflowOptions withQueue(String queue) {
     return new StartWorkflowOptions(
-        this.workflowId, this.timeout, this.deadline, queue, this.deduplicationId, this.priority);
+        this.workflowId,
+        this.timeout,
+        this.deadline,
+        queue,
+        this.deduplicationId,
+        this.priority,
+        this.queuePartitionKey);
   }
 
   /** Produces a new StartWorkflowOptions that assigns the started workflow to a queue */
@@ -132,7 +155,8 @@ public StartWorkflowOptions withDeduplicationId(String deduplicationId) {
         this.deadline,
         this.queueName,
         deduplicationId,
-        this.priority);
+        this.priority,
+        this.queuePartitionKey);
   }
 
   /**
@@ -146,7 +170,20 @@ public StartWorkflowOptions withPriority(Integer priority) {
         this.deadline,
         this.queueName,
         this.deduplicationId,
-        priority);
+        priority,
+        this.queuePartitionKey);
+  }
+
+  /** Produces a new StartWorkflowOptions that assigns a queue partition key */
+  public StartWorkflowOptions withQueuePartitionKey(String queuePartitionKey) {
+    return new StartWorkflowOptions(
+        this.workflowId,
+        this.timeout,
+        this.deadline,
+        this.queueName,
+        this.deduplicationId,
+        this.priority,
+        queuePartitionKey);
   }
 
   /** Get the assigned workflow ID, replacing empty with null */

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

@@ -9,15 +9,20 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.net.InetSocketAddress;
+import java.time.Duration;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
+import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.module.SimpleModule;
+import com.fasterxml.jackson.databind.ser.std.StdSerializer;
 import com.sun.net.httpserver.HttpExchange;
 import com.sun.net.httpserver.HttpHandler;
 import com.sun.net.httpserver.HttpServer;
@@ -28,6 +33,25 @@ public class AdminServer implements AutoCloseable {
   private static final Logger logger = LoggerFactory.getLogger(AdminServer.class);
   private static final ObjectMapper mapper = new ObjectMapper();
 
+  static class DurationSerializer extends StdSerializer<Duration> {
+
+    public DurationSerializer() {
+      super(Duration.class);
+    }
+
+    @Override
+    public void serialize(Duration value, JsonGenerator gen, SerializerProvider provider)
+        throws IOException {
+      gen.writeNumber(value.toMillis() / 1000.0);
+    }
+  }
+
+  static {
+    var module = new SimpleModule();
+    module.addSerializer(Duration.class, new DurationSerializer());
+    mapper.registerModule(module);
+  }
+
   private HttpServer server;
   private final SystemDatabase systemDatabase;
   private final DBOSExecutor dbosExecutor;

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

@@ -67,7 +67,7 @@ public WorkflowsOutput(WorkflowStatus status) {
         status.deadlineEpochMs() == null ? null : status.deadlineEpochMs().toString();
     this.DeduplicationID = status.deduplicationId();
     this.Priority = Objects.requireNonNullElse(status.priority(), 0).toString();
-    this.QueuePartitionKey = status.partitionKey();
+    this.QueuePartitionKey = status.queuePartitionKey();
     this.ForkedFrom = status.forkedFrom();
   }
 }