Add a Quit RPC so that we can conveniently shut down the workers from the driver.

8 files changed, 55 insertions, 4 deletions

diff --git a/src/proto/grpc/testing/control.proto b/src/proto/grpc/testing/control.proto
index 2f352e652f..d135cb4d32 100644
--- a/src/proto/grpc/testing/control.proto
+++ b/src/proto/grpc/testing/control.proto
@@ -150,3 +150,6 @@ message ServerStatus {
   // Number of cores on the server. See gpr_cpu_num_cores.
   int32 cores = 3;
 }
+
+message Void {
+}
diff --git a/src/proto/grpc/testing/services.proto b/src/proto/grpc/testing/services.proto
index af285ceab8..57cd9ecf76 100644
--- a/src/proto/grpc/testing/services.proto
+++ b/src/proto/grpc/testing/services.proto
@@ -62,4 +62,7 @@ service WorkerService {
   // and once the shutdown has finished, the OK status is sent to terminate
   // this RPC.
   rpc RunClient(stream ClientArgs) returns (stream ClientStatus);
+
+  // Quit this worker
+  rpc QuitWorker(Void) returns (Void);
 }
diff --git a/test/cpp/qps/driver.cc b/test/cpp/qps/driver.cc
index 490156aec2..370a3a834a 100644
--- a/test/cpp/qps/driver.cc
+++ b/test/cpp/qps/driver.cc
@@ -283,5 +283,18 @@ std::unique_ptr<ScenarioResult> RunScenario(
   delete[] servers;
   return result;
 }
+
+void RunQuit() {
+  // Get client, server lists
+  auto workers = get_hosts("QPS_WORKERS");
+  for (size_t i = 0; i < workers.size(); i++) {
+    auto stub = WorkerService::NewStub(
+       CreateChannel(workers[i], InsecureChannelCredentials()));
+    Void dummy;
+    grpc::ClientContext ctx;
+    GPR_ASSERT(stub->QuitWorker(&ctx, dummy, &dummy).ok());
+  }
+}
+
 }  // namespace testing
 }  // namespace grpc
diff --git a/test/cpp/qps/driver.h b/test/cpp/qps/driver.h
index 2a7cf805e5..4b2b400c0c 100644
--- a/test/cpp/qps/driver.h
+++ b/test/cpp/qps/driver.h
@@ -70,6 +70,7 @@ std::unique_ptr<ScenarioResult> RunScenario(
     const grpc::testing::ServerConfig& server_config, size_t num_servers,
     int warmup_seconds, int benchmark_seconds, int spawn_local_worker_count);
 
+void RunQuit();
 }  // namespace testing
 }  // namespace grpc
 
diff --git a/test/cpp/qps/qps_driver.cc b/test/cpp/qps/qps_driver.cc
index aa3cb68821..1fe37b1667 100644
--- a/test/cpp/qps/qps_driver.cc
+++ b/test/cpp/qps/qps_driver.cc
@@ -77,6 +77,8 @@ DEFINE_double(pareto_alpha, -1.0, "Pareto alpha value");
 
 DEFINE_bool(secure_test, false, "Run a secure test");
 
+DEFINE_bool(quit, false, "Quit the workers");
+
 using grpc::testing::ClientConfig;
 using grpc::testing::ServerConfig;
 using grpc::testing::ClientType;
@@ -89,6 +91,11 @@ namespace grpc {
 namespace testing {
 
 static void QpsDriver() {
+  if (FLAGS_quit) {
+    RunQuit();
+    return;
+  }
+
   RpcType rpc_type;
   GPR_ASSERT(RpcType_Parse(FLAGS_rpc_type, &rpc_type));
 
diff --git a/test/cpp/qps/qps_worker.cc b/test/cpp/qps/qps_worker.cc
index 6316605aaf..f33b13b5b2 100644
--- a/test/cpp/qps/qps_worker.cc
+++ b/test/cpp/qps/qps_worker.cc
@@ -107,8 +107,8 @@ static std::unique_ptr<Server> CreateServer(const ServerConfig& config) {
 
 class WorkerServiceImpl GRPC_FINAL : public WorkerService::Service {
  public:
-  explicit WorkerServiceImpl(int server_port)
-      : acquired_(false), server_port_(server_port) {}
+  WorkerServiceImpl(int server_port, QpsWorker *worker)
+    : acquired_(false), server_port_(server_port), worker_(worker) {}
 
   Status RunClient(ServerContext* ctx,
                    ServerReaderWriter<ClientStatus, ClientArgs>* stream)
@@ -138,6 +138,16 @@ class WorkerServiceImpl GRPC_FINAL : public WorkerService::Service {
     return ret;
   }
 
+  Status QuitWorker(ServerContext *ctx, const Void*, Void*) GRPC_OVERRIDE {
+    InstanceGuard g(this);
+    if (!g.Acquired()) {
+      return Status(StatusCode::RESOURCE_EXHAUSTED, "");
+    }
+
+    worker_->MarkDone();
+    return Status::OK;
+  }
+  
  private:
   // Protect against multiple clients using this worker at once.
   class InstanceGuard {
@@ -248,10 +258,12 @@ class WorkerServiceImpl GRPC_FINAL : public WorkerService::Service {
   std::mutex mu_;
   bool acquired_;
   int server_port_;
+  QpsWorker *worker_;
 };
 
 QpsWorker::QpsWorker(int driver_port, int server_port) {
-  impl_.reset(new WorkerServiceImpl(server_port));
+  impl_.reset(new WorkerServiceImpl(server_port, this));
+  gpr_atm_rel_store(&done_, static_cast<gpr_atm>(0));
 
   char* server_address = NULL;
   gpr_join_host_port(&server_address, "::", driver_port);
@@ -267,5 +279,11 @@ QpsWorker::QpsWorker(int driver_port, int server_port) {
 
 QpsWorker::~QpsWorker() {}
 
+bool QpsWorker::Done() const {
+  return (gpr_atm_acq_load(&done_) != static_cast<gpr_atm>(0));
+}
+void QpsWorker::MarkDone() {
+  gpr_atm_rel_store(&done_, static_cast<gpr_atm>(1));
+}
 }  // namespace testing
 }  // namespace grpc
diff --git a/test/cpp/qps/qps_worker.h b/test/cpp/qps/qps_worker.h
index 27de69fa65..f14a5c95ad 100644
--- a/test/cpp/qps/qps_worker.h
+++ b/test/cpp/qps/qps_worker.h
@@ -36,6 +36,8 @@
 
 #include <memory>
 
+#include <grpc/support/atm.h>
+
 namespace grpc {
 
 class Server;
@@ -49,9 +51,13 @@ class QpsWorker {
   explicit QpsWorker(int driver_port, int server_port = 0);
   ~QpsWorker();
 
+  bool Done() const;
+  void MarkDone();
  private:
   std::unique_ptr<WorkerServiceImpl> impl_;
   std::unique_ptr<Server> server_;
+
+  gpr_atm done_;
 };
 
 }  // namespace testing
diff --git a/test/cpp/qps/worker.cc b/test/cpp/qps/worker.cc
index a1e73e9abe..f42cfe3255 100644
--- a/test/cpp/qps/worker.cc
+++ b/test/cpp/qps/worker.cc
@@ -56,7 +56,7 @@ namespace testing {
 static void RunServer() {
   QpsWorker worker(FLAGS_driver_port, FLAGS_server_port);
 
-  while (!got_sigint) {
+  while (!got_sigint && !worker.Done()) {
     gpr_sleep_until(gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
                                  gpr_time_from_seconds(5, GPR_TIMESPAN)));
   }