MCS connection scaling interop tests for Java (#12651)

Max Concurrent Streams connection scaling interop test client and server for Java.

Author: kannanjgithub

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"),
+  MAX_CONCURRENT_STREAMS_CONNECTION_SCALING("max concurrent streaming connection scaling");
 
   private final String description;
 

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

@@ -563,7 +563,11 @@ private void runTest(TestCases testCase) throws Exception {
         tester.testOrcaOob();
         break;
       }
-
+      
+      case MAX_CONCURRENT_STREAMS_CONNECTION_SCALING: {
+        tester.testMcs();
+        break;
+      }
       default:
         throw new IllegalArgumentException("Unknown test case: " + testCase);
     }
@@ -596,6 +600,7 @@ private ClientInterceptor maybeCreateAdditionalMetadataInterceptor(
   }
 
   private class Tester extends AbstractInteropTest {
+
     @Override
     protected ManagedChannelBuilder<?> createChannelBuilder() {
       boolean useGeneric = false;
@@ -979,31 +984,17 @@ public void testOrcaOob() throws Exception {
           .build();
 
       final int retryLimit = 5;
-      BlockingQueue<Object> queue = new LinkedBlockingQueue<>();
       final Object lastItem = new Object();
+      StreamingOutputCallResponseObserver streamingOutputCallResponseObserver =
+          new StreamingOutputCallResponseObserver(lastItem);
       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 +1007,8 @@ public void onCompleted() {
       streamObserver.onNext(StreamingOutputCallRequest.newBuilder()
           .setOrcaOobReport(answer2)
           .addResponseParameters(ResponseParameters.newBuilder().setSize(1).build()).build());
-      assertThat(queue.take()).isInstanceOf(StreamingOutputCallResponse.class);
+      assertThat(streamingOutputCallResponseObserver.take())
+          .isInstanceOf(StreamingOutputCallResponse.class);
 
       for (i = 0; i < retryLimit; i++) {
         Thread.sleep(1000);
@@ -1027,7 +1019,7 @@ public void onCompleted() {
       }
       assertThat(i).isLessThan(retryLimit);
       streamObserver.onCompleted();
-      assertThat(queue.take()).isSameInstanceAs(lastItem);
+      assertThat(streamingOutputCallResponseObserver.take()).isSameInstanceAs(lastItem);
     }
 
     @Override
@@ -1054,6 +1046,85 @@ protected ServerBuilder<?> getHandshakerServerBuilder() {
     protected int operationTimeoutMillis() {
       return 15000;
     }
+
+    class StreamingOutputCallResponseObserver implements
+        StreamObserver<StreamingOutputCallResponse> {
+      private final Object lastItem;
+      private final BlockingQueue<Object> queue = new LinkedBlockingQueue<>();
+
+      public StreamingOutputCallResponseObserver(Object lastItem) {
+        this.lastItem = lastItem;
+      }
+
+      @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);
+      }
+
+      Object take() throws InterruptedException {
+        return queue.take();
+      }
+    }
+
+    public void testMcs() throws Exception {
+      final Object lastItem = new Object();
+      StreamingOutputCallResponseObserver responseObserver1 =
+          new StreamingOutputCallResponseObserver(lastItem);
+      StreamObserver<StreamingOutputCallRequest> streamObserver1 =
+          asyncStub.fullDuplexCall(responseObserver1);
+      StreamingOutputCallRequest request = StreamingOutputCallRequest.newBuilder()
+          .addResponseParameters(ResponseParameters.newBuilder()
+              .setFillPeerSocketAddress(
+                  Messages.BoolValue.newBuilder().setValue(true).build())
+              .build())
+          .build();
+      streamObserver1.onNext(request);
+      Object responseObj = responseObserver1.take();
+      StreamingOutputCallResponse callResponse = (StreamingOutputCallResponse) responseObj;
+      String clientSocketAddressInCall1 = callResponse.getPeerSocketAddress();
+      assertThat(clientSocketAddressInCall1).isNotEmpty();
+
+      StreamingOutputCallResponseObserver responseObserver2 =
+          new StreamingOutputCallResponseObserver(lastItem);
+      StreamObserver<StreamingOutputCallRequest> streamObserver2 =
+          asyncStub.fullDuplexCall(responseObserver2);
+      streamObserver2.onNext(request);
+      callResponse = (StreamingOutputCallResponse) responseObserver2.take();
+      String clientSocketAddressInCall2 = callResponse.getPeerSocketAddress();
+
+      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(lastItem);
+      StreamObserver<StreamingOutputCallRequest> streamObserver3 =
+          asyncStub.fullDuplexCall(responseObserver3);
+      streamObserver3.onNext(request);
+      callResponse = (StreamingOutputCallResponse) responseObserver3.take();
+      String clientSocketAddressInCall3 = callResponse.getPeerSocketAddress();
+
+      // This assertion is currently failing because connection scaling when MCS limit has been
+      // reached is not yet implemented in gRPC Java.
+      assertThat(clientSocketAddressInCall3).isNotEqualTo(clientSocketAddressInCall1);
+
+      streamObserver1.onCompleted();
+      streamObserver2.onCompleted();
+      streamObserver3.onCompleted();
+      assertThat(responseObserver1.take()).isSameInstanceAs(lastItem);
+      assertThat(responseObserver2.take()).isSameInstanceAs(lastItem);
+      assertThat(responseObserver3.take()).isSameInstanceAs(lastItem);
+    }
   }
 
   private static String validTestCasesHelpText() {

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

@@ -16,13 +16,18 @@
 
 package io.grpc.testing.integration;
 
+import static io.grpc.Grpc.TRANSPORT_ATTR_REMOTE_ADDR;
+
 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 +47,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;
@@ -59,8 +65,8 @@
  * sent in response streams.
  */
 public class TestServiceImpl implements io.grpc.BindableService, AsyncService {
+  static final 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;
@@ -441,7 +447,13 @@ public Queue<Chunk> toChunkQueue(StreamingOutputCallRequest request) {
     Queue<Chunk> chunkQueue = new ArrayDeque<>();
     int offset = 0;
     for (ResponseParameters params : request.getResponseParametersList()) {
-      chunkQueue.add(new Chunk(params.getIntervalUs(), offset, params.getSize()));
+      String peerSocketAddress = null;
+      if (params.getFillPeerSocketAddress().getValue()) {
+        SocketAddress peerAddress = PEER_ADDRESS_CONTEXT_KEY.get();
+        peerSocketAddress = peerAddress != null ? peerAddress.toString() : "";
+      }
+      chunkQueue.add(
+          new Chunk(params.getIntervalUs(), offset, params.getSize(), peerSocketAddress));
 
       // Increment the offset past this chunk. Buffer need to be circular.
       offset = (offset + params.getSize()) % compressableBuffer.size();
@@ -459,11 +471,17 @@ private class Chunk {
     private final int delayMicroseconds;
     private final int offset;
     private final int length;
+    private final String peerSocketAddress;
 
     public Chunk(int delayMicroseconds, int offset, int length) {
+      this(delayMicroseconds, offset, length, null);
+    }
+
+    public Chunk(int delayMicroseconds, int offset, int length, String peerSocketAddress) {
       this.delayMicroseconds = delayMicroseconds;
       this.offset = offset;
       this.length = length;
+      this.peerSocketAddress = peerSocketAddress;
     }
 
     /**
@@ -472,10 +490,15 @@ public Chunk(int delayMicroseconds, int offset, int length) {
     private StreamingOutputCallResponse toResponse() {
       StreamingOutputCallResponse.Builder responseBuilder =
           StreamingOutputCallResponse.newBuilder();
-      ByteString payload = generatePayload(compressableBuffer, offset, length);
-      responseBuilder.setPayload(
-          Payload.newBuilder()
-              .setBody(payload));
+      if (length > 0) {
+        ByteString payload = generatePayload(compressableBuffer, offset, length);
+        responseBuilder.setPayload(
+            Payload.newBuilder()
+                .setBody(payload));
+      }
+      if (peerSocketAddress != null) {
+        responseBuilder.setPeerSocketAddress(peerSocketAddress);
+      }
       return responseBuilder.build();
     }
   }
@@ -505,7 +528,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 AddPeerAddressToContextInterceptor());
   }
 
   /**
@@ -540,6 +564,22 @@ public void close(Status status, Metadata trailers) {
     };
   }
 
+  static class AddPeerAddressToContextInterceptor 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 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 int mcsLimit = -1;
 
   private ScheduledExecutorService executor;
   private Server server;
@@ -118,6 +119,10 @@ void parseArgs(String[] args) {
           usage = true;
           break;
         }
+      } else if ("max_concurrent_streams_limit".equals(key)) {
+        mcsLimit = Integer.parseInt(value);
+        // TODO: Make Netty server builder usable for IPV6 as well (not limited to MCS handling)
+        addressType = Util.AddressType.IPV4; // To use NettyServerBuilder
       } else {
         System.err.println("Unknown argument: " + key);
         usage = true;
@@ -141,6 +146,8 @@ void parseArgs(String[] args) {
               + "\n                        for testing. Only effective when --use_alts=true."
               + "\n  --address_type=IPV4|IPV6|IPV4_IPV6"
               + "\n                        What type of addresses to listen on. Default IPV4_IPV6"
+              + "\n  --max_concurrent_streams_limit=LIMIT"
+              + "\n                        Set the maximum concurrent streams limit"
       );
       System.exit(1);
     }
@@ -186,6 +193,9 @@ void start() throws Exception {
         if (v4Address != null && !v4Address.equals(localV4Address)) {
           ((NettyServerBuilder) serverBuilder).addListenAddress(v4Address);
         }
+        if (mcsLimit != -1) {
+          ((NettyServerBuilder) serverBuilder).maxConcurrentCallsPerConnection(mcsLimit);
+        }
         break;
       case IPV6:
         List<SocketAddress> v6Addresses = Util.getV6Addresses(port);

interop-testing/src/main/proto/grpc/testing/messages.proto

@@ -159,6 +159,10 @@ message ResponseParameters {
   // implement the full compression tests by introspecting the call to verify
   // the response's compression status.
   BoolValue compressed = 3;
+
+  // Whether to request the server to send the requesting peer's socket
+  // address in the response.
+  BoolValue fill_peer_socket_address = 4;
 }
 
 // Server-streaming request.
@@ -186,6 +190,9 @@ message StreamingOutputCallRequest {
 message StreamingOutputCallResponse {
   // Payload to increase response size.
   Payload payload = 1;
+
+  // The peer's socket address if requested.
+  string peer_socket_address = 2;
 }
 
 // For reconnect interop test only.

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",
+      "max_concurrent_streams_connection_scaling",
     };
 
     // additional test cases