HBASE-26942 cache region locations when getAllRegionLocations

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

@@ -649,7 +649,7 @@ private void removeLocationFromCache(HRegionLocation loc) {
     }
   }
 
-  private void addLocationToCache(HRegionLocation loc) {
+  void addLocationToCache(HRegionLocation loc) {
     addToCache(getTableCache(loc.getRegion().getTable()), createRegionLocations(loc));
   }
 

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

@@ -18,6 +18,7 @@
 package org.apache.hadoop.hbase.client;
 
 import static org.apache.hadoop.hbase.trace.TraceUtil.tracedFuture;
+import static org.apache.hadoop.hbase.util.FutureUtils.addListener;
 
 import java.util.Arrays;
 import java.util.List;
@@ -61,8 +62,22 @@ public CompletableFuture<List<HRegionLocation>> getAllRegionLocations() {
         return conn.registry.getMetaRegionLocations()
           .thenApply(locs -> Arrays.asList(locs.getRegionLocations()));
       }
-      return ClientMetaTableAccessor
-        .getTableHRegionLocations(conn.getTable(TableName.META_TABLE_NAME), tableName);
+      CompletableFuture<List<HRegionLocation>> future = new CompletableFuture<>();
+      addListener(ClientMetaTableAccessor.getTableHRegionLocations(conn
+          .getTable(TableName.META_TABLE_NAME), tableName), (locs, error) -> {
+          if (error != null) {
+            future.completeExceptionally(error);
+            return;
+          }
+          try {
+            locs.forEach(loc -> conn.getLocator()
+              .getNonMetaRegionLocator().addLocationToCache(loc));
+            future.complete(locs);
+          } catch (Exception e) {
+            future.completeExceptionally(e);
+          }
+        });
+      return future;
     }, getClass().getSimpleName() + ".getAllRegionLocations");
   }
 

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

@@ -25,9 +25,13 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 import java.io.IOException;
+import java.net.SocketAddress;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
@@ -36,6 +40,7 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.stream.IntStream;
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.DoNotRetryIOException;
 import org.apache.hadoop.hbase.HBaseClassTestRule;
 import org.apache.hadoop.hbase.HBaseTestingUtil;
 import org.apache.hadoop.hbase.HRegionLocation;
@@ -64,6 +69,8 @@
 
 import org.apache.hbase.thirdparty.com.google.common.io.Closeables;
 
+import org.apache.hadoop.hbase.shaded.protobuf.generated.ClientProtos.ClientService;
+
 @Category({ MediumTests.class, ClientTests.class })
 @RunWith(Parameterized.class)
 public class TestAsyncNonMetaRegionLocator {
@@ -460,4 +467,58 @@ public void testConcurrentUpdateCachedLocationOnError() throws Exception {
     IntStream.range(0, 100).parallel()
       .forEach(i -> locator.updateCachedLocationOnError(loc, new NotServingRegionException()));
   }
+
+  @Test
+  public void testCacheLocationWhenGetAllLocations() throws Exception {
+    createMultiRegionTable();
+    AsyncConnectionImpl conn = (AsyncConnectionImpl)
+      ConnectionFactory.createAsyncConnection(TEST_UTIL.getConfiguration()).get();
+    conn.getRegionLocator(TABLE_NAME).getAllRegionLocations().get();
+    List<RegionInfo> regions = TEST_UTIL.getAdmin().getRegions(TABLE_NAME);
+    for (RegionInfo region : regions) {
+      assertNotNull(conn.getLocator().getRegionLocationInCache(TABLE_NAME, region.getStartKey()));
+    }
+  }
+
+  @Test
+  public void testCacheLocationExceptionallyWhenGetAllLocations() throws Exception {
+    createMultiRegionTable();
+    Configuration config = new Configuration(TEST_UTIL.getConfiguration());
+    config.setClass(ConnectionFactory.HBASE_CLIENT_ASYNC_CONNECTION_IMPL,
+      BadAsyncConnectionImpl.class, AsyncConnection.class);
+    AsyncConnectionImpl clientConn = (AsyncConnectionImpl)
+      ConnectionFactory.createAsyncConnection(config).get();
+
+    Exception ex = null;
+    try {
+      clientConn.getRegionLocator(TABLE_NAME).getAllRegionLocations().get();
+    } catch (Exception e) {
+      ex = e;
+    }
+
+    assertNotNull(ex);
+    assertTrue(ex instanceof ExecutionException);
+    assertTrue(ex.getCause() != null && ex.getCause() instanceof RegionLocationResultException);
+
+    List<RegionInfo> regions = TEST_UTIL.getAdmin().getRegions(TABLE_NAME);
+    for (RegionInfo region : regions) {
+      assertNull(clientConn.getLocator()
+        .getRegionLocationInCache(TABLE_NAME, region.getStartKey()));
+    }
+  }
+
+  private static final class RegionLocationResultException extends DoNotRetryIOException {}
+
+  private static final class BadAsyncConnectionImpl extends AsyncConnectionImpl {
+
+    public BadAsyncConnectionImpl(Configuration conf, ConnectionRegistry registry,
+        String clusterId, SocketAddress localAddress, User user) {
+      super(conf, registry, clusterId, localAddress, user);
+    }
+
+    @Override
+    ClientService.Interface getRegionServerStub(ServerName serverName) throws IOException {
+      throw new RegionLocationResultException();
+    }
+  }
 }