From 0eefd8b20981d90ca453bfa4794b253ba6364654 Mon Sep 17 00:00:00 2001
From: Esun Kim <veblush@google.com>
Date: Fri, 12 Jun 2020 14:22:34 -0700
Subject: [PATCH] Finish client first completely for client-lead tests

---
 test/cpp/qps/driver.cc | 42 ++++++++++++++++++++++++++++++++----------
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/test/cpp/qps/driver.cc b/test/cpp/qps/driver.cc
index deac6c43cbe..59640ebb5cc 100644
--- a/test/cpp/qps/driver.cc
+++ b/test/cpp/qps/driver.cc
@@ -492,6 +492,10 @@ std::unique_ptr<ScenarioResult> RunScenario(
   Histogram merged_latencies;
   std::unordered_map<int, int64_t> merged_statuses;
 
+  // For the case where clients leads the test, it's going to finish
+  // clients first and wait until it's completely done for generaous termination.
+  bool client_finish_first = (client_config.rpc_type() != STREAMING_FROM_SERVER);
+
   gpr_log(GPR_INFO, "Finishing clients");
   for (size_t i = 0; i < num_clients; i++) {
     auto client = &clients[i];
@@ -504,16 +508,19 @@ std::unique_ptr<ScenarioResult> RunScenario(
       GPR_ASSERT(false);
     }
   }
-  gpr_log(GPR_INFO, "Finishing servers");
-  for (size_t i = 0; i < num_servers; i++) {
-    auto server = &servers[i];
-    if (!server->stream->Write(server_mark)) {
-      gpr_log(GPR_ERROR, "Couldn't write mark to server %zu", i);
-      GPR_ASSERT(false);
-    }
-    if (!server->stream->WritesDone()) {
-      gpr_log(GPR_ERROR, "Failed WritesDone for server %zu", i);
-      GPR_ASSERT(false);
+
+  if (!client_finish_first) {
+    gpr_log(GPR_INFO, "Finishing servers");
+    for (size_t i = 0; i < num_servers; i++) {
+      auto server = &servers[i];
+      if (!server->stream->Write(server_mark)) {
+        gpr_log(GPR_ERROR, "Couldn't write mark to server %zu", i);
+        GPR_ASSERT(false);
+      }
+      if (!server->stream->WritesDone()) {
+        gpr_log(GPR_ERROR, "Failed WritesDone for server %zu", i);
+        GPR_ASSERT(false);
+      }
     }
   }
 
@@ -559,6 +566,21 @@ std::unique_ptr<ScenarioResult> RunScenario(
     rrc->set_count(it->second);
   }
 
+  if (client_finish_first) {
+    gpr_log(GPR_INFO, "Finishing servers");
+    for (size_t i = 0; i < num_servers; i++) {
+      auto server = &servers[i];
+      if (!server->stream->Write(server_mark)) {
+        gpr_log(GPR_ERROR, "Couldn't write mark to server %zu", i);
+        GPR_ASSERT(false);
+      }
+      if (!server->stream->WritesDone()) {
+        gpr_log(GPR_ERROR, "Failed WritesDone for server %zu", i);
+        GPR_ASSERT(false);
+      }
+    }
+  }
+
   for (size_t i = 0; i < num_servers; i++) {
     auto server = &servers[i];
     // Read the server final status