fix error with partial count returned

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs

@@ -218,7 +218,7 @@ protected abstract string ContinuationToken
         /// <summary>
         /// Gets a value indicating whether the context still has more results.
         /// </summary>
-        private bool HasMoreResults => this.itemProducerForest.Count != 0 && this.CurrentItemProducerTree().HasMoreResults;
+        private bool HasMoreResults => this.FailureResponse != null || (this.itemProducerForest.Count != 0 && this.CurrentItemProducerTree().HasMoreResults);
 
         /// <summary>
         /// Gets the number of documents we can still buffer.
@@ -336,7 +336,7 @@ protected async Task<bool> MoveNextHelperAsync(ItemProducerTree itemProducerTree
                 this.FailureResponse = moveNextResponse.failureResponse;
             }
 
-            return !moveNextResponse.successfullyMovedNext;
+            return moveNextResponse.successfullyMovedNext;
         }
 
         /// <summary>
@@ -353,15 +353,21 @@ public override async Task<QueryResponseCore> DrainAsync(int maxElements, Cancel
             if (this.FailureResponse != null)
             {
                 this.Stop();
-                return this.FailureResponse.Value;
+
+                QueryResponseCore failure = this.FailureResponse.Value;
+                this.FailureResponse = null;
+                return failure;
             }
 
             // Drain the results. If there is no results and a failure then return the failure.
             IReadOnlyList<CosmosElement> results = await this.InternalDrainAsync(maxElements, cancellationToken);
             if ((results == null || results.Count == 0) && this.FailureResponse != null)
             {
                 this.Stop();
-                return this.FailureResponse.Value;
+                QueryResponseCore failure = this.FailureResponse.Value;
+                this.FailureResponse = null;
+                return failure;
+
             }
 
             string continuation = this.ContinuationToken;

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs

@@ -209,7 +209,8 @@ public override async Task<IReadOnlyList<CosmosElement>> InternalDrainAsync(int
             ////  2) <i, j> always come before <i, k> where j < k
 
             List<CosmosElement> results = new List<CosmosElement>();
-            while (!this.IsDone && results.Count < maxElements)
+            bool isSuccessToMoveNext = true;
+            while (!this.IsDone && results.Count < maxElements && isSuccessToMoveNext)
             {
                 // Only drain from the highest priority document producer 
                 // We need to pop and push back the document producer tree, since the priority changes according to the sort order.
@@ -234,10 +235,7 @@ public override async Task<IReadOnlyList<CosmosElement>> InternalDrainAsync(int
 
                 this.previousRid = orderByQueryResult.Rid;
 
-                if (await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken))
-                {
-                    break;
-                }
+                isSuccessToMoveNext = await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken);
 
                 this.PushCurrentItemProducerTree(currentItemProducerTree);
             }

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosParallelItemQueryExecutionContext.cs

@@ -155,7 +155,7 @@ public override async Task<IReadOnlyList<CosmosElement>> InternalDrainAsync(int
             for (int i = 0; i < Math.Min(itemsLeftInCurrentPage, maxElements); i++)
             {
                 results.Add(currentItemProducerTree.Current);
-                if (await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken))
+                if (!await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken))
                 {
                     break;
                 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockPartitionResponse.cs

@@ -15,6 +15,8 @@ namespace Microsoft.Azure.Cosmos.Tests
     /// </summary>
     internal class MockPartitionResponse
     {
+        internal const int MessageWithToManyRequestFailure = -1;
+
         public PartitionKeyRange PartitionKeyRange { get; set; }
 
         /// <summary>
@@ -24,6 +26,7 @@ internal class MockPartitionResponse
         /// </summary>
         /// <remarks>
         /// Empty int[] represent an empty page
+        /// -1 represents a 429 failure
         /// </remarks>
         public List<int[]> MessagesWithItemIndex { get; set; } = new List<int[]>();
 
@@ -46,21 +49,29 @@ public int GetTotalItemCount()
             int totalItemCount = 0;
             if(this.MessagesWithItemIndex != null)
             {
-                foreach (var message in this.MessagesWithItemIndex)
+                foreach (int[] message in this.MessagesWithItemIndex)
                 {
-                    totalItemCount += message.Length;
+                    if (!IsFailurePage(message))
+                    {
+                        totalItemCount += message.Length;
+                    }
                 }
             }
 
             if(this.Split != null)
             {
-                foreach(var partitionResponse in this.Split)
+                foreach(MockPartitionResponse partitionResponse in this.Split)
                 {
                     totalItemCount += partitionResponse.GetTotalItemCount();
                 }
             }
 
             return totalItemCount;
         }
+
+        public static bool IsFailurePage(int[] page)
+        {
+            return page != null && page.Length == 1 && page[0] < 0;
+        }
     }
 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockQueryFactory.cs

@@ -17,6 +17,7 @@ namespace Microsoft.Azure.Cosmos.Tests
     internal static class MockQueryFactory
     {
         public static readonly int[] EmptyPage = new int[] { };
+        public static readonly int[] FailureToManyRequests = new int[] { MockPartitionResponse.MessageWithToManyRequestFailure };
         public static readonly string DefaultDatabaseRid = ResourceId.NewDatabaseId(3810641).ToString();
         public static readonly string DefaultCollectionRid = ResourceId.NewDocumentCollectionId(DefaultDatabaseRid, 1376573569).ToString();
         public static readonly SqlQuerySpec DefaultQuerySpec = new SqlQuerySpec("SELECT * FROM C ");
@@ -92,37 +93,53 @@ private static IList<ToDoItem> GenerateAndMockResponseHelper(
                 string previousContinuationToken = initContinuationToken;
 
                 // Loop through each message inside the partition
-                List<int[]> messages = partitionAndMessages.MessagesWithItemIndex;
-                int messagesCount = messages == null ? 0 : messages.Count;
-                int lastMessageIndex = messagesCount - 1;
-                for (int i = 0; i < messagesCount; i++)
+                List<int[]> pages = partitionAndMessages.MessagesWithItemIndex;
+                int pagesCount = pages == null ? 0 : pages.Count;
+                int lastPageIndex = pagesCount - 1;
+                for (int i = 0; i < pagesCount; i++)
                 {
-                    int[] message = partitionAndMessages.MessagesWithItemIndex[i];
-
+                    int[] itemsInPage = partitionAndMessages.MessagesWithItemIndex[i];
                     string newContinuationToken = null;
+                    QueryResponseCore queryResponse;
 
-                    List<ToDoItem> currentPageItems = new List<ToDoItem>();
-                    // Null represents an empty page
-                    if (message != null)
+                    bool isFailureResponse = itemsInPage.Length == 1 && itemsInPage[0] < 0;
+                    if (isFailureResponse)
                     {
-                        foreach (int itemPosition in message)
+                        if(itemsInPage[0] == MockPartitionResponse.MessageWithToManyRequestFailure)
+                        {
+                            queryResponse = QueryResponseMessageFactory.CreateFailureToManyRequestResponse();
+                        }
+                        else
                         {
-                            currentPageItems.Add(allItemsOrdered[itemPosition]);
+                            throw new ArgumentException($"Unknown mocked failure response {itemsInPage[0]}");
                         }
                     }
-
-                    // Last message should have null continuation token
-                    // Split means it's not the last message for this PK range
-                    if (i != lastMessageIndex || partitionAndMessages.HasSplit)
+                    else
                     {
-                        newContinuationToken = Guid.NewGuid().ToString();
-                    }
+                        List<ToDoItem> currentPageItems = new List<ToDoItem>();
+
+                        // Null represents an empty page. Page with a single negative value represents a failure response
+                        if (itemsInPage != null)
+                        {
+                            foreach (int itemPosition in itemsInPage)
+                            {
+                                currentPageItems.Add(allItemsOrdered[itemPosition]);
+                            }
+                        }
+
+                        // Last message should have null continuation token
+                        // Split means it's not the last message for this PK range
+                        if (i != lastPageIndex || partitionAndMessages.HasSplit)
+                        {
+                            newContinuationToken = Guid.NewGuid().ToString();
+                        }
 
-                    QueryResponseCore queryResponse = QueryResponseMessageFactory.CreateQueryResponse(
-                        currentPageItems,
-                        isOrderByQuery,
-                        newContinuationToken,
-                        containerRid);
+                        queryResponse = QueryResponseMessageFactory.CreateQueryResponse(
+                            currentPageItems,
+                            isOrderByQuery,
+                            newContinuationToken,
+                            containerRid);
+                    }
 
                     mockQueryClient.Setup(x =>
                       x.ExecuteItemQueryAsync(
@@ -260,6 +277,22 @@ public static MockPartitionResponse[] CreateDefaultSplit(
             };
         }
 
+        public static List<MockPartitionResponse[]> GetFailureScenarios()
+        {
+            return new List<MockPartitionResponse[]>
+            {
+                CreateDefaultResponse(
+                    MockQueryFactory.FailureToManyRequests),
+                CreateDefaultResponse(
+                    new int[] { 0 },
+                    MockQueryFactory.FailureToManyRequests),
+                CreateDefaultResponse(
+                    new int[] { 0 },
+                    new int[] { 1 },
+                    MockQueryFactory.FailureToManyRequests),
+            };
+        }
+
         public static List<MockPartitionResponse[]> GetSplitScenarios()
         {
             List<MockPartitionResponse[]> allSplitScenario = new List<MockPartitionResponse[]>();