[Internal] Query: Fixes occasional hang while querying using partial partition key against a sub-partitioned container

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs

@@ -10,27 +10,31 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel
     using Microsoft.Azure.Cosmos.Pagination;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination;
+    using Microsoft.Azure.Cosmos.Query.Core.QueryClient;
     using Microsoft.Azure.Cosmos.Tracing;
 
     internal sealed class QueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator<QueryPage, QueryState>
     {
         private readonly IQueryDataSource queryDataSource;
         private readonly SqlQuerySpec sqlQuerySpec;
         private readonly QueryPaginationOptions queryPaginationOptions;
+        private readonly ContainerQueryProperties containerQueryProperties;
         private readonly Cosmos.PartitionKey? partitionKey;
 
         public QueryPartitionRangePageAsyncEnumerator(
             IQueryDataSource queryDataSource,
             SqlQuerySpec sqlQuerySpec,
             FeedRangeState<QueryState> feedRangeState,
             Cosmos.PartitionKey? partitionKey,
-            QueryPaginationOptions queryPaginationOptions)
+            QueryPaginationOptions queryPaginationOptions,
+            ContainerQueryProperties containerQueryProperties)
             : base(feedRangeState)
         {
             this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource));
             this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec));
             this.queryPaginationOptions = queryPaginationOptions;
             this.partitionKey = partitionKey;
+            this.containerQueryProperties = containerQueryProperties;
         }
 
         public override ValueTask DisposeAsync() => default;
@@ -42,15 +46,103 @@ protected override Task<TryCatch<QueryPage>> GetNextPageAsync(ITrace trace, Canc
                 throw new ArgumentNullException(nameof(trace));
             }
 
-            // We sadly need to check the partition key, since a user can set a partition key in the request options with a different continuation token.
-            // In the future the partition filtering and continuation information needs to be a tightly bounded contract (like cross feed range state).
-            FeedRangeInternal feedRange = this.partitionKey.HasValue ? new FeedRangePartitionKey(this.partitionKey.Value) : this.FeedRangeState.FeedRange;
+            FeedRangeInternal feedRange = this.LimitFeedRangeToSinglePartition();
             return this.queryDataSource.MonadicQueryAsync(
               sqlQuerySpec: this.sqlQuerySpec,
               feedRangeState: new FeedRangeState<QueryState>(feedRange, this.FeedRangeState.State),
               queryPaginationOptions: this.queryPaginationOptions,
               trace: trace,
               cancellationToken);
         }
+
+        /// <summary>
+        /// Updates the FeedRange to limit the scope of this enumerator to single physical partition.
+        /// Generally speaking, a subpartitioned container can experience split partition at any level of hierarchical partition key.
+        /// This could cause a situation where more than one physical partition contains the data for a partial partition key.
+        /// Currently, enumerator instantiation does not honor physical partition boundary and allocates entire epk range which could spans across multiple physical partitions to the enumerator.
+        /// Since such an epk range does not exist at the container level, Service generates a GoneException.
+        /// This method restrics the range of each container by shrinking the ends of the range so that they do not span across physical partition.
+        /// </summary>
+        private FeedRangeInternal LimitFeedRangeToSinglePartition()
+        {
+            // We sadly need to check the partition key, since a user can set a partition key in the request options with a different continuation token.
+            // In the future the partition filtering and continuation information needs to be a tightly bounded contract (like cross feed range state).
+            FeedRangeInternal feedRange = this.FeedRangeState.FeedRange;
+            if (this.partitionKey.HasValue)
+            {
+                // ISSUE-HACK-adityasa-3/25/2024 - We should not update the original feed range inside this class.
+                // Instead we should guarantee that when enumerator is instantiated it is limited to a single physical partition.
+                // Ultimately we should remove enumerator's dependency on PartitionKey.
+                if ((this.containerQueryProperties.PartitionKeyDefinition.Paths.Count > 1) &&
+                    (this.partitionKey.Value.InternalKey.Components.Count != this.containerQueryProperties.PartitionKeyDefinition.Paths.Count) &&
+                    (feedRange is FeedRangeEpk feedRangeEpk))
+                {
+                    if (this.containerQueryProperties.EffectiveRangesForPartitionKey == null ||
+                        this.containerQueryProperties.EffectiveRangesForPartitionKey.Count == 0)
+                    {
+                        throw new InvalidOperationException(
+                            "EffectiveRangesForPartitionKey should be populated when PK is specified in request options.");
+                    }
+
+                    foreach (Documents.Routing.Range<String> epkForPartitionKey in
+                        this.containerQueryProperties.EffectiveRangesForPartitionKey)
+                    {
+                        if (Documents.Routing.Range<String>.CheckOverlapping(
+                                feedRangeEpk.Range,
+                                epkForPartitionKey))
+                        {
+                            if (!feedRangeEpk.Range.Equals(epkForPartitionKey))
+                            {
+                                String overlappingMin;
+                                bool minInclusive;
+                                String overlappingMax;
+                                bool maxInclusive;
+
+                                if (Documents.Routing.Range<String>.MinComparer.Instance.Compare(
+                                        epkForPartitionKey,
+                                        feedRangeEpk.Range) < 0)
+                                {
+                                    overlappingMin = feedRangeEpk.Range.Min;
+                                    minInclusive = feedRangeEpk.Range.IsMinInclusive;
+                                }
+                                else
+                                {
+                                    overlappingMin = epkForPartitionKey.Min;
+                                    minInclusive = epkForPartitionKey.IsMinInclusive;
+                                }
+
+                                if (Documents.Routing.Range<String>.MaxComparer.Instance.Compare(
+                                        epkForPartitionKey,
+                                        feedRangeEpk.Range) > 0)
+                                {
+                                    overlappingMax = feedRangeEpk.Range.Max;
+                                    maxInclusive = feedRangeEpk.Range.IsMaxInclusive;
+                                }
+                                else
+                                {
+                                    overlappingMax = epkForPartitionKey.Max;
+                                    maxInclusive = epkForPartitionKey.IsMaxInclusive;
+                                }
+
+                                feedRange = new FeedRangeEpk(
+                                    new Documents.Routing.Range<String>(
+                                        overlappingMin,
+                                        overlappingMax,
+                                        minInclusive,
+                                        maxInclusive));
+                            }
+
+                            break;
+                        }
+                    }
+                }
+                else
+                {
+                    feedRange = new FeedRangePartitionKey(this.partitionKey.Value);
+                }
+            }
+
+            return feedRange;
+        }
     }
 }