Skip to content

Commit 38fff07

Browse files
ananth7592ananth7592
committed
Refactor relevant uses of Range.Min/MaxComparer to use the Range.LengthAwareMin/Max comparers in the v3 codebase
1 parent 7bcccd7 commit 38fff07

10 files changed

Lines changed: 222 additions & 71 deletions

File tree

‎Microsoft.Azure.Cosmos/src/DocumentClient.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1146,7 +1146,7 @@ private async Task InitializeCachesAsync(string databaseName, DocumentCollection
11461146
true,
11471147
false),
11481148
NoOpTrace.Singleton,
1149-
resolvedCollection.PartitionKey);
1149+
resolvedCollection?.PartitionKey);
11501150

11511151
// In Gateway mode, AddressCache is null
11521152
if (this.AddressResolver != null)

‎Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs‎

Lines changed: 3 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -112,17 +112,12 @@ await GatewayStoreModel.ApplySessionTokenAsync(
112112
CancellationToken.None,
113113
NoOpTrace.Singleton);
114114

115-
PartitionKeyDefinition partitionKeyDefinition = null;
116-
if (collection != null)
117-
{
118-
partitionKeyDefinition = collection.PartitionKey;
119-
}
120115
(bool isSuccess, PartitionKeyRange partitionKeyRange) = await TryResolvePartitionKeyRangeAsync(
121116
request: request,
122117
sessionContainer: this.sessionContainer,
123118
partitionKeyRangeCache: this.partitionKeyRangeCache,
124119
clientCollectionCache: this.clientCollectionCache,
125-
partitionKeyDefinition: partitionKeyDefinition,
120+
partitionKeyDefinition: collection.PartitionKey,
126121
refreshCache: false);
127122

128123
request.RequestContext.ResolvedPartitionKeyRange = partitionKeyRange;
@@ -318,17 +313,11 @@ internal async Task CaptureSessionTokenAndHandleSplitAsync(
318313
CancellationToken.None,
319314
NoOpTrace.Singleton);
320315

321-
PartitionKeyDefinition partitionKeyDefinition = null;
322-
if (collection != null)
323-
{
324-
partitionKeyDefinition = collection.PartitionKey;
325-
}
326-
327316
await this.partitionKeyRangeCache.TryGetPartitionKeyRangeByIdAsync(
328317
request.RequestContext.ResolvedCollectionRid,
329318
partitionKeyRangeInResponse,
330319
NoOpTrace.Singleton,
331-
partitionKeyDefinition,
320+
collection?.PartitionKey,
332321
forceRefresh: true);
333322
}
334323
}
@@ -427,18 +416,12 @@ internal static async Task<Tuple<bool, string>> TryResolveSessionTokenAsync(
427416
CancellationToken.None,
428417
NoOpTrace.Singleton);
429418

430-
PartitionKeyDefinition partitionKeyDefinition = null;
431-
if (collection != null)
432-
{
433-
partitionKeyDefinition = collection.PartitionKey;
434-
}
435-
436419
(bool isSuccess, PartitionKeyRange partitionKeyRange) = await TryResolvePartitionKeyRangeAsync(
437420
request: request,
438421
sessionContainer: sessionContainer,
439422
partitionKeyRangeCache: partitionKeyRangeCache,
440423
clientCollectionCache: clientCollectionCache,
441-
partitionKeyDefinition: partitionKeyDefinition,
424+
partitionKeyDefinition: collection?.PartitionKey,
442425
refreshCache: false);
443426

444427
if (isSuccess && sessionContainer is SessionContainer gatewaySessionContainer)

‎Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/HierarchicalPartitionUtils.cs‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition
66
{
77
using System;
8+
using System.Collections.Generic;
89
using Microsoft.Azure.Cosmos.Query.Core.QueryClient;
910

1011
internal static class HierarchicalPartitionUtils
@@ -50,7 +51,11 @@ public static FeedRangeInternal LimitFeedRangeToSinglePartition(PartitionKey? pa
5051
String overlappingMax;
5152
bool maxInclusive;
5253

53-
if (Documents.Routing.Range<String>.MinComparer.Instance.Compare(
54+
bool isLengthAwareComparisonEnabled = ConfigurationManager.IsLengthAwareRangeComparatorEnabled();
55+
56+
IComparer<Documents.Routing.Range<string>> minComparer = GetComparer(isMinComparer: true, useLengthAwareComparison: isLengthAwareComparisonEnabled);
57+
58+
if (minComparer.Compare(
5459
epkForPartitionKey,
5560
feedRangeEpk.Range) < 0)
5661
{
@@ -63,7 +68,9 @@ public static FeedRangeInternal LimitFeedRangeToSinglePartition(PartitionKey? pa
6368
minInclusive = epkForPartitionKey.IsMinInclusive;
6469
}
6570

66-
if (Documents.Routing.Range<String>.MaxComparer.Instance.Compare(
71+
IComparer<Documents.Routing.Range<string>> maxComparer = GetComparer(isMinComparer: false, useLengthAwareComparison: isLengthAwareComparisonEnabled);
72+
73+
if (maxComparer.Compare(
6774
epkForPartitionKey,
6875
feedRangeEpk.Range) > 0)
6976
{
@@ -96,5 +103,12 @@ public static FeedRangeInternal LimitFeedRangeToSinglePartition(PartitionKey? pa
96103

97104
return feedRange;
98105
}
106+
107+
private static IComparer<Documents.Routing.Range<string>> GetComparer(bool isMinComparer, bool useLengthAwareComparison)
108+
{
109+
return isMinComparer
110+
? (useLengthAwareComparison ? Documents.Routing.Range<string>.LengthAwareMinComparer.Instance : Documents.Routing.Range<string>.MinComparer.Instance)
111+
: (useLengthAwareComparison ? Documents.Routing.Range<string>.LengthAwareMaxComparer.Instance : Documents.Routing.Range<string>.MaxComparer.Instance);
112+
}
99113
}
100114
}

‎Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextFactory.cs‎

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -144,19 +144,13 @@ internal static async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesA
144144
{
145145
List<PartitionKeyRange> targetRanges = null;
146146
if (!string.IsNullOrEmpty(feedOptions.PartitionKeyRangeId))
147-
{
148-
PartitionKeyDefinition partitionKeyDefinition = null;
149-
if (collection != null)
150-
{
151-
partitionKeyDefinition = collection.PartitionKey;
152-
}
153-
147+
{
154148
targetRanges = new List<PartitionKeyRange>()
155149
{
156150
await queryExecutionContext.GetTargetPartitionKeyRangeByIdAsync(
157151
collection.ResourceId,
158152
feedOptions.PartitionKeyRangeId,
159-
partitionKeyDefinition)
153+
collection.PartitionKey)
160154
};
161155
}
162156
else if (feedOptions.PartitionKey != null)

‎Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,7 @@ public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesA
266266
{
267267
IRoutingMapProvider routingMapProvider = await this.GetRoutingMapProviderAsync();
268268

269-
List<PartitionKeyRange> ranges = await routingMapProvider.TryGetOverlappingRangesAsync(collectionResourceId, providedRanges, getPKRangesTrace, partitionKeyDefinition, false);
269+
List<PartitionKeyRange> ranges = await routingMapProvider.TryGetOverlappingRangesAsync(collectionResourceId, providedRanges, getPKRangesTrace, partitionKeyDefinition);
270270
if (ranges == null && PathsHelper.IsNameBased(resourceLink))
271271
{
272272
// Refresh the cache and don't try to re-resolve collection as it is not clear what already

‎Microsoft.Azure.Cosmos/src/Routing/CollectionRoutingMap.cs‎

Lines changed: 11 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ namespace Microsoft.Azure.Cosmos.Routing
2020
/// </summary>
2121
internal sealed class CollectionRoutingMap
2222
{
23-
private const string useLengthAwareComparatorEnvVar = "UseLengthAwareRangeComparator";
2423
/// <summary>
2524
/// Partition key range id to partition address and range.
2625
/// </summary>
@@ -34,19 +33,6 @@ internal sealed class CollectionRoutingMap
3433

3534
internal int HighestNonOfflinePkRangeId { get; private set; }
3635

37-
public CollectionRoutingMap(
38-
CollectionRoutingMap collectionRoutingMap,
39-
string changeFeedNextIfNoneMatch)
40-
{
41-
this.rangeById = new Dictionary<string, Tuple<PartitionKeyRange, ServiceIdentity>>(collectionRoutingMap.rangeById);
42-
this.orderedPartitionKeyRanges = new List<PartitionKeyRange>(collectionRoutingMap.orderedPartitionKeyRanges);
43-
this.orderedRanges = new List<Range<string>>(collectionRoutingMap.orderedRanges);
44-
this.goneRanges = new HashSet<string>(collectionRoutingMap.goneRanges);
45-
this.HighestNonOfflinePkRangeId = collectionRoutingMap.HighestNonOfflinePkRangeId;
46-
this.CollectionUniqueId = collectionRoutingMap.CollectionUniqueId;
47-
this.ChangeFeedNextIfNoneMatch = changeFeedNextIfNoneMatch;
48-
}
49-
5036
private CollectionRoutingMap(
5137
Dictionary<string, Tuple<PartitionKeyRange, ServiceIdentity>> rangeById,
5238
List<PartitionKeyRange> orderedPartitionKeyRanges,
@@ -143,12 +129,11 @@ public IReadOnlyList<PartitionKeyRange> GetOverlappingRanges(IReadOnlyList<Range
143129

144130
SortedList<string, PartitionKeyRange> partitionRanges = new SortedList<string, PartitionKeyRange>();
145131

146-
string useLengthAwareRangeComparator = Environment.GetEnvironmentVariable(useLengthAwareComparatorEnvVar);
147-
bool lengthAwareComparison = !string.IsNullOrEmpty(useLengthAwareRangeComparator) && useLengthAwareRangeComparator.Equals(bool.FalseString, StringComparison.OrdinalIgnoreCase);
132+
bool isLengthAwareComparisonEnabled = ConfigurationManager.IsLengthAwareRangeComparatorEnabled();
148133

149134
bool useLengthAwareComparison = false;
150135
// Enable length-aware comparison for provided partition EPK ranges only when the container uses a hierarchical partition key (MultiHash).
151-
if (!lengthAwareComparison && this.partitionKeyDefinition != null && this.partitionKeyDefinition.Kind == PartitionKind.MultiHash)
136+
if (isLengthAwareComparisonEnabled && this.partitionKeyDefinition != null && this.partitionKeyDefinition.Kind == PartitionKind.MultiHash && this.partitionKeyDefinition.Paths.Count > 1)
152137
{
153138
useLengthAwareComparison = true;
154139
}
@@ -157,23 +142,14 @@ public IReadOnlyList<PartitionKeyRange> GetOverlappingRanges(IReadOnlyList<Range
157142
// Then within that two positions, check for overlapping partition key ranges
158143
foreach (Range<string> providedRange in providedPartitionKeyRanges)
159144
{
160-
IComparer<Range<string>> minComparer = Range<string>.MinComparer.Instance;
161-
if (useLengthAwareComparison)
162-
{
163-
minComparer = Range<string>.LengthAwareMinComparer.Instance;
164-
}
165-
145+
IComparer<Range<string>> minComparer = this.GetComparer(isMinComparer: true, useLengthAwareComparison: useLengthAwareComparison);
166146
int minIndex = this.orderedRanges.BinarySearch(providedRange, minComparer);
167147
if (minIndex < 0)
168148
{
169149
minIndex = Math.Max(0, (~minIndex) - 1);
170150
}
171151

172-
IComparer<Range<string>> maxComparer = Range<string>.MaxComparer.Instance;
173-
if (useLengthAwareComparison)
174-
{
175-
maxComparer = Range<string>.LengthAwareMaxComparer.Instance;
176-
}
152+
IComparer<Range<string>> maxComparer = this.GetComparer(isMinComparer: false, useLengthAwareComparison: useLengthAwareComparison);
177153

178154
int maxIndex = this.orderedRanges.BinarySearch(providedRange, maxComparer);
179155
if (maxIndex < 0)
@@ -193,6 +169,13 @@ public IReadOnlyList<PartitionKeyRange> GetOverlappingRanges(IReadOnlyList<Range
193169
return new ReadOnlyCollection<PartitionKeyRange>(partitionRanges.Values);
194170
}
195171

172+
private IComparer<Range<string>> GetComparer(bool isMinComparer, bool useLengthAwareComparison)
173+
{
174+
return isMinComparer
175+
? (useLengthAwareComparison ? Range<string>.LengthAwareMinComparer.Instance : Range<string>.MinComparer.Instance)
176+
: (useLengthAwareComparison ? Range<string>.LengthAwareMaxComparer.Instance : Range<string>.MaxComparer.Instance);
177+
}
178+
196179
public PartitionKeyRange GetRangeByEffectivePartitionKey(string effectivePartitionKeyValue)
197180
{
198181
if (string.CompareOrdinal(effectivePartitionKeyValue, PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey) >= 0)

‎Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs‎

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ namespace Microsoft.Azure.Cosmos.Routing
1313
using System.Net;
1414
using System.Threading.Tasks;
1515
using Microsoft.Azure.Cosmos.Core.Trace;
16-
using Microsoft.Azure.Cosmos.Query.Core;
1716
using Microsoft.Azure.Cosmos.Query.Core.Monads;
1817
using Microsoft.Azure.Cosmos.Query.Core.QueryPlan;
1918
using Microsoft.Azure.Cosmos.Tracing;
@@ -232,11 +231,21 @@ await routingMapProvider.TryGetRangeByEffectivePartitionKeyAsync(
232231
PartitionKeyRange lastPartitionKeyRange = partitionKeyRanges[partitionKeyRanges.Count - 1];
233232

234233
return new ResolvedRangeInfo(lastPartitionKeyRange, suppliedTokens);
235-
}
234+
}
235+
236+
bool isLengthAwareComparisonEnabled = ConfigurationManager.IsLengthAwareRangeComparatorEnabled();
237+
bool useLengthAwareComparison = false;
238+
// Enable length-aware comparison for provided partition EPK ranges only when the container uses a hierarchical partition key (MultiHash).
239+
if (isLengthAwareComparisonEnabled && partitionKeyDefinition != null && partitionKeyDefinition.Kind == PartitionKind.MultiHash && partitionKeyDefinition.Paths.Count > 1)
240+
{
241+
useLengthAwareComparison = true;
242+
}
243+
244+
IComparer<Range<string>> minComparer = this.GetComparer(isMinComparer: true, useLengthAwareComparison);
236245

237246
Range<string> minimumRange = PartitionRoutingHelper.Min(
238247
providedPartitionKeyRanges,
239-
Range<string>.MinComparer.Instance);
248+
minComparer);
240249

241250
return new ResolvedRangeInfo(
242251
await routingMapProvider.TryGetRangeByEffectivePartitionKeyAsync(collectionRid, minimumRange.Min, trace, partitionKeyDefinition),
@@ -352,7 +361,15 @@ public virtual async Task<bool> TryAddPartitionKeyRangeToContinuationTokenAsync(
352361

353362
// We only need to get the next range if we have to
354363
if (string.IsNullOrEmpty(backendResponseHeaders[HttpConstants.HttpHeaders.Continuation]))
355-
{
364+
{
365+
bool isLengthAwareComparisonEnabled = ConfigurationManager.IsLengthAwareRangeComparatorEnabled();
366+
bool useLengthAwareComparison = false;
367+
// Enable length-aware comparison for provided partition EPK ranges only when the container uses a hierarchical partition key (MultiHash).
368+
if (isLengthAwareComparisonEnabled && partitionKeyDefinition != null && partitionKeyDefinition.Kind == PartitionKind.MultiHash && partitionKeyDefinition.Paths.Count > 1)
369+
{
370+
useLengthAwareComparison = true;
371+
}
372+
356373
if (direction == RntdbEnumerationDirection.Reverse)
357374
{
358375
rangeToUse = PartitionRoutingHelper.MinBefore(
@@ -361,14 +378,15 @@ public virtual async Task<bool> TryAddPartitionKeyRangeToContinuationTokenAsync(
361378
providedPartitionKeyRanges.Single(),
362379
trace, partitionKeyDefinition,
363380
false)).ToList(),
364-
currentRange);
381+
currentRange,
382+
this.GetComparer(isMinComparer: true, useLengthAwareComparison: useLengthAwareComparison));
365383
}
366384
else
367385
{
368386
Range<string> nextProvidedRange = PartitionRoutingHelper.MinAfter(
369387
providedPartitionKeyRanges,
370388
currentRange.ToRange(),
371-
Range<string>.MaxComparer.Instance);
389+
this.GetComparer(isMinComparer: false, useLengthAwareComparison: useLengthAwareComparison));
372390

373391
if (nextProvidedRange == null)
374392
{
@@ -554,14 +572,14 @@ private static T MinAfter<T>(IReadOnlyList<T> values, T minValue, IComparer<T> c
554572
return min;
555573
}
556574

557-
private static PartitionKeyRange MinBefore(IReadOnlyList<PartitionKeyRange> values, PartitionKeyRange minValue)
575+
private static PartitionKeyRange MinBefore(IReadOnlyList<PartitionKeyRange> values, PartitionKeyRange minValue,
576+
IComparer<Range<string>> comparer)
558577
{
559578
if (values.Count == 0)
560579
{
561580
throw new ArgumentException(nameof(values));
562581
}
563582

564-
IComparer<Range<string>> comparer = Range<string>.MinComparer.Instance;
565583
PartitionKeyRange min = null;
566584
foreach (PartitionKeyRange value in values)
567585
{
@@ -572,6 +590,13 @@ private static PartitionKeyRange MinBefore(IReadOnlyList<PartitionKeyRange> valu
572590
}
573591

574592
return min;
593+
}
594+
595+
private IComparer<Range<string>> GetComparer(bool isMinComparer, bool useLengthAwareComparison)
596+
{
597+
return isMinComparer
598+
? (useLengthAwareComparison ? Range<string>.LengthAwareMinComparer.Instance : Range<string>.MinComparer.Instance)
599+
: (useLengthAwareComparison ? Range<string>.LengthAwareMaxComparer.Instance : Range<string>.MaxComparer.Instance);
575600
}
576601

577602
public readonly struct ResolvedRangeInfo

0 commit comments

Comments
 (0)