3 * Copyright 2020 gRPC authors.
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
9 * http://www.apache.org/licenses/LICENSE-2.0
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
19 #include <grpcpp/grpcpp.h>
20 #include <grpcpp/server.h>
21 #include <grpcpp/server_builder.h>
22 #include <grpcpp/server_context.h>
26 #include <condition_variable>
36 #include "absl/algorithm/container.h"
37 #include "absl/flags/flag.h"
38 #include "absl/strings/str_split.h"
39 #include "src/core/lib/channel/status_util.h"
40 #include "src/core/lib/gpr/env.h"
41 #include "src/proto/grpc/testing/empty.pb.h"
42 #include "src/proto/grpc/testing/messages.pb.h"
43 #include "src/proto/grpc/testing/test.grpc.pb.h"
44 #include "test/core/util/test_config.h"
45 #include "test/cpp/util/test_config.h"
47 ABSL_FLAG(bool, fail_on_failed_rpc, false,
48 "Fail client if any RPCs fail after first successful RPC.");
49 ABSL_FLAG(int32_t, num_channels, 1, "Number of channels.");
50 ABSL_FLAG(bool, print_response, false, "Write RPC response to stdout.");
51 ABSL_FLAG(int32_t, qps, 1, "Qps per channel.");
52 // TODO(Capstan): Consider using absl::Duration
53 ABSL_FLAG(int32_t, rpc_timeout_sec, 30, "Per RPC timeout seconds.");
54 ABSL_FLAG(std::string, server, "localhost:50051", "Address of server.");
55 ABSL_FLAG(int32_t, stats_port, 50052,
56 "Port to expose peer distribution stats service.");
57 ABSL_FLAG(std::string, rpc, "UnaryCall",
58 "a comma separated list of rpc methods.");
59 ABSL_FLAG(std::string, metadata, "", "metadata to send with the RPC.");
60 ABSL_FLAG(std::string, expect_status, "OK",
61 "RPC status for the test RPC to be considered successful");
64 using grpc::ClientAsyncResponseReader;
65 using grpc::ClientContext;
66 using grpc::CompletionQueue;
68 using grpc::ServerBuilder;
69 using grpc::ServerContext;
71 using grpc::testing::ClientConfigureRequest;
72 using grpc::testing::ClientConfigureRequest_RpcType_Name;
73 using grpc::testing::ClientConfigureResponse;
74 using grpc::testing::Empty;
75 using grpc::testing::LoadBalancerAccumulatedStatsRequest;
76 using grpc::testing::LoadBalancerAccumulatedStatsResponse;
77 using grpc::testing::LoadBalancerStatsRequest;
78 using grpc::testing::LoadBalancerStatsResponse;
79 using grpc::testing::LoadBalancerStatsService;
80 using grpc::testing::SimpleRequest;
81 using grpc::testing::SimpleResponse;
82 using grpc::testing::TestService;
83 using grpc::testing::XdsUpdateClientConfigureService;
85 class XdsStatsWatcher;
87 struct StatsWatchers {
88 // Unique ID for each outgoing RPC
89 int global_request_id = 0;
90 // Unique ID for each outgoing RPC by RPC method type
91 std::map<int, int> global_request_id_by_type;
92 // Stores a set of watchers that should be notified upon outgoing RPC
94 std::set<XdsStatsWatcher*> watchers;
95 // Global watcher for accumululated stats.
96 XdsStatsWatcher* global_watcher;
97 // Mutex for global_request_id and watchers
100 // Whether at least one RPC has succeeded, indicating xDS resolution completed.
101 std::atomic<bool> one_rpc_succeeded(false);
102 // RPC configuration detailing how RPC should be sent.
104 ClientConfigureRequest::RpcType type;
105 std::vector<std::pair<std::string, std::string>> metadata;
108 struct RpcConfigurationsQueue {
109 // A queue of RPC configurations detailing how RPCs should be sent.
110 std::deque<std::vector<RpcConfig>> rpc_configs_queue;
111 // Mutex for rpc_configs_queue
112 std::mutex mu_rpc_configs_queue;
114 struct AsyncClientCall {
115 Empty empty_response;
116 SimpleResponse simple_response;
117 ClientContext context;
119 int saved_request_id;
120 ClientConfigureRequest::RpcType rpc_type;
121 std::unique_ptr<ClientAsyncResponseReader<Empty>> empty_response_reader;
122 std::unique_ptr<ClientAsyncResponseReader<SimpleResponse>>
123 simple_response_reader;
126 /** Records the remote peer distribution for a given range of RPCs. */
127 class XdsStatsWatcher {
129 XdsStatsWatcher(int start_id, int end_id)
130 : start_id_(start_id), end_id_(end_id), rpcs_needed_(end_id - start_id) {}
132 // Upon the completion of an RPC, we will look at the request_id, the
133 // rpc_type, and the peer the RPC was sent to in order to count
134 // this RPC into the right stats bin.
135 void RpcCompleted(AsyncClientCall* call, const std::string& peer) {
136 // We count RPCs for global watcher or if the request_id falls into the
137 // watcher's interested range of request ids.
138 if ((start_id_ == 0 && end_id_ == 0) ||
139 (start_id_ <= call->saved_request_id &&
140 call->saved_request_id < end_id_)) {
142 std::lock_guard<std::mutex> lock(m_);
145 ++no_remote_peer_by_type_[call->rpc_type];
147 // RPC is counted into both per-peer bin and per-method-per-peer bin.
148 rpcs_by_peer_[peer]++;
149 rpcs_by_type_[call->rpc_type][peer]++;
152 // Report accumulated stats.
153 auto& stats_per_method = *accumulated_stats_.mutable_stats_per_method();
155 stats_per_method[ClientConfigureRequest_RpcType_Name(
157 auto& result = *method_stat.mutable_result();
158 grpc_status_code code =
159 static_cast<grpc_status_code>(call->status.error_code());
160 auto& num_rpcs = result[code];
162 auto rpcs_started = method_stat.rpcs_started();
163 method_stat.set_rpcs_started(++rpcs_started);
169 void WaitForRpcStatsResponse(LoadBalancerStatsResponse* response,
171 std::unique_lock<std::mutex> lock(m_);
172 cv_.wait_for(lock, std::chrono::seconds(timeout_sec),
173 [this] { return rpcs_needed_ == 0; });
174 response->mutable_rpcs_by_peer()->insert(rpcs_by_peer_.begin(),
175 rpcs_by_peer_.end());
176 auto& response_rpcs_by_method = *response->mutable_rpcs_by_method();
177 for (const auto& rpc_by_type : rpcs_by_type_) {
178 std::string method_name;
179 if (rpc_by_type.first == ClientConfigureRequest::EMPTY_CALL) {
180 method_name = "EmptyCall";
181 } else if (rpc_by_type.first == ClientConfigureRequest::UNARY_CALL) {
182 method_name = "UnaryCall";
186 // TODO(@donnadionne): When the test runner changes to accept EMPTY_CALL
187 // and UNARY_CALL we will just use the name of the enum instead of the
188 // method_name variable.
189 auto& response_rpc_by_method = response_rpcs_by_method[method_name];
190 auto& response_rpcs_by_peer =
191 *response_rpc_by_method.mutable_rpcs_by_peer();
192 for (const auto& rpc_by_peer : rpc_by_type.second) {
193 auto& response_rpc_by_peer = response_rpcs_by_peer[rpc_by_peer.first];
194 response_rpc_by_peer = rpc_by_peer.second;
197 response->set_num_failures(no_remote_peer_ + rpcs_needed_);
200 void GetCurrentRpcStats(LoadBalancerAccumulatedStatsResponse* response,
201 StatsWatchers* stats_watchers) {
202 std::unique_lock<std::mutex> lock(m_);
203 response->CopyFrom(accumulated_stats_);
204 // TODO(@donnadionne): delete deprecated stats below when the test is no
205 // longer using them.
206 auto& response_rpcs_started_by_method =
207 *response->mutable_num_rpcs_started_by_method();
208 auto& response_rpcs_succeeded_by_method =
209 *response->mutable_num_rpcs_succeeded_by_method();
210 auto& response_rpcs_failed_by_method =
211 *response->mutable_num_rpcs_failed_by_method();
212 for (const auto& rpc_by_type : rpcs_by_type_) {
213 auto total_succeeded = 0;
214 for (const auto& rpc_by_peer : rpc_by_type.second) {
215 total_succeeded += rpc_by_peer.second;
217 response_rpcs_succeeded_by_method[ClientConfigureRequest_RpcType_Name(
218 rpc_by_type.first)] = total_succeeded;
219 response_rpcs_started_by_method[ClientConfigureRequest_RpcType_Name(
220 rpc_by_type.first)] =
221 stats_watchers->global_request_id_by_type[rpc_by_type.first];
222 response_rpcs_failed_by_method[ClientConfigureRequest_RpcType_Name(
223 rpc_by_type.first)] = no_remote_peer_by_type_[rpc_by_type.first];
231 int no_remote_peer_ = 0;
232 std::map<int, int> no_remote_peer_by_type_;
233 // A map of stats keyed by peer name.
234 std::map<std::string, int> rpcs_by_peer_;
235 // A two-level map of stats keyed at top level by RPC method and second level
237 std::map<int, std::map<std::string, int>> rpcs_by_type_;
238 // Storing accumulated stats in the response proto format.
239 LoadBalancerAccumulatedStatsResponse accumulated_stats_;
241 std::condition_variable cv_;
246 TestClient(const std::shared_ptr<Channel>& channel,
247 StatsWatchers* stats_watchers)
248 : stub_(TestService::NewStub(channel)), stats_watchers_(stats_watchers) {}
250 void AsyncUnaryCall(const RpcConfig& config) {
251 SimpleResponse response;
252 int saved_request_id;
254 std::lock_guard<std::mutex> lock(stats_watchers_->mu);
255 saved_request_id = ++stats_watchers_->global_request_id;
257 ->global_request_id_by_type[ClientConfigureRequest::UNARY_CALL];
259 std::chrono::system_clock::time_point deadline =
260 std::chrono::system_clock::now() +
261 std::chrono::seconds(config.timeout_sec != 0
263 : absl::GetFlag(FLAGS_rpc_timeout_sec));
264 AsyncClientCall* call = new AsyncClientCall;
265 for (const auto& data : config.metadata) {
266 call->context.AddMetadata(data.first, data.second);
267 // TODO(@donnadionne): move deadline to separate proto.
268 if (data.first == "rpc-behavior" && data.second == "keep-open") {
270 std::chrono::system_clock::now() + std::chrono::seconds(INT_MAX);
273 call->context.set_deadline(deadline);
274 call->saved_request_id = saved_request_id;
275 call->rpc_type = ClientConfigureRequest::UNARY_CALL;
276 call->simple_response_reader = stub_->PrepareAsyncUnaryCall(
277 &call->context, SimpleRequest::default_instance(), &cq_);
278 call->simple_response_reader->StartCall();
279 call->simple_response_reader->Finish(&call->simple_response, &call->status,
283 void AsyncEmptyCall(const RpcConfig& config) {
285 int saved_request_id;
287 std::lock_guard<std::mutex> lock(stats_watchers_->mu);
288 saved_request_id = ++stats_watchers_->global_request_id;
290 ->global_request_id_by_type[ClientConfigureRequest::EMPTY_CALL];
292 std::chrono::system_clock::time_point deadline =
293 std::chrono::system_clock::now() +
294 std::chrono::seconds(config.timeout_sec != 0
296 : absl::GetFlag(FLAGS_rpc_timeout_sec));
297 AsyncClientCall* call = new AsyncClientCall;
298 for (const auto& data : config.metadata) {
299 call->context.AddMetadata(data.first, data.second);
300 // TODO(@donnadionne): move deadline to separate proto.
301 if (data.first == "rpc-behavior" && data.second == "keep-open") {
303 std::chrono::system_clock::now() + std::chrono::seconds(INT_MAX);
306 call->context.set_deadline(deadline);
307 call->saved_request_id = saved_request_id;
308 call->rpc_type = ClientConfigureRequest::EMPTY_CALL;
309 call->empty_response_reader = stub_->PrepareAsyncEmptyCall(
310 &call->context, Empty::default_instance(), &cq_);
311 call->empty_response_reader->StartCall();
312 call->empty_response_reader->Finish(&call->empty_response, &call->status,
316 void AsyncCompleteRpc() {
319 while (cq_.Next(&got_tag, &ok)) {
320 AsyncClientCall* call = static_cast<AsyncClientCall*>(got_tag);
323 std::lock_guard<std::mutex> lock(stats_watchers_->mu);
324 auto server_initial_metadata = call->context.GetServerInitialMetadata();
325 auto metadata_hostname =
326 call->context.GetServerInitialMetadata().find("hostname");
327 std::string hostname =
328 metadata_hostname != call->context.GetServerInitialMetadata().end()
329 ? std::string(metadata_hostname->second.data(),
330 metadata_hostname->second.length())
331 : call->simple_response.hostname();
332 for (auto watcher : stats_watchers_->watchers) {
333 watcher->RpcCompleted(call, hostname);
337 if (!RpcStatusCheckSuccess(call)) {
338 if (absl::GetFlag(FLAGS_print_response) ||
339 absl::GetFlag(FLAGS_fail_on_failed_rpc)) {
340 std::cout << "RPC failed: " << call->status.error_code() << ": "
341 << call->status.error_message() << std::endl;
343 if (absl::GetFlag(FLAGS_fail_on_failed_rpc) &&
344 one_rpc_succeeded.load()) {
348 if (absl::GetFlag(FLAGS_print_response)) {
349 auto metadata_hostname =
350 call->context.GetServerInitialMetadata().find("hostname");
351 std::string hostname =
353 call->context.GetServerInitialMetadata().end()
354 ? std::string(metadata_hostname->second.data(),
355 metadata_hostname->second.length())
356 : call->simple_response.hostname();
357 std::cout << "Greeting: Hello world, this is " << hostname
358 << ", from " << call->context.peer() << std::endl;
360 one_rpc_succeeded = true;
368 static bool RpcStatusCheckSuccess(AsyncClientCall* call) {
369 // Determine RPC success based on expected status.
370 grpc_status_code code;
371 GPR_ASSERT(grpc_status_code_from_string(
372 absl::GetFlag(FLAGS_expect_status).c_str(), &code));
373 return code == static_cast<grpc_status_code>(call->status.error_code());
376 std::unique_ptr<TestService::Stub> stub_;
377 StatsWatchers* stats_watchers_;
381 class LoadBalancerStatsServiceImpl : public LoadBalancerStatsService::Service {
383 explicit LoadBalancerStatsServiceImpl(StatsWatchers* stats_watchers)
384 : stats_watchers_(stats_watchers) {}
386 Status GetClientStats(ServerContext* /*context*/,
387 const LoadBalancerStatsRequest* request,
388 LoadBalancerStatsResponse* response) override {
391 XdsStatsWatcher* watcher;
393 std::lock_guard<std::mutex> lock(stats_watchers_->mu);
394 start_id = stats_watchers_->global_request_id + 1;
395 end_id = start_id + request->num_rpcs();
396 watcher = new XdsStatsWatcher(start_id, end_id);
397 stats_watchers_->watchers.insert(watcher);
399 watcher->WaitForRpcStatsResponse(response, request->timeout_sec());
401 std::lock_guard<std::mutex> lock(stats_watchers_->mu);
402 stats_watchers_->watchers.erase(watcher);
408 Status GetClientAccumulatedStats(
409 ServerContext* /*context*/,
410 const LoadBalancerAccumulatedStatsRequest* /*request*/,
411 LoadBalancerAccumulatedStatsResponse* response) override {
412 std::lock_guard<std::mutex> lock(stats_watchers_->mu);
413 stats_watchers_->global_watcher->GetCurrentRpcStats(response,
419 StatsWatchers* stats_watchers_;
422 class XdsUpdateClientConfigureServiceImpl
423 : public XdsUpdateClientConfigureService::Service {
425 explicit XdsUpdateClientConfigureServiceImpl(
426 RpcConfigurationsQueue* rpc_configs_queue)
427 : rpc_configs_queue_(rpc_configs_queue) {}
429 Status Configure(ServerContext* /*context*/,
430 const ClientConfigureRequest* request,
431 ClientConfigureResponse* /*response*/) override {
432 std::map<int, std::vector<std::pair<std::string, std::string>>>
434 for (const auto& data : request->metadata()) {
435 metadata_map[data.type()].push_back({data.key(), data.value()});
437 std::vector<RpcConfig> configs;
438 for (const auto& rpc : request->types()) {
440 config.timeout_sec = request->timeout_sec();
441 config.type = static_cast<ClientConfigureRequest::RpcType>(rpc);
442 auto metadata_iter = metadata_map.find(rpc);
443 if (metadata_iter != metadata_map.end()) {
444 config.metadata = metadata_iter->second;
446 configs.push_back(std::move(config));
449 std::lock_guard<std::mutex> lock(
450 rpc_configs_queue_->mu_rpc_configs_queue);
451 rpc_configs_queue_->rpc_configs_queue.emplace_back(std::move(configs));
457 RpcConfigurationsQueue* rpc_configs_queue_;
460 void RunTestLoop(std::chrono::duration<double> duration_per_query,
461 StatsWatchers* stats_watchers,
462 RpcConfigurationsQueue* rpc_configs_queue) {
463 TestClient client(grpc::CreateChannel(absl::GetFlag(FLAGS_server),
464 grpc::InsecureChannelCredentials()),
466 std::chrono::time_point<std::chrono::system_clock> start =
467 std::chrono::system_clock::now();
468 std::chrono::duration<double> elapsed;
470 std::thread thread = std::thread(&TestClient::AsyncCompleteRpc, &client);
472 std::vector<RpcConfig> configs;
475 std::lock_guard<std::mutex> lockk(
476 rpc_configs_queue->mu_rpc_configs_queue);
477 if (!rpc_configs_queue->rpc_configs_queue.empty()) {
478 configs = std::move(rpc_configs_queue->rpc_configs_queue.front());
479 rpc_configs_queue->rpc_configs_queue.pop_front();
483 elapsed = std::chrono::system_clock::now() - start;
484 if (elapsed > duration_per_query) {
485 start = std::chrono::system_clock::now();
486 for (const auto& config : configs) {
487 if (config.type == ClientConfigureRequest::EMPTY_CALL) {
488 client.AsyncEmptyCall(config);
489 } else if (config.type == ClientConfigureRequest::UNARY_CALL) {
490 client.AsyncUnaryCall(config);
500 void RunServer(const int port, StatsWatchers* stats_watchers,
501 RpcConfigurationsQueue* rpc_configs_queue) {
502 GPR_ASSERT(port != 0);
503 std::ostringstream server_address;
504 server_address << "0.0.0.0:" << port;
506 LoadBalancerStatsServiceImpl stats_service(stats_watchers);
507 XdsUpdateClientConfigureServiceImpl client_config_service(rpc_configs_queue);
509 ServerBuilder builder;
510 builder.RegisterService(&stats_service);
511 builder.RegisterService(&client_config_service);
512 builder.AddListeningPort(server_address.str(),
513 grpc::InsecureServerCredentials());
514 std::unique_ptr<Server> server(builder.BuildAndStart());
515 gpr_log(GPR_DEBUG, "Server listening on %s", server_address.str().c_str());
520 void BuildRpcConfigsFromFlags(RpcConfigurationsQueue* rpc_configs_queue) {
521 // Store Metadata like
522 // "EmptyCall:key1:value1,UnaryCall:key1:value1,UnaryCall:key2:value2" into a
523 // map where the key is the RPC method and value is a vector of key:value
524 // pairs. {EmptyCall, [{key1,value1}],
525 // UnaryCall, [{key1,value1}, {key2,value2}]}
526 std::vector<std::string> rpc_metadata =
527 absl::StrSplit(absl::GetFlag(FLAGS_metadata), ',', absl::SkipEmpty());
528 std::map<int, std::vector<std::pair<std::string, std::string>>> metadata_map;
529 for (auto& data : rpc_metadata) {
530 std::vector<std::string> metadata =
531 absl::StrSplit(data, ':', absl::SkipEmpty());
532 GPR_ASSERT(metadata.size() == 3);
533 if (metadata[0] == "EmptyCall") {
534 metadata_map[ClientConfigureRequest::EMPTY_CALL].push_back(
535 {metadata[1], metadata[2]});
536 } else if (metadata[0] == "UnaryCall") {
537 metadata_map[ClientConfigureRequest::UNARY_CALL].push_back(
538 {metadata[1], metadata[2]});
543 std::vector<RpcConfig> configs;
544 std::vector<std::string> rpc_methods =
545 absl::StrSplit(absl::GetFlag(FLAGS_rpc), ',', absl::SkipEmpty());
546 for (const std::string& rpc_method : rpc_methods) {
548 if (rpc_method == "EmptyCall") {
549 config.type = ClientConfigureRequest::EMPTY_CALL;
550 } else if (rpc_method == "UnaryCall") {
551 config.type = ClientConfigureRequest::UNARY_CALL;
555 auto metadata_iter = metadata_map.find(config.type);
556 if (metadata_iter != metadata_map.end()) {
557 config.metadata = metadata_iter->second;
559 configs.push_back(std::move(config));
562 std::lock_guard<std::mutex> lock(rpc_configs_queue->mu_rpc_configs_queue);
563 rpc_configs_queue->rpc_configs_queue.emplace_back(std::move(configs));
567 int main(int argc, char** argv) {
568 grpc::testing::TestEnvironment env(argc, argv);
569 grpc::testing::InitTest(&argc, &argv, true);
570 // Validate the expect_status flag.
571 grpc_status_code code;
572 GPR_ASSERT(grpc_status_code_from_string(
573 absl::GetFlag(FLAGS_expect_status).c_str(), &code));
574 StatsWatchers stats_watchers;
575 RpcConfigurationsQueue rpc_config_queue;
578 std::lock_guard<std::mutex> lock(stats_watchers.mu);
579 stats_watchers.global_watcher = new XdsStatsWatcher(0, 0);
580 stats_watchers.watchers.insert(stats_watchers.global_watcher);
583 BuildRpcConfigsFromFlags(&rpc_config_queue);
585 std::chrono::duration<double> duration_per_query =
586 std::chrono::nanoseconds(std::chrono::seconds(1)) /
587 absl::GetFlag(FLAGS_qps);
589 std::vector<std::thread> test_threads;
590 test_threads.reserve(absl::GetFlag(FLAGS_num_channels));
591 for (int i = 0; i < absl::GetFlag(FLAGS_num_channels); i++) {
592 test_threads.emplace_back(std::thread(&RunTestLoop, duration_per_query,
593 &stats_watchers, &rpc_config_queue));
596 RunServer(absl::GetFlag(FLAGS_stats_port), &stats_watchers,
599 for (auto it = test_threads.begin(); it != test_threads.end(); it++) {