Merge pull request #132 from nimf/use-existing-keys

Bind keys on the fly.
GoogleCloudPlatform · May 24, 2022 · 9ec2aa5 · 9ec2aa5
2 parents 9b056e9 + 712560a
commit 9ec2aa5
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/grpc-gcp/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java b/grpc-gcp/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java
@@ -822,7 +822,12 @@ protected ChannelRef getChannelRef(@Nullable String key) {
       return pickLeastBusyChannel(/* forFallback= */ false);
     }
     ChannelRef mappedChannel = affinityKeyToChannelRef.get(key);
-    if (mappedChannel == null || !fallbackEnabled) {
+    if (mappedChannel == null) {
+      ChannelRef channelRef = pickLeastBusyChannel(/*forFallback= */ false);
+      bind(channelRef, Collections.singletonList(key));
+      return channelRef;
+    }
+    if (!fallbackEnabled) {
       return mappedChannel;
     }
     // Look up if the channelRef is not ready.

diff --git a/grpc-gcp/src/test/java/com/google/cloud/grpc/GcpManagedChannelTest.java b/grpc-gcp/src/test/java/com/google/cloud/grpc/GcpManagedChannelTest.java
@@ -382,7 +382,7 @@ public void testGetChannelRefMaxSize() {
   public void testBindUnbindKey() {
     // Initialize the channel and bind the key, check the affinity count.
     ChannelRef cf1 = gcpChannel.new ChannelRef(builder.build(), 1, 0, 5);
-    ChannelRef cf2 = gcpChannel.new ChannelRef(builder.build(), 1, 0, 4);
+    ChannelRef cf2 = gcpChannel.new ChannelRef(builder.build(), 2, 0, 4);
     gcpChannel.channelRefs.add(cf1);
     gcpChannel.channelRefs.add(cf2);
     gcpChannel.bind(cf1, Collections.singletonList("key1"));
@@ -418,6 +418,27 @@ public void testBindUnbindKey() {
     assertEquals(0, gcpChannel.affinityKeyToChannelRef.size());
   }
 
+  @Test
+  public void testUsingKeyWithoutBinding() {
+    // Initialize the channel and bind the key, check the affinity count.
+    ChannelRef cf1 = gcpChannel.new ChannelRef(builder.build(), 1, 0, 5);
+    ChannelRef cf2 = gcpChannel.new ChannelRef(builder.build(), 2, 0, 4);
+    gcpChannel.channelRefs.add(cf1);
+    gcpChannel.channelRefs.add(cf2);
+
+    final String key = "non-binded-key";
+    ChannelRef channelRef = gcpChannel.getChannelRef(key);
+    // Should bind on the fly to the least busy channel, which is 2.
+    assertThat(channelRef.getId()).isEqualTo(2);
+
+    cf1.activeStreamsCountDecr(System.nanoTime(), Status.OK, true);
+    cf1.activeStreamsCountDecr(System.nanoTime(), Status.OK, true);
+    channelRef = gcpChannel.getChannelRef(key);
+    // Even after channel 1 now has less active streams (3) the channel 2 is still mapped for the
+    // same key.
+    assertThat(channelRef.getId()).isEqualTo(2);
+  }
+
   @Test
   public void testGetKeysFromRequest() {
     String expected = "thisisaname";