diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index 92a984a015..92cafdc52c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/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()); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsConstants.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsConstants.java index 78ed689cc3..652db2cf92 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsConstants.java +++ b/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> GRPC_METRICS = diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/OutstandingRpcsMetricTracker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/OutstandingRpcsMetricTracker.java new file mode 100644 index 0000000000..25885275ce --- /dev/null +++ b/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 + 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 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); + } + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelInsight.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelInsight.java new file mode 100644 index 0000000000..dc496eaf43 --- /dev/null +++ b/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(); +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelInsightsProvider.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelInsightsProvider.java new file mode 100644 index 0000000000..3353016c14 --- /dev/null +++ b/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 getChannelInfos(); +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java index 173722f2f4..199eee5b33 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java @@ -57,7 +57,7 @@ *

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()); @@ -196,7 +196,7 @@ private int pickEntryIndexLeastInFlight() { for (int i = 0; i < localEntries.size(); i++) { Entry entry = localEntries.get(i); - int rpcs = entry.outstandingRpcs.get(); + int rpcs = entry.totalOutstandingRpcs(); if (rpcs < minRpcs) { minRpcs = rpcs; candidates.clear(); @@ -222,7 +222,7 @@ private int pickEntryIndexPowerOfTwoLeastInFlight() { Entry entry1 = localEntries.get(choice1); Entry entry2 = localEntries.get(choice2); - return entry1.outstandingRpcs.get() < entry2.outstandingRpcs.get() ? choice1 : choice2; + return entry1.totalOutstandingRpcs() < entry2.totalOutstandingRpcs() ? choice1 : choice2; } Channel getChannel(int index) { @@ -471,7 +471,7 @@ void refresh() { * Get and retain a Channel Entry. The returned Entry will have its rpc count incremented, * preventing it from getting recycled. */ - private Entry getRetainedEntry(int affinity) { + private Entry getRetainedEntry(int affinity, boolean isStreaming) { // If an entry is not retainable, that usually means that it's about to be replaced and if we // retry we should get a new useable entry. // The maximum number of concurrent calls to this method for any given time span is at most 2, @@ -479,7 +479,7 @@ private Entry getRetainedEntry(int affinity) { // code evolving for (int i = 0; i < 5; i++) { Entry entry = getEntry(affinity); - if (entry.retain()) { + if (entry.retain(isStreaming)) { return entry; } } @@ -507,8 +507,14 @@ private Entry getEntry(int affinity) { return localEntries.get(index); } + /** Gets the current list of BigtableChannelInsight objects. */ + @Override + public List 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; /** @@ -525,9 +531,12 @@ static class Entry { * outstanding RPCs has to happen when the ClientCall is closed or the ClientCall failed to * start. */ - @VisibleForTesting final AtomicInteger outstandingRpcs = new AtomicInteger(0); + @VisibleForTesting final AtomicInteger outstandingUnaryRpcs = new AtomicInteger(0); + + @VisibleForTesting final AtomicInteger outstandingStreamingRpcs = new AtomicInteger(0); - private final AtomicInteger maxOutstanding = new AtomicInteger(); + private final AtomicInteger maxOutstandingUnaryRpcs = new AtomicInteger(); + private final AtomicInteger maxOutstandingStreamingRpcs = new AtomicInteger(); /** Queue storing the last 5 minutes of probe results */ @VisibleForTesting @@ -555,8 +564,17 @@ ManagedChannel getManagedChannel() { return this.channel; } + @VisibleForTesting + int totalOutstandingRpcs() { + return outstandingUnaryRpcs.get() + outstandingStreamingRpcs.get(); + } + int getAndResetMaxOutstanding() { - return maxOutstanding.getAndSet(outstandingRpcs.get()); + int currentUnary = outstandingUnaryRpcs.get(); + int currentStreaming = outstandingStreamingRpcs.get(); + int prevMaxUnary = maxOutstandingUnaryRpcs.getAndSet(currentUnary); + int prevMaxStreaming = maxOutstandingStreamingRpcs.getAndSet(currentStreaming); + return prevMaxStreaming + prevMaxUnary; } /** @@ -565,19 +583,16 @@ int getAndResetMaxOutstanding() { * channel has been successfully retained and it is the responsibility of the caller to release * it. */ - private boolean retain() { - // register desire to start RPC - int currentOutstanding = outstandingRpcs.incrementAndGet(); - - // Rough bookkeeping - int prevMax = maxOutstanding.get(); - if (currentOutstanding > prevMax) { - maxOutstanding.incrementAndGet(); - } - + @VisibleForTesting + boolean retain(boolean isStreaming) { + AtomicInteger counter = isStreaming ? outstandingStreamingRpcs : outstandingUnaryRpcs; + AtomicInteger maxCounter = + isStreaming ? maxOutstandingStreamingRpcs : maxOutstandingUnaryRpcs; + int currentOutstanding = counter.incrementAndGet(); + maxCounter.accumulateAndGet(currentOutstanding, Math::max); // abort if the channel is closing if (shutdownRequested.get()) { - release(); + release(isStreaming); return false; } return true; @@ -587,15 +602,17 @@ private boolean retain() { * Notify the channel that the number of outstanding RPCs has decreased. If shutdown has been * previously requested, this method will shutdown the channel if its the last outstanding RPC. */ - private void release() { - int newCount = outstandingRpcs.decrementAndGet(); + void release(boolean isStreaming) { + AtomicInteger counter = isStreaming ? outstandingStreamingRpcs : outstandingUnaryRpcs; + int newCount = counter.decrementAndGet(); if (newCount < 0) { LOG.log(Level.WARNING, "Bug! Reference count is negative (" + newCount + ")!"); } - // Must check outstandingRpcs after shutdownRequested (in reverse order of retain()) to ensure + // Must check toalOutstandingRpcs after shutdownRequested (in reverse order of retain()) to + // ensure // mutual exclusion. - if (shutdownRequested.get() && outstandingRpcs.get() == 0) { + if (shutdownRequested.get() && totalOutstandingRpcs() == 0) { shutdown(); } } @@ -606,7 +623,7 @@ private void release() { */ private void requestShutdown() { shutdownRequested.set(true); - if (outstandingRpcs.get() == 0) { + if (totalOutstandingRpcs() == 0) { shutdown(); } } @@ -617,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. */ @@ -635,8 +663,10 @@ public String authority() { @Override public ClientCall newCall( MethodDescriptor methodDescriptor, CallOptions callOptions) { - Entry entry = getRetainedEntry(index); - return new ReleasingClientCall<>(entry.channel.newCall(methodDescriptor, callOptions), entry); + boolean isStreaming = methodDescriptor.getType() != MethodDescriptor.MethodType.UNARY; + Entry entry = getRetainedEntry(index, isStreaming); + return new ReleasingClientCall<>( + entry.channel.newCall(methodDescriptor, callOptions), entry, isStreaming); } } @@ -644,12 +674,14 @@ public ClientCall newCall( static class ReleasingClientCall extends SimpleForwardingClientCall { @Nullable private CancellationException cancellationException; final Entry entry; + private final boolean isStreaming; private final AtomicBoolean wasClosed = new AtomicBoolean(); private final AtomicBoolean wasReleased = new AtomicBoolean(); - public ReleasingClientCall(ClientCall delegate, Entry entry) { + public ReleasingClientCall(ClientCall delegate, Entry entry, boolean isStreaming) { super(delegate); this.entry = entry; + this.isStreaming = isStreaming; } @Override @@ -673,7 +705,7 @@ public void onClose(Status status, Metadata trailers) { super.onClose(status, trailers); } finally { if (wasReleased.compareAndSet(false, true)) { - entry.release(); + entry.release(isStreaming); } else { LOG.log( Level.WARNING, @@ -687,7 +719,7 @@ public void onClose(Status status, Metadata trailers) { } catch (Exception e) { // In case start failed, make sure to release if (wasReleased.compareAndSet(false, true)) { - entry.release(); + entry.release(isStreaming); } else { LOG.log( Level.WARNING, diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java index ba18994619..a53604f753 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java +++ b/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 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); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettingsTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettingsTest.java index 28d5a43738..4235f0c814 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettingsTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettingsTest.java @@ -19,6 +19,7 @@ import com.google.api.gax.grpc.ChannelPoolSettings; import com.google.common.collect.ImmutableSet; +import io.grpc.ManagedChannel; import java.lang.reflect.Method; import java.lang.reflect.Modifier; import java.util.Arrays; @@ -27,6 +28,7 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.Mockito; @RunWith(JUnit4.class) public class BigtableChannelPoolSettingsTest { @@ -48,6 +50,44 @@ public void testToBigtableChannelPoolSettingsAllFieldsSetCopiesCorrectly() throw assertSettingsCopiedCorrectly(originalSettings, copiedSettings); } + @Test + public void testEntryRetainReleaseByType() { + ManagedChannel mockChannel = Mockito.mock(ManagedChannel.class); + BigtableChannelPool.Entry entry = new BigtableChannelPool.Entry(mockChannel); + + // Test Unary + assertThat(entry.retain(false)).isTrue(); // Unary + assertThat(entry.outstandingUnaryRpcs.get()).isEqualTo(1); + assertThat(entry.outstandingStreamingRpcs.get()).isEqualTo(0); + assertThat(entry.totalOutstandingRpcs()).isEqualTo(1); + // Test Unary release + entry.release(false); + assertThat(entry.outstandingUnaryRpcs.get()).isEqualTo(0); + assertThat(entry.outstandingStreamingRpcs.get()).isEqualTo(0); + assertThat(entry.totalOutstandingRpcs()).isEqualTo(0); + + // Test Streaming + assertThat(entry.retain(true)).isTrue(); // Streaming + assertThat(entry.outstandingUnaryRpcs.get()).isEqualTo(0); + assertThat(entry.outstandingStreamingRpcs.get()).isEqualTo(1); + assertThat(entry.totalOutstandingRpcs()).isEqualTo(1); + // Test Streaming again + assertThat(entry.retain(true)).isTrue(); // Streaming again + assertThat(entry.outstandingStreamingRpcs.get()).isEqualTo(2); + assertThat(entry.outstandingUnaryRpcs.get()).isEqualTo(0); + assertThat(entry.totalOutstandingRpcs()).isEqualTo(2); + + entry.release(true); + assertThat(entry.outstandingStreamingRpcs.get()).isEqualTo(1); + assertThat(entry.outstandingUnaryRpcs.get()).isEqualTo(0); + assertThat(entry.totalOutstandingRpcs()).isEqualTo(1); + + entry.release(true); + assertThat(entry.outstandingStreamingRpcs.get()).isEqualTo(0); + assertThat(entry.outstandingUnaryRpcs.get()).isEqualTo(0); + assertThat(entry.totalOutstandingRpcs()).isEqualTo(0); + } + @Test public void testToBigtableChannelPoolSettingsDefaultValuesCopiesCorrectly() throws Exception { ChannelPoolSettings originalSettings = ChannelPoolSettings.builder().build();