#include <list>
#include <memory>
#include <mutex>
#include <sstream>
#include <string>
#include <thread>
#include <utility>
#include <vector>
#include <grpc/grpc.h>
#include <grpc/support/cpu.h>
#include <grpc/support/log.h>
#include <grpcpp/alarm.h>
#include <grpcpp/channel.h>
#include <grpcpp/client_context.h>
#include "src/proto/grpc/testing/benchmark_service.grpc.pb.h"
#include "test/cpp/qps/client.h"
#include "test/cpp/qps/usage_timer.h"
namespace grpc {
namespace testing {
struct CallbackClientRpcContext {
CallbackClientRpcContext(BenchmarkService::Stub* stub) : stub_(stub) {}
~CallbackClientRpcContext() {}
SimpleResponse response_;
ClientContext context_;
Alarm alarm_;
BenchmarkService::Stub* stub_;
};
static std::unique_ptr<BenchmarkService::Stub> BenchmarkStubCreator(
const std::shared_ptr<Channel>& ch) {
return BenchmarkService::NewStub(ch);
}
class CallbackClient
: public ClientImpl<BenchmarkService::Stub, SimpleRequest> {
public:
CallbackClient(const ClientConfig& config)
: ClientImpl<BenchmarkService::Stub, SimpleRequest>(
config, BenchmarkStubCreator) {
num_threads_ = NumThreads(config);
rpcs_done_ = 0;
SetupLoadTest(config, num_threads_);
total_outstanding_rpcs_ =
config.client_channels() * config.outstanding_rpcs_per_channel();
}
virtual ~CallbackClient() {}
protected:
size_t num_threads_;
size_t total_outstanding_rpcs_;
std::mutex shutdown_mu_;
std::condition_variable shutdown_cv_;
size_t rpcs_done_;
std::vector<std::unique_ptr<CallbackClientRpcContext>> ctx_;
virtual void InitThreadFuncImpl(size_t thread_idx) = 0;
virtual bool ThreadFuncImpl(Thread* t, size_t thread_idx) = 0;
void ThreadFunc(size_t thread_idx, Thread* t) override {
InitThreadFuncImpl(thread_idx);
ThreadFuncImpl(t, thread_idx);
}
virtual void ScheduleRpc(Thread* t, size_t thread_idx,
size_t ctx_vector_idx) = 0;
void NotifyMainThreadOfThreadCompletion() {
std::lock_guard<std::mutex> l(shutdown_mu_);
rpcs_done_++;
if (rpcs_done_ == total_outstanding_rpcs_) {
shutdown_cv_.notify_one();
}
}
private:
int NumThreads(const ClientConfig& config) {
int num_threads = config.async_client_threads();
if (num_threads <= 0) { num_threads = cores_;
gpr_log(GPR_INFO, "Sizing callback client to %d threads", num_threads);
}
return num_threads;
}
void DestroyMultithreading() final {
std::unique_lock<std::mutex> l(shutdown_mu_);
while (rpcs_done_ != total_outstanding_rpcs_) {
shutdown_cv_.wait(l);
}
EndThreads();
}
};
class CallbackUnaryClient final : public CallbackClient {
public:
CallbackUnaryClient(const ClientConfig& config) : CallbackClient(config) {
for (int ch = 0; ch < config.client_channels(); ch++) {
for (int i = 0; i < config.outstanding_rpcs_per_channel(); i++) {
ctx_.emplace_back(
new CallbackClientRpcContext(channels_[ch].get_stub()));
}
}
StartThreads(num_threads_);
}
~CallbackUnaryClient() {}
protected:
bool ThreadFuncImpl(Thread* t, size_t thread_idx) override {
for (size_t vector_idx = thread_idx; vector_idx < total_outstanding_rpcs_;
vector_idx += num_threads_) {
ScheduleRpc(t, thread_idx, vector_idx);
}
return true;
}
void InitThreadFuncImpl(size_t thread_idx) override { return; }
private:
void ScheduleRpc(Thread* t, size_t thread_idx, size_t vector_idx) override {
if (!closed_loop_) {
gpr_timespec next_issue_time = NextIssueTime(thread_idx);
ctx_[vector_idx]->alarm_.experimental().Set(
next_issue_time, [this, t, thread_idx, vector_idx](bool ok) {
IssueUnaryCallbackRpc(t, thread_idx, vector_idx);
});
} else {
IssueUnaryCallbackRpc(t, thread_idx, vector_idx);
}
}
void IssueUnaryCallbackRpc(Thread* t, size_t thread_idx, size_t vector_idx) {
GPR_TIMER_SCOPE("CallbackUnaryClient::ThreadFunc", 0);
double start = UsageTimer::Now();
ctx_[vector_idx]->stub_->experimental_async()->UnaryCall(
(&ctx_[vector_idx]->context_), &request_, &ctx_[vector_idx]->response_,
[this, t, thread_idx, start, vector_idx](grpc::Status s) {
HistogramEntry entry;
if (s.ok()) {
entry.set_value((UsageTimer::Now() - start) * 1e9);
}
entry.set_status(s.error_code());
t->UpdateHistogram(&entry);
if (ThreadCompleted() || !s.ok()) {
NotifyMainThreadOfThreadCompletion();
} else {
ctx_[vector_idx].reset(
new CallbackClientRpcContext(ctx_[vector_idx]->stub_));
ScheduleRpc(t, thread_idx, vector_idx);
}
});
}
};
std::unique_ptr<Client> CreateCallbackClient(const ClientConfig& config) {
switch (config.rpc_type()) {
case UNARY:
return std::unique_ptr<Client>(new CallbackUnaryClient(config));
case STREAMING:
case STREAMING_FROM_CLIENT:
case STREAMING_FROM_SERVER:
case STREAMING_BOTH_WAYS:
assert(false);
return nullptr;
default:
assert(false);
return nullptr;
}
}
} }