Internal Query: Adds adoption of pagination lib

Microsoft.Azure.Cosmos/src/ChangeFeed/StandByFeedContinuationToken.cs

@@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos.ChangeFeed
     using System.Diagnostics;
     using System.Linq;
     using System.Threading.Tasks;
-    using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens;
+    using Microsoft.Azure.Cosmos.Routing;
     using Newtonsoft.Json;
 
     /// <summary>

Microsoft.Azure.Cosmos/src/ChangeFeed/StandByFeedIteratorCore.cs

@@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos.ChangeFeed
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.CosmosElements;
-    using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens;
+    using Microsoft.Azure.Cosmos.Query;
     using Microsoft.Azure.Cosmos.Routing;
 
     /// <summary>

Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs

@@ -68,7 +68,7 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
 
             this.retryContext = null;
             // Received Connection error (HttpRequestException), initiate the endpoint rediscovery
-            if (exception is HttpRequestException)
+            if (exception is HttpRequestException httpException)
             {
                 DefaultTrace.TraceWarning("Endpoint not reachable. Refresh cache and retry");
                 return await this.ShouldRetryOnEndpointFailureAsync(this.isReadRequest, false);

Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs

@@ -10,8 +10,8 @@ namespace Microsoft.Azure.Cosmos.CosmosElements
     using System.Collections.Generic;
     using System.Linq;
     using Microsoft.Azure.Cosmos.Json;
-    using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
+    using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct;
 
 #if INTERNAL
 #pragma warning disable CS1591 // Missing XML comment for publicly visible type or member

Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs

@@ -10,8 +10,8 @@ namespace Microsoft.Azure.Cosmos.CosmosElements
     using System.Collections.Generic;
     using System.Linq;
     using Microsoft.Azure.Cosmos.Json;
-    using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
+    using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct;
 
 #if INTERNAL
 #pragma warning disable CS1591 // Missing XML comment for publicly visible type or member

Microsoft.Azure.Cosmos/src/FeedRange/Continuations/FeedRangeCompositeContinuation.cs

@@ -10,8 +10,8 @@ namespace Microsoft.Azure.Cosmos
     using System.Net;
     using System.Threading;
     using System.Threading.Tasks;
-    using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
+    using Microsoft.Azure.Cosmos.Routing;
     using Newtonsoft.Json;
 
     /// <summary>

Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEpk.cs

@@ -59,9 +59,11 @@ public override Task<TResult> AcceptAsync<TResult>(
             return visitor.VisitAsync(this, cancellationToken);
         }
 
-        public override string ToString()
-        {
-            return this.Range.ToString();
-        }
+        public override Task<TResult> AcceptAsync<TResult, TArg>(
+            IFeedRangeAsyncVisitor<TResult, TArg> visitor,
+            TArg argument,
+            CancellationToken cancellationToken) => visitor.VisitAsync(this, argument, cancellationToken);
+
+        public override string ToString() => this.Range.ToString();
     }
 }

Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeInternal.cs

@@ -29,6 +29,11 @@ public abstract Task<IEnumerable<string>> GetPartitionKeyRangesAsync(
 
         public abstract Task<TResult> AcceptAsync<TResult>(IFeedRangeAsyncVisitor<TResult> visitor, CancellationToken cancellationToken = default);
 
+        public abstract Task<TResult> AcceptAsync<TResult, TArg>(
+            IFeedRangeAsyncVisitor<TResult, TArg> visitor,
+            TArg argument,
+            CancellationToken cancellationToken);
+
         public abstract override string ToString();
 
         public override string ToJsonString()

Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKey.cs

@@ -57,9 +57,11 @@ public override Task<TResult> AcceptAsync<TResult>(
             return visitor.VisitAsync(this, cancellationToken);
         }
 
-        public override string ToString()
-        {
-            return this.PartitionKey.InternalKey.ToJsonString();
-        }
+        public override Task<TResult> AcceptAsync<TResult, TArg>(
+           IFeedRangeAsyncVisitor<TResult, TArg> visitor,
+           TArg argument,
+           CancellationToken cancellationToken) => visitor.VisitAsync(this, argument, cancellationToken);
+
+        public override string ToString() => this.PartitionKey.InternalKey.ToJsonString();
     }
 }

Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKeyRange.cs

@@ -85,9 +85,11 @@ public override Task<TResult> AcceptAsync<TResult>(
             return visitor.VisitAsync(this, cancellationToken);
         }
 
-        public override string ToString()
-        {
-            return this.PartitionKeyRangeId;
-        }
+        public override Task<TResult> AcceptAsync<TResult, TArg>(
+           IFeedRangeAsyncVisitor<TResult, TArg> visitor,
+           TArg argument,
+           CancellationToken cancellationToken) => visitor.VisitAsync(this, argument, cancellationToken);
+
+        public override string ToString() => this.PartitionKeyRangeId;
     }
 }

Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/IFeedRangeAsyncVisitor{TResult,TArg}.cs

@@ -0,0 +1,18 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System.Threading;
+    using System.Threading.Tasks;
+
+    internal interface IFeedRangeAsyncVisitor<TResult, TArg>
+    {
+        public abstract Task<TResult> VisitAsync(FeedRangePartitionKey feedRange, TArg argument, CancellationToken cancellationToken);
+
+        public abstract Task<TResult> VisitAsync(FeedRangePartitionKeyRange feedRange, TArg argument, CancellationToken cancellationToken);
+
+        public abstract Task<TResult> VisitAsync(FeedRangeEpk feedRange, TArg argument, CancellationToken cancellationToken);
+    }
+}

Microsoft.Azure.Cosmos/src/Handler/PartitionKeyRangeHandler.cs

@@ -11,7 +11,6 @@ namespace Microsoft.Azure.Cosmos.Handlers
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Common;
-    using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens;
     using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Routing;

Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs

@@ -79,7 +79,9 @@ internal CosmosQueryResponseMessageHeaders CloneKnownProperties(
         internal static CosmosQueryResponseMessageHeaders ConvertToQueryHeaders(
             Headers sourceHeaders,
             ResourceType resourceType,
-            string containerRid)
+            string containerRid,
+            int? substatusCode = null,
+            string activityId = null)
         {
             if (sourceHeaders == null)
             {
@@ -98,11 +100,11 @@ internal static CosmosQueryResponseMessageHeaders ConvertToQueryHeaders(
             {
                 RequestCharge = sourceHeaders.RequestCharge,
                 ContentLength = sourceHeaders.ContentLength,
-                ActivityId = sourceHeaders.ActivityId,
+                ActivityId = sourceHeaders.ActivityId ?? activityId,
                 ETag = sourceHeaders.ETag,
                 Location = sourceHeaders.Location,
                 RetryAfterLiteral = sourceHeaders.RetryAfterLiteral,
-                SubStatusCodeLiteral = sourceHeaders.SubStatusCodeLiteral,
+                SubStatusCodeLiteral = sourceHeaders.SubStatusCodeLiteral ?? (substatusCode.HasValue ? substatusCode.Value.ToString() : null),
                 ContentType = sourceHeaders.ContentType,
                 QueryMetricsText = sourceHeaders.QueryMetricsText
             };

Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs

@@ -0,0 +1,50 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Pagination
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Query.Core.Monads;
+
+    internal sealed class BufferedPartitionRangePageAsyncEnumerator<TPage, TState> : PartitionRangePageAsyncEnumerator<TPage, TState>, IPrefetcher
+        where TPage : Page<TState>
+        where TState : State
+    {
+        private readonly PartitionRangePageAsyncEnumerator<TPage, TState> enumerator;
+        private TryCatch<TPage>? bufferedPage;
+
+        public BufferedPartitionRangePageAsyncEnumerator(PartitionRangePageAsyncEnumerator<TPage, TState> enumerator, CancellationToken cancellationToken)
+            : base(enumerator.Range, cancellationToken, enumerator.State)
+        {
+            this.enumerator = enumerator ?? throw new ArgumentNullException(nameof(enumerator));
+        }
+
+        public override ValueTask DisposeAsync() => this.enumerator.DisposeAsync();
+
+        protected override async Task<TryCatch<TPage>> GetNextPageAsync(CancellationToken cancellationToken)
+        {
+            await this.PrefetchAsync(cancellationToken);
+
+            // Serve from the buffered page first.
+            TryCatch<TPage> returnValue = this.bufferedPage.Value;
+            this.bufferedPage = null;
+            return returnValue;
+        }
+
+        public async ValueTask PrefetchAsync(CancellationToken cancellationToken)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            if (this.bufferedPage.HasValue)
+            {
+                return;
+            }
+
+            await this.enumerator.MoveNextAsync();
+            this.bufferedPage = this.enumerator.Current;
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs

@@ -17,17 +17,20 @@ internal sealed class CrossPartitionRangePageAsyncEnumerable<TPage, TState> : IA
         private readonly CreatePartitionRangePageAsyncEnumerator<TPage, TState> createPartitionRangeEnumerator;
         private readonly IComparer<PartitionRangePageAsyncEnumerator<TPage, TState>> comparer;
         private readonly IFeedRangeProvider feedRangeProvider;
+        private readonly int maxConcurrency;
 
         public CrossPartitionRangePageAsyncEnumerable(
             IFeedRangeProvider feedRangeProvider,
             CreatePartitionRangePageAsyncEnumerator<TPage, TState> createPartitionRangeEnumerator,
             IComparer<PartitionRangePageAsyncEnumerator<TPage, TState>> comparer,
+            int maxConcurrency,
             CrossPartitionState<TState> state = default)
         {
             this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer));
             this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator));
             this.comparer = comparer ?? throw new ArgumentNullException(nameof(comparer));
             this.state = state;
+            this.maxConcurrency = maxConcurrency < 0 ? throw new ArgumentOutOfRangeException(nameof(maxConcurrency)) : maxConcurrency;
         }
 
         public IAsyncEnumerator<TryCatch<CrossPartitionPage<TPage, TState>>> GetAsyncEnumerator(CancellationToken cancellationToken = default)
@@ -38,6 +41,8 @@ public IAsyncEnumerator<TryCatch<CrossPartitionPage<TPage, TState>>> GetAsyncEnu
                 this.feedRangeProvider,
                 this.createPartitionRangeEnumerator,
                 this.comparer,
+                this.maxConcurrency,
+                cancellationToken,
                 this.state);
         }
     }

Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs

@@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Pagination
 {
     using System;
     using System.Collections.Generic;
+    using System.Linq;
     using System.Net;
     using System.Threading;
     using System.Threading.Tasks;
@@ -24,21 +25,25 @@ internal sealed class CrossPartitionRangePageAsyncEnumerator<TPage, TState> : IA
         private readonly IFeedRangeProvider feedRangeProvider;
         private readonly CreatePartitionRangePageAsyncEnumerator<TPage, TState> createPartitionRangeEnumerator;
         private readonly AsyncLazy<PriorityQueue<PartitionRangePageAsyncEnumerator<TPage, TState>>> lazyEnumerators;
+        private CancellationToken cancellationToken;
 
         public CrossPartitionRangePageAsyncEnumerator(
             IFeedRangeProvider feedRangeProvider,
             CreatePartitionRangePageAsyncEnumerator<TPage, TState> createPartitionRangeEnumerator,
             IComparer<PartitionRangePageAsyncEnumerator<TPage, TState>> comparer,
+            int? maxConcurrency,
+            CancellationToken cancellationToken,
             CrossPartitionState<TState> state = default)
         {
-            this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider));
-            this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator));
-
             if (comparer == null)
             {
                 throw new ArgumentNullException(nameof(comparer));
             }
 
+            this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider));
+            this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator));
+            this.cancellationToken = cancellationToken;
+
             this.lazyEnumerators = new AsyncLazy<PriorityQueue<PartitionRangePageAsyncEnumerator<TPage, TState>>>(async (CancellationToken token) =>
             {
                 IReadOnlyList<(PartitionKeyRange, TState)> rangeAndStates;
@@ -60,13 +65,23 @@ public CrossPartitionRangePageAsyncEnumerator(
                     rangeAndStates = rangesAndStatesBuilder;
                 }
 
-                PriorityQueue<PartitionRangePageAsyncEnumerator<TPage, TState>> enumerators = new PriorityQueue<PartitionRangePageAsyncEnumerator<TPage, TState>>(comparer);
-                foreach ((PartitionKeyRange range, TState rangeState) in rangeAndStates)
+                List<BufferedPartitionRangePageAsyncEnumerator<TPage, TState>> bufferedEnumerators = rangeAndStates
+                    .Select(rangeAndState =>
+                    {
+                        PartitionRangePageAsyncEnumerator<TPage, TState> enumerator = createPartitionRangeEnumerator(rangeAndState.Item1, rangeAndState.Item2);
+                        BufferedPartitionRangePageAsyncEnumerator<TPage, TState> bufferedEnumerator = new BufferedPartitionRangePageAsyncEnumerator<TPage, TState>(enumerator, cancellationToken);
+                        return bufferedEnumerator;
+                    })
+                    .ToList();
+
+                if (maxConcurrency.HasValue)
                 {
-                    PartitionRangePageAsyncEnumerator<TPage, TState> enumerator = createPartitionRangeEnumerator(range, rangeState);
-                    enumerators.Enqueue(enumerator);
+                    await ParallelPrefetch.PrefetchInParallelAsync(bufferedEnumerators, maxConcurrency.Value, token);
                 }
 
+                PriorityQueue<PartitionRangePageAsyncEnumerator<TPage, TState>> enumerators = new PriorityQueue<PartitionRangePageAsyncEnumerator<TPage, TState>>(
+                    bufferedEnumerators,
+                    comparer);
                 return enumerators;
             });
         }
@@ -75,7 +90,9 @@ public CrossPartitionRangePageAsyncEnumerator(
 
         public async ValueTask<bool> MoveNextAsync()
         {
-            PriorityQueue<PartitionRangePageAsyncEnumerator<TPage, TState>> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default);
+            this.cancellationToken.ThrowIfCancellationRequested();
+
+            PriorityQueue<PartitionRangePageAsyncEnumerator<TPage, TState>> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: this.cancellationToken);
             if (enumerators.Count == 0)
             {
                 return false;
@@ -103,7 +120,7 @@ public async ValueTask<bool> MoveNextAsync()
                     // Handle split
                     IEnumerable<PartitionKeyRange> childRanges = await this.feedRangeProvider.GetChildRangeAsync(
                         currentPaginator.Range,
-                        cancellationToken: default);
+                        cancellationToken: this.cancellationToken);
                     foreach (PartitionKeyRange childRange in childRanges)
                     {
                         PartitionRangePageAsyncEnumerator<TPage, TState> childPaginator = this.createPartitionRangeEnumerator(
@@ -120,18 +137,18 @@ public async ValueTask<bool> MoveNextAsync()
                 {
                     throw new NotImplementedException();
                 }
-            }
 
-            if (currentPaginator.State != null)
-            {
+                // Just enqueue the paginator and the user can decide if they want to retry.
                 enumerators.Enqueue(currentPaginator);
+
+                this.Current = TryCatch<CrossPartitionPage<TPage, TState>>.FromException(currentPaginator.Current.Exception);
+                return true;
             }
 
-            TryCatch<TPage> backendPage = currentPaginator.Current;
-            if (backendPage.Failed)
+            if (currentPaginator.State != default)
             {
-                this.Current = TryCatch<CrossPartitionPage<TPage, TState>>.FromException(backendPage.Exception);
-                return true;
+                // Don't enqueue the paginator otherwise it's an infinite loop.
+                enumerators.Enqueue(currentPaginator);
             }
 
             CrossPartitionState<TState> crossPartitionState;
@@ -151,7 +168,7 @@ public async ValueTask<bool> MoveNextAsync()
             }
 
             this.Current = TryCatch<CrossPartitionPage<TPage, TState>>.FromResult(
-                new CrossPartitionPage<TPage, TState>(backendPage.Result, crossPartitionState));
+                new CrossPartitionPage<TPage, TState>(currentPaginator.Current.Result, crossPartitionState));
             return true;
         }
 
@@ -161,6 +178,11 @@ public ValueTask DisposeAsync()
             return default;
         }
 
+        public void SetCancellationToken(CancellationToken cancellationToken)
+        {
+            this.cancellationToken = cancellationToken;
+        }
+
         private static bool IsSplitException(Exception exeception)
         {
             return exeception is CosmosException cosmosException