Change Feed Processor: Fixes duplicate lease docs crashing load balancer on partition split

When the lease container is partitioned by /partitionKey, PartitionSynchronizerCore.HandlePartitionGoneAsync creates child lease documents using Guid.NewGuid().ToString() as the partitionKey value via DocumentServiceLeaseManagerCosmos.CreateLeaseIfNotExistAsync. Because each split retry (host restart mid-split, concurrent hosts, or error retries) picks a new Guid, TryCreateItemAsync's per-partition-key id-uniqueness check never catches cross-PK duplicates. Once duplicates exist, EqualPartitionsBalancingStrategy.CategorizeLeases throws from Dictionary.Add on every balance tick, blocking all lease acquisition for the container until the extra documents are manually deleted (IcM 768856224).

Two fixes:

1. EqualPartitionsBalancingStrategy.CategorizeLeases: tolerate duplicate CurrentLeaseToken entries (keep the first, log a warning with the conflicting ids/PKs and remediation pointer). This unblocks the load balancer even when duplicates already exist in a customer's lease container.

2. PartitionSynchronizerCore.HandlePartitionGoneAsync (both DocumentServiceLeaseCore and DocumentServiceLeaseCoreEpk overloads): pre-query existing leases once and skip CreateLeaseIfNotExistAsync for ranges whose child lease token is already present, mirroring the dedup in CreateLeasesAsync.

Tests added:

- EqualPartitionsBalancingStrategyTests.CalculateLeasesToTake_DuplicateLeaseTokens_DoesNotThrow

- PartitionSynchronizerCoreTests.HandlePartitionGoneAsync_PKRangeBasedLease_Split_DoesNotCreateDuplicateChildLeases

- PartitionSynchronizerCoreTests.HandlePartitionGoneAsync_EpkBasedLease_Split_DoesNotCreateDuplicateChildLeases

- PartitionSynchronizerCoreTests.HandlePartitionGoneAsync_PKRangeBasedLease_Split_CreatesOnlyMissingChildLeases

- Microsoft.Azure.Cosmos.EmulatorTests.ChangeFeed.DuplicateLeaseRegressionTests (emulator-based end-to-end regression)

Existing HandlePartitionGoneAsync tests were updated to mock GetAllLeasesAsync (via a CreateEmptyLeaseContainer helper) since the synchronizer now consults the lease container before creating children.

Out of scope: long-term deterministic partition-key derivation from LeaseToken. That change is backward-incompatible for existing lease containers (customers have GUID-based PKs on all current lease docs) and will need a migration story; deferred to follow-up work.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: NaluTripician

Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/Bootstrapping/PartitionSynchronizerCore.cs

@@ -84,29 +84,75 @@ public override async Task CreateMissingLeasesAsync()
                 throw new InvalidOperationException();
             }
 
+            // Pre-check existing leases so that retried split handling on the same host (e.g. a retry
+            // after a transient failure within this call) does not create duplicate child lease
+            // documents. This matters especially when the lease container is partitioned by
+            // /partitionKey: each retry picks a new random Guid partition-key value, so the
+            // per-partition-key id uniqueness check on TryCreateItemAsync does not catch cross-PK
+            // duplicates. Duplicates would later crash EqualPartitionsBalancingStrategy.CategorizeLeases
+            // and block all lease acquisition.
+            //
+            // Note: this is a TOCTOU pre-check, not a hard guarantee. Two hosts racing into
+            // HandlePartitionGoneAsync for the same parent lease can still both observe the children
+            // as missing and both proceed to create. That residual race is handled by the
+            // duplicate-tolerant logic in EqualPartitionsBalancingStrategy.CategorizeLeases, which
+            // keeps the first lease seen per token and logs a warning instead of throwing.
+            IReadOnlyList<DocumentServiceLease> existingLeases = await this.leaseContainer.GetAllLeasesAsync().ConfigureAwait(false);
+
+            HashSet<string> existingPkRangeLeaseTokens = new HashSet<string>(
+                existingLeases.Where(l => l is DocumentServiceLeaseCore).Select(l => l.CurrentLeaseToken),
+                StringComparer.Ordinal);
+            HashSet<string> existingEpkLeaseTokens = new HashSet<string>(
+                existingLeases.Where(l => l is DocumentServiceLeaseCoreEpk).Select(l => l.CurrentLeaseToken),
+                StringComparer.Ordinal);
+
             return lease switch
             {
-                DocumentServiceLeaseCoreEpk feedRangeBaseLease => await this.HandlePartitionGoneAsync(leaseToken, lastContinuationToken, feedRangeBaseLease, overlappingRanges),
-                _ => await this.HandlePartitionGoneAsync(leaseToken, lastContinuationToken, (DocumentServiceLeaseCore)lease, overlappingRanges)
+                DocumentServiceLeaseCoreEpk feedRangeBaseLease => await this.HandlePartitionGoneAsync(leaseToken, lastContinuationToken, feedRangeBaseLease, overlappingRanges, existingEpkLeaseTokens),
+                _ => await this.HandlePartitionGoneAsync(leaseToken, lastContinuationToken, (DocumentServiceLeaseCore)lease, overlappingRanges, existingPkRangeLeaseTokens, existingEpkLeaseTokens)
             };
         }
 
         /// <summary>
         /// Handles splits and merges for partition based leases.
         /// </summary>
+        /// <remarks>
+        /// When every child lease for a split is already present, this method returns an empty
+        /// sequence and <c>shouldDeleteGoneLease = true</c>. The caller
+        /// (<see cref="FeedManagement.PartitionControllerCore"/>) therefore deletes the parent
+        /// lease but does not start observers for the existing children on this host; those leases
+        /// are picked up by the next load-balancer tick in
+        /// <see cref="FeedManagement.EqualPartitionsBalancingStrategy.SelectLeasesToTake"/>. This
+        /// is intentional — the existing children may already be owned by another host, and the
+        /// previous behaviour of eagerly grabbing them would have propagated this host's parent
+        /// lease properties onto leases it did not create.
+        /// </remarks>
         private async Task<(IEnumerable<DocumentServiceLease>, bool)> HandlePartitionGoneAsync(
             string leaseToken,
             string lastContinuationToken,
             DocumentServiceLeaseCore partitionBasedLease,
-            IReadOnlyList<PartitionKeyRange> overlappingRanges)
+            IReadOnlyList<PartitionKeyRange> overlappingRanges,
+            HashSet<string> existingPkRangeLeaseTokens,
+            HashSet<string> existingEpkLeaseTokens)
         {
             ConcurrentQueue<DocumentServiceLease> newLeases = new ConcurrentQueue<DocumentServiceLease>();
             if (overlappingRanges.Count > 1)
             {
-                // Split: More than two children
+                // Split: More than two children. The children this overload creates are PK-range
+                // based (DocumentServiceLeaseCore with LeaseToken == PartitionKeyRange.Id), so the
+                // dedup compares against the PK-range token set. This is deliberately a stricter
+                // match than the heuristic used by CreateLeasesAsync (which also treats overlapping
+                // EPK leases as "covered") because at this point the children's partition-key-range
+                // ids are known exactly from the split result.
                 await overlappingRanges.ForEachAsync(
                     async addedRange =>
                     {
+                        if (existingPkRangeLeaseTokens.Contains(addedRange.Id))
+                        {
+                            DefaultTrace.TraceInformation("Skipping creation of child lease for range '{0}'; a lease with that token already exists.", addedRange.Id);
+                            return;
+                        }
+
                         DocumentServiceLease newLease = await this.leaseManager.CreateLeaseIfNotExistAsync(addedRange, lastContinuationToken);
                         if (newLease != null)
                         {
@@ -119,14 +165,27 @@ await overlappingRanges.ForEachAsync(
             }
             else
             {
-                // Merge: 1 children, multiple ranges merged into 1
+                // Merge: 1 children, multiple ranges merged into 1. The new lease created for a
+                // merged range is EPK-based (DocumentServiceLeaseCoreEpk) with
+                // LeaseToken == "{Min}-{Max}", so the dedup must compare against the EPK token set,
+                // not the PK-range set, even though we entered through the PK-range overload.
                 PartitionKeyRange mergedRange = overlappingRanges[0];
                 DefaultTrace.TraceInformation("Lease {0} merged into {1}", leaseToken, mergedRange.Id);
 
-                DocumentServiceLease newLease = await this.leaseManager.CreateLeaseIfNotExistAsync((FeedRangeEpk)partitionBasedLease.FeedRange, lastContinuationToken);
-                if (newLease != null)
+                FeedRangeEpk mergedFeedRange = (FeedRangeEpk)partitionBasedLease.FeedRange;
+                string mergedLeaseToken = $"{mergedFeedRange.Range.Min}-{mergedFeedRange.Range.Max}";
+
+                if (existingEpkLeaseTokens.Contains(mergedLeaseToken))
                 {
-                    newLeases.Enqueue(newLease);
+                    DefaultTrace.TraceInformation("Skipping creation of merged child lease for range '{0}'; a lease with that token already exists.", mergedLeaseToken);
+                }
+                else
+                {
+                    DocumentServiceLease newLease = await this.leaseManager.CreateLeaseIfNotExistAsync(mergedFeedRange, lastContinuationToken);
+                    if (newLease != null)
+                    {
+                        newLeases.Enqueue(newLease);
+                    }
                 }
             }
 
@@ -136,11 +195,17 @@ await overlappingRanges.ForEachAsync(
         /// <summary>
         /// Handles splits and merges for feed range based leases.
         /// </summary>
+        /// <remarks>
+        /// See the remarks on the PK-range overload for the behaviour when every child lease is
+        /// already present: the method returns an empty sequence; acquisition of the existing
+        /// children is deferred to the next load-balancer tick.
+        /// </remarks>
         private async Task<(IEnumerable<DocumentServiceLease>, bool)> HandlePartitionGoneAsync(
             string leaseToken,
             string lastContinuationToken,
             DocumentServiceLeaseCoreEpk feedRangeBasedLease,
-            IReadOnlyList<PartitionKeyRange> overlappingRanges)
+            IReadOnlyList<PartitionKeyRange> overlappingRanges,
+            HashSet<string> existingEpkLeaseTokens)
         {
             List<DocumentServiceLease> newLeases = new List<DocumentServiceLease>();
             if (overlappingRanges.Count > 1)
@@ -155,21 +220,37 @@ await overlappingRanges.ForEachAsync(
                 {
                     Documents.Routing.Range<string> partitionRange = overlappingRanges[i].ToRange();
                     Documents.Routing.Range<string> mergedRange = new Documents.Routing.Range<string>(min, partitionRange.Max, true, false);
-                    DocumentServiceLease newLease = await this.leaseManager.CreateLeaseIfNotExistAsync(new FeedRangeEpk(mergedRange), lastContinuationToken);
-                    if (newLease != null)
+                    string childLeaseToken = $"{mergedRange.Min}-{mergedRange.Max}";
+                    if (existingEpkLeaseTokens.Contains(childLeaseToken))
+                    {
+                        DefaultTrace.TraceInformation("Skipping creation of child EPK lease '{0}'; a lease with that token already exists.", childLeaseToken);
+                    }
+                    else
                     {
-                        newLeases.Add(newLease);
+                        DocumentServiceLease newLease = await this.leaseManager.CreateLeaseIfNotExistAsync(new FeedRangeEpk(mergedRange), lastContinuationToken);
+                        if (newLease != null)
+                        {
+                            newLeases.Add(newLease);
+                        }
                     }
 
                     min = partitionRange.Max;
                 }
 
                 // Add the last range with the original max and the last min from the split
                 Documents.Routing.Range<string> lastRangeAfterSplit = new Documents.Routing.Range<string>(min, max, true, false);
-                DocumentServiceLease lastLease = await this.leaseManager.CreateLeaseIfNotExistAsync(new FeedRangeEpk(lastRangeAfterSplit), lastContinuationToken);
-                if (lastLease != null)
+                string lastChildLeaseToken = $"{lastRangeAfterSplit.Min}-{lastRangeAfterSplit.Max}";
+                if (existingEpkLeaseTokens.Contains(lastChildLeaseToken))
+                {
+                    DefaultTrace.TraceInformation("Skipping creation of child EPK lease '{0}'; a lease with that token already exists.", lastChildLeaseToken);
+                }
+                else
                 {
-                    newLeases.Add(lastLease);
+                    DocumentServiceLease lastLease = await this.leaseManager.CreateLeaseIfNotExistAsync(new FeedRangeEpk(lastRangeAfterSplit), lastContinuationToken);
+                    if (lastLease != null)
+                    {
+                        newLeases.Add(lastLease);
+                    }
                 }
 
                 DefaultTrace.TraceInformation("Lease {0} split into {1}", leaseToken, string.Join(", ", newLeases.Select(l => l.CurrentLeaseToken)));

Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/FeedManagement/EqualPartitionsBalancingStrategy.cs

@@ -138,7 +138,28 @@ private void CategorizeLeases(
             {
                 Debug.Assert(lease.CurrentLeaseToken != null, "TakeLeasesAsync: lease.CurrentLeaseToken cannot be null.");
 
+                // Defensive: duplicate lease documents for the same lease token can exist when the lease
+                // container is partitioned by /partitionKey and split handling re-ran with a different
+                // Guid-based partition key (see PartitionSynchronizerCore.HandlePartitionGoneAsync).
+                // Use TryAdd so the load balancer keeps the first lease seen and continues functioning
+                // instead of throwing and blocking lease acquisition for the whole container. The
+                // duplicate documents must be deleted manually (or the lease container recreated with
+                // /id partitioning) to fully resolve the condition.
+                if (allPartitions.ContainsKey(lease.CurrentLeaseToken))
+                {
+                    DocumentServiceLease existing = allPartitions[lease.CurrentLeaseToken];
+                    DefaultTrace.TraceWarning(
+                        "Duplicate lease document detected for lease token '{0}'. Keeping lease with id '{1}' (partitionKey '{2}') and ignoring duplicate with id '{3}' (partitionKey '{4}'). To fully resolve, delete the duplicate lease document(s) (same id, different partitionKey values) from the lease container; if feasible, migrate the lease container to /id partitioning which avoids this condition entirely.",
+                        lease.CurrentLeaseToken,
+                        existing.Id,
+                        existing.PartitionKey,
+                        lease.Id,
+                        lease.PartitionKey);
+                    continue;
+                }
+
                 allPartitions.Add(lease.CurrentLeaseToken, lease);
+
                 if (string.IsNullOrWhiteSpace(lease.Owner) || this.IsExpired(lease))
                 {
                     DefaultTrace.TraceVerbose("Found unused or expired lease: {0}", lease);