feat(AzureStorage): support orchestrator-only and activity-only worke…

src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs

@@ -276,11 +276,13 @@ public BehaviorOnContinueAsNew EventBehaviourForContinueAsNew
         }
 
         // We always leave the dispatcher counts at one unless we can find a customer workload that requires more.
+        // A count of zero disables the corresponding dispatcher entirely, which is how a worker is restricted
+        // to only orchestrations or only activities via <see cref="WorkerDispatchMode"/>.
         /// <inheritdoc />
-        public int TaskActivityDispatcherCount { get; } = 1;
+        public int TaskActivityDispatcherCount => this.settings.WorkerDispatchMode == WorkerDispatchMode.Orchestrator ? 0 : 1;
 
         /// <inheritdoc />
-        public int TaskOrchestrationDispatcherCount { get; } = 1;
+        public int TaskOrchestrationDispatcherCount => this.settings.WorkerDispatchMode == WorkerDispatchMode.Activity ? 0 : 1;
 
         #region IEntityOrchestrationService
 
@@ -462,7 +464,12 @@ public async Task StartAsync()
             this.shutdownSource = new CancellationTokenSource();
             this.statsLoop = Task.Run(() => this.ReportStatsLoop(this.shutdownSource.Token));
 
-            await this.appLeaseManager.StartAsync();
+            // Activity-only workers never lease a control-queue partition, so the partition manager
+            // (started via the app lease manager) is skipped entirely.
+            if (this.settings.WorkerDispatchMode != WorkerDispatchMode.Activity)
+            {
+                await this.appLeaseManager.StartAsync();
+            }
 
             this.isStarted = true;
         }
@@ -487,7 +494,12 @@ public async Task StopAsync(bool isForced)
                 this.orchestrationSessionManager.AbortAllSessions();
             }
 
-            await this.appLeaseManager.StopAsync();
+            // Symmetric with StartAsync: activity-only workers never started the partition manager.
+            if (this.settings.WorkerDispatchMode != WorkerDispatchMode.Activity)
+            {
+                await this.appLeaseManager.StopAsync();
+            }
+
             this.isStarted = false;
         }
 
@@ -690,6 +702,20 @@ static TaskHubInfo GetTaskHubInfo(string taskHub, int partitionCount)
 
         #endregion
 
+        // When a dispatcher is disabled for this worker's WorkerDispatchMode, its fetch loop (if it runs at all)
+        // should idle rather than spin. This waits for the normal polling interval and swallows cancellation.
+        async Task BackoffAsync(CancellationToken cancellationToken)
+        {
+            try
+            {
+                await Task.Delay(this.settings.MaxQueuePollingInterval, cancellationToken);
+            }
+            catch (OperationCanceledException)
+            {
+                // Shutting down; fall through to return no work item.
+            }
+        }
+
         #region Orchestration Work Item Methods
         /// <inheritdoc />
         public Task<TaskOrchestrationWorkItem> LockNextTaskOrchestrationWorkItemAsync(
@@ -705,6 +731,13 @@ public Task<TaskOrchestrationWorkItem> LockNextTaskOrchestrationWorkItemAsync(
 
         async Task<TaskOrchestrationWorkItem> LockNextTaskOrchestrationWorkItemAsync(bool entitiesOnly, CancellationToken cancellationToken)
         {
+            // Safety net: an activity-only worker must never dispatch orchestrations or entities.
+            if (this.settings.WorkerDispatchMode == WorkerDispatchMode.Activity)
+            {
+                await this.BackoffAsync(cancellationToken);
+                return null;
+            }
+
             Guid traceActivityId = StartNewLogicalTraceScope(useExisting: true);
 
             await this.EnsureTaskHubAsync();
@@ -1540,6 +1573,13 @@ public async Task<TaskActivityWorkItem> LockNextTaskActivityWorkItem(
             TimeSpan receiveTimeout,
             CancellationToken cancellationToken)
         {
+            // Safety net: an orchestrator-only worker must never dequeue the work-item queue.
+            if (this.settings.WorkerDispatchMode == WorkerDispatchMode.Orchestrator)
+            {
+                await this.BackoffAsync(cancellationToken);
+                return null;
+            }
+
             await this.EnsureTaskHubAsync();
 
             using (var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, this.shutdownSource.Token))

src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs

@@ -295,6 +295,13 @@ internal LogHelper Logger
         /// </summary>
         public QueueClientMessageEncoding QueueClientMessageEncoding { get; set; } = QueueClientMessageEncoding.UTF8;
 
+        /// <summary>
+        /// Gets or sets which kinds of work this worker instance dispatches, allowing orchestrations
+        /// and activities to be scaled independently across separate deployments on the same task hub.
+        /// The default is <see cref="WorkerDispatchMode.Both"/>, which preserves the historical behavior.
+        /// </summary>
+        public WorkerDispatchMode WorkerDispatchMode { get; set; } = WorkerDispatchMode.Both;
+
         /// <summary>
         /// When true, an etag is used when attempting to make instance table updates upon completing an orchestration work item.
         /// </summary>

src/DurableTask.AzureStorage/WorkerDispatchMode.cs

@@ -0,0 +1,44 @@
+//  ----------------------------------------------------------------------------------
+//  Copyright Microsoft Corporation
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//  http://www.apache.org/licenses/LICENSE-2.0
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//  ----------------------------------------------------------------------------------
+
+namespace DurableTask.AzureStorage
+{
+    /// <summary>
+    /// Specifies which kinds of work a single worker instance dispatches, enabling orchestrations
+    /// and activities to be scaled independently across separate worker deployments on the same task hub.
+    /// </summary>
+    /// <remarks>
+    /// This is a single, self-contained value so it can be surfaced from a single environment variable
+    /// or configuration setting without any code changes.
+    /// </remarks>
+    public enum WorkerDispatchMode
+    {
+        /// <summary>
+        /// The worker dispatches both orchestrations (including entities) and activities.
+        /// This is the default and matches the historical behavior.
+        /// </summary>
+        Both = 0,
+
+        /// <summary>
+        /// The worker leases control-queue partitions and dispatches only orchestrations and entities.
+        /// It does not dequeue the work-item queue, so it runs no activities.
+        /// </summary>
+        Orchestrator = 1,
+
+        /// <summary>
+        /// The worker dispatches only activities from the work-item queue.
+        /// It does not lease any control-queue partition, so it runs no orchestrations or entities.
+        /// </summary>
+        Activity = 2,
+    }
+}

src/DurableTask.Core/TaskHubWorker.cs

@@ -296,6 +296,18 @@ public async Task<TaskHubWorker> StartAsync()
 
                 this.logHelper.TaskHubWorkerStarting();
                 var sw = Stopwatch.StartNew();
+
+                // Read the dispatcher counts here (rather than caching them in the constructor) so a backend
+                // that derives them from mutable settings is honored if they changed after construction.
+                // A count of zero disables that dispatcher for this worker; entities ride with orchestrations.
+                bool dispatchOrchestrations = this.orchestrationService.TaskOrchestrationDispatcherCount > 0;
+                bool dispatchActivities = this.orchestrationService.TaskActivityDispatcherCount > 0;
+                bool dispatchEntities = this.dispatchEntitiesSeparately && dispatchOrchestrations;
+
+                // Dispatcher objects are always constructed so the public TaskOrchestrationDispatcher/
+                // TaskActivityDispatcher properties stay non-null for callers that configure them after
+                // StartAsync (e.g. IncludeDetails). Whether each one actually polls is gated below on
+                // StartAsync, and a disabled dispatcher (dispatcher count 0) starts zero fetch loops.
                 this.orchestrationDispatcher = new TaskOrchestrationDispatcher(
                     this.orchestrationService,
                     this.orchestrationManager,
@@ -324,10 +336,18 @@ public async Task<TaskHubWorker> StartAsync()
                 }
 
                 await this.orchestrationService.StartAsync();
-                await this.orchestrationDispatcher.StartAsync();
-                await this.activityDispatcher.StartAsync();
 
-                if (this.dispatchEntitiesSeparately)
+                if (dispatchOrchestrations)
+                {
+                    await this.orchestrationDispatcher.StartAsync();
+                }
+
+                if (dispatchActivities)
+                {
+                    await this.activityDispatcher.StartAsync();
+                }
+
+                if (dispatchEntities)
                 {
                     await this.entityDispatcher.StartAsync();
                 }

test/DurableTask.AzureStorage.Tests/DecoupledWorkerDispatchTests.cs

@@ -0,0 +1,187 @@
+//  ----------------------------------------------------------------------------------
+//  Copyright Microsoft Corporation
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//  http://www.apache.org/licenses/LICENSE-2.0
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//  ----------------------------------------------------------------------------------
+
+namespace DurableTask.AzureStorage.Tests
+{
+    using DurableTask.AzureStorage;
+    using DurableTask.Core;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using System;
+    using System.Linq;
+    using System.Threading.Tasks;
+
+    [TestClass]
+    public class DecoupledWorkerDispatchTests
+    {
+        static readonly string TestConnectionString = TestHelpers.GetTestStorageAccountConnectionString();
+
+        // ---------------------------------------------------------------------------------------
+        // Unit tests: the WorkerDispatchMode setting maps onto the IOrchestrationService dispatcher
+        // counts, which is the signal TaskHubWorker uses to gate each dispatcher. No storage needed.
+        // ---------------------------------------------------------------------------------------
+
+        [TestMethod]
+        public void DefaultMode_IsBoth_AndRunsBothDispatchers()
+        {
+            var service = CreateService("DispatchModeDefault", mode: null);
+
+            Assert.AreEqual(WorkerDispatchMode.Both, new AzureStorageOrchestrationServiceSettings().WorkerDispatchMode);
+            Assert.AreEqual(1, service.TaskOrchestrationDispatcherCount);
+            Assert.AreEqual(1, service.TaskActivityDispatcherCount);
+        }
+
+        [TestMethod]
+        public void BothMode_RunsBothDispatchers()
+        {
+            var service = CreateService("DispatchModeBoth", WorkerDispatchMode.Both);
+
+            Assert.AreEqual(1, service.TaskOrchestrationDispatcherCount);
+            Assert.AreEqual(1, service.TaskActivityDispatcherCount);
+        }
+
+        [TestMethod]
+        public void OrchestratorMode_DisablesActivityDispatcher()
+        {
+            var service = CreateService("DispatchModeOrch", WorkerDispatchMode.Orchestrator);
+
+            Assert.AreEqual(1, service.TaskOrchestrationDispatcherCount);
+            Assert.AreEqual(0, service.TaskActivityDispatcherCount);
+        }
+
+        [TestMethod]
+        public void ActivityMode_DisablesOrchestrationDispatcher()
+        {
+            var service = CreateService("DispatchModeAct", WorkerDispatchMode.Activity);
+
+            Assert.AreEqual(0, service.TaskOrchestrationDispatcherCount);
+            Assert.AreEqual(1, service.TaskActivityDispatcherCount);
+        }
+
+        static AzureStorageOrchestrationService CreateService(string taskHub, WorkerDispatchMode? mode)
+        {
+            var settings = new AzureStorageOrchestrationServiceSettings
+            {
+                TaskHubName = TestHelpers.GetTestTaskHubName() + taskHub,
+                StorageAccountClientProvider = new StorageAccountClientProvider(TestConnectionString),
+            };
+
+            if (mode.HasValue)
+            {
+                settings.WorkerDispatchMode = mode.Value;
+            }
+
+            return new AzureStorageOrchestrationService(settings);
+        }
+
+        // ---------------------------------------------------------------------------------------
+        // Integration test (Azurite): two workers on the SAME task hub, one Orchestrator-only and
+        // one Activity-only. The orchestration (which calls an activity) must run to completion, and
+        // the roles must be decoupled: the orchestrator worker never dequeues the work-item queue and
+        // the activity worker never leases a control-queue partition.
+        // ---------------------------------------------------------------------------------------
+
+        [TestMethod]
+        public async Task DecoupledWorkers_RunOrchestrationToCompletion()
+        {
+            string taskHub = TestHelpers.GetTestTaskHubName() + "Decoupled";
+            const string input = "world";
+
+            var orchestratorSettings = new AzureStorageOrchestrationServiceSettings
+            {
+                TaskHubName = taskHub,
+                StorageAccountClientProvider = new StorageAccountClientProvider(TestConnectionString),
+                WorkerDispatchMode = WorkerDispatchMode.Orchestrator,
+                WorkerId = "orchestrator-worker",
+            };
+
+            var activitySettings = new AzureStorageOrchestrationServiceSettings
+            {
+                TaskHubName = taskHub,
+                StorageAccountClientProvider = new StorageAccountClientProvider(TestConnectionString),
+                WorkerDispatchMode = WorkerDispatchMode.Activity,
+                WorkerId = "activity-worker",
+            };
+
+            var orchestratorService = new AzureStorageOrchestrationService(orchestratorSettings);
+            var activityService = new AzureStorageOrchestrationService(activitySettings);
+
+            // Sanity check on the mode gating before starting anything.
+            Assert.AreEqual(1, orchestratorService.TaskOrchestrationDispatcherCount);
+            Assert.AreEqual(0, orchestratorService.TaskActivityDispatcherCount);
+            Assert.AreEqual(0, activityService.TaskOrchestrationDispatcherCount);
+            Assert.AreEqual(1, activityService.TaskActivityDispatcherCount);
+
+            // Both deployments ship the same code, so both workers register the same orchestration and
+            // activity. Each worker only dispatches the kind of work allowed by its mode.
+            var orchestratorWorker = new TaskHubWorker(orchestratorService);
+            orchestratorWorker.AddTaskOrchestrations(typeof(DecoupledHelloOrchestrator));
+            orchestratorWorker.AddTaskActivities(typeof(DecoupledHello));
+
+            var activityWorker = new TaskHubWorker(activityService);
+            activityWorker.AddTaskOrchestrations(typeof(DecoupledHelloOrchestrator));
+            activityWorker.AddTaskActivities(typeof(DecoupledHello));
+
+            var client = new TaskHubClient(orchestratorService);
+
+            await orchestratorWorker.StartAsync();
+            await activityWorker.StartAsync();
+
+            try
+            {
+                OrchestrationInstance instance = await client.CreateOrchestrationInstanceAsync(
+                    typeof(DecoupledHelloOrchestrator), input);
+
+                OrchestrationState state = await client.WaitForOrchestrationAsync(instance, TimeSpan.FromSeconds(90));
+
+                // The orchestrator-only worker cannot dequeue activities and the activity-only worker
+                // cannot dequeue orchestrations, so completion proves the two roles cooperated.
+                Assert.IsNotNull(state);
+                Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus);
+                Assert.AreEqual($"\"Hello, {input}!\"", state.Output);
+
+                // The activity worker must never have leased a control-queue partition...
+                Assert.AreEqual(0, activityService.OwnedControlQueues.Count(), "Activity-only worker should not lease any control-queue partition.");
+                // ...while the orchestrator worker owns at least one partition.
+                await TestHelpers.WaitFor(() => orchestratorService.OwnedControlQueues.Any(), TimeSpan.FromSeconds(30));
+                Assert.IsTrue(orchestratorService.OwnedControlQueues.Any(), "Orchestrator-only worker should lease control-queue partitions.");
+            }
+            finally
+            {
+                await orchestratorWorker.StopAsync(isForced: true);
+                await activityWorker.StopAsync(isForced: true);
+                await orchestratorService.DeleteAsync();
+            }
+        }
+
+        internal class DecoupledHelloOrchestrator : TaskOrchestration<string, string>
+        {
+            public override Task<string> RunTask(OrchestrationContext context, string input)
+            {
+                return context.ScheduleTask<string>(typeof(DecoupledHello), input);
+            }
+        }
+
+        internal class DecoupledHello : TaskActivity<string, string>
+        {
+            protected override string Execute(TaskContext context, string input)
+            {
+                if (string.IsNullOrEmpty(input))
+                {
+                    throw new ArgumentNullException(nameof(input));
+                }
+
+                return $"Hello, {input}!";
+            }
+        }
+    }
+}