Interceptor logic.

Author: kannanjgithub

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

@@ -60,8 +60,7 @@ public enum TestCases {
   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"),
-  MCS("max concurrent streaming"),
-  MCSSS("mcs server streaming");
+  MCS_CS("max concurrent streaming connection scaling");
 
   private final String description;
 

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

@@ -17,7 +17,7 @@
 package io.grpc.testing.integration;
 
 import static com.google.common.truth.Truth.assertThat;
-import static io.grpc.testing.integration.TestCases.MCS;
+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;
@@ -576,16 +576,10 @@ private void runTest(TestCases testCase) throws Exception {
         break;
       }
 
-      case MCS: {
+      case MCS_CS: {
         tester.testMcs();
         break;
       }
-
-      case MCSSS: {
-        tester.testMcs_serverStreaming();
-        break;
-      }
-
       default:
         throw new IllegalArgumentException("Unknown test case: " + testCase);
     }
@@ -622,8 +616,8 @@ private class Tester extends AbstractInteropTest {
 
     @Override
     protected ManagedChannelBuilder<?> createChannelBuilder() {
-      boolean useSubchannelMetricsSink = testCase.equals(MCS.toString());
-      boolean useGeneric = testCase.equals(MCS.toString())? true : false;
+      boolean useSubchannelMetricsSink = testCase.equals(MCS_CS.toString());
+      boolean useGeneric = testCase.equals(MCS_CS.toString())? true : false;
       ChannelCredentials channelCredentials;
       if (customCredentialsType != null) {
         useGeneric = true; // Retain old behavior; avoids erroring if incompatible
@@ -683,7 +677,7 @@ protected ManagedChannelBuilder<?> createChannelBuilder() {
         if (serverHostOverride != null) {
           channelBuilder.overrideAuthority(serverHostOverride);
         }
-        if (testCase.equals(MCS.toString())) {
+        if (testCase.equals(MCS_CS.toString())) {
           channelBuilder.disableServiceConfigLookUp();
           try {
             @SuppressWarnings("unchecked")
@@ -1101,32 +1095,37 @@ Object take() throws InterruptedException {
     }
 
     public void testMcs() throws Exception {
-      StreamingOutputCallResponseObserver responseObserver1 = new StreamingOutputCallResponseObserver();
+      StreamingOutputCallResponseObserver responseObserver1 =
+          new StreamingOutputCallResponseObserver();
       StreamObserver<StreamingOutputCallRequest> streamObserver1 =
           asyncStub.fullDuplexCall(responseObserver1);
-      streamObserver1.onNext(StreamingOutputCallRequest.newBuilder()
-          .addResponseParameters(ResponseParameters.newBuilder().setSize(1).build()).build());
-      assertThat(responseObserver1.take()).isInstanceOf(StreamingOutputCallResponse.class);
+      StreamingOutputCallRequest request = StreamingOutputCallRequest.newBuilder()
+          .setPayload(Payload.newBuilder().setBody(
+              ByteString.copyFrom(MCS_CS.description().getBytes())).build()).build();
+      streamObserver1.onNext(request);
+      Object responseObj = responseObserver1.take();
+      StreamingOutputCallResponse callResponse = (StreamingOutputCallResponse) responseObj;
+      String clientSocketAddressInCall1 = new String(callResponse.getPayload().getBody().toByteArray());
 
       StreamingOutputCallResponseObserver responseObserver2 = new StreamingOutputCallResponseObserver();
       StreamObserver<StreamingOutputCallRequest> streamObserver2 =
           asyncStub.fullDuplexCall(responseObserver2);
-      streamObserver2.onNext(StreamingOutputCallRequest.newBuilder()
-          .addResponseParameters(ResponseParameters.newBuilder().setSize(1).build()).build());
-      assertThat(responseObserver2.take()).isInstanceOf(StreamingOutputCallResponse.class);
+      streamObserver2.onNext(request);
+      callResponse = (StreamingOutputCallResponse) responseObserver2.take();
+      String clientSocketAddressInCall2 = new String(callResponse.getPayload().getBody().toByteArray());
 
-      assertThat(fakeMetricsSink.openConnectionCount).isEqualTo(1);
+      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(StreamingOutputCallRequest.newBuilder()
-          .addResponseParameters(ResponseParameters.newBuilder().setSize(1).build()).build());
-      assertThat(responseObserver3.take()).isInstanceOf(StreamingOutputCallResponse.class);
+      streamObserver3.onNext(request);
+      callResponse = (StreamingOutputCallResponse) responseObserver3.take();
+      String clientSocketAddressInCall3 = new String(callResponse.getPayload().getBody().toByteArray());
 
-      assertThat(fakeMetricsSink.openConnectionCount).isEqualTo(2);
+      assertThat(clientSocketAddressInCall3).isNotEqualTo(clientSocketAddressInCall1);
     }
 
     public void testMcs_serverStreaming() throws Exception {

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

@@ -16,13 +16,19 @@
 
 package io.grpc.testing.integration;
 
+import static io.grpc.Grpc.TRANSPORT_ATTR_REMOTE_ADDR;
+import static io.grpc.testing.integration.TestCases.MCS_CS;
+
 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.Attributes;
+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 +48,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;
@@ -55,12 +62,14 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
+import io.grpc.Context;
 
 /**
  * Implementation of the business logic for the TestService. Uses an executor to schedule chunks
  * 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;
@@ -235,6 +244,16 @@ public void onNext(StreamingOutputCallRequest request) {
               .asRuntimeException());
           return;
         }
+        if (new String(request.getPayload().getBody().toByteArray()).equals(MCS_CS.description())) {
+          SocketAddress peerAddress = PEER_ADDRESS_CONTEXT_KEY.get();
+          ByteString payload = ByteString.copyFrom(peerAddress.toString().getBytes());
+          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,25 @@ 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 {
+
+        // Continue the call processing within the new context
+        // return newContext.call(() -> next.startCall(call, headers));
+        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.
    */