From 33748f7dae188b1f9f253e87f0ecfb7b11974c53 Mon Sep 17 00:00:00 2001
From: Jake Willey <jawilley@microsoft.com>
Date: Tue, 22 Oct 2019 10:04:03 -0700
Subject: [PATCH] Updated the failure handling

---
 .../CosmosCrossPartitionQueryExecutionContext.cs       |  2 +-
 .../CosmosOrderByItemQueryExecutionContext.cs          | 10 ++++------
 .../CosmosParallelItemQueryExecutionContext.cs         |  2 +-
 3 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs
index aa6841c603..6fe01206db 100644
--- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs
+++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs
@@ -336,7 +336,7 @@ protected async Task<bool> MoveNextHelperAsync(ItemProducerTree itemProducerTree
                 this.FailureResponse = moveNextResponse.failureResponse;
             }
 
-            return !moveNextResponse.successfullyMovedNext;
+            return moveNextResponse.successfullyMovedNext;
         }
 
         /// <summary>
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs
index 159347ecfe..8ed3b4101e 100644
--- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs
+++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs
@@ -106,7 +106,7 @@ protected override string ContinuationToken
             // With this information we have captured the progress for all partitions in a single continuation token.
             get
             {
-                if (this.IsDone || this.FailureResponse != null)
+                if (this.IsDone)
                 {
                     return null;
                 }
@@ -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);
             }
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosParallelItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosParallelItemQueryExecutionContext.cs
index 682c029a2e..270b587958 100644
--- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosParallelItemQueryExecutionContext.cs
+++ b/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;
                 }