diff --git a/src/proto/grpc/testing/control.proto b/src/proto/grpc/testing/control.proto index 2f352e652f3..d135cb4d323 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 af285ceab81..57cd9ecf76b 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 490156aec2c..370a3a834a3 100644 --- a/test/cpp/qps/driver.cc +++ b/test/cpp/qps/driver.cc @@ -283,5 +283,18 @@ std::unique_ptr 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 2a7cf805e57..4b2b400c0ce 100644 --- a/test/cpp/qps/driver.h +++ b/test/cpp/qps/driver.h @@ -70,6 +70,7 @@ std::unique_ptr 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 aa3cb688214..1fe37b1667c 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 6316605aafd..f33b13b5b24 100644 --- a/test/cpp/qps/qps_worker.cc +++ b/test/cpp/qps/qps_worker.cc @@ -107,8 +107,8 @@ static std::unique_ptr 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* 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(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(0)); +} +void QpsWorker::MarkDone() { + gpr_atm_rel_store(&done_, static_cast(1)); +} } // namespace testing } // namespace grpc diff --git a/test/cpp/qps/qps_worker.h b/test/cpp/qps/qps_worker.h index 27de69fa65a..f14a5c95ade 100644 --- a/test/cpp/qps/qps_worker.h +++ b/test/cpp/qps/qps_worker.h @@ -36,6 +36,8 @@ #include +#include + 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 impl_; std::unique_ptr server_; + + gpr_atm done_; }; } // namespace testing diff --git a/test/cpp/qps/worker.cc b/test/cpp/qps/worker.cc index a1e73e9abe9..f42cfe3255a 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))); }