From 89c02dc0cc5997df17ff240ecec23f801a81e09d Mon Sep 17 00:00:00 2001
From: zhangduo <zhangduo@apache.org>
Date: Sat, 2 Feb 2019 20:31:13 +0800
Subject: [PATCH] HBASE-21828 Make sure we do not return CompletionException
 when locating region

---
 .../hbase/client/AsyncRegionLocator.java      |  4 ++-
 ...meout.java => TestAsyncRegionLocator.java} | 30 +++++++++++++++----
 2 files changed, 28 insertions(+), 6 deletions(-)
 rename hbase-server/src/test/java/org/apache/hadoop/hbase/client/{TestAsyncRegionLocatorTimeout.java => TestAsyncRegionLocator.java} (81%)

diff --git a/hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncRegionLocator.java b/hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncRegionLocator.java
index 82187617141..cc2a21b633a 100644
--- a/hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncRegionLocator.java
+++ b/hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncRegionLocator.java
@@ -28,6 +28,7 @@ import org.apache.hadoop.hbase.RegionLocations;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.exceptions.TimeoutIOException;
 import org.apache.hadoop.hbase.util.Bytes;
+import org.apache.hadoop.hbase.util.FutureUtils;
 import org.apache.yetus.audience.InterfaceAudience;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -66,12 +67,13 @@ class AsyncRegionLocator {
       }
       future.completeExceptionally(new TimeoutIOException(timeoutMsg.get()));
     }, timeoutNs, TimeUnit.NANOSECONDS);
-    return future.whenComplete((loc, error) -> {
+    FutureUtils.addListener(future, (loc, error) -> {
       if (error != null && error.getClass() != TimeoutIOException.class) {
         // cancel timeout task if we are not completed by it.
         timeoutTask.cancel();
       }
     });
+    return future;
   }
 
   private boolean isMeta(TableName tableName) {
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAsyncRegionLocatorTimeout.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAsyncRegionLocator.java
similarity index 81%
rename from hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAsyncRegionLocatorTimeout.java
rename to hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAsyncRegionLocator.java
index 758aa303fca..0b51ce87235 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAsyncRegionLocatorTimeout.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAsyncRegionLocator.java
@@ -28,14 +28,17 @@ import static org.junit.Assert.fail;
 
 import java.io.IOException;
 import java.util.Optional;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.commons.io.IOUtils;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.HBaseClassTestRule;
 import org.apache.hadoop.hbase.HBaseTestingUtility;
 import org.apache.hadoop.hbase.HRegionLocation;
 import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.TableNotFoundException;
 import org.apache.hadoop.hbase.coprocessor.ObserverContext;
 import org.apache.hadoop.hbase.coprocessor.RegionCoprocessor;
 import org.apache.hadoop.hbase.coprocessor.RegionCoprocessorEnvironment;
@@ -53,11 +56,11 @@ import org.junit.Test;
 import org.junit.experimental.categories.Category;
 
 @Category({ MediumTests.class, ClientTests.class })
-public class TestAsyncRegionLocatorTimeout {
+public class TestAsyncRegionLocator {
 
   @ClassRule
   public static final HBaseClassTestRule CLASS_RULE =
-      HBaseClassTestRule.forClass(TestAsyncRegionLocatorTimeout.class);
+    HBaseClassTestRule.forClass(TestAsyncRegionLocator.class);
 
   private static final HBaseTestingUtility TEST_UTIL = new HBaseTestingUtility();
 
@@ -96,7 +99,7 @@ public class TestAsyncRegionLocatorTimeout {
     TEST_UTIL.waitTableAvailable(TABLE_NAME);
     AsyncRegistry registry = AsyncRegistryFactory.getRegistry(TEST_UTIL.getConfiguration());
     CONN = new AsyncConnectionImpl(TEST_UTIL.getConfiguration(), registry,
-        registry.getClusterId().get(), User.getCurrent());
+      registry.getClusterId().get(), User.getCurrent());
     LOCATOR = CONN.getLocator();
   }
 
@@ -107,7 +110,7 @@ public class TestAsyncRegionLocatorTimeout {
   }
 
   @Test
-  public void test() throws InterruptedException, ExecutionException {
+  public void testTimeout() throws InterruptedException, ExecutionException {
     SLEEP_MS = 1000;
     long startNs = System.nanoTime();
     try {
@@ -115,7 +118,6 @@ public class TestAsyncRegionLocatorTimeout {
         TimeUnit.MILLISECONDS.toNanos(500)).get();
       fail();
     } catch (ExecutionException e) {
-      e.printStackTrace();
       assertThat(e.getCause(), instanceOf(TimeoutIOException.class));
     }
     long costMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNs);
@@ -129,4 +131,22 @@ public class TestAsyncRegionLocatorTimeout {
     assertEquals(loc.getServerName(),
       TEST_UTIL.getHBaseCluster().getRegionServer(0).getServerName());
   }
+
+  @Test
+  public void testNoCompletionException() {
+    // make sure that we do not get CompletionException
+    SLEEP_MS = 0;
+    AtomicReference<Throwable> errorHolder = new AtomicReference<>();
+    try {
+      LOCATOR.getRegionLocation(TableName.valueOf("NotExist"), EMPTY_START_ROW,
+        RegionLocateType.CURRENT, TimeUnit.SECONDS.toNanos(1))
+        .whenComplete((r, e) -> errorHolder.set(e)).join();
+      fail();
+    } catch (CompletionException e) {
+      // join will return a CompletionException, which is OK
+      assertThat(e.getCause(), instanceOf(TableNotFoundException.class));
+    }
+    // but we need to make sure that we do not get a CompletionException in the callback
+    assertThat(errorHolder.get(), instanceOf(TableNotFoundException.class));
+  }
 }