Get instance id for desired control-queue(s)

Test/DurableTask.AzureStorage.Tests/TestPartitionIndex.cs

@@ -0,0 +1,93 @@
+//  ----------------------------------------------------------------------------------
+//  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 System;
+    using System.Collections.Generic;
+    using System.Data;
+    using System.Threading;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+    [TestClass]
+    public class TestPartitionIndex
+    {
+        private AzureStorageOrchestrationService azureStorageOrchestrationService;
+        private AzureStorageOrchestrationServiceSettings settings;
+        private int partitionCount = 4;
+        private Dictionary<string, int> controlQueueNumberToNameMap;
+        private CancellationTokenSource cancellationTokenSource;
+        private const string TaskHub = "taskHubName";
+
+        [TestInitialize]
+        public void Initialize()
+        {
+            cancellationTokenSource = new CancellationTokenSource();
+
+            settings = new AzureStorageOrchestrationServiceSettings()
+            {
+                StorageConnectionString = TestHelpers.GetTestStorageAccountConnectionString(),
+                TaskHubName = TaskHub,
+                PartitionCount = partitionCount
+            };
+
+            azureStorageOrchestrationService = new AzureStorageOrchestrationService(settings);
+
+            controlQueueNumberToNameMap = new Dictionary<string, int>();
+
+            for (int i = 0; i < partitionCount; i++)
+            {
+                var controlQueueName = AzureStorageOrchestrationService.GetControlQueueName(settings.TaskHubName, i);
+                controlQueueNumberToNameMap[controlQueueName] = i;
+            }
+        }
+
+        [TestMethod]
+        [DataRow(20, false)]
+        [DataRow(20, true)]
+        public void GetPartitionIndexTest(int maxInstanceIdCount, bool enableExplicitPartitionPlacement)
+        {
+            settings.EnableExplicitPartitionPlacement = enableExplicitPartitionPlacement;
+
+            for (uint instanceIdSuffix = 0; instanceIdSuffix < settings.PartitionCount * 4; instanceIdSuffix++)
+            {
+                Dictionary<uint, int> indexNumberToCount = new Dictionary<uint, int>();
+
+                for (uint indexCount = 0; indexCount < settings.PartitionCount; indexCount++)
+                {
+                    indexNumberToCount[indexCount] = 0;
+                }
+
+                for (int instanceCount = 0; instanceCount < maxInstanceIdCount; instanceCount++)
+                {
+                    var instanceIdPrefix = Guid.NewGuid().ToString();
+
+                    var instanceId = $"{instanceIdPrefix}!{instanceIdSuffix}";
+
+                    var partitionIndex = azureStorageOrchestrationService.GetPartitionIndex(instanceId);
+
+                    indexNumberToCount[partitionIndex]++;
+                }
+
+                if (enableExplicitPartitionPlacement)
+                {
+                    Assert.AreEqual(indexNumberToCount[(uint)(instanceIdSuffix % settings.PartitionCount)], maxInstanceIdCount);
+                }
+                else
+                {
+                    Assert.AreNotEqual(indexNumberToCount[(uint)(instanceIdSuffix % settings.PartitionCount)], maxInstanceIdCount);
+                }
+            }
+        }
+    }
+}

src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs

@@ -2048,9 +2048,9 @@ public Task<string> DownloadBlobAsync(string blobUri)
 
         // TODO: Change this to a sticky assignment so that partition count changes can
         //       be supported: https://github.com/Azure/azure-functions-durable-extension/issues/1
-        async Task<ControlQueue?> GetControlQueueAsync(string instanceId)
+        internal async Task<ControlQueue?> GetControlQueueAsync(string instanceId)
         {
-            uint partitionIndex = Fnv1aHashHelper.ComputeHash(instanceId) % (uint)this.settings.PartitionCount;
+            uint partitionIndex = GetPartitionIndex(instanceId);
             string queueName = GetControlQueueName(this.settings.TaskHubName, (int)partitionIndex);
 
             ControlQueue cachedQueue;
@@ -2075,6 +2075,28 @@ public Task<string> DownloadBlobAsync(string blobUri)
             return cachedQueue;
         }
 
+        internal uint GetPartitionIndex(string instanceId)
+        {
+            uint totalPartitions = (uint)this.settings.PartitionCount;
+
+            int placementSeparatorPosition = instanceId.LastIndexOf('!');
+
+            // if the instance id ends with !nnn, where nnn is an unsigned number, it indicates explicit partition placement
+            if (
+                this.settings.EnableExplicitPartitionPlacement
+                && placementSeparatorPosition != -1
+                && uint.TryParse(instanceId.Substring(placementSeparatorPosition + 1), out uint index))
+            {
+                var partitionId = index % totalPartitions;
+                return (uint)partitionId;
+            }
+            else
+            {
+                return Fnv1aHashHelper.ComputeHash(instanceId) % totalPartitions;
+
+            }
+        }
+
         /// <summary>
         /// Disposes of the current object.
         /// </summary>

src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs

@@ -318,5 +318,11 @@ internal LogHelper Logger
         /// Consumers that require separate dispatch (such as the new out-of-proc v2 SDKs) must set this to true.
         /// </summary>
         public bool UseSeparateQueueForEntityWorkItems { get; set; } = false;
+
+        /// <summary>
+        /// Enabled explicit placement of instance to parition id.
+        /// if the instance id ends with !nnn, where nnn is an unsigned number, it indicates explicit partition placement
+        /// </summary>
+        public bool EnableExplicitPartitionPlacement { get; set; } = false;
     }
 }