fix"

Author: sushanb

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

@@ -557,8 +557,8 @@ static class Entry {
     ManagedChannel getManagedChannel() {
       return this.channel;
     }
-
-    private int totalOutstandingRpcs() {
+    @VisibleForTesting
+    int totalOutstandingRpcs() {
       return outstandingUnaryRpcs.get() + outstandingStreamingRpcs.get();
     }
 
@@ -576,7 +576,8 @@ int getAndResetMaxOutstanding() {
      * channel has been successfully retained and it is the responsibility of the caller to release
      * it.
      */
-    private boolean retain(boolean isStreaming) {
+    @VisibleForTesting
+    boolean retain(boolean isStreaming) {
       AtomicInteger counter = isStreaming ? outstandingStreamingRpcs : outstandingUnaryRpcs;
       AtomicInteger maxCounter = isStreaming ? maxOutstandingStreamingRpcs : maxOutstandingUnaryRpcs;
       int currentOutstanding = counter.incrementAndGet();
@@ -593,7 +594,7 @@ private boolean retain(boolean isStreaming) {
      * 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(boolean isStreaming) {
+    void release(boolean isStreaming) {
       AtomicInteger counter = isStreaming ? outstandingStreamingRpcs : outstandingUnaryRpcs;
       int newCount = counter.decrementAndGet();
       if (newCount < 0) {

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettingsTest.java

@@ -24,9 +24,12 @@
 import java.util.Arrays;
 import java.util.Set;
 import java.util.stream.Collectors;
+
+import io.grpc.ManagedChannel;
 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 +51,45 @@ 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();