feat: support the bare minimum of RegionMetrics to support RegionSize… (#4342)

…Calculator in hbase 2x

Change-Id: I4d3c9176966772a0f6dac5dd4494142dbe49791a

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigtable-hbase/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> ☕️

If you write sample code, please follow the [samples format](
https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).

Author: igorbernstein2

bigtable-hbase-2.x-parent/bigtable-hbase-2.x/pom.xml

@@ -101,6 +101,12 @@ limitations under the License.
       <version>${junit.version}</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.truth</groupId>
+      <artifactId>truth</artifactId>
+      <version>${truth.version}</version>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>

bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/BigtableAdmin.java

@@ -39,6 +39,8 @@
 import org.apache.hadoop.hbase.ClusterStatus;
 import org.apache.hadoop.hbase.HRegionInfo;
 import org.apache.hadoop.hbase.HTableDescriptor;
+import org.apache.hadoop.hbase.RegionMetrics;
+import org.apache.hadoop.hbase.ServerName;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.AbstractBigtableAdmin;
 import org.apache.hadoop.hbase.client.AbstractBigtableConnection;
@@ -54,7 +56,6 @@
  */
 @InternalApi("For internal usage only")
 public abstract class BigtableAdmin extends AbstractBigtableAdmin {
-
   private final BigtableAsyncAdmin asyncAdmin;
 
   public BigtableAdmin(AbstractBigtableConnection connection) throws IOException {
@@ -295,6 +296,12 @@ public ClusterStatus getClusterStatus() throws IOException {
         -1);
   }
 
+  @Override
+  public List<RegionMetrics> getRegionMetrics(ServerName serverName, TableName tableName)
+      throws IOException {
+    return FutureUtil.unwrap(asyncAdmin.getRegionMetrics(serverName, tableName));
+  }
+
   private static Class<? extends BigtableAdmin> adminClass = null;
 
   /**

bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/BigtableAsyncAdmin.java

@@ -25,6 +25,7 @@
 import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest;
 import com.google.cloud.bigtable.admin.v2.models.RestoreTableRequest;
 import com.google.cloud.bigtable.admin.v2.models.Table;
+import com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.cloud.bigtable.hbase.BigtableOptionsFactory;
 import com.google.cloud.bigtable.hbase.util.Logger;
 import com.google.cloud.bigtable.hbase.util.ModifyTableBuilder;
@@ -35,9 +36,13 @@
 import com.google.common.base.Strings;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
+import com.google.protobuf.ByteString;
 import io.grpc.Status;
 import java.io.IOException;
+import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -55,6 +60,10 @@
 import net.bytebuddy.matcher.ElementMatchers;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.HTableDescriptor;
+import org.apache.hadoop.hbase.RegionMetrics;
+import org.apache.hadoop.hbase.ServerName;
+import org.apache.hadoop.hbase.Size;
+import org.apache.hadoop.hbase.Size.Unit;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.TableNotDisabledException;
 import org.apache.hadoop.hbase.TableNotEnabledException;
@@ -555,6 +564,36 @@ public CompletableFuture<List<RegionInfo>> getRegions(TableName tableName) {
         });
   }
 
+  @Override
+  public CompletableFuture<List<RegionMetrics>> getRegionMetrics(
+      ServerName ignored, TableName tableName) {
+    // TODO: implement caching
+    CompletableFuture<List<KeyOffset>> keyOffsetsFuture =
+        toCompletableFuture(
+            asyncConnection
+                .getBigtableApi()
+                .getDataClient()
+                .sampleRowKeysAsync(tableName.getNameAsString()));
+
+    return keyOffsetsFuture.thenApply(
+        keyOffsets -> {
+          long now = System.currentTimeMillis();
+          List<RegionMetrics> metrics = new ArrayList<>();
+          ByteString lastKey = ByteString.EMPTY;
+          long lastOffset = 0;
+          for (KeyOffset keyOffset : keyOffsets) {
+            byte[] regionName =
+                RegionInfo.createRegionName(
+                    tableName, lastKey.toByteArray(), Bytes.toBytes(now), false);
+            metrics.add(
+                BasicRegionMetrics.create(regionName, keyOffset.getOffsetBytes() - lastOffset));
+            lastKey = keyOffset.getKey();
+            lastOffset = keyOffset.getOffsetBytes();
+          }
+          return metrics;
+        });
+  }
+
   @Override
   public CompletableFuture<Void> snapshot(SnapshotDescription snapshot) {
     Objects.requireNonNull(snapshot);
@@ -601,4 +640,67 @@ public static BigtableAsyncAdmin createInstance(CommonConnection connection) thr
       throw new IOException(e);
     }
   }
+
+  private static Class<? extends RegionMetrics> regionMetricsClass = null;
+
+  private static synchronized Class<? extends RegionMetrics> getRegionMetricsSubclass()
+      throws NoSuchMethodException {
+    if (regionMetricsClass == null) {
+      regionMetricsClass =
+          new ByteBuddy()
+              .subclass(BasicRegionMetrics.class)
+              .name(BasicRegionMetrics.class.getName() + "Impl")
+              .method(ElementMatchers.isAbstract())
+              .intercept(InvocationHandlerAdapter.of(new UnsupportedOperationsHandler()))
+              .make()
+              .load(BigtableAsyncAdmin.class.getClassLoader())
+              .getLoaded();
+    }
+    return regionMetricsClass;
+  }
+
+  public abstract static class BasicRegionMetrics implements RegionMetrics {
+    private final byte[] regionName;
+    private final long size;
+
+    static RegionMetrics create(byte[] regionName, long size) {
+      try {
+        return getRegionMetricsSubclass()
+            .getConstructor(byte[].class, long.class)
+            .newInstance(regionName, size);
+      } catch (NoSuchMethodException
+          | InstantiationException
+          | IllegalAccessException
+          | InvocationTargetException e) {
+        throw new IllegalStateException("Failed to instantiate RegionMetrics subclass", e);
+      }
+    }
+
+    public BasicRegionMetrics(byte[] regionName, long size) {
+      this.regionName = regionName;
+      this.size = size;
+    }
+
+    @Override
+    public byte[] getRegionName() {
+      return regionName;
+    }
+
+    @Override
+    public int getStoreFileCount() {
+      return 1;
+    }
+
+    @Override
+    public Size getStoreFileSize() {
+      return new Size(size, Unit.BYTE);
+    }
+  }
+  /** Handler for unsupported operations for generating Admin class at runtime. */
+  public static class UnsupportedOperationsHandler implements InvocationHandler {
+    @Override
+    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+      throw new UnsupportedOperationException(method.getName());
+    }
+  }
 }

bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/test/java/com/google/cloud/bigtable/BigtableAsyncAdminTest.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.when;
+
+import com.google.api.core.SettableApiFuture;
+import com.google.cloud.bigtable.data.v2.models.KeyOffset;
+import com.google.cloud.bigtable.hbase.wrappers.AdminClientWrapper;
+import com.google.cloud.bigtable.hbase.wrappers.BigtableApi;
+import com.google.cloud.bigtable.hbase.wrappers.BigtableHBaseSettings;
+import com.google.cloud.bigtable.hbase.wrappers.DataClientWrapper;
+import com.google.cloud.bigtable.hbase2_x.BigtableAsyncAdmin;
+import com.google.protobuf.ByteString;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.RegionMetrics;
+import org.apache.hadoop.hbase.ServerName;
+import org.apache.hadoop.hbase.Size;
+import org.apache.hadoop.hbase.Size.Unit;
+import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.client.CommonConnection;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+@RunWith(JUnit4.class)
+public class BigtableAsyncAdminTest {
+  @Rule public final MockitoRule mockitoRule = MockitoJUnit.rule();
+
+  @Mock CommonConnection connection;
+  @Mock BigtableHBaseSettings settings;
+  @Mock BigtableApi bigtableApi;
+  @Mock AdminClientWrapper adminClientWrapper;
+  @Mock DataClientWrapper dataClientWrapper;
+
+  private BigtableAsyncAdmin admin;
+
+  @Before
+  public void setUp() throws Exception {
+    when(connection.getBigtableSettings()).thenReturn(settings);
+    when(connection.getBigtableApi()).thenReturn(bigtableApi);
+    when(bigtableApi.getAdminClient()).thenReturn(adminClientWrapper);
+    when(bigtableApi.getDataClient()).thenReturn(dataClientWrapper);
+    when(connection.getDisabledTables()).thenReturn(new HashSet<>());
+    when(connection.getConfiguration()).thenReturn(new Configuration(false));
+
+    admin = BigtableAsyncAdmin.createInstance(connection);
+  }
+
+  @Test
+  public void testRegionMetricsFileSize() throws ExecutionException, InterruptedException {
+    String tableId = "fake-table";
+    SettableApiFuture<List<KeyOffset>> resultFuture = SettableApiFuture.create();
+    resultFuture.set(
+        Arrays.asList(
+            KeyOffset.create(ByteString.copyFromUtf8("a"), 100),
+            KeyOffset.create(ByteString.copyFromUtf8("b"), 250),
+            KeyOffset.create(ByteString.copyFromUtf8(""), 401)));
+    when(dataClientWrapper.sampleRowKeysAsync(tableId)).thenReturn(resultFuture);
+    List<RegionMetrics> results =
+        admin
+            .getRegionMetrics(ServerName.valueOf("host", 123, 0), TableName.valueOf("fake-table"))
+            .get();
+    assertThat(results).hasSize(3);
+    assertThat(results.get(0).getStoreFileSize()).isEqualTo(new Size(100, Unit.BYTE));
+    assertThat(results.get(1).getStoreFileSize()).isEqualTo(new Size(150, Unit.BYTE));
+    assertThat(results.get(2).getStoreFileSize()).isEqualTo(new Size(151, Unit.BYTE));
+  }
+}