Merge branch 'main' into vabachu/entities

Author: bachuv

CHANGELOG.md

@@ -1,4 +1,5 @@
 ## Unreleased
+* Add work item filtering support for `DurableTaskGrpcWorker` to enable worker-side filtering of orchestration and activity work items ([#275](https://github.com/microsoft/durabletask-java/pull/275)) 
 * Add support for calls to HTTP endpoints ([#271](https://github.com/microsoft/durabletask-java/pull/271))
 * Add getSuspendPostUri and getResumePostUri getters to HttpManagementPayload ([#264](https://github.com/microsoft/durabletask-java/pull/264))
 

client/src/main/java/com/microsoft/durabletask/DurableTaskGrpcWorker.java

@@ -48,12 +48,15 @@ public final class DurableTaskGrpcWorker implements AutoCloseable {
     private final DurableTaskGrpcWorkerVersioningOptions versioningOptions;
     private final int maxConcurrentEntityWorkItems;
     private final ExecutorService workItemExecutor;
+
+    private final WorkItemFilter workItemFilter;
+    private final GetWorkItemsRequest getWorkItemsRequest;
     private final PayloadHelper payloadHelper;
     private final int chunkSizeBytes;
 
     private final TaskHubSidecarServiceBlockingStub sidecarClient;
 
-    DurableTaskGrpcWorker(DurableTaskGrpcWorkerBuilder builder) {
+    DurableTaskGrpcWorker(DurableTaskGrpcWorkerBuilder builder, WorkItemFilter workItemFilter) {
         this.orchestrationFactories.putAll(builder.orchestrationFactories);
         this.activityFactories.putAll(builder.activityFactories);
         this.entityFactories.putAll(builder.entityFactories);
@@ -94,6 +97,8 @@ public final class DurableTaskGrpcWorker implements AutoCloseable {
                     return t;
                 },
                 new ThreadPoolExecutor.CallerRunsPolicy());
+        this.workItemFilter = workItemFilter;
+        this.getWorkItemsRequest = buildGetWorkItemsRequest();
         this.payloadHelper = builder.payloadStore != null
             ? new PayloadHelper(builder.payloadStore, builder.largePayloadOptions)
             : null;
@@ -173,16 +178,7 @@ public void startAndBlock() {
         // TODO: How do we interrupt manually?
         while (true) {
             try {
-                GetWorkItemsRequest.Builder requestBuilder = GetWorkItemsRequest.newBuilder();
-                if (!this.entityFactories.isEmpty()) {
-                    // Signal to the sidecar that this worker can handle entity work items
-                    requestBuilder.setMaxConcurrentEntityWorkItems(this.maxConcurrentEntityWorkItems);
-                }
-                if (this.payloadHelper != null) {
-                    requestBuilder.addCapabilities(WorkerCapability.WORKER_CAPABILITY_LARGE_PAYLOADS);
-                }
-                GetWorkItemsRequest getWorkItemsRequest = requestBuilder.build();
-                Iterator<WorkItem> workItemStream = this.sidecarClient.getWorkItems(getWorkItemsRequest);
+                Iterator<WorkItem> workItemStream = this.sidecarClient.getWorkItems(this.getWorkItemsRequest);
                 while (workItemStream.hasNext()) {
                     WorkItem workItem = workItemStream.next();
                     RequestCase requestType = workItem.getRequestCase();
@@ -605,6 +601,47 @@ public void stop() {
         this.close();
     }
 
+    /**
+     * Returns the work item filter configured for this worker, or {@code null} if none.
+     */
+    WorkItemFilter getWorkItemFilter() {
+        return this.workItemFilter;
+    }
+
+    private GetWorkItemsRequest buildGetWorkItemsRequest() {
+        GetWorkItemsRequest.Builder builder = GetWorkItemsRequest.newBuilder();
+        if (!this.entityFactories.isEmpty()) {
+            // Signal to the sidecar that this worker can handle entity work items
+            builder.setMaxConcurrentEntityWorkItems(this.maxConcurrentEntityWorkItems);
+        }
+        if (this.workItemFilter != null) {
+            builder.setWorkItemFilters(toProtoWorkItemFilters(this.workItemFilter));
+        }
+        if (this.payloadHelper != null) {
+            builder.addCapabilities(WorkerCapability.WORKER_CAPABILITY_LARGE_PAYLOADS);
+        }
+        return builder.build();
+    }
+
+    static WorkItemFilters toProtoWorkItemFilters(WorkItemFilter filter) {
+        WorkItemFilters.Builder builder = WorkItemFilters.newBuilder();
+        for (WorkItemFilter.OrchestrationFilter orch : filter.getOrchestrations()) {
+            com.microsoft.durabletask.implementation.protobuf.OrchestratorService.OrchestrationFilter.Builder orchBuilder =
+                    com.microsoft.durabletask.implementation.protobuf.OrchestratorService.OrchestrationFilter.newBuilder()
+                            .setName(orch.getName());
+            orchBuilder.addAllVersions(orch.getVersions());
+            builder.addOrchestrations(orchBuilder.build());
+        }
+        for (WorkItemFilter.ActivityFilter activity : filter.getActivities()) {
+            com.microsoft.durabletask.implementation.protobuf.OrchestratorService.ActivityFilter.Builder actBuilder =
+                    com.microsoft.durabletask.implementation.protobuf.OrchestratorService.ActivityFilter.newBuilder()
+                            .setName(activity.getName());
+            actBuilder.addAllVersions(activity.getVersions());
+            builder.addActivities(actBuilder.build());
+        }
+        return builder.build();
+    }
+    
     /**
      * Sends an orchestrator response, chunking it if it exceeds the configured chunk size.
      */

client/src/main/java/com/microsoft/durabletask/DurableTaskGrpcWorkerBuilder.java

@@ -6,7 +6,10 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Locale;
 
 /**
@@ -42,6 +45,8 @@ public final class DurableTaskGrpcWorkerBuilder {
     DurableTaskGrpcWorkerVersioningOptions versioningOptions;
     int maxConcurrentEntityWorkItems = 1;
     int maxWorkItemThreads;
+    private WorkItemFilter workItemFilter;
+    private boolean autoGenerateWorkItemFilters;
     PayloadStore payloadStore;
     LargePayloadOptions largePayloadOptions;
     int chunkSizeBytes = DEFAULT_CHUNK_SIZE_BYTES;
@@ -309,6 +314,24 @@ public DurableTaskGrpcWorkerBuilder useVersioning(DurableTaskGrpcWorkerVersionin
         return this;
     }
 
+    /**
+     * Sets explicit work item filters for this worker. When set, only work items matching the filters
+     * will be dispatched to this worker by the backend.
+     * <p>
+     * Work item filtering can improve efficiency in multi-worker deployments by ensuring each worker
+     * only receives work items it can handle. However, if an orchestration calls a task type
+     * (e.g., an activity or sub-orchestrator) that is not registered with any connected worker,
+     * the call may hang indefinitely instead of failing with an error.
+     *
+     * @param workItemFilter the work item filter to use, or {@code null} to disable filtering
+     * @return this builder object
+     */
+    public DurableTaskGrpcWorkerBuilder useWorkItemFilters(WorkItemFilter workItemFilter) {
+        this.workItemFilter = workItemFilter;
+        this.autoGenerateWorkItemFilters = false;
+        return this;
+    }
+
     /**
      * Enables large payload externalization with default options.
      * <p>
@@ -346,6 +369,27 @@ public DurableTaskGrpcWorkerBuilder useExternalizedPayloads(PayloadStore payload
         return this;
     }
 
+    /**
+     * Enables automatic work item filtering by generating filters from the registered
+     * orchestrations and activities. When enabled, the backend will only dispatch work items
+     * for registered orchestrations and activities to this worker.
+     * <p>
+     * Work item filtering can improve efficiency in multi-worker deployments by ensuring each worker
+     * only receives work items it can handle. However, if an orchestration calls a task type
+     * (e.g., an activity or sub-orchestrator) that is not registered with any connected worker,
+     * the call may hang indefinitely instead of failing with an error.
+     * <p>
+     * Only use this method when all task types referenced by orchestrations are guaranteed to be
+     * registered with at least one connected worker.
+     *
+     * @return this builder object
+     */
+    public DurableTaskGrpcWorkerBuilder useWorkItemFilters() {
+        this.autoGenerateWorkItemFilters = true;
+        this.workItemFilter = null;
+        return this;
+    }
+
     /**
      * Sets the maximum size in bytes for a single orchestrator response chunk sent over gRPC.
      * Responses larger than this will be automatically split into multiple chunks.
@@ -381,6 +425,31 @@ int getChunkSizeBytes() {
      * @return a new {@link DurableTaskGrpcWorker} object
      */
     public DurableTaskGrpcWorker build() {
-        return new DurableTaskGrpcWorker(this);
+        WorkItemFilter resolvedFilter = this.autoGenerateWorkItemFilters
+                ? buildAutoWorkItemFilter()
+                : this.workItemFilter;
+        return new DurableTaskGrpcWorker(this, resolvedFilter);
+    }
+
+    private WorkItemFilter buildAutoWorkItemFilter() {
+        List<String> versions = Collections.emptyList();
+        if (this.versioningOptions != null
+                && this.versioningOptions.getMatchStrategy() == DurableTaskGrpcWorkerVersioningOptions.VersionMatchStrategy.STRICT
+                && this.versioningOptions.getVersion() != null) {
+            versions = Collections.singletonList(this.versioningOptions.getVersion());
+        }
+
+        WorkItemFilter.Builder builder = WorkItemFilter.newBuilder();
+        List<String> orchestrationNames = new ArrayList<>(this.orchestrationFactories.keySet());
+        Collections.sort(orchestrationNames);
+        for (String name : orchestrationNames) {
+            builder.addOrchestration(name, versions);
+        }
+        List<String> activityNames = new ArrayList<>(this.activityFactories.keySet());
+        Collections.sort(activityNames);
+        for (String name : activityNames) {
+            builder.addActivity(name, versions);
+        }
+        return builder.build();
     }
 }