fix: per-connection metrics issue when using a different Bigtable project (#2143)

rkaregar · web-flow · commit 8dbd680abc1a · 2024-03-05T12:37:56.000-05:00
* fix: per-connection metric issue when using a different Bigtable project.

* Improve annotation.
diff --git a/google-cloud-bigtable-stats/src/main/java/com/google/cloud/bigtable/stats/BigtableCreateTimeSeriesExporter.java b/google-cloud-bigtable-stats/src/main/java/com/google/cloud/bigtable/stats/BigtableCreateTimeSeriesExporter.java
@@ -58,7 +58,7 @@ public void export(Collection<Metric> metrics) {
                   Collectors.groupingBy(
                       timeSeries ->
                           BigtableStackdriverExportUtils.getProjectId(
-                              metric.getMetricDescriptor(), timeSeries),
+                              metric.getMetricDescriptor(), timeSeries, gceOrGkeMonitoredResource),
                       Collectors.mapping(
                           timeSeries ->
                               BigtableStackdriverExportUtils.convertTimeSeries(
diff --git a/google-cloud-bigtable-stats/src/main/java/com/google/cloud/bigtable/stats/BigtableStackdriverExportUtils.java b/google-cloud-bigtable-stats/src/main/java/com/google/cloud/bigtable/stats/BigtableStackdriverExportUtils.java
@@ -22,6 +22,7 @@
 import com.google.api.Metric;
 import com.google.api.MetricDescriptor.MetricKind;
 import com.google.api.MonitoredResource;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
 import com.google.monitoring.v3.TimeInterval;
@@ -55,6 +56,10 @@
 
 class BigtableStackdriverExportUtils {
   private static final String BIGTABLE_RESOURCE_TYPE = "bigtable_client_raw";
+
+  @VisibleForTesting static final String GCE_RESOURCE_TYPE = "gce_instance";
+  @VisibleForTesting static final String GKE_RESOURCE_TYPE = "k8s_container";
+  @VisibleForTesting static final String GCE_OR_GKE_PROJECT_ID_KEY = "project_id";
   private static final Logger logger =
       Logger.getLogger(BigtableStackdriverExportUtils.class.getName());
 
@@ -209,7 +214,19 @@ private static com.google.monitoring.v3.TimeSeries.Builder setupBuilderForGceOrG
     return builder;
   }
 
-  static String getProjectId(MetricDescriptor metricDescriptor, TimeSeries timeSeries) {
+  static String getProjectId(
+      MetricDescriptor metricDescriptor,
+      TimeSeries timeSeries,
+      MonitoredResource gceOrGkeMonitoredResource) {
+    if (isBigtableTableMetric(metricDescriptor)) {
+      return getProjectIdForBigtableTableResource(metricDescriptor, timeSeries);
+    } else {
+      return getProjectIdForGceOrGkeResource(gceOrGkeMonitoredResource);
+    }
+  }
+
+  static String getProjectIdForBigtableTableResource(
+      MetricDescriptor metricDescriptor, TimeSeries timeSeries) {
     List<LabelKey> labelKeys = metricDescriptor.getLabelKeys();
     List<LabelValue> labelValues = timeSeries.getLabelValues();
     for (int i = 0; i < labelKeys.size(); i++) {
@@ -220,6 +237,15 @@ static String getProjectId(MetricDescriptor metricDescriptor, TimeSeries timeSer
     throw new IllegalStateException("Can't find project id for the current timeseries");
   }
 
+  static String getProjectIdForGceOrGkeResource(MonitoredResource gceOrGkeMonitoredResource) {
+    if (!gceOrGkeMonitoredResource.getType().equals(GCE_RESOURCE_TYPE)
+        && !gceOrGkeMonitoredResource.getType().equals(GKE_RESOURCE_TYPE)) {
+      throw new IllegalStateException(
+          "Expected GCE or GKE resource type, but found " + gceOrGkeMonitoredResource);
+    }
+    return gceOrGkeMonitoredResource.getLabelsOrThrow(GCE_OR_GKE_PROJECT_ID_KEY);
+  }
+
   static String getDefaultTaskValue() {
     // Something like '<pid>@<hostname>'
     final String jvmName = ManagementFactory.getRuntimeMXBean().getName();
diff --git a/google-cloud-bigtable-stats/src/test/java/com/google/cloud/bigtable/stats/BigtableCreateTimeSeriesExporterTest.java b/google-cloud-bigtable-stats/src/test/java/com/google/cloud/bigtable/stats/BigtableCreateTimeSeriesExporterTest.java
@@ -61,6 +61,8 @@ public class BigtableCreateTimeSeriesExporterTest {
   private static final String bigtableZone = "us-east-1";
   private static final String bigtableCluster = "cluster-1";
   private static final String clientName = "client-name";
+  private static final String gceProjectId = "fake-gce-project";
+  private static final String gkeProjectId = "fake-gke-project";
 
   @Rule public final MockitoRule mockitoRule = MockitoJUnit.rule();
 
@@ -121,6 +123,7 @@ public void testTimeSeriesForMetricWithBigtableResource() {
 
     CreateTimeSeriesRequest request = argumentCaptor.getValue();
 
+    assertThat(request.getName()).isEqualTo("projects/" + bigtableProjectId);
     assertThat(request.getTimeSeriesList()).hasSize(1);
 
     com.google.monitoring.v3.TimeSeries timeSeries = request.getTimeSeriesList().get(0);
@@ -148,8 +151,9 @@ public void testTimeSeriesForMetricWithGceResource() {
         new BigtableCreateTimeSeriesExporter(
             fakeMetricServiceClient,
             MonitoredResource.newBuilder()
-                .setType("gce-instance")
-                .putLabels("some-gce-key", "some-gce-value")
+                .setType(BigtableStackdriverExportUtils.GCE_RESOURCE_TYPE)
+                .putLabels(BigtableStackdriverExportUtils.GCE_OR_GKE_PROJECT_ID_KEY, gceProjectId)
+                .putLabels("another-gce-key", "another-gce-value")
                 .build());
     ArgumentCaptor<CreateTimeSeriesRequest> argumentCaptor =
         ArgumentCaptor.forClass(CreateTimeSeriesRequest.class);
@@ -197,12 +201,17 @@ public void testTimeSeriesForMetricWithGceResource() {
 
     CreateTimeSeriesRequest request = argumentCaptor.getValue();
 
+    assertThat(request.getName()).isEqualTo("projects/" + gceProjectId);
     assertThat(request.getTimeSeriesList()).hasSize(1);
 
     com.google.monitoring.v3.TimeSeries timeSeries = request.getTimeSeriesList().get(0);
 
     assertThat(timeSeries.getResource().getLabelsMap())
-        .containsExactly("some-gce-key", "some-gce-value");
+        .containsExactly(
+            BigtableStackdriverExportUtils.GCE_OR_GKE_PROJECT_ID_KEY,
+            gceProjectId,
+            "another-gce-key",
+            "another-gce-value");
 
     assertThat(timeSeries.getMetric().getLabelsMap()).hasSize(5);
     assertThat(timeSeries.getMetric().getLabelsMap())
@@ -225,8 +234,9 @@ public void testTimeSeriesForMetricWithGkeResource() {
         new BigtableCreateTimeSeriesExporter(
             fakeMetricServiceClient,
             MonitoredResource.newBuilder()
-                .setType("gke-container")
-                .putLabels("some-gke-key", "some-gke-value")
+                .setType(BigtableStackdriverExportUtils.GKE_RESOURCE_TYPE)
+                .putLabels(BigtableStackdriverExportUtils.GCE_OR_GKE_PROJECT_ID_KEY, gkeProjectId)
+                .putLabels("another-gke-key", "another-gke-value")
                 .build());
     ArgumentCaptor<CreateTimeSeriesRequest> argumentCaptor =
         ArgumentCaptor.forClass(CreateTimeSeriesRequest.class);
@@ -275,12 +285,17 @@ public void testTimeSeriesForMetricWithGkeResource() {
 
     CreateTimeSeriesRequest request = argumentCaptor.getValue();
 
+    assertThat(request.getName()).isEqualTo("projects/" + gkeProjectId);
     assertThat(request.getTimeSeriesList()).hasSize(1);
 
     com.google.monitoring.v3.TimeSeries timeSeries = request.getTimeSeriesList().get(0);
 
     assertThat(timeSeries.getResource().getLabelsMap())
-        .containsExactly("some-gke-key", "some-gke-value");
+        .containsExactly(
+            BigtableStackdriverExportUtils.GCE_OR_GKE_PROJECT_ID_KEY,
+            gkeProjectId,
+            "another-gke-key",
+            "another-gke-value");
 
     assertThat(timeSeries.getMetric().getLabelsMap()).hasSize(5);
     assertThat(timeSeries.getMetric().getLabelsMap())
