HBASE-26648 Improve fidelity of RegionLocator spans (addendum)

Addendum extends the test case to cover both `RpcConnectionRegistry` and `ZKConnectionRegistry`.

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

@@ -34,7 +34,9 @@ import java.util.concurrent.TimeUnit;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.ConnectionRule;
 import org.apache.hadoop.hbase.HBaseClassTestRule;
+import org.apache.hadoop.hbase.HBaseConfiguration;
 import org.apache.hadoop.hbase.HBaseTestingUtility;
+import org.apache.hadoop.hbase.HConstants;
 import org.apache.hadoop.hbase.HRegionLocation;
 import org.apache.hadoop.hbase.MatcherPredicate;
 import org.apache.hadoop.hbase.MiniClusterRule;
@@ -54,39 +56,61 @@ import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
+import org.junit.experimental.runners.Enclosed;
 import org.junit.rules.ExternalResource;
 import org.junit.rules.RuleChain;
+import org.junit.rules.TestName;
 import org.junit.rules.TestRule;
+import org.junit.runner.RunWith;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-@Category({ MediumTests.class, ClientTests.class })
+@RunWith(Enclosed.class)
 public class TestAsyncMetaRegionLocator {
   private static final Logger logger = LoggerFactory.getLogger(TestAsyncMetaRegionLocator.class);
 
-  @ClassRule
-  public static final HBaseClassTestRule CLASS_RULE =
-    HBaseClassTestRule.forClass(TestAsyncMetaRegionLocator.class);
-
-  private static final OpenTelemetryClassRule otelClassRule = OpenTelemetryClassRule.create();
-  private static final MiniClusterRule miniClusterRule = MiniClusterRule.newBuilder()
-    .setMiniClusterOption(StartMiniClusterOption.builder().numWorkers(3).build()).build();
-  private static final ConnectionRule connectionRule =
-    ConnectionRule.createAsyncConnectionRule(miniClusterRule::createAsyncConnection);
-
   private static final class Setup extends ExternalResource {
+
+    private final MiniClusterRule miniClusterRule;
+    private final ConnectionRule connectionRule;
+
+    private boolean initialized = false;
+    private HBaseTestingUtility testUtil;
+    private AsyncMetaRegionLocator locator;
     private ConnectionRegistry registry;
 
+    public Setup(final ConnectionRule connectionRule, final MiniClusterRule miniClusterRule) {
+      this.connectionRule = connectionRule;
+      this.miniClusterRule = miniClusterRule;
+    }
+
+    public HBaseTestingUtility getTestingUtility() {
+      assertInitialized();
+      return testUtil;
+    }
+
+    public AsyncMetaRegionLocator getLocator() {
+      assertInitialized();
+      return locator;
+    }
+
+    private void assertInitialized() {
+      if (!initialized) {
+        throw new IllegalStateException("before method has not been called.");
+      }
+    }
+
     @Override
     protected void before() throws Throwable {
       final AsyncAdmin admin = connectionRule.getAsyncConnection().getAdmin();
-      TEST_UTIL = miniClusterRule.getTestingUtility();
+      testUtil = miniClusterRule.getTestingUtility();
       HBaseTestingUtility.setReplicas(admin, TableName.META_TABLE_NAME, 3);
-      TEST_UTIL.waitUntilNoRegionsInTransition();
+      testUtil.waitUntilNoRegionsInTransition();
-      registry = ConnectionRegistryFactory.getRegistry(TEST_UTIL.getConfiguration());
+      registry = ConnectionRegistryFactory.getRegistry(testUtil.getConfiguration());
-      RegionReplicaTestHelper.waitUntilAllMetaReplicasAreReady(TEST_UTIL, registry);
+      RegionReplicaTestHelper.waitUntilAllMetaReplicasAreReady(testUtil, registry);
       admin.balancerSwitch(false).get();
-      LOCATOR = new AsyncMetaRegionLocator(registry);
+      locator = new AsyncMetaRegionLocator(registry);
+      initialized = true;
     }
 
     @Override
@@ -95,62 +119,127 @@ public class TestAsyncMetaRegionLocator {
     }
   }
 
-  @ClassRule
+  public static abstract class AbstractBase {
-  public static final TestRule classRule = RuleChain.outerRule(otelClassRule)
+    private final OpenTelemetryClassRule otelClassRule = OpenTelemetryClassRule.create();
-    .around(miniClusterRule).around(connectionRule).around(new Setup());
+    private final MiniClusterRule miniClusterRule;
+    private final Setup setup;
 
-  private static HBaseTestingUtility TEST_UTIL;
+    protected Matcher<SpanData> parentSpanMatcher;
-  private static AsyncMetaRegionLocator LOCATOR;
+    protected List<SpanData> spans;
+    protected Matcher<SpanData> registryGetMetaRegionLocationsMatcher;
 
-  @Rule
+    @Rule
-  public final OpenTelemetryTestRule otelTestRule = new OpenTelemetryTestRule(otelClassRule);
+    public final TestRule classRule;
 
-  @Test
+    @Rule
-  public void test() throws Exception {
+    public final OpenTelemetryTestRule otelTestRule = new OpenTelemetryTestRule(otelClassRule);
-    TraceUtil.trace(() -> {
-      try {
-        testLocator(miniClusterRule.getTestingUtility(), TableName.META_TABLE_NAME, new Locator() {
-          @Override
-          public void updateCachedLocationOnError(HRegionLocation loc, Throwable error) {
-            LOCATOR.updateCachedLocationOnError(loc, error);
-          }
 
-          @Override
+    @Rule
-          public RegionLocations getRegionLocations(TableName tableName, int replicaId,
+    public TestName testName = new TestName();
-            boolean reload) throws Exception {
-            return LOCATOR.getRegionLocations(replicaId, reload).get();
-          }
-        });
-      } catch (Exception e) {
-        throw new RuntimeException(e);
-      }
-    }, "test");
 
-    final Configuration conf = TEST_UTIL.getConfiguration();
+    public AbstractBase() {
-    final Matcher<SpanData> parentSpanMatcher = allOf(hasName("test"), hasEnded());
+      miniClusterRule = MiniClusterRule.newBuilder()
-    Waiter.waitFor(conf, TimeUnit.SECONDS.toMillis(5),
+        .setMiniClusterOption(StartMiniClusterOption.builder().numWorkers(3).build())
-      new MatcherPredicate<>(otelClassRule::getSpans, hasItem(parentSpanMatcher)));
+        .setConfiguration(() -> {
-    final List<SpanData> spans = otelClassRule.getSpans();
+          final Configuration conf = HBaseConfiguration.create();
-    if (logger.isDebugEnabled()) {
+          conf.setClass(HConstants.CLIENT_CONNECTION_REGISTRY_IMPL_CONF_KEY,
-      StringTraceRenderer renderer = new StringTraceRenderer(spans);
+            getConnectionRegistryClass(), ConnectionRegistry.class);
-      renderer.render(logger::debug);
+          return conf;
+        }).build();
+      final ConnectionRule connectionRule =
+        ConnectionRule.createAsyncConnectionRule(miniClusterRule::createAsyncConnection);
+      setup = new Setup(connectionRule, miniClusterRule);
+      classRule = RuleChain.outerRule(otelClassRule).around(miniClusterRule).around(connectionRule)
+        .around(setup);
     }
 
-    assertThat(spans, hasItem(parentSpanMatcher));
+    protected abstract Class<? extends ConnectionRegistry> getConnectionRegistryClass();
-    final SpanData parentSpan =
-      spans.stream().filter(parentSpanMatcher::matches).findAny().orElseThrow(AssertionError::new);
 
-    final Matcher<SpanData> registryGetMetaRegionLocationsMatcher =
+    @Test
-      allOf(hasName(endsWith("ConnectionRegistry.getMetaRegionLocations")),
+    public void test() throws Exception {
-        hasParentSpanId(parentSpan), hasKind(SpanKind.INTERNAL), hasEnded());
+      final AsyncMetaRegionLocator locator = setup.getLocator();
-    assertThat(spans, hasItem(registryGetMetaRegionLocationsMatcher));
+      final HBaseTestingUtility testUtil = setup.getTestingUtility();
-    final SpanData registry_getMetaRegionLocationsSpan =
+
-      spans.stream().filter(registryGetMetaRegionLocationsMatcher::matches).findAny()
+      TraceUtil.trace(() -> {
+        try {
+          testLocator(miniClusterRule.getTestingUtility(), TableName.META_TABLE_NAME,
+            new Locator() {
+              @Override
+              public void updateCachedLocationOnError(HRegionLocation loc, Throwable error) {
+                locator.updateCachedLocationOnError(loc, error);
+              }
+
+              @Override
+              public RegionLocations getRegionLocations(TableName tableName, int replicaId,
+                boolean reload) throws Exception {
+                return locator.getRegionLocations(replicaId, reload).get();
+              }
+            });
+        } catch (Exception e) {
+          throw new RuntimeException(e);
+        }
+      }, testName.getMethodName());
+
+      final Configuration conf = testUtil.getConfiguration();
+      parentSpanMatcher = allOf(hasName(testName.getMethodName()), hasEnded());
+      Waiter.waitFor(conf, TimeUnit.SECONDS.toMillis(5),
+        new MatcherPredicate<>(otelClassRule::getSpans, hasItem(parentSpanMatcher)));
+      spans = otelClassRule.getSpans();
+      if (logger.isDebugEnabled()) {
+        StringTraceRenderer renderer = new StringTraceRenderer(spans);
+        renderer.render(logger::debug);
+      }
+      assertThat(spans, hasItem(parentSpanMatcher));
+      final SpanData parentSpan = spans.stream().filter(parentSpanMatcher::matches).findAny()
         .orElseThrow(AssertionError::new);
 
-    final Matcher<SpanData> clientGetMetaRegionLocationsMatcher =
+      registryGetMetaRegionLocationsMatcher =
-      allOf(hasName(endsWith("ClientMetaService/GetMetaRegionLocations")),
+        allOf(hasName(endsWith("ConnectionRegistry.getMetaRegionLocations")),
+          hasParentSpanId(parentSpan), hasKind(SpanKind.INTERNAL), hasEnded());
+      assertThat(spans, hasItem(registryGetMetaRegionLocationsMatcher));
+    }
+  }
+
+  /**
+   * Test covers when client is configured with {@link ZKConnectionRegistry}.
+   */
+  @Category({ MediumTests.class, ClientTests.class })
+  public static class TestZKConnectionRegistry extends AbstractBase {
+    @ClassRule
+    public static final HBaseClassTestRule CLASS_RULE =
+      HBaseClassTestRule.forClass(TestZKConnectionRegistry.class);
+
+    @Override
+    protected Class<? extends ConnectionRegistry> getConnectionRegistryClass() {
+      return ZKConnectionRegistry.class;
+    }
+  }
+
+  /**
+   * Test covers when client is configured with {@link RpcConnectionRegistry}.
+   */
+  @Category({ MediumTests.class, ClientTests.class })
+  public static class TestRpcConnectionRegistry extends AbstractBase {
+    @ClassRule
+    public static final HBaseClassTestRule CLASS_RULE =
+      HBaseClassTestRule.forClass(TestRpcConnectionRegistry.class);
+
+    @Override
+    protected Class<? extends ConnectionRegistry> getConnectionRegistryClass() {
+      return RpcConnectionRegistry.class;
+    }
+
+    @Test
+    @Override
+    public void test() throws Exception {
+      super.test();
+      final SpanData registry_getMetaRegionLocationsSpan =
+        spans.stream().filter(registryGetMetaRegionLocationsMatcher::matches).findAny()
+          .orElseThrow(AssertionError::new);
+      final Matcher<SpanData> clientGetMetaRegionLocationsMatcher = allOf(
+        hasName(endsWith("ClientMetaService/GetMetaRegionLocations")),
         hasParentSpanId(registry_getMetaRegionLocationsSpan), hasKind(SpanKind.CLIENT), hasEnded());
-    assertThat(spans, hasItem(clientGetMetaRegionLocationsMatcher));
+      assertThat(spans, hasItem(clientGetMetaRegionLocationsMatcher));
+    }
   }
 }