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,21 @@ 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));
+          } finally {
+            future.complete(locs);
+          }
+        });
+      return future;
     }, getClass().getSimpleName() + ".getAllRegionLocations");
   }
 

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

@@ -25,7 +25,10 @@
 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.assertThrows;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -61,6 +64,7 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import org.junit.runners.Parameterized.Parameter;
+import org.mockito.Mockito;
 
 import org.apache.hbase.thirdparty.com.google.common.io.Closeables;
 
@@ -460,4 +464,38 @@ 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();
+    AsyncConnectionImpl conn = (AsyncConnectionImpl)
+      ConnectionFactory.createAsyncConnection(TEST_UTIL.getConfiguration()).get();
+    AsyncTable<AdvancedScanResultConsumer> metaTable = conn.getTable(TableName.META_TABLE_NAME);
+    AsyncTable<AdvancedScanResultConsumer> spyMetaTable = Mockito.spy(metaTable);
+    Mockito.doThrow(new MockedBadScanResultException()).when(spyMetaTable)
+      .scan(Mockito.any(Scan.class), Mockito.any(AdvancedScanResultConsumer.class));
+    AsyncConnectionImpl spyConn = Mockito.spy(conn);
+    Mockito.doReturn(spyMetaTable).when(spyConn).getTable(TableName.META_TABLE_NAME);
+    assertThrows(MockedBadScanResultException.class, () ->
+      spyConn.getRegionLocator(TABLE_NAME).getAllRegionLocations().get());
+    List<RegionInfo> regions = TEST_UTIL.getAdmin().getRegions(TABLE_NAME);
+    for (RegionInfo region : regions) {
+      assertNull(conn.getLocator().getRegionLocationInCache(TABLE_NAME, region.getStartKey()));
+    }
+  }
+
+  /** This is used to mock that getting all region locations completes exceptionally. */
+  private static final class MockedBadScanResultException extends RuntimeException {}
 }