tools/rpc_replay/rpc_replay.cpp - brpc - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you under the Apache License, Version 2.0 (the
 // "License"); you may not use this file except in compliance
 // with the License.  You may obtain a copy of the License at
 //
 //   http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing,
 // software distributed under the License is distributed on an
 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.


 #include <gflags/gflags.h>
 #include <butil/logging.h>
 #include <butil/time.h>
 #include <butil/macros.h>
 #include <butil/file_util.h>
 #include <bvar/bvar.h>
 #include <bthread/bthread.h>
 #include <brpc/channel.h>
 #include <brpc/server.h>
 #include <brpc/rpc_dump.h>
 #include <brpc/serialized_request.h>
 #include <brpc/nshead_message.h>
 #include <brpc/details/http_message.h>
 #include "brpc/options.pb.h"
 #include "info_thread.h"

 DEFINE_string(dir, "", "The directory of dumped requests");
 DEFINE_int32(times, 1, "Repeat replaying for so many times");
 DEFINE_int32(qps, 0, "Limit QPS if this flag is positive");
 DEFINE_int32(thread_num, 0, "Number of threads for replaying");
 DEFINE_bool(use_bthread, true, "Use bthread to replay");
 DEFINE_string(connection_type, "", "Connection type, choose automatically "
               "according to protocol by default");
 DEFINE_string(server, "0.0.0.0:8002", "IP Address of server");
 DEFINE_string(load_balancer, "", "The algorithm for load balancing");
 DEFINE_int32(timeout_ms, 100, "RPC timeout in milliseconds");
 DEFINE_int32(max_retry, 3, "Maximum retry times");
 DEFINE_int32(dummy_port, 8899, "Port of dummy server(to monitor replaying)");
 DEFINE_string(http_host, "", "Host field for http protocol");

 bvar::LatencyRecorder g_latency_recorder("rpc_replay");
 bvar::Adder<int64_t> g_error_count("rpc_replay_error_count");
 bvar::Adder<int64_t> g_sent_count;

 // Include channels for all protocols that support both client and server.
 class ChannelGroup {
 public:
     int Init();

     ~ChannelGroup();

     // Get channel by protocol type.
     brpc::Channel* channel(brpc::ProtocolType type) {
         if ((size_t)type < _chans.size()) {
             return _chans[(size_t)type];
         }
         return NULL;
     }

 private:
     std::vector<brpc::Channel*> _chans;
 };

 int ChannelGroup::Init() {
     {
         // force global initialization of rpc.
         brpc::Channel dummy_channel;
     }
     std::vector<std::pair<brpc::ProtocolType, brpc::Protocol> > protocols;
     brpc::ListProtocols(&protocols);
     size_t max_protocol_size = 0;
     for (size_t i = 0; i < protocols.size(); ++i) {
         max_protocol_size = std::max(max_protocol_size,
                                      (size_t)protocols[i].first);
     }
     _chans.resize(max_protocol_size + 1);
     for (size_t i = 0; i < protocols.size(); ++i) {
         const brpc::ProtocolType protocol_type = protocols[i].first;
         const brpc::Protocol protocol = protocols[i].second;
         brpc::ChannelOptions options;
         options.protocol = protocol_type;
         options.connection_type = FLAGS_connection_type;
         options.timeout_ms = FLAGS_timeout_ms/*milliseconds*/;
         options.max_retry = FLAGS_max_retry;
         if ((options.connection_type == brpc::CONNECTION_TYPE_UNKNOWN ||
             options.connection_type & protocol.supported_connection_type) &&
             protocol.support_client() &&
             protocol.support_server()) {
             brpc::Channel* chan = new brpc::Channel;
             if (chan->Init(FLAGS_server.c_str(), FLAGS_load_balancer.c_str(),
                         &options) != 0) {
                 LOG(ERROR) << "Fail to initialize channel";
                 delete chan;
                 return -1;
             }
             _chans[protocol_type] = chan;
         }
     }
     return 0;
 }

 ChannelGroup::~ChannelGroup() {
     for (size_t i = 0; i < _chans.size(); ++i) {
         delete _chans[i];
     }
     _chans.clear();
 }

 static void handle_response(brpc::Controller* cntl, int64_t start_time,
                             bool sleep_on_error/*note*/) {
     // TODO(gejun): some bthreads are starved when new bthreads are created
     // continuously, which happens when server is down and RPC keeps failing.
     // Sleep a while on error to avoid that now.
     const int64_t end_time = butil::gettimeofday_us();
     const int64_t elp = end_time - start_time;
     if (!cntl->Failed()) {
         g_latency_recorder << elp;
     } else {
         g_error_count << 1;
         if (sleep_on_error) {
             bthread_usleep(10000);
         }
     }
     delete cntl;
 }

 butil::atomic<int> g_thread_offset(0);

 static void* replay_thread(void* arg) {
     ChannelGroup* chan_group = static_cast<ChannelGroup*>(arg);
     const int thread_offset = g_thread_offset.fetch_add(1, butil::memory_order_relaxed);
     double req_rate = FLAGS_qps / (double)FLAGS_thread_num;
     brpc::SerializedRequest req;
     brpc::NsheadMessage nshead_req;
     int64_t last_expected_time = butil::monotonic_time_ns();
     const int64_t interval = (int64_t) (1000000000L / req_rate);
     // the max tolerant delay between end_time and expected_time. 10ms or 10 intervals
     int64_t max_tolerant_delay = std::max((int64_t) 10000000L, 10 * interval);
     for (int i = 0; !brpc::IsAskedToQuit() && i < FLAGS_times; ++i) {
         brpc::SampleIterator it(FLAGS_dir);
         int j = 0;
         for (brpc::SampledRequest* sample = it.Next();
              !brpc::IsAskedToQuit() && sample != NULL; sample = it.Next(), ++j) {
             std::unique_ptr<brpc::SampledRequest> sample_guard(sample);
             if ((j % FLAGS_thread_num) != thread_offset) {
                 continue;
             }
             brpc::Channel* chan =
                 chan_group->channel(sample->meta.protocol_type());
             if (chan == NULL) {
                 LOG(ERROR) << "No channel on protocol="
                            << sample->meta.protocol_type();
                 continue;
             }

             brpc::Controller* cntl = new brpc::Controller;
             req.Clear();

             google::protobuf::Message* req_ptr = &req;
             cntl->reset_sampled_request(sample_guard.release());
             if (sample->meta.protocol_type() == brpc::PROTOCOL_HTTP) {
                 brpc::HttpMessage http_message;
                 http_message.ParseFromIOBuf(sample->request);
                 cntl->http_request().Swap(http_message.header());
                 if (!FLAGS_http_host.empty()) {
                     // reset Host in header
                     cntl->http_request().SetHeader("Host", FLAGS_http_host);
                 }
                 cntl->request_attachment() = http_message.body().movable();
                 req_ptr = NULL;
             } else if (sample->meta.protocol_type() == brpc::PROTOCOL_NSHEAD) {
                 nshead_req.Clear();
                 memcpy(&nshead_req.head, sample->meta.nshead().c_str(), sample->meta.nshead().length());
                 nshead_req.body = sample->request;
                 req_ptr = &nshead_req;
             } else if (sample->meta.attachment_size() > 0) {
                 sample->request.cutn(
                     &req.serialized_data(),
                     sample->request.size() - sample->meta.attachment_size());
                 cntl->request_attachment() = sample->request.movable();
             } else {
                 req.serialized_data() = sample->request.movable();
             }
             g_sent_count << 1;
             const int64_t start_time = butil::gettimeofday_us();
             if (FLAGS_qps <= 0) {
                 chan->CallMethod(NULL/*use rpc_dump_context in cntl instead*/,
                         cntl, req_ptr, NULL/*ignore response*/, NULL);
                 handle_response(cntl, start_time, true);
             } else {
                 google::protobuf::Closure* done =
                     brpc::NewCallback(handle_response, cntl, start_time, false);
                 chan->CallMethod(NULL/*use rpc_dump_context in cntl instead*/,
                         cntl, req_ptr, NULL/*ignore response*/, done);
                 int64_t end_time = butil::monotonic_time_ns();
                 int64_t expected_time = last_expected_time + interval;
                 if (end_time < expected_time) {
                     usleep((expected_time - end_time)/1000);
                 }
                 if (end_time - expected_time > max_tolerant_delay) {
                     expected_time = end_time;
                 }
                 last_expected_time = expected_time;
             }
         }
     }
     return NULL;
 }

 int main(int argc, char* argv[]) {
     // Parse gflags. We recommend you to use gflags as well.
     GFLAGS_NAMESPACE::ParseCommandLineFlags(&argc, &argv, true);

     if (FLAGS_dir.empty() ||
         !butil::DirectoryExists(butil::FilePath(FLAGS_dir))) {
         LOG(ERROR) << "--dir=<dir-of-dumped-files> is required";
         return -1;
     }

     if (FLAGS_dummy_port >= 0) {
         brpc::StartDummyServerAt(FLAGS_dummy_port);
     }

     ChannelGroup chan_group;
     if (chan_group.Init() != 0) {
         LOG(ERROR) << "Fail to init ChannelGroup";
         return -1;
     }

     if (FLAGS_thread_num <= 0) {
         if (FLAGS_qps <= 0) { // unlimited qps
             FLAGS_thread_num = 50;
         } else {
             FLAGS_thread_num = FLAGS_qps / 10000;
             if (FLAGS_thread_num < 1) {
                 FLAGS_thread_num = 1;
             }
             if (FLAGS_thread_num > 50) {
                 FLAGS_thread_num = 50;
             }
         }
     }

     const int rate_limit_per_thread = 1000000;
     int req_rate_per_thread = FLAGS_qps / FLAGS_thread_num;
     if (req_rate_per_thread > rate_limit_per_thread) {
         LOG(ERROR) << "req_rate: " << (int64_t) req_rate_per_thread << " is too large in one thread. The rate limit is "
                 <<  rate_limit_per_thread << " in one thread";
         return -1;
     }

     std::vector<bthread_t> bids;
     std::vector<pthread_t> pids;
     if (!FLAGS_use_bthread) {
         pids.resize(FLAGS_thread_num);
         for (int i = 0; i < FLAGS_thread_num; ++i) {
             if (pthread_create(&pids[i], NULL, replay_thread, &chan_group) != 0) {
                 LOG(ERROR) << "Fail to create pthread";
                 return -1;
             }
         }
     } else {
         bids.resize(FLAGS_thread_num);
         for (int i = 0; i < FLAGS_thread_num; ++i) {
             if (bthread_start_background(
                     &bids[i], NULL, replay_thread, &chan_group) != 0) {
                 LOG(ERROR) << "Fail to create bthread";
                 return -1;
             }
         }
     }
     brpc::InfoThread info_thr;
     brpc::InfoThreadOptions info_thr_opt;
     info_thr_opt.latency_recorder = &g_latency_recorder;
     info_thr_opt.error_count = &g_error_count;
     info_thr_opt.sent_count = &g_sent_count;

     if (!info_thr.start(info_thr_opt)) {
         LOG(ERROR) << "Fail to create info_thread";
         return -1;
     }

     for (int i = 0; i < FLAGS_thread_num; ++i) {
         if (!FLAGS_use_bthread) {
             pthread_join(pids[i], NULL);
         } else {
             bthread_join(bids[i], NULL);
         }
     }
     info_thr.stop();

     return 0;
 }
	// Licensed to the Apache Software Foundation (ASF) under one
	// or more contributor license agreements. See the NOTICE file
	// distributed with this work for additional information
	// regarding copyright ownership. The ASF licenses this file
	// to you under the Apache License, Version 2.0 (the
	// "License"); you may not use this file except in compliance
	// with the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing,
	// software distributed under the License is distributed on an
	// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	// KIND, either express or implied. See the License for the
	// specific language governing permissions and limitations
	// under the License.


	#include <gflags/gflags.h>
	#include <butil/logging.h>
	#include <butil/time.h>
	#include <butil/macros.h>
	#include <butil/file_util.h>
	#include <bvar/bvar.h>
	#include <bthread/bthread.h>
	#include <brpc/channel.h>
	#include <brpc/server.h>
	#include <brpc/rpc_dump.h>
	#include <brpc/serialized_request.h>
	#include <brpc/nshead_message.h>
	#include <brpc/details/http_message.h>
	#include "brpc/options.pb.h"
	#include "info_thread.h"

	DEFINE_string(dir, "", "The directory of dumped requests");
	DEFINE_int32(times, 1, "Repeat replaying for so many times");
	DEFINE_int32(qps, 0, "Limit QPS if this flag is positive");
	DEFINE_int32(thread_num, 0, "Number of threads for replaying");
	DEFINE_bool(use_bthread, true, "Use bthread to replay");
	DEFINE_string(connection_type, "", "Connection type, choose automatically "
	"according to protocol by default");
	DEFINE_string(server, "0.0.0.0:8002", "IP Address of server");
	DEFINE_string(load_balancer, "", "The algorithm for load balancing");
	DEFINE_int32(timeout_ms, 100, "RPC timeout in milliseconds");
	DEFINE_int32(max_retry, 3, "Maximum retry times");
	DEFINE_int32(dummy_port, 8899, "Port of dummy server(to monitor replaying)");
	DEFINE_string(http_host, "", "Host field for http protocol");

	bvar::LatencyRecorder g_latency_recorder("rpc_replay");
	bvar::Adder<int64_t> g_error_count("rpc_replay_error_count");
	bvar::Adder<int64_t> g_sent_count;

	// Include channels for all protocols that support both client and server.
	class ChannelGroup {
	public:
	int Init();

	~ChannelGroup();

	// Get channel by protocol type.
	brpc::Channel* channel(brpc::ProtocolType type) {
	if ((size_t)type < _chans.size()) {
	return _chans[(size_t)type];
	}
	return NULL;
	}

	private:
	std::vector<brpc::Channel*> _chans;
	};

	int ChannelGroup::Init() {
	{
	// force global initialization of rpc.
	brpc::Channel dummy_channel;
	}
	std::vector<std::pair<brpc::ProtocolType, brpc::Protocol> > protocols;
	brpc::ListProtocols(&protocols);
	size_t max_protocol_size = 0;
	for (size_t i = 0; i < protocols.size(); ++i) {
	max_protocol_size = std::max(max_protocol_size,
	(size_t)protocols[i].first);
	}
	_chans.resize(max_protocol_size + 1);
	for (size_t i = 0; i < protocols.size(); ++i) {
	const brpc::ProtocolType protocol_type = protocols[i].first;
	const brpc::Protocol protocol = protocols[i].second;
	brpc::ChannelOptions options;
	options.protocol = protocol_type;
	options.connection_type = FLAGS_connection_type;
	options.timeout_ms = FLAGS_timeout_ms/milliseconds/;
	options.max_retry = FLAGS_max_retry;
	if ((options.connection_type == brpc::CONNECTION_TYPE_UNKNOWN \|\|
	options.connection_type & protocol.supported_connection_type) &&
	protocol.support_client() &&
	protocol.support_server()) {
	brpc::Channel* chan = new brpc::Channel;
	if (chan->Init(FLAGS_server.c_str(), FLAGS_load_balancer.c_str(),
	&options) != 0) {
	LOG(ERROR) << "Fail to initialize channel";
	delete chan;
	return -1;
	}
	_chans[protocol_type] = chan;
	}
	}
	return 0;
	}

	ChannelGroup::~ChannelGroup() {
	for (size_t i = 0; i < _chans.size(); ++i) {
	delete _chans[i];
	}
	_chans.clear();
	}

	static void handle_response(brpc::Controller* cntl, int64_t start_time,
	bool sleep_on_error/note/) {
	// TODO(gejun): some bthreads are starved when new bthreads are created
	// continuously, which happens when server is down and RPC keeps failing.
	// Sleep a while on error to avoid that now.
	const int64_t end_time = butil::gettimeofday_us();
	const int64_t elp = end_time - start_time;
	if (!cntl->Failed()) {
	g_latency_recorder << elp;
	} else {
	g_error_count << 1;
	if (sleep_on_error) {
	bthread_usleep(10000);
	}
	}
	delete cntl;
	}

	butil::atomic<int> g_thread_offset(0);

	static void* replay_thread(void* arg) {
	ChannelGroup* chan_group = static_cast<ChannelGroup*>(arg);
	const int thread_offset = g_thread_offset.fetch_add(1, butil::memory_order_relaxed);
	double req_rate = FLAGS_qps / (double)FLAGS_thread_num;
	brpc::SerializedRequest req;
	brpc::NsheadMessage nshead_req;
	int64_t last_expected_time = butil::monotonic_time_ns();
	const int64_t interval = (int64_t) (1000000000L / req_rate);
	// the max tolerant delay between end_time and expected_time. 10ms or 10 intervals
	int64_t max_tolerant_delay = std::max((int64_t) 10000000L, 10 * interval);
	for (int i = 0; !brpc::IsAskedToQuit() && i < FLAGS_times; ++i) {
	brpc::SampleIterator it(FLAGS_dir);
	int j = 0;
	for (brpc::SampledRequest* sample = it.Next();
	!brpc::IsAskedToQuit() && sample != NULL; sample = it.Next(), ++j) {
	std::unique_ptr<brpc::SampledRequest> sample_guard(sample);
	if ((j % FLAGS_thread_num) != thread_offset) {
	continue;
	}
	brpc::Channel* chan =
	chan_group->channel(sample->meta.protocol_type());
	if (chan == NULL) {
	LOG(ERROR) << "No channel on protocol="
	<< sample->meta.protocol_type();
	continue;
	}

	brpc::Controller* cntl = new brpc::Controller;
	req.Clear();

	google::protobuf::Message* req_ptr = &req;
	cntl->reset_sampled_request(sample_guard.release());
	if (sample->meta.protocol_type() == brpc::PROTOCOL_HTTP) {
	brpc::HttpMessage http_message;
	http_message.ParseFromIOBuf(sample->request);
	cntl->http_request().Swap(http_message.header());
	if (!FLAGS_http_host.empty()) {
	// reset Host in header
	cntl->http_request().SetHeader("Host", FLAGS_http_host);
	}
	cntl->request_attachment() = http_message.body().movable();
	req_ptr = NULL;
	} else if (sample->meta.protocol_type() == brpc::PROTOCOL_NSHEAD) {
	nshead_req.Clear();
	memcpy(&nshead_req.head, sample->meta.nshead().c_str(), sample->meta.nshead().length());
	nshead_req.body = sample->request;
	req_ptr = &nshead_req;
	} else if (sample->meta.attachment_size() > 0) {
	sample->request.cutn(
	&req.serialized_data(),
	sample->request.size() - sample->meta.attachment_size());
	cntl->request_attachment() = sample->request.movable();
	} else {
	req.serialized_data() = sample->request.movable();
	}
	g_sent_count << 1;
	const int64_t start_time = butil::gettimeofday_us();
	if (FLAGS_qps <= 0) {
	chan->CallMethod(NULL/use rpc_dump_context in cntl instead/,
	cntl, req_ptr, NULL/ignore response/, NULL);
	handle_response(cntl, start_time, true);
	} else {
	google::protobuf::Closure* done =
	brpc::NewCallback(handle_response, cntl, start_time, false);
	chan->CallMethod(NULL/use rpc_dump_context in cntl instead/,
	cntl, req_ptr, NULL/ignore response/, done);
	int64_t end_time = butil::monotonic_time_ns();
	int64_t expected_time = last_expected_time + interval;
	if (end_time < expected_time) {
	usleep((expected_time - end_time)/1000);
	}
	if (end_time - expected_time > max_tolerant_delay) {
	expected_time = end_time;
	}
	last_expected_time = expected_time;
	}
	}
	}
	return NULL;
	}

	int main(int argc, char* argv[]) {
	// Parse gflags. We recommend you to use gflags as well.
	GFLAGS_NAMESPACE::ParseCommandLineFlags(&argc, &argv, true);

	if (FLAGS_dir.empty() \|\|
	!butil::DirectoryExists(butil::FilePath(FLAGS_dir))) {
	LOG(ERROR) << "--dir=<dir-of-dumped-files> is required";
	return -1;
	}

	if (FLAGS_dummy_port >= 0) {
	brpc::StartDummyServerAt(FLAGS_dummy_port);
	}

	ChannelGroup chan_group;
	if (chan_group.Init() != 0) {
	LOG(ERROR) << "Fail to init ChannelGroup";
	return -1;
	}

	if (FLAGS_thread_num <= 0) {
	if (FLAGS_qps <= 0) { // unlimited qps
	FLAGS_thread_num = 50;
	} else {
	FLAGS_thread_num = FLAGS_qps / 10000;
	if (FLAGS_thread_num < 1) {
	FLAGS_thread_num = 1;
	}
	if (FLAGS_thread_num > 50) {
	FLAGS_thread_num = 50;
	}
	}
	}

	const int rate_limit_per_thread = 1000000;
	int req_rate_per_thread = FLAGS_qps / FLAGS_thread_num;
	if (req_rate_per_thread > rate_limit_per_thread) {
	LOG(ERROR) << "req_rate: " << (int64_t) req_rate_per_thread << " is too large in one thread. The rate limit is "
	<< rate_limit_per_thread << " in one thread";
	return -1;
	}

	std::vector<bthread_t> bids;
	std::vector<pthread_t> pids;
	if (!FLAGS_use_bthread) {
	pids.resize(FLAGS_thread_num);
	for (int i = 0; i < FLAGS_thread_num; ++i) {
	if (pthread_create(&pids[i], NULL, replay_thread, &chan_group) != 0) {
	LOG(ERROR) << "Fail to create pthread";
	return -1;
	}
	}
	} else {
	bids.resize(FLAGS_thread_num);
	for (int i = 0; i < FLAGS_thread_num; ++i) {
	if (bthread_start_background(
	&bids[i], NULL, replay_thread, &chan_group) != 0) {
	LOG(ERROR) << "Fail to create bthread";
	return -1;
	}
	}
	}
	brpc::InfoThread info_thr;
	brpc::InfoThreadOptions info_thr_opt;
	info_thr_opt.latency_recorder = &g_latency_recorder;
	info_thr_opt.error_count = &g_error_count;
	info_thr_opt.sent_count = &g_sent_count;

	if (!info_thr.start(info_thr_opt)) {
	LOG(ERROR) << "Fail to create info_thread";
	return -1;
	}

	for (int i = 0; i < FLAGS_thread_num; ++i) {
	if (!FLAGS_use_bthread) {
	pthread_join(pids[i], NULL);
	} else {
	bthread_join(bids[i], NULL);
	}
	}
	info_thr.stop();

	return 0;
	}