ChangeFeedProcessor: Fixes lease de-duplication for /partitionKey-partitioned lease containers

Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/LeaseManagement/DocumentServiceLeaseManagerCosmos.cs

@@ -31,6 +31,8 @@ internal sealed class DocumentServiceLeaseManagerCosmos : DocumentServiceLeaseMa
         private readonly AsyncLazy<TryCatch<string>> lazyContainerRid;
         private PartitionKeyRangeCache partitionKeyRangeCache;
 
+        internal static bool IsChangeFeedLeaseIdAsPartitionKeyEnabled = ConfigurationManager.IsChangeFeedLeaseIdAsPartitionKeyEnabled();
+
         public DocumentServiceLeaseManagerCosmos(
             ContainerInternal monitoredContainer,
             ContainerInternal leaseContainer,
@@ -128,7 +130,9 @@ public override Task<DocumentServiceLease> CreateLeaseIfNotExistAsync(
                 Mode = this.GetChangeFeedMode()
             };
 
-            this.requestOptionsFactory.AddPartitionKeyIfNeeded((string pk) => documentServiceLease.LeasePartitionKey = pk, Guid.NewGuid().ToString());
+            this.requestOptionsFactory.AddPartitionKeyIfNeeded(
+                (string pk) => documentServiceLease.LeasePartitionKey = pk,
+                DocumentServiceLeaseManagerCosmos.GetLeasePartitionKeyValue(documentServiceLease.LeaseId));
 
             return this.TryCreateDocumentServiceLeaseAsync(documentServiceLease);
         }
@@ -153,11 +157,20 @@ public override Task<DocumentServiceLease> CreateLeaseIfNotExistAsync(
                 Mode = this.GetChangeFeedMode()
             };
 
-            this.requestOptionsFactory.AddPartitionKeyIfNeeded((string pk) => documentServiceLease.LeasePartitionKey = pk, Guid.NewGuid().ToString());
+            this.requestOptionsFactory.AddPartitionKeyIfNeeded(
+                (string pk) => documentServiceLease.LeasePartitionKey = pk,
+                DocumentServiceLeaseManagerCosmos.GetLeasePartitionKeyValue(documentServiceLease.LeaseId));
 
             return this.TryCreateDocumentServiceLeaseAsync(documentServiceLease);
         }
 
+        private static string GetLeasePartitionKeyValue(string leaseId)
+        {
+            return DocumentServiceLeaseManagerCosmos.IsChangeFeedLeaseIdAsPartitionKeyEnabled
+                ? leaseId
+                : Guid.NewGuid().ToString();
+        }
+
         private string GetChangeFeedMode()
         {
             return this.options.Mode == ChangeFeedMode.AllVersionsAndDeletes

Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs

@@ -146,6 +146,11 @@ internal static class ConfigurationManager
         /// </summary>
         internal static readonly string Http2KeepAlivePingTimeoutInSeconds = "AZURE_COSMOS_HTTP2_KEEPALIVE_PING_TIMEOUT_IN_SECONDS";
 
+        /// <summary>
+        /// Environment variable name to enable deterministic lease-id partition key values for Change Feed lease creation.
+        /// </summary>
+        internal static readonly string ChangeFeedLeaseIdAsPartitionKeyEnabled = "AZURE_COSMOS_CHANGE_FEED_LEASE_ID_AS_PARTITION_KEY_ENABLED";
+
         public static T GetEnvironmentVariable<T>(string variable, T defaultValue)
         {
             string value = Environment.GetEnvironmentVariable(variable);
@@ -212,7 +217,19 @@ public static bool IsThinClientEnabled(
                     .GetEnvironmentVariable(
                         variable: ConfigurationManager.ThinClientModeEnabled,
                         defaultValue: defaultValue);
-        }
+        }
+
+        /// <summary>
+        /// Gets the boolean value indicating whether Change Feed lease creation should use lease id as the partition key value.
+        /// </summary>
+        /// <returns>A boolean flag indicating if deterministic lease-id partition key behavior is enabled.</returns>
+        public static bool IsChangeFeedLeaseIdAsPartitionKeyEnabled()
+        {
+            return ConfigurationManager
+                    .GetEnvironmentVariable(
+                        variable: ConfigurationManager.ChangeFeedLeaseIdAsPartitionKeyEnabled,
+                        defaultValue: true);
+        }
 
         /// <summary>
         /// Gets the AAD scope value to override.

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ChangeFeed/GremlinSmokeTests.cs

@@ -8,10 +8,12 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests.ChangeFeed
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.ChangeFeed.LeaseManagement;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
 
     [TestClass]
     [TestCategory("ChangeFeed")]
+    [DoNotParallelize]
     public class GremlinSmokeTests : BaseChangeFeedClientHelper
     {
         private Container Container;
@@ -75,39 +77,64 @@ public async Task WritesTriggerDelegate_WithLeaseContainer()
 
         }
 
-        [TestMethod]
-        public async Task Schema_DefaultsToHavingPartitionKey()
+        [DataTestMethod]
+        [DataRow(true, DisplayName = "Default (deterministic PK): partition key equals lease id")]
+        [DataRow(false, DisplayName = "Opt-out (legacy GUID PK): partition key differs from lease id")]
+        public async Task Schema_DefaultsToHavingPartitionKey(bool isChangeFeedLeaseIdAsPartitionKeyEnabled)
         {
-            ChangeFeedProcessor processor = this.Container
-                .GetChangeFeedProcessorBuilder("test", (IReadOnlyCollection<TestClass> docs, CancellationToken token) => Task.CompletedTask)
-                .WithInstanceName("random")
-                .WithLeaseContainer(this.LeaseContainer).Build();
-
-            await processor.StartAsync();
-            // Letting processor initialize
-            await Task.Delay(BaseChangeFeedClientHelper.ChangeFeedSetupTime);
+            DocumentServiceLeaseManagerCosmos.IsChangeFeedLeaseIdAsPartitionKeyEnabled = isChangeFeedLeaseIdAsPartitionKeyEnabled;
 
-            // Verify that leases have the partitionKey attribute
-            using FeedIterator<dynamic> iterator = this.LeaseContainer.GetItemQueryIterator<dynamic>();
-            while (iterator.HasMoreResults)
+            try
             {
-                FeedResponse<dynamic> page = await iterator.ReadNextAsync();
-                foreach (dynamic lease in page)
+                ChangeFeedProcessor processor = this.Container
+                    .GetChangeFeedProcessorBuilder("test", (IReadOnlyCollection<TestClass> docs, CancellationToken token) => Task.CompletedTask)
+                    .WithInstanceName("random")
+                    .WithLeaseContainer(this.LeaseContainer).Build();
+
+                await processor.StartAsync();
+                // Letting processor initialize
+                await Task.Delay(BaseChangeFeedClientHelper.ChangeFeedSetupTime);
+
+                // Verify that leases have the partitionKey attribute
+                using FeedIterator<dynamic> iterator = this.LeaseContainer.GetItemQueryIterator<dynamic>();
+                while (iterator.HasMoreResults)
                 {
-                    string leaseId = lease.id;
-                    Assert.IsNotNull(lease.partitionKey);
-                    if (leaseId.Contains(".info") || leaseId.Contains(".lock"))
+                    FeedResponse<dynamic> page = await iterator.ReadNextAsync();
+                    foreach (dynamic lease in page)
                     {
-                        // These are the store initialization marks
-                        continue;
+                        string leaseId = lease.id;
+                        Assert.IsNotNull(lease.partitionKey);
+                        if (leaseId.Contains(".info") || leaseId.Contains(".lock"))
+                        {
+                            // These are the store initialization marks
+                            continue;
+                        }
+
+                        if (isChangeFeedLeaseIdAsPartitionKeyEnabled)
+                        {
+                            Assert.AreEqual((string)lease.id, (string)lease.partitionKey,
+                                "Deterministic mode: lease partition key must equal lease id for dedup invariant.");
+                        }
+                        else
+                        {
+                            Assert.AreNotEqual((string)lease.id, (string)lease.partitionKey,
+                                "Legacy mode: lease partition key should be a random GUID, not the lease id.");
+                            Assert.IsTrue(System.Guid.TryParse((string)lease.partitionKey, out _),
+                                "Legacy mode: partition key should be a valid GUID.");
+                        }
+
+                        Assert.IsNotNull(lease.LeaseToken);
+                        Assert.IsNull(lease.PartitionId);
                     }
-
-                    Assert.IsNotNull(lease.LeaseToken);
-                    Assert.IsNull(lease.PartitionId);
                 }
-            }
 
-            await processor.StopAsync();
+                await processor.StopAsync();
+            }
+            finally
+            {
+                DocumentServiceLeaseManagerCosmos.IsChangeFeedLeaseIdAsPartitionKeyEnabled =
+                    ConfigurationManager.IsChangeFeedLeaseIdAsPartitionKeyEnabled();
+            }
         }
 
         public class TestClass