diff options
Diffstat (limited to 'test/cpp/qps/qps_worker.cc')
-rw-r--r-- | test/cpp/qps/qps_worker.cc | 24 |
1 files changed, 21 insertions, 3 deletions
diff --git a/test/cpp/qps/qps_worker.cc b/test/cpp/qps/qps_worker.cc index 7e9e05f7ec..9442017ddf 100644 --- a/test/cpp/qps/qps_worker.cc +++ b/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 |