From 9538455e0f390466811a07a6ae45a95524dfd85c Mon Sep 17 00:00:00 2001
From: Liangliang Gu <marsishandsome@gmail.com>
Date: Mon, 1 Jul 2019 16:33:40 +0800
Subject: [PATCH] fix count error, if advanceNextResponse is empty, we should
 read next region (#878)

---
 .../tikv/operation/iterator/DAGIterator.java  | 20 +++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/tikv-client/src/main/java/com/pingcap/tikv/operation/iterator/DAGIterator.java b/tikv-client/src/main/java/com/pingcap/tikv/operation/iterator/DAGIterator.java
index df6c9f10eb..5d885fb638 100644
--- a/tikv-client/src/main/java/com/pingcap/tikv/operation/iterator/DAGIterator.java
+++ b/tikv-client/src/main/java/com/pingcap/tikv/operation/iterator/DAGIterator.java
@@ -123,11 +123,27 @@ private boolean advanceNextResponse() {
     return true;
   }
 
+  /**
+   * chunk maybe empty while there is still data transmitting from TiKV. In this case, {@code
+   * readNextRegionChunks} cannot just returns false because the iterator thinks there is no data to
+   * process. This while loop ensures we can drain all possible data transmitting from TiKV.
+   *
+   * @return
+   */
   private boolean readNextRegionChunks() {
-    if (eof || regionTasks == null || taskIndex >= regionTasks.size()) {
-      return false;
+    while (hasNextRegionTask()) {
+      if (doReadNextRegionChunks()) {
+        return true;
+      }
     }
+    return false;
+  }
+
+  private boolean hasNextRegionTask() {
+    return !(eof || regionTasks == null || taskIndex >= regionTasks.size());
+  }
 
+  private boolean doReadNextRegionChunks() {
     try {
       switch (pushDownType) {
         case STREAMING: