From 42b32c64f567a5477ad00d01c41bf3cf5b3db409 Mon Sep 17 00:00:00 2001
From: apolcyn <apolcyn@google.com>
Date: Fri, 4 Mar 2022 12:18:07 -0800
Subject: [PATCH] interop-testing: Add a configurable warmup phase to fallback
 test client

---
 .../integration/GrpclbFallbackTestClient.java | 41 ++++++++++++++++---
 1 file changed, 35 insertions(+), 6 deletions(-)

diff --git a/interop-testing/src/main/java/io/grpc/testing/integration/GrpclbFallbackTestClient.java b/interop-testing/src/main/java/io/grpc/testing/integration/GrpclbFallbackTestClient.java
index 52c9e8238b..8eea59e895 100644
--- a/interop-testing/src/main/java/io/grpc/testing/integration/GrpclbFallbackTestClient.java
+++ b/interop-testing/src/main/java/io/grpc/testing/integration/GrpclbFallbackTestClient.java
@@ -75,6 +75,7 @@ public final class GrpclbFallbackTestClient {
   private String customCredentialsType;
   private String testCase;
   private Boolean skipNetCmd = false;
+  private int numWarmupRpcs;
 
   private ManagedChannel channel;
   private TestServiceGrpc.TestServiceBlockingStub blockingStub;
@@ -111,6 +112,8 @@ public final class GrpclbFallbackTestClient {
         customCredentialsType = value;
       } else if ("skip_net_cmd".equals(key)) {
         skipNetCmd = Boolean.valueOf(value);
+      } else if ("num_warmup_rpcs".equals(key)) {
+        numWarmupRpcs = Integer.valueOf(value);
       } else {
         System.err.println("Unknown argument: " + key);
         usage = true;
@@ -136,6 +139,10 @@ public final class GrpclbFallbackTestClient {
           + "shell command to allow setting the net config outside of the test "
           + "client. Default: "
           + c.skipNetCmd
+          + "\n  --num_warmup_rpcs                     Number of RPCs to perform "
+          + "on a separate warmup channel before the actual test runs (each warmup "
+          + "RPC uses a 1 second deadline). Default: "
+          + c.numWarmupRpcs
           + "\n  --test_case=TEST_CASE        Test case to run. Valid options are:"
           + "\n      fast_fallback_before_startup : fallback before LB connection"
           + "\n      fast_fallback_after_startup : fallback after startup due to "
@@ -197,14 +204,15 @@ public final class GrpclbFallbackTestClient {
     assertEquals(0, exitCode);
   }
 
-  private GrpclbRouteType doRpcAndGetPath(Deadline deadline) {
+  private GrpclbRouteType doRpcAndGetPath(
+      TestServiceGrpc.TestServiceBlockingStub stub, Deadline deadline) {
     logger.info("doRpcAndGetPath deadline: " + deadline);
     final SimpleRequest request = SimpleRequest.newBuilder()
         .setFillGrpclbRouteType(true)
         .build();
     GrpclbRouteType result = GrpclbRouteType.GRPCLB_ROUTE_TYPE_UNKNOWN;
     try {
-      SimpleResponse response = blockingStub
+      SimpleResponse response = stub
           .withDeadline(deadline)
           .unaryCall(request);
       result = response.getGrpclbRouteType();
@@ -226,7 +234,7 @@ public final class GrpclbFallbackTestClient {
     boolean fallBack = false;
     while (!fallbackDeadline.isExpired()) {
       GrpclbRouteType grpclbRouteType = doRpcAndGetPath(
-          Deadline.after(1, TimeUnit.SECONDS));
+          blockingStub, Deadline.after(1, TimeUnit.SECONDS));
       if (grpclbRouteType == GrpclbRouteType.GRPCLB_ROUTE_TYPE_BACKEND) {
         throw new AssertionError("Got grpclb route type backend. Backends are "
             + "supposed to be unreachable, so this test is broken");
@@ -247,7 +255,7 @@ public final class GrpclbFallbackTestClient {
     for (int i = 0; i < 30; i++) {
       assertEquals(
           GrpclbRouteType.GRPCLB_ROUTE_TYPE_FALLBACK,
-          doRpcAndGetPath(Deadline.after(20, TimeUnit.SECONDS)));
+          doRpcAndGetPath(blockingStub, Deadline.after(20, TimeUnit.SECONDS)));
       Thread.sleep(1000);
     }
   }
@@ -270,7 +278,7 @@ public final class GrpclbFallbackTestClient {
     initStub();
     assertEquals(
         GrpclbRouteType.GRPCLB_ROUTE_TYPE_BACKEND,
-        doRpcAndGetPath(Deadline.after(20, TimeUnit.SECONDS)));
+        doRpcAndGetPath(blockingStub, Deadline.after(20, TimeUnit.SECONDS)));
     runShellCmd(unrouteLbAndBackendAddrsCmd);
     final Deadline fallbackDeadline = Deadline.after(40, TimeUnit.SECONDS);
     waitForFallbackAndDoRpcs(fallbackDeadline);
@@ -280,13 +288,34 @@ public final class GrpclbFallbackTestClient {
     initStub();
     assertEquals(
         GrpclbRouteType.GRPCLB_ROUTE_TYPE_BACKEND,
-        doRpcAndGetPath(Deadline.after(20, TimeUnit.SECONDS)));
+        doRpcAndGetPath(blockingStub, Deadline.after(20, TimeUnit.SECONDS)));
     runShellCmd(blackholeLbAndBackendAddrsCmd);
     final Deadline fallbackDeadline = Deadline.after(40, TimeUnit.SECONDS);
     waitForFallbackAndDoRpcs(fallbackDeadline);
   }
 
+  // The purpose of this warmup method is to get potentially expensive one-per-process
+  // initialization out of the way, so that we can use aggressive timeouts in the actual
+  // test cases. Note that the warmup phase is done using a separate channel from the
+  // actual test cases, so that we don't affect the states of LB policies in the channel
+  // of the actual test case.
+  private void warmup() throws Exception {
+    logger.info("Begin warmup, performing " + numWarmupRpcs + " RPCs on the warmup channel");
+    ManagedChannel channel = createChannel();
+    TestServiceGrpc.TestServiceBlockingStub stub = TestServiceGrpc.newBlockingStub(channel);
+    for (int i = 0; i < numWarmupRpcs; i++) {
+      doRpcAndGetPath(stub, Deadline.after(1, TimeUnit.SECONDS));
+    }
+    try {
+      channel.shutdownNow();
+      channel.awaitTermination(1, TimeUnit.SECONDS);
+    } catch (Exception ex) {
+      throw new RuntimeException(ex);
+    }
+  }
+
   private void run() throws Exception {
+    warmup();
     logger.info("Begin test case: " + testCase);
     if (testCase.equals("fast_fallback_before_startup")) {
       runFastFallbackBeforeStartup();