diff --git a/src/proto/grpc/testing/control.proto b/src/proto/grpc/testing/control.proto index 2f352e652f323945a613be6a5bb233a8c47b71df..d135cb4d3236fffc1f7b03bcc185d36324d6f604 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 af285ceab8130056f02a91bd4945c422a85c4251..57cd9ecf76b85cb78d52da7c06be483c72ed3323 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 490156aec2cec0d9cf19927bbe050c040b2b9f58..370a3a834a32415a9aa3b9be09df3f0236c4db04 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 2a7cf805e57d70263b816a6b8c6cc761dfb309ca..4b2b400c0ce9ec3aa10e30ce0b47f83cbb478b48 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 aa3cb6882146ab771b3d00036e2ecb00ed006ffb..1fe37b1667c2028f6340b964606e7b70cdf44c3c 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 6316605aafd15bef66cf14c9a0858aad157bb714..f33b13b5b24f1dd01425accf65ae190986899c7b 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 27de69fa65a3896ac544f3998fdbe15eb667d408..f14a5c95adedd94b7fb5c4101dc0d5b1b79d5d61 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 a1e73e9abe921a8ce197e866b6643c77342fa510..f42cfe3255a399c9a67d8b8c56445826b8894962 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))); }