Merge pull request #5008 from vjpai/worker_quit

Provide an RPC to quit the performance benchmarking workers

src/node/performance/worker_service_impl.js

@@ -37,6 +37,11 @@ var os = require('os');
 var BenchmarkClient = require('./benchmark_client');
 var BenchmarkServer = require('./benchmark_server');
 
+exports.quitWorker = function quitWorker(call, callback) {
+  callback(null, {});
+  process.exit(0);
+}
+
 exports.runClient = function runClient(call) {
   var client;
   call.on('data', function(request) {

src/proto/grpc/testing/control.proto

@@ -166,3 +166,6 @@ message CoreResponse {
   // Number of cores available on the server
   int32 cores = 1;
 }
+
+message Void {
+}

src/proto/grpc/testing/services.proto

@@ -65,4 +65,7 @@ service WorkerService {
 
   // Just return the core count - unary call
   rpc CoreCount(CoreRequest) returns (CoreResponse);
+
+  // Quit this worker
+  rpc QuitWorker(Void) returns (Void);
 }

test/cpp/qps/driver.cc

@@ -384,5 +384,18 @@ std::unique_ptr<ScenarioResult> RunScenario(
   delete[] servers;
   return result;
 }
+
+void RunQuit() {
+  // Get client, server lists
+  auto workers = get_workers("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

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
 

test/cpp/qps/qps_driver.cc

@@ -78,6 +78,8 @@ DEFINE_int32(client_core_limit, -1, "Limit on client cores to use");
 
 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;
@@ -90,6 +92,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));
 

test/cpp/qps/qps_worker.cc

@@ -103,8 +103,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)
@@ -140,6 +140,16 @@ class WorkerServiceImpl GRPC_FINAL : public WorkerService::Service {
     return Status::OK;
   }
 
+  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 {
@@ -250,10 +260,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);
@@ -269,5 +281,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

test/cpp/qps/qps_worker.h

@@ -36,6 +36,8 @@
 
 #include <memory>
 
+#include <grpc/support/atm.h>
+
 namespace grpc {
 
 class Server;
@@ -49,9 +51,14 @@ 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

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)));
   }