add outstanding metrics tracker

Author: sushanb

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java

@@ -26,12 +26,7 @@
 import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
 import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
 import com.google.cloud.bigtable.data.v2.internal.JwtCredentialsWithAudience;
-import com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants;
-import com.google.cloud.bigtable.data.v2.stub.metrics.CustomOpenTelemetryMetricsProvider;
-import com.google.cloud.bigtable.data.v2.stub.metrics.DefaultMetricsProvider;
-import com.google.cloud.bigtable.data.v2.stub.metrics.ErrorCountPerConnectionMetricTracker;
-import com.google.cloud.bigtable.data.v2.stub.metrics.MetricsProvider;
-import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider;
+import com.google.cloud.bigtable.data.v2.stub.metrics.*;
 import com.google.cloud.bigtable.gaxx.grpc.BigtableTransportChannelProvider;
 import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer;
 import io.grpc.ManagedChannelBuilder;
@@ -98,7 +93,7 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
 
     @Nullable OpenTelemetrySdk internalOtel = null;
     @Nullable ErrorCountPerConnectionMetricTracker errorCountPerConnectionMetricTracker = null;
-
+    @Nullable OutstandingRpcsMetricTracker outstandingRpcsMetricTracker = null;
     // Internal metrics are scoped to the connections, so we need a mutable transportProvider,
     // otherwise there is
     // no reason to build the internal OtelProvider
@@ -111,6 +106,8 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
         errorCountPerConnectionMetricTracker =
             setupPerConnectionErrorTracer(builder, transportProvider, internalOtel);
 
+        outstandingRpcsMetricTracker = new OutstandingRpcsMetricTracker(internalOtel, "LB_POLICY");
+
         // Configure grpc metrics
         configureGrpcOtel(transportProvider, internalOtel);
       }
@@ -137,7 +134,9 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
 
       BigtableTransportChannelProvider btTransportProvider =
           BigtableTransportChannelProvider.create(
-              (InstantiatingGrpcChannelProvider) transportProvider.build(), channelPrimer);
+              (InstantiatingGrpcChannelProvider) transportProvider.build(),
+              channelPrimer,
+              outstandingRpcsMetricTracker);
 
       builder.setTransportChannelProvider(btTransportProvider);
     }
@@ -149,6 +148,10 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
           clientContext.getExecutor());
     }
 
+    if (outstandingRpcsMetricTracker != null) {
+      outstandingRpcsMetricTracker.start(clientContext.getExecutor());
+    }
+
     return new BigtableClientContext(
         clientContext, openTelemetry, internalOtel, settings.getMetricsProvider());
   }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsConstants.java

@@ -70,6 +70,7 @@ public class BuiltinMetricsConstants {
   static final String REMAINING_DEADLINE_NAME = "remaining_deadline";
   static final String CLIENT_BLOCKING_LATENCIES_NAME = "throttling_latencies";
   static final String PER_CONNECTION_ERROR_COUNT_NAME = "per_connection_error_count";
+  static final String OUTSTANDING_RPCS_PER_CHANNEL_NAME = "connection_pool/outstanding_rpc";
 
   // Start allow list of metrics that will be exported as internal
   public static final Map<String, Set<String>> GRPC_METRICS =

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/OutstandingRpcsMetricTracker.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2025 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
+ *
+ *     https://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.data.v2.stub.metrics;
+
+import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.METER_NAME;
+import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.OUTSTANDING_RPCS_PER_CHANNEL_NAME;
+
+import com.google.cloud.bigtable.gaxx.grpc.BigtableChannelInsight;
+import com.google.cloud.bigtable.gaxx.grpc.BigtableChannelInsightsProvider;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.metrics.LongHistogram;
+import io.opentelemetry.api.metrics.Meter;
+import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class OutstandingRpcsMetricTracker implements Runnable {
+  private static final int SAMPLING_PERIOD_SECONDS = 60;
+  private final LongHistogram outstandingRpcsHistogram;
+  private final Attributes baseAttributes;
+  private final AtomicReference<BigtableChannelInsightsProvider>
+      bigtableChannelInsightsProviderRef = new AtomicReference<>();
+
+  public OutstandingRpcsMetricTracker(OpenTelemetry openTelemetry, String lbPolicy) {
+    Meter meter = openTelemetry.getMeter(METER_NAME);
+    this.outstandingRpcsHistogram =
+        meter
+            .histogramBuilder(OUTSTANDING_RPCS_PER_CHANNEL_NAME)
+            .ofLongs()
+            .setDescription(
+                "A distribution of the number of outstanding RPCs per connection in the client pool, sampled periodically.")
+            .setUnit("1")
+            .build();
+
+    this.baseAttributes =
+        Attributes.builder().put("transport_type", "grpc").put("lb_policy", lbPolicy).build();
+  }
+
+  /**
+   * Registers the provider for the channel pool entries. This should be called by the component
+   * that creates the BigtableChannelPool.
+   */
+  public void registerChannelInsightsProvider(
+      BigtableChannelInsightsProvider channelInsightsProvider) {
+    this.bigtableChannelInsightsProviderRef.set(channelInsightsProvider);
+  }
+
+  /** Starts the periodic collection. */
+  public ScheduledFuture<?> start(ScheduledExecutorService scheduler) {
+    return scheduler.scheduleAtFixedRate(
+        this, SAMPLING_PERIOD_SECONDS, SAMPLING_PERIOD_SECONDS, TimeUnit.SECONDS);
+  }
+
+  @Override
+  public void run() {
+    BigtableChannelInsightsProvider channelInsightsProvider =
+        bigtableChannelInsightsProviderRef.get();
+    if (channelInsightsProvider == null) {
+      return; // Not registered yet
+    }
+    List<? extends BigtableChannelInsight> channelInsights =
+        channelInsightsProvider.getChannelInfos();
+    if (channelInsights == null || channelInsights.isEmpty()) {
+      return;
+    }
+    for (BigtableChannelInsight info : channelInsights) {
+      long currentOutstandingUnaryRpcs = info.getOutstandingStreamingRpcs();
+      long currentOutstandingStreamingRpcs = info.getOutstandingStreamingRpcs();
+      // Record outstanding unary RPCs with streaming=false
+      Attributes unaryAttributes = baseAttributes.toBuilder().put("streaming", false).build();
+      outstandingRpcsHistogram.record(currentOutstandingUnaryRpcs, unaryAttributes);
+
+      // Record outstanding streaming RPCs with streaming=true
+      Attributes streamingAttributes = baseAttributes.toBuilder().put("streaming", true).build();
+      outstandingRpcsHistogram.record(currentOutstandingStreamingRpcs, streamingAttributes);
+    }
+  }
+}

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelInsight.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2025 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
+ *
+ *     https://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.gaxx.grpc;
+
+import com.google.api.core.InternalApi;
+
+/** Provides insights about a single channel in the channel pool. */
+@InternalApi
+public interface BigtableChannelInsight {
+  /** Gets the current number of outstanding Unary RPCs on this channel. */
+  int getOutstandingUnaryRpcs();
+
+  int getOutstandingStreamingRpcs();
+}

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelInsightsProvider.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2025 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
+ *
+ *     https://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.gaxx.grpc;
+
+import com.google.api.core.InternalApi;
+import java.util.List;
+
+@InternalApi
+@FunctionalInterface
+public interface BigtableChannelInsightsProvider {
+  /** Gets the current list of BigtableChannelInsight objects. */
+  List<? extends BigtableChannelInsight> getChannelInfos();
+}

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java

@@ -57,7 +57,7 @@
  * <p>Internal API
  */
 @InternalApi
-public class BigtableChannelPool extends ManagedChannel {
+public class BigtableChannelPool extends ManagedChannel implements BigtableChannelInsightsProvider {
   @VisibleForTesting
   static final Logger LOG = Logger.getLogger(BigtableChannelPool.class.getName());
 
@@ -507,8 +507,14 @@ private Entry getEntry(int affinity) {
     return localEntries.get(index);
   }
 
+  /** Gets the current list of BigtableChannelInsight objects. */
+  @Override
+  public List<? extends BigtableChannelInsight> getChannelInfos() {
+    return entries.get();
+  }
+
   /** Bundles a gRPC {@link ManagedChannel} with some usage accounting. */
-  static class Entry {
+  static class Entry implements BigtableChannelInsight {
     private final ManagedChannel channel;
 
     /**
@@ -628,6 +634,17 @@ private void shutdown() {
         channel.shutdown();
       }
     }
+
+    /** Gets the current number of outstanding Unary RPCs on this channel. */
+    @Override
+    public int getOutstandingUnaryRpcs() {
+      return outstandingUnaryRpcs.get();
+    }
+
+    @Override
+    public int getOutstandingStreamingRpcs() {
+      return 0;
+    }
   }
 
   /** Thin wrapper to ensure that new calls are properly reference counted. */

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java

@@ -23,12 +23,14 @@
 import com.google.api.gax.rpc.TransportChannel;
 import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.auth.Credentials;
+import com.google.cloud.bigtable.data.v2.stub.metrics.OutstandingRpcsMetricTracker;
 import com.google.common.base.Preconditions;
 import io.grpc.ManagedChannel;
 import java.io.IOException;
 import java.util.Map;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
+import javax.annotation.Nullable;
 
 /**
  * An instance of TransportChannelProvider that provides a TransportChannel through a supplied
@@ -39,12 +41,15 @@ public final class BigtableTransportChannelProvider implements TransportChannelP
 
   private final InstantiatingGrpcChannelProvider delegate;
   private final ChannelPrimer channelPrimer;
+  @Nullable private final OutstandingRpcsMetricTracker outstandingRpcsMetricTracker;
 
   private BigtableTransportChannelProvider(
       InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider,
-      ChannelPrimer channelPrimer) {
+      ChannelPrimer channelPrimer,
+      OutstandingRpcsMetricTracker outstandingRpcsMetricTracker) {
     delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider);
     this.channelPrimer = channelPrimer;
+    this.outstandingRpcsMetricTracker = outstandingRpcsMetricTracker;
   }
 
   @Override
@@ -66,7 +71,8 @@ public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService ex
   public BigtableTransportChannelProvider withExecutor(Executor executor) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor);
-    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
+    return new BigtableTransportChannelProvider(
+        newChannelProvider, channelPrimer, outstandingRpcsMetricTracker);
   }
 
   @Override
@@ -78,7 +84,8 @@ public boolean needsHeaders() {
   public BigtableTransportChannelProvider withHeaders(Map<String, String> headers) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers);
-    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
+    return new BigtableTransportChannelProvider(
+        newChannelProvider, channelPrimer, outstandingRpcsMetricTracker);
   }
 
   @Override
@@ -90,7 +97,8 @@ public boolean needsEndpoint() {
   public TransportChannelProvider withEndpoint(String endpoint) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint);
-    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
+    return new BigtableTransportChannelProvider(
+        newChannelProvider, channelPrimer, outstandingRpcsMetricTracker);
   }
 
   @Deprecated
@@ -104,7 +112,8 @@ public boolean acceptsPoolSize() {
   public TransportChannelProvider withPoolSize(int size) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withPoolSize(size);
-    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
+    return new BigtableTransportChannelProvider(
+        newChannelProvider, channelPrimer, outstandingRpcsMetricTracker);
   }
 
   /** Expected to only be called once when BigtableClientContext is created */
@@ -136,6 +145,10 @@ public TransportChannel getTransportChannel() throws IOException {
     BigtableChannelPool btChannelPool =
         BigtableChannelPool.create(btPoolSettings, channelFactory, channelPrimer);
 
+    if (outstandingRpcsMetricTracker != null) {
+      outstandingRpcsMetricTracker.registerChannelInsightsProvider(btChannelPool::getChannelInfos);
+    }
+
     return GrpcTransportChannel.create(btChannelPool);
   }
 
@@ -153,13 +166,16 @@ public boolean needsCredentials() {
   public TransportChannelProvider withCredentials(Credentials credentials) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withCredentials(credentials);
-    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
+    return new BigtableTransportChannelProvider(
+        newChannelProvider, channelPrimer, outstandingRpcsMetricTracker);
   }
 
   /** Creates a BigtableTransportChannelProvider. */
   public static BigtableTransportChannelProvider create(
       InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider,
-      ChannelPrimer channelPrimer) {
-    return new BigtableTransportChannelProvider(instantiatingGrpcChannelProvider, channelPrimer);
+      ChannelPrimer channelPrimer,
+      OutstandingRpcsMetricTracker outstandingRpcsMetricTracke) {
+    return new BigtableTransportChannelProvider(
+        instantiatingGrpcChannelProvider, channelPrimer, outstandingRpcsMetricTracke);
   }
 }