HBASE-17723 ClientAsyncPrefetchScanner may end prematurely when the size of the cache is one

hbase-client/src/main/java/org/apache/hadoop/hbase/client/ClientAsyncPrefetchScanner.java

@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.client;
 
+import com.google.common.annotations.VisibleForTesting;
 import static org.apache.hadoop.hbase.client.ConnectionUtils.calcEstimatedSize;
 
 import java.io.IOException;
@@ -26,6 +27,7 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.TableName;
@@ -62,6 +64,8 @@ public class ClientAsyncPrefetchScanner extends ClientSimpleScanner {
   private AtomicBoolean prefetchRunning;
   // an attribute for synchronizing close between scanner and prefetch threads
   private AtomicLong closingThreadId;
+  // used for testing
+  private Consumer<Boolean> prefetchListener;
   private static final int NO_THREAD = -1;
 
   public ClientAsyncPrefetchScanner(Configuration configuration, Scan scan, TableName name,
@@ -72,6 +76,11 @@ public class ClientAsyncPrefetchScanner extends ClientSimpleScanner {
         replicaCallTimeoutMicroSecondScan);
   }
 
+  @VisibleForTesting
+  void setPrefetchListener(Consumer<Boolean> prefetchListener) {
+    this.prefetchListener = prefetchListener;
+  }
+
   @Override
   protected void initCache() {
     // concurrent cache
@@ -88,17 +97,21 @@ public class ClientAsyncPrefetchScanner extends ClientSimpleScanner {
   public Result next() throws IOException {
 
     try {
+      boolean hasExecutedPrefetch = false;
+      do {
         handleException();
 
         // If the scanner is closed and there's nothing left in the cache, next is a no-op.
         if (getCacheCount() == 0 && this.closed) {
           return null;
         }
+
         if (prefetchCondition()) {
           // run prefetch in the background only if no prefetch is already running
           if (!isPrefetchRunning()) {
             if (prefetchRunning.compareAndSet(false, true)) {
               getPool().execute(prefetchRunnable);
+              hasExecutedPrefetch = true;
             }
           }
         }
@@ -116,6 +129,7 @@ public class ClientAsyncPrefetchScanner extends ClientSimpleScanner {
         if (getCacheCount() > 0) {
           return pollCache();
         }
+      } while (!hasExecutedPrefetch);
 
       // if we exhausted this scanner before calling close, write out the scan metrics
       writeScanMetrics();
@@ -219,11 +233,16 @@ public class ClientAsyncPrefetchScanner extends ClientSimpleScanner {
 
     @Override
     public void run() {
+      boolean succeed = false;
       try {
         loadCache();
+        succeed = true;
       } catch (Exception e) {
         exceptionsQueue.add(e);
       } finally {
+        if (prefetchListener != null) {
+          prefetchListener.accept(succeed);
+        }
         prefetchRunning.set(false);
         if(closed) {
           if (closingThreadId.compareAndSet(NO_THREAD, Thread.currentThread().getId())) {

hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestScannersFromClientSide.java

@@ -22,6 +22,7 @@ import static org.junit.Assert.assertTrue;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 import java.util.stream.IntStream;
 
 import org.apache.commons.logging.Log;
@@ -656,7 +657,9 @@ public class TestScannersFromClientSide {
     testAsyncScanner(TableName.valueOf(name.getMethodName()),
       2,
       3,
-      10);
+      10,
+      -1,
+      null);
   }
 
   @Test
@@ -664,11 +667,28 @@ public class TestScannersFromClientSide {
     testAsyncScanner(TableName.valueOf(name.getMethodName()),
       30000,
       1,
-      1);
+      1,
+      -1,
+      null);
+  }
+
+  @Test
+  public void testAsyncScannerWithoutCaching() throws Exception {
+    testAsyncScanner(TableName.valueOf(name.getMethodName()),
+      5,
+      1,
+      1,
+      1,
+      (b) -> {
+        try {
+          TimeUnit.MILLISECONDS.sleep(500);
+        } catch (InterruptedException ex) {
+        }
+      });
   }
 
   private void testAsyncScanner(TableName table, int rowNumber, int familyNumber,
-      int qualifierNumber) throws Exception {
+      int qualifierNumber, int caching, Consumer<Boolean> listener) throws Exception {
     assert rowNumber > 0;
     assert familyNumber > 0;
     assert qualifierNumber > 0;
@@ -707,11 +727,18 @@ public class TestScannersFromClientSide {
 
     Scan scan = new Scan();
     scan.setAsyncPrefetch(true);
-    ResultScanner scanner = ht.getScanner(scan);
+    if (caching > 0) {
+      scan.setCaching(caching);
+    }
+    try (ResultScanner scanner = ht.getScanner(scan)) {
+      assertTrue("Not instance of async scanner",scanner instanceof ClientAsyncPrefetchScanner);
+      ((ClientAsyncPrefetchScanner) scanner).setPrefetchListener(listener);
       List<Cell> kvListScan = new ArrayList<>();
       Result result;
       boolean first = true;
+      int actualRows = 0;
       while ((result = scanner.next()) != null) {
+        ++actualRows;
         // waiting for cache. see HBASE-17376
         if (first) {
           TimeUnit.SECONDS.sleep(1);
@@ -721,9 +748,12 @@ public class TestScannersFromClientSide {
           kvListScan.add(kv);
         }
       }
+      assertEquals(rowNumber, actualRows);
+      // These cells may have different rows but it is ok. The Result#getRow
+      // isn't used in the verifyResult()
       result = Result.create(kvListScan);
-    assertTrue("Not instance of async scanner",scanner instanceof ClientAsyncPrefetchScanner);
       verifyResult(result, kvListExp, toLog, "Testing async scan");
+    }
     TEST_UTIL.deleteTable(table);
   }