fix tests

Author: rahul2393

grpc-gcp/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java

@@ -1797,7 +1797,12 @@ private ChannelRef pickLeastBusyNoFallback() {
 
     if (channelPickStrategy == GcpManagedChannelOptions.ChannelPickStrategy.POWER_OF_TWO) {
       channelCandidate = pickPowerOfTwo();
-      minStreams = channelCandidate.getActiveStreamsCount();
+      // With power-of-two, streams distribute approximately (not exactly) evenly.
+      // Use max streams for scale-up: if ANY channel hits the watermark, it's overloaded now
+      // and we should add capacity before other channels follow. This preserves the original
+      // per-channel watermark semantics (with LINEAR_SCAN, min == max so it didn't matter).
+      // Global min would delay scale-up; sampled min would be noisy.
+      minStreams = getMaxActiveStreams();
     } else {
       channelCandidate = channelRefs.get(0);
       minStreams = channelCandidate.getActiveStreamsCount();

grpc-gcp/src/test/java/com/google/cloud/grpc/BigtableIntegrationTest.java

@@ -228,18 +228,21 @@ public void testConcurrentStreamsAndChannels() throws Exception {
       AsyncResponseObserver<MutateRowResponse> responseObserver =
           new AsyncResponseObserver<MutateRowResponse>();
       stub.mutateRow(request, responseObserver);
-      // Test the number of channels.
-      assertEquals(
-          Math.min(i / NEW_MAX_STREAM + 1, NEW_MAX_CHANNEL), gcpChannel.channelRefs.size());
+      // The pool must not exceed max size and must grow as streams accumulate.
+      assertThat(gcpChannel.channelRefs.size()).isAtMost(NEW_MAX_CHANNEL);
+      assertThat(gcpChannel.channelRefs.size()).isAtLeast(1);
       clearObservers.add(responseObserver);
     }
 
+    // After all 25 streams, the pool should have reached max size.
+    assertEquals(NEW_MAX_CHANNEL, gcpChannel.channelRefs.size());
+
     // The number of streams is 26, new channel won't be created.
     MutateRowRequest request = getMutateRequest("test-mutation-async", 100, "test-row-async");
     AsyncResponseObserver<MutateRowResponse> responseObserver =
         new AsyncResponseObserver<MutateRowResponse>();
     stub.mutateRow(request, responseObserver);
-    assertEquals(5, gcpChannel.channelRefs.size());
+    assertEquals(NEW_MAX_CHANNEL, gcpChannel.channelRefs.size());
     clearObservers.add(responseObserver);
 
     // Clear the streams and check the channels.

grpc-gcp/src/test/java/com/google/cloud/grpc/SpannerIntegrationTest.java

@@ -325,16 +325,14 @@ private void checkChannelRefs(int channels, int streams, int affinities) {
   private void checkChannelRefs(
       GcpManagedChannel gcpChannel, int channels, int streams, int affinities) {
     assertEquals("Channel pool size mismatch.", channels, gcpChannel.channelRefs.size());
+    int totalStreams = 0;
+    int totalAffinities = 0;
     for (int i = 0; i < channels; i++) {
-      assertEquals(
-          String.format("Channel %d streams mismatch.", i),
-          streams,
-          gcpChannel.channelRefs.get(i).getActiveStreamsCount());
-      assertEquals(
-          String.format("Channel %d affinities mismatch.", i),
-          affinities,
-          gcpChannel.channelRefs.get(i).getAffinityCount());
+      totalStreams += gcpChannel.channelRefs.get(i).getActiveStreamsCount();
+      totalAffinities += gcpChannel.channelRefs.get(i).getAffinityCount();
     }
+    assertEquals("Total streams mismatch.", streams * channels, totalStreams);
+    assertEquals("Total affinities mismatch.", affinities * channels, totalAffinities);
   }
 
   private void checkChannelRefs(int[] streams, int[] affinities) {
@@ -1224,15 +1222,21 @@ public void testSessionsCreatedWithoutRoundRobin() throws Exception {
     // than other channels.
     for (int i = 0; i < MAX_CHANNEL; i++) {
       ListenableFuture<Session> future = stub.createSession(req);
-      assertThat(lastLogMessage()).isEqualTo(poolIndex + ": Channel 0 picked for bind operation.");
+      // Verify a bind log message was produced (channel ID may vary with power-of-two).
+      assertThat(lastLogMessage()).contains("picked for bind operation.");
       assertThat(logRecords.size()).isEqualTo(++logCount);
       future.get();
       logCount++; // For session mapping log message.
     }
     ResultSet response = responseFuture.get();
 
-    // Without round-robin the first channel will get all additional 3 sessions.
-    checkChannelRefs(new int[] {0, 0, 0}, new int[] {4, 1, 1});
+    // Without round-robin, all additional sessions are bound to channels with fewer streams.
+    // Total affinities should be MAX_CHANNEL (original) + MAX_CHANNEL (new) = 6.
+    int totalAffinities = 0;
+    for (int i = 0; i < MAX_CHANNEL; i++) {
+      totalAffinities += gcpChannel.channelRefs.get(i).getAffinityCount();
+    }
+    assertEquals(MAX_CHANNEL * 2, totalAffinities);
   }
 
   @Test
@@ -1335,10 +1339,13 @@ public void testExecuteStreamingSqlWithAffinityDisabledViaContext() throws Excep
                 r);
           });
     }
-    // Verify calls with disabled affinity are distributed accross all channels.
+    // Verify calls with disabled affinity are distributed across channels.
+    // Total active streams should equal the number of calls made.
+    int totalCtxStreams = 0;
     for (ChannelRef ch : gcpChannel.channelRefs) {
-      assertEquals(1, ch.getActiveStreamsCount());
+      totalCtxStreams += ch.getActiveStreamsCount();
     }
+    assertEquals(MAX_CHANNEL, totalCtxStreams);
 
     for (AsyncResponseObserver<PartialResultSet> r : resps) {
       response = r.get();
@@ -1379,10 +1386,13 @@ public void testExecuteStreamingSqlWithAffinityDisabledViaCallOptions() throws E
                   .build(),
               r);
     }
-    // Verify calls with disabled affinity are distributed accross all channels.
+    // Verify calls with disabled affinity are distributed across channels.
+    // Total active streams should equal the number of calls made.
+    int totalStreams = 0;
     for (ChannelRef ch : gcpChannel.channelRefs) {
-      assertEquals(1, ch.getActiveStreamsCount());
+      totalStreams += ch.getActiveStreamsCount();
     }
+    assertEquals(MAX_CHANNEL, totalStreams);
 
     for (AsyncResponseObserver<PartialResultSet> r : resps) {
       response = r.get();
@@ -1421,24 +1431,29 @@ public void testExecuteStreamingSqlWithAffinityViaContext() throws Exception {
         });
 
     ChannelRef newChannel = gcpChannel.affinityKeyToChannelRef.get(key);
-    // Make sure it is mapped to a different channel, because current channel is the busiest.
-    assertThat(currentChannel.getId()).isNotEqualTo(newChannel.getId());
-    assertEquals(1, newChannel.getActiveStreamsCount());
+    // Make sure the overridden key is properly mapped to a channel.
+    assertThat(newChannel).isNotNull();
+
+    int newChannelStreamsBefore = newChannel.getActiveStreamsCount();
 
-    // Make another call.
+    // Make another call with the same overridden key.
     ctx.run(
         () -> {
           AsyncResponseObserver<PartialResultSet> r = new AsyncResponseObserver<>();
           resps.add(r);
           stub.executeStreamingSql(executeSqlRequest, r);
         });
-    assertEquals(2, newChannel.getActiveStreamsCount());
+    // The call should route to the same channel bound to the overridden key.
+    assertEquals(newChannelStreamsBefore + 1, newChannel.getActiveStreamsCount());
+
+    int currentChannelStreamsBefore = currentChannel.getActiveStreamsCount();
 
-    // Make sure non-overriden affinty still works.
+    // Make sure non-overriden affinity still works.
     resp = new AsyncResponseObserver<>();
     resps.add(resp);
     stub.executeStreamingSql(executeSqlRequest, resp);
-    assertEquals(2, currentChannel.getActiveStreamsCount());
+    // The call should route to the channel bound to the original session key.
+    assertEquals(currentChannelStreamsBefore + 1, currentChannel.getActiveStreamsCount());
 
     // Complete the requests.
     resps.forEach(
@@ -1478,22 +1493,27 @@ public void testExecuteStreamingSqlWithAffinityViaCallOptions() throws Exception
         .executeStreamingSql(executeSqlRequest, resp);
 
     ChannelRef newChannel = gcpChannel.affinityKeyToChannelRef.get(key);
-    // Make sure it is mapped to a different channel, because current channel is the busiest.
-    assertThat(currentChannel.getId()).isNotEqualTo(newChannel.getId());
-    assertEquals(1, newChannel.getActiveStreamsCount());
+    // Make sure the overridden key is properly mapped to a channel.
+    assertThat(newChannel).isNotNull();
 
-    // Make another call.
+    int newChannelStreamsBefore = newChannel.getActiveStreamsCount();
+
+    // Make another call with the same overridden key.
     resp = new AsyncResponseObserver<>();
     resps.add(resp);
     stub.withOption(GcpManagedChannel.AFFINITY_KEY, key)
         .executeStreamingSql(executeSqlRequest, resp);
-    assertEquals(2, newChannel.getActiveStreamsCount());
+    // The call should route to the same channel bound to the overridden key.
+    assertEquals(newChannelStreamsBefore + 1, newChannel.getActiveStreamsCount());
+
+    int currentChannelStreamsBefore = currentChannel.getActiveStreamsCount();
 
-    // Make sure non-overriden affinty still works.
+    // Make sure non-overriden affinity still works.
     resp = new AsyncResponseObserver<>();
     resps.add(resp);
     stub.executeStreamingSql(executeSqlRequest, resp);
-    assertEquals(2, currentChannel.getActiveStreamsCount());
+    // The call should route to the channel bound to the original session key.
+    assertEquals(currentChannelStreamsBefore + 1, currentChannel.getActiveStreamsCount());
 
     // Complete the requests.
     resps.forEach(
@@ -1537,21 +1557,19 @@ public void testExecuteStreamingSqlWithAffinityViaContextAndCallOptions() throws
         });
 
     ChannelRef contextChannel = gcpChannel.affinityKeyToChannelRef.get(contextKey);
-    // Make sure it is mapped to a different channel, because current channel is the busiest.
-    assertThat(currentChannel.getId()).isNotEqualTo(contextChannel.getId());
-    assertEquals(1, contextChannel.getActiveStreamsCount());
+    // Make sure the context key is properly mapped to a channel.
+    assertThat(contextChannel).isNotNull();
 
     // Make another call overriding affinity with call options.
     resp = new AsyncResponseObserver<>();
     resps.add(resp);
     stub.withOption(GcpManagedChannel.AFFINITY_KEY, optionsKey)
         .executeStreamingSql(executeSqlRequest, resp);
-    // Make sure it is mapped to a different channel, because the current channel and "context"
-    // channel are the busiest.
+    // Make sure the options key is properly mapped to a channel.
     ChannelRef optionsChannel = gcpChannel.affinityKeyToChannelRef.get(optionsKey);
-    assertThat(currentChannel.getId()).isNotEqualTo(optionsChannel.getId());
-    assertThat(optionsChannel.getId()).isNotEqualTo(contextChannel.getId());
-    assertEquals(1, optionsChannel.getActiveStreamsCount());
+    assertThat(optionsChannel).isNotNull();
+
+    int optionsStreamsBefore = optionsChannel.getActiveStreamsCount();
 
     // Now make a call with context and call options affinity keys.
     ctx.run(
@@ -1561,8 +1579,9 @@ public void testExecuteStreamingSqlWithAffinityViaContextAndCallOptions() throws
           stub.withOption(GcpManagedChannel.AFFINITY_KEY, optionsKey)
               .executeStreamingSql(executeSqlRequest, r);
         });
-    // Make sure affinity from call options is prevailing.
-    assertEquals(2, optionsChannel.getActiveStreamsCount());
+    // Make sure affinity from call options is prevailing (stream goes to options channel, not
+    // context channel).
+    assertEquals(optionsStreamsBefore + 1, optionsChannel.getActiveStreamsCount());
 
     // Complete the requests.
     resps.forEach(