MCS connection scaling interop tests for Java

android-interop-testing/build.gradle

@@ -17,6 +17,7 @@ android {
                 srcDirs += "${projectDir}/../interop-testing/src/main/java/"
                 setIncludes(["io/grpc/android/integrationtest/**",
                              "io/grpc/testing/integration/AbstractInteropTest.java",
+                             "io/grpc/testing/integration/TestCases.java",
                              "io/grpc/testing/integration/TestServiceImpl.java",
                              "io/grpc/testing/integration/Util.java"])
             }

interop-testing/src/main/java/io/grpc/testing/integration/TestCases.java

@@ -59,7 +59,8 @@ public enum TestCases {
   RPC_SOAK("sends 'soak_iterations' large_unary rpcs in a loop, each on the same channel"),
   CHANNEL_SOAK("sends 'soak_iterations' large_unary rpcs in a loop, each on a new channel"),
   ORCA_PER_RPC("report backend metrics per query"),
-  ORCA_OOB("report backend metrics out-of-band");
+  ORCA_OOB("report backend metrics out-of-band"),
+  MCS_CS("max concurrent streaming connection scaling");
 
   private final String description;
 

interop-testing/src/main/java/io/grpc/testing/integration/TestServiceClient.java

@@ -17,6 +17,7 @@
 package io.grpc.testing.integration;
 
 import static com.google.common.truth.Truth.assertThat;
+import static io.grpc.testing.integration.TestCases.MCS_CS;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -70,6 +71,7 @@
 import io.grpc.testing.integration.Messages.TestOrcaReport;
 import java.io.File;
 import java.io.FileInputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.Charset;
 import java.util.Arrays;
@@ -563,7 +565,11 @@ private void runTest(TestCases testCase) throws Exception {
         tester.testOrcaOob();
         break;
       }
-
+
+      case MCS_CS: {
+        tester.testMcs();
+        break;
+      }
       default:
         throw new IllegalArgumentException("Unknown test case: " + testCase);
     }
@@ -596,9 +602,10 @@ private ClientInterceptor maybeCreateAdditionalMetadataInterceptor(
   }
 
   private class Tester extends AbstractInteropTest {
+
     @Override
     protected ManagedChannelBuilder<?> createChannelBuilder() {
-      boolean useGeneric = false;
+      boolean useGeneric = testCase.equals(MCS_CS.toString()) ? true : false;
       ChannelCredentials channelCredentials;
       if (customCredentialsType != null) {
         useGeneric = true; // Retain old behavior; avoids erroring if incompatible
@@ -658,7 +665,17 @@ protected ManagedChannelBuilder<?> createChannelBuilder() {
         if (serverHostOverride != null) {
           channelBuilder.overrideAuthority(serverHostOverride);
         }
-        if (serviceConfig != null) {
+        if (testCase.equals(MCS_CS.toString())) {
+          channelBuilder.disableServiceConfigLookUp();
+          try {
+            @SuppressWarnings("unchecked")
+            Map<String, ?> serviceConfigMap = (Map<String, ?>) JsonParser.parse(
+                "{\"connection_scaling\":{\"max_connections_per_subchannel\": 2}}");
+            channelBuilder.defaultServiceConfig(serviceConfigMap);
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        } else if (serviceConfig != null) {
           channelBuilder.disableServiceConfigLookUp();
           channelBuilder.defaultServiceConfig(serviceConfig);
         }
@@ -979,31 +996,16 @@ public void testOrcaOob() throws Exception {
           .build();
 
       final int retryLimit = 5;
-      BlockingQueue<Object> queue = new LinkedBlockingQueue<>();
-      final Object lastItem = new Object();
+      StreamingOutputCallResponseObserver streamingOutputCallResponseObserver =
+          new StreamingOutputCallResponseObserver();
       StreamObserver<StreamingOutputCallRequest> streamObserver =
-          asyncStub.fullDuplexCall(new StreamObserver<StreamingOutputCallResponse>() {
-
-            @Override
-            public void onNext(StreamingOutputCallResponse value) {
-              queue.add(value);
-            }
-
-            @Override
-            public void onError(Throwable t) {
-              queue.add(t);
-            }
-
-            @Override
-            public void onCompleted() {
-              queue.add(lastItem);
-            }
-          });
+          asyncStub.fullDuplexCall(streamingOutputCallResponseObserver);
 
       streamObserver.onNext(StreamingOutputCallRequest.newBuilder()
           .setOrcaOobReport(answer)
           .addResponseParameters(ResponseParameters.newBuilder().setSize(1).build()).build());
-      assertThat(queue.take()).isInstanceOf(StreamingOutputCallResponse.class);
+      assertThat(streamingOutputCallResponseObserver.take())
+          .isInstanceOf(StreamingOutputCallResponse.class);
       int i = 0;
       for (; i < retryLimit; i++) {
         Thread.sleep(1000);
@@ -1016,7 +1018,7 @@ public void onCompleted() {
       streamObserver.onNext(StreamingOutputCallRequest.newBuilder()
           .setOrcaOobReport(answer2)
           .addResponseParameters(ResponseParameters.newBuilder().setSize(1).build()).build());
-      assertThat(queue.take()).isInstanceOf(StreamingOutputCallResponse.class);
+      assertThat(streamingOutputCallResponseObserver.isCompleted).isTrue();
 
       for (i = 0; i < retryLimit; i++) {
         Thread.sleep(1000);
@@ -1026,8 +1028,6 @@ public void onCompleted() {
         }
       }
       assertThat(i).isLessThan(retryLimit);
-      streamObserver.onCompleted();
-      assertThat(queue.take()).isSameInstanceAs(lastItem);
     }
 
     @Override
@@ -1054,6 +1054,78 @@ protected ServerBuilder<?> getHandshakerServerBuilder() {
     protected int operationTimeoutMillis() {
       return 15000;
     }
+
+    class StreamingOutputCallResponseObserver implements
+        StreamObserver<StreamingOutputCallResponse> {
+      private final BlockingQueue<Object> queue = new LinkedBlockingQueue<>();
+      private volatile boolean isCompleted = true;
+
+      @Override
+      public void onNext(StreamingOutputCallResponse value) {
+        queue.add(value);
+      }
+
+      @Override
+      public void onError(Throwable t) {
+        queue.add(t);
+      }
+
+      @Override
+      public void onCompleted() {
+        isCompleted = true;
+      }
+
+      Object take() throws InterruptedException {
+        return queue.take();
+      }
+    }
+
+    public void testMcs() throws Exception {
+      StreamingOutputCallResponseObserver responseObserver1 =
+          new StreamingOutputCallResponseObserver();
+      StreamObserver<StreamingOutputCallRequest> streamObserver1 =
+          asyncStub.fullDuplexCall(responseObserver1);
+      StreamingOutputCallRequest request = StreamingOutputCallRequest.newBuilder()
+          .setPayload(Payload.newBuilder().setBody(
+              ByteString.copyFromUtf8(MCS_CS.description())).build()).build();
+      streamObserver1.onNext(request);
+      Object responseObj = responseObserver1.take();
+      StreamingOutputCallResponse callResponse = (StreamingOutputCallResponse) responseObj;
+      String clientSocketAddressInCall1 = new String(callResponse.getPayload().getBody()
+          .toByteArray(), UTF_8);
+      assertThat(clientSocketAddressInCall1).isNotEmpty();
+
+      StreamingOutputCallResponseObserver responseObserver2 =
+          new StreamingOutputCallResponseObserver();
+      StreamObserver<StreamingOutputCallRequest> streamObserver2 =
+          asyncStub.fullDuplexCall(responseObserver2);
+      streamObserver2.onNext(request);
+      callResponse = (StreamingOutputCallResponse) responseObserver2.take();
+      String clientSocketAddressInCall2 =
+          new String(callResponse.getPayload().getBody().toByteArray(), UTF_8);
+
+      assertThat(clientSocketAddressInCall1).isEqualTo(clientSocketAddressInCall2);
+
+      // The first connection is at max rpc call count of 2, so the 3rd rpc will cause a new
+      // connection to be created in the same subchannel and not get queued.
+      StreamingOutputCallResponseObserver responseObserver3 =
+          new StreamingOutputCallResponseObserver();
+      StreamObserver<StreamingOutputCallRequest> streamObserver3 =
+          asyncStub.fullDuplexCall(responseObserver3);
+      streamObserver3.onNext(request);
+      callResponse = (StreamingOutputCallResponse) responseObserver3.take();
+      String clientSocketAddressInCall3 =
+          new String(callResponse.getPayload().getBody().toByteArray(), UTF_8);
+
+      assertThat(clientSocketAddressInCall3).isNotEqualTo(clientSocketAddressInCall1);
+
+      streamObserver1.onCompleted();
+      assertThat(responseObserver1.isCompleted).isTrue();
+      streamObserver2.onCompleted();
+      assertThat(responseObserver2.isCompleted).isTrue();
+      streamObserver3.onCompleted();
+      assertThat(responseObserver3.isCompleted).isTrue();
+    }
   }
 
   private static String validTestCasesHelpText() {

interop-testing/src/main/java/io/grpc/testing/integration/TestServiceImpl.java

@@ -16,13 +16,20 @@
 
 package io.grpc.testing.integration;
 
+import static io.grpc.Grpc.TRANSPORT_ATTR_REMOTE_ADDR;
+import static io.grpc.testing.integration.TestCases.MCS_CS;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Queues;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import com.google.protobuf.ByteString;
+import io.grpc.Context;
+import io.grpc.Contexts;
 import io.grpc.ForwardingServerCall.SimpleForwardingServerCall;
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
+import io.grpc.ServerCall.Listener;
 import io.grpc.ServerCallHandler;
 import io.grpc.ServerInterceptor;
 import io.grpc.Status;
@@ -42,6 +49,7 @@
 import io.grpc.testing.integration.Messages.StreamingOutputCallResponse;
 import io.grpc.testing.integration.Messages.TestOrcaReport;
 import io.grpc.testing.integration.TestServiceGrpc.AsyncService;
+import java.net.SocketAddress;
 import java.util.ArrayDeque;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -61,8 +69,8 @@
  * sent in response streams.
  */
 public class TestServiceImpl implements io.grpc.BindableService, AsyncService {
+  static Context.Key<SocketAddress> PEER_ADDRESS_CONTEXT_KEY = Context.key("peer-address");
   private final Random random = new Random();
-
   private final ScheduledExecutorService executor;
   private final ByteString compressableBuffer;
   private final MetricRecorder metricRecorder;
@@ -235,6 +243,17 @@ public void onNext(StreamingOutputCallRequest request) {
               .asRuntimeException());
           return;
         }
+        if (new String(request.getPayload().getBody().toByteArray(), UTF_8)
+            .equals(MCS_CS.description())) {
+          SocketAddress peerAddress = PEER_ADDRESS_CONTEXT_KEY.get();
+          ByteString payload = ByteString.copyFromUtf8(peerAddress.toString());
+          StreamingOutputCallResponse.Builder responseBuilder =
+              StreamingOutputCallResponse.newBuilder();
+          responseBuilder.setPayload(
+              Payload.newBuilder()
+                  .setBody(payload));
+          responseObserver.onNext(responseBuilder.build());
+        }
         dispatcher.enqueue(toChunkQueue(request));
       }
 
@@ -507,7 +526,8 @@ public static List<ServerInterceptor> interceptors() {
     return Arrays.asList(
         echoRequestHeadersInterceptor(Util.METADATA_KEY),
         echoRequestMetadataInHeaders(Util.ECHO_INITIAL_METADATA_KEY),
-        echoRequestMetadataInTrailers(Util.ECHO_TRAILING_METADATA_KEY));
+        echoRequestMetadataInTrailers(Util.ECHO_TRAILING_METADATA_KEY),
+        new McsScalingTestcaseInterceptor());
   }
 
   /**
@@ -539,6 +559,22 @@ public void close(Status status, Metadata trailers) {
     };
   }
 
+  static class McsScalingTestcaseInterceptor implements ServerInterceptor {
+    @Override
+    public <ReqT, RespT> Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call,
+        Metadata headers, ServerCallHandler<ReqT, RespT> next) {
+      SocketAddress peerAddress = call.getAttributes().get(TRANSPORT_ATTR_REMOTE_ADDR);
+
+      // Create a new context with the peer address value
+      Context newContext = Context.current().withValue(PEER_ADDRESS_CONTEXT_KEY, peerAddress);
+      try {
+        return Contexts.interceptCall(newContext, call, headers, next);
+      } catch (Exception ex) {
+        throw new RuntimeException(ex);
+      }
+    }
+  }
+
   /**
    * Echoes request headers with the specified key(s) from a client into response headers only.
    */

interop-testing/src/main/java/io/grpc/testing/integration/TestServiceServer.java

@@ -75,6 +75,7 @@ public void run() {
   private int port = 8080;
   private boolean useTls = true;
   private boolean useAlts = false;
+  private boolean useMcs = false;
 
   private ScheduledExecutorService executor;
   private Server server;
@@ -118,6 +119,9 @@ void parseArgs(String[] args) {
           usage = true;
           break;
         }
+      } else if ("use_mcs".equals(key)) {
+        useMcs = Boolean.parseBoolean(value);
+        addressType = Util.AddressType.IPV4; // To use NettyServerBuilder
       } else {
         System.err.println("Unknown argument: " + key);
         usage = true;
@@ -186,6 +190,9 @@ void start() throws Exception {
         if (v4Address != null && !v4Address.equals(localV4Address)) {
           ((NettyServerBuilder) serverBuilder).addListenAddress(v4Address);
         }
+        if (useMcs) {
+          ((NettyServerBuilder) serverBuilder).maxConcurrentCallsPerConnection(2);
+        }
         break;
       case IPV6:
         List<SocketAddress> v6Addresses = Util.getV6Addresses(port);

interop-testing/src/test/java/io/grpc/testing/integration/TestCasesTest.java

@@ -67,7 +67,8 @@ public void testCaseNamesShouldMapToEnums() {
       "cancel_after_first_response",
       "timeout_on_sleeping_server",
       "orca_per_rpc",
-      "orca_oob"
+      "orca_oob",
+      "mcs_cs",
     };
 
     // additional test cases