[thirdparty/pdns.git] / pdns / dnstcpbench.cc

/*
    PowerDNS Versatile Database Driven Nameserver
    Copyright (C) 2002-2013  PowerDNS.COM BV

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License version 2
    as published by the Free Software Foundation
    
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
*/
#include <boost/accumulators/statistics/median.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/accumulators.hpp>

#include <boost/accumulators/statistics.hpp>

#include "dnsparser.hh"
#include "sstuff.hh"
#include "misc.hh"
#include "dnswriter.hh"
#include "dnsrecords.hh"
#include "statbag.hh"
#include <netinet/tcp.h>
#include <boost/array.hpp>
#include <boost/program_options.hpp>
#include <boost/foreach.hpp>

StatBag S;
namespace po = boost::program_options;

po::variables_map g_vm;
bool g_verbose;
bool g_onlyTCP;
bool g_tcpNoDelay;
unsigned int g_timeoutMsec;
AtomicCounter g_networkErrors, g_otherErrors, g_OK, g_truncates, g_authAnswers, g_timeOuts;
ComboAddress g_dest;

unsigned int makeUsec(const struct timeval& tv)
{
  return 1000000*tv.tv_sec + tv.tv_usec;
}

/* On Linux, run echo 1 > /proc/sys/net/ipv4/tcp_tw_recycle 
   to prevent running out of free TCP ports */

struct BenchQuery
{
  BenchQuery(const std::string& qname_, uint16_t qtype_) : qname(qname_), qtype(qtype_), udpUsec(0), tcpUsec(0), answerSecond(0) {}
  BenchQuery(){}
  std::string qname;
  uint16_t qtype;
  uint32_t udpUsec, tcpUsec;
  time_t answerSecond;
};

void doQuery(BenchQuery* q)
try
{
  vector<uint8_t> packet;
  DNSPacketWriter pw(packet, q->qname, q->qtype);
  int res;
  string reply;

  struct timeval tv, now;
  gettimeofday(&tv, 0);

  if(!g_onlyTCP) {
    Socket udpsock((AddressFamily)g_dest.sin4.sin_family, Datagram);
    
    udpsock.sendTo(string((char*)&*packet.begin(), (char*)&*packet.end()), g_dest);
    ComboAddress origin;
    res = waitForData(udpsock.getHandle(), 0, 1000 * g_timeoutMsec);
    if(res < 0)
      throw NetworkError("Error waiting for response");
    if(!res) {
      g_timeOuts++;
      return;
    }

    udpsock.recvFrom(reply, origin);

    gettimeofday(&now, 0);
    q->udpUsec = makeUsec(now - tv);
    tv=now;

    MOADNSParser mdp(reply);
    if(!mdp.d_header.tc)
      return;
    g_truncates++;
  }

  Socket sock((AddressFamily)g_dest.sin4.sin_family, Stream);
  int tmp=1;
  if(setsockopt(sock.getHandle(),SOL_SOCKET,SO_REUSEADDR,(char*)&tmp,sizeof tmp)<0) 
    throw runtime_error("Unable to set socket reuse: "+string(strerror(errno)));
    
  if(g_tcpNoDelay && setsockopt(sock.getHandle(), IPPROTO_TCP, TCP_NODELAY,(char*)&tmp,sizeof tmp)<0) 
    throw runtime_error("Unable to set socket no delay: "+string(strerror(errno)));

  sock.connect(g_dest);
  uint16_t len = htons(packet.size());
  string tcppacket((char*)& len, 2);
  tcppacket.append((char*)&*packet.begin(), (char*)&*packet.end());

  sock.writen(tcppacket);

  res = waitForData(sock.getHandle(), 0, 1000 * g_timeoutMsec);
  if(res < 0)
    throw NetworkError("Error waiting for response");
  if(!res) {
    g_timeOuts++;
    return;
  }
  
  if(sock.read((char *) &len, 2) != 2)
    throw AhuException("tcp read failed");
  
  len=ntohs(len);
  char *creply = new char[len];
  int n=0;
  int numread;
  while(n<len) {
    numread=sock.read(creply+n, len-n);
    if(numread<0)
      throw AhuException("tcp read failed");
    n+=numread;
  }
  
  reply=string(creply, len);
  delete[] creply;
  
  gettimeofday(&now, 0);
  q->tcpUsec = makeUsec(now - tv);
  q->answerSecond = now.tv_sec;

  MOADNSParser mdp(reply);
  //  cout<<"Had correct TCP/IP response, "<<mdp.d_answers.size()<<" answers, aabit="<<mdp.d_header.aa<<endl;
  if(mdp.d_header.aa)
    g_authAnswers++;
  g_OK++;
}
catch(NetworkError& ne)
{
  cerr<<"Network error: "<<ne.what()<<endl;
  g_networkErrors++;
}
catch(...)
{
  g_otherErrors++;
}

/* read queries from stdin, put in vector
   launch n worker threads, each picks a query using AtomicCounter
   If a worker reaches the end of its queue, it stops */

AtomicCounter g_pos;

vector<BenchQuery> g_queries;

static void* worker(void*)
{
  for(;;) {
    unsigned int pos = g_pos++; 
    if(pos >= g_queries.size())
      break;

    doQuery(&g_queries[pos]); // this is safe as long as nobody *inserts* to g_queries
  }
  return 0;
}

int main(int argc, char** argv)
try
{
  po::options_description desc("Allowed options"), hidden, alloptions;
  desc.add_options()
    ("help,h", "produce help message")
    ("verbose,v", "be verbose")
    ("udp-first,u", "try UDP first")
    ("file,f", po::value<string>(), "source file - if not specified, defaults to stdin")
    ("tcp-no-delay", po::value<bool>()->default_value(true), "use TCP_NODELAY socket option")
    ("timeout-msec", po::value<int>()->default_value(10), "wait for this amount of milliseconds for an answer")
    ("workers", po::value<int>()->default_value(100), "number of parallel workers");

  hidden.add_options()
    ("remote-host", po::value<string>(), "remote-host")
    ("remote-port", po::value<int>()->default_value(53), "remote-port");
  alloptions.add(desc).add(hidden); 

  po::positional_options_description p;
  p.add("remote-host", 1);
  p.add("remote-port", 1);

  po::store(po::command_line_parser(argc, argv).options(alloptions).positional(p).run(), g_vm);
  po::notify(g_vm);
  
  if(g_vm.count("help")) {
    cout << desc<<endl;
    exit(EXIT_SUCCESS);
  }
  g_tcpNoDelay = g_vm["tcp-no-delay"].as<bool>();

  g_onlyTCP = !g_vm.count("udp-first");
  g_verbose = g_vm.count("verbose");
  g_timeoutMsec = g_vm["timeout-msec"].as<int>();

  reportAllTypes();

  if(g_vm["remote-host"].empty()) {
    cerr<<"Syntax: tcpbench remote [port] < queries"<<endl;
    cerr<<"Where queries is one query per line, format: qname qtype, just 1 space"<<endl;
    cerr<<desc<<endl;
    exit(EXIT_FAILURE);
  }

  g_dest = ComboAddress(g_vm["remote-host"].as<string>().c_str(), g_vm["remote-port"].as<int>());

  unsigned int numworkers=g_vm["workers"].as<int>();
  
  if(g_verbose) {
    cout<<"Sending queries to: "<<g_dest.toStringWithPort()<<endl;
    cout<<"Attempting UDP first: " << (g_onlyTCP ? "no" : "yes") <<endl;
    cout<<"Timeout: "<< g_timeoutMsec<<"msec"<<endl;
    cout << "Using TCP_NODELAY: "<<g_tcpNoDelay<<endl;
  }


  pthread_t workers[numworkers];

  FILE* fp;
  if(!g_vm.count("file"))
    fp=fdopen(0, "r");
  else {
    fp=fopen(g_vm["file"].as<string>().c_str(), "r");
    if(!fp)
      unixDie("Unable to open "+g_vm["file"].as<string>()+" for input");
  }
  pair<string, string> q;
  string line;
  while(stringfgets(fp, line)) {
    trim_right(line);
    q=splitField(line, ' ');
    g_queries.push_back(BenchQuery(q.first, DNSRecordContent::TypeToNumber(q.second)));
  }
  fclose(fp);
    
  for(unsigned int n = 0; n < numworkers; ++n) {
    pthread_create(&workers[n], 0, worker, 0);
  }
  for(unsigned int n = 0; n < numworkers; ++n) {
    void* status;
    pthread_join(workers[n], &status);
  }
  
  using namespace boost::accumulators;
  typedef accumulator_set<
    unsigned int
    , stats<boost::accumulators::tag::median(with_p_square_quantile),
      boost::accumulators::tag::mean(immediate)
    >
  > acc_t;

  acc_t udpspeeds, tcpspeeds, qps;
  
  typedef map<time_t, uint32_t> counts_t;
  counts_t counts;

  BOOST_FOREACH(const BenchQuery& bq, g_queries) {
    counts[bq.answerSecond]++;
    udpspeeds(bq.udpUsec);
    tcpspeeds(bq.tcpUsec);
  }

  BOOST_FOREACH(const counts_t::value_type& val, counts) {
    qps(val.second);
  }

  cout<<"Average qps: "<<mean(qps)<<", median qps: "<<median(qps)<<endl;
  cout<<"Average UDP latency: "<<mean(udpspeeds)<<"usec, median: "<<median(udpspeeds)<<"usec"<<endl;
  cout<<"Average TCP latency: "<<mean(tcpspeeds)<<"usec, median: "<<median(tcpspeeds)<<"usec"<<endl;

  cout<<"OK: "<<g_OK<<", network errors: "<<g_networkErrors<<", other errors: "<<g_otherErrors<<endl;
  cout<<"Timeouts: "<<g_timeOuts<<endl;
  cout<<"Truncateds: "<<g_truncates<<", auth answers: "<<g_authAnswers<<endl;
}
catch(std::exception &e)
{
  cerr<<"Fatal: "<<e.what()<<endl;
}
Commit	Line	Data
22f4bd59	1	/*
	2	PowerDNS Versatile Database Driven Nameserver
	3	Copyright (C) 2002-2013 PowerDNS.COM BV
	4
	5	This program is free software; you can redistribute it and/or modify
	6	it under the terms of the GNU General Public License version 2
	7	as published by the Free Software Foundation
	8
	9	This program is distributed in the hope that it will be useful,
	10	but WITHOUT ANY WARRANTY; without even the implied warranty of
	11	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
	12	GNU General Public License for more details.
	13
	14	You should have received a copy of the GNU General Public License
	15	along with this program; if not, write to the Free Software
	16	Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
	17	*/
4106b16e	18	#include <boost/accumulators/statistics/median.hpp>
	19	#include <boost/accumulators/statistics/mean.hpp>
	20	#include <boost/accumulators/accumulators.hpp>
	21
	22	#include <boost/accumulators/statistics.hpp>
	23
a51b52c9	24	#include "dnsparser.hh"
	25	#include "sstuff.hh"
	26	#include "misc.hh"
	27	#include "dnswriter.hh"
	28	#include "dnsrecords.hh"
	29	#include "statbag.hh"
a8f90a6a	30	#include <netinet/tcp.h>
a51b52c9	31	#include <boost/array.hpp>
36658df5	32	#include <boost/program_options.hpp>
4106b16e	33	#include <boost/foreach.hpp>
a51b52c9	34
36658df5	35	StatBag S;
36658df5	36	namespace po = boost::program_options;
4106b16e	37
36658df5	38	po::variables_map g_vm;
36658df5	39	bool g_verbose;
a51b52c9	40	bool g_onlyTCP;
a8f90a6a	41	bool g_tcpNoDelay;
36658df5	42	unsigned int g_timeoutMsec;
36658df5	43	AtomicCounter g_networkErrors, g_otherErrors, g_OK, g_truncates, g_authAnswers, g_timeOuts;
1d74012c	44	ComboAddress g_dest;
f7896b52	45
4106b16e	46	unsigned int makeUsec(const struct timeval& tv)
	47	{
	48	return 1000000*tv.tv_sec + tv.tv_usec;
	49	}
	50
22f4bd59	51	/* On Linux, run echo 1 > /proc/sys/net/ipv4/tcp_tw_recycle
	52	to prevent running out of free TCP ports */
	53
	54	struct BenchQuery
	55	{
4106b16e	56	BenchQuery(const std::string& qname_, uint16_t qtype_) : qname(qname_), qtype(qtype_), udpUsec(0), tcpUsec(0), answerSecond(0) {}
22f4bd59	57	BenchQuery(){}
	58	std::string qname;
	59	uint16_t qtype;
4106b16e	60	uint32_t udpUsec, tcpUsec;
4106b16e	61	time_t answerSecond;
22f4bd59	62	};
f7896b52	63
22f4bd59	64	void doQuery(BenchQuery* q)
f7896b52	65	try
a51b52c9	66	{
a51b52c9	67	vector<uint8_t> packet;
1d74012c	68	DNSPacketWriter pw(packet, q->qname, q->qtype);
36658df5	69	int res;
a51b52c9	70	string reply;
a51b52c9	71
4106b16e	72	struct timeval tv, now;
	73	gettimeofday(&tv, 0);
	74
a51b52c9	75	if(!g_onlyTCP) {
1d74012c	76	Socket udpsock((AddressFamily)g_dest.sin4.sin_family, Datagram);
a51b52c9	77
1d74012c	78	udpsock.sendTo(string((char)&packet.begin(), (char)&packet.end()), g_dest);
a51b52c9	79	ComboAddress origin;
36658df5	80	res = waitForData(udpsock.getHandle(), 0, 1000 * g_timeoutMsec);
	81	if(res < 0)
	82	throw NetworkError("Error waiting for response");
	83	if(!res) {
	84	g_timeOuts++;
	85	return;
	86	}
	87
a51b52c9	88	udpsock.recvFrom(reply, origin);
4106b16e	89
	90	gettimeofday(&now, 0);
	91	q->udpUsec = makeUsec(now - tv);
	92	tv=now;
	93
a51b52c9	94	MOADNSParser mdp(reply);
	95	if(!mdp.d_header.tc)
	96	return;
f7896b52	97	g_truncates++;
a51b52c9	98	}
a51b52c9	99
1d74012c	100	Socket sock((AddressFamily)g_dest.sin4.sin_family, Stream);
f7896b52	101	int tmp=1;
	102	if(setsockopt(sock.getHandle(),SOL_SOCKET,SO_REUSEADDR,(char*)&tmp,sizeof tmp)<0)
	103	throw runtime_error("Unable to set socket reuse: "+string(strerror(errno)));
a8f90a6a	104
	105	if(g_tcpNoDelay && setsockopt(sock.getHandle(), IPPROTO_TCP, TCP_NODELAY,(char*)&tmp,sizeof tmp)<0)
	106	throw runtime_error("Unable to set socket no delay: "+string(strerror(errno)));
a51b52c9	107
1d74012c	108	sock.connect(g_dest);
36658df5	109	uint16_t len = htons(packet.size());
	110	string tcppacket((char*)& len, 2);
	111	tcppacket.append((char)&packet.begin(), (char)&packet.end());
	112
	113	sock.writen(tcppacket);
	114
	115	res = waitForData(sock.getHandle(), 0, 1000 * g_timeoutMsec);
	116	if(res < 0)
	117	throw NetworkError("Error waiting for response");
	118	if(!res) {
	119	g_timeOuts++;
	120	return;
	121	}
a51b52c9	122
	123	if(sock.read((char *) &len, 2) != 2)
	124	throw AhuException("tcp read failed");
	125
	126	len=ntohs(len);
	127	char *creply = new char[len];
	128	int n=0;
	129	int numread;
	130	while(n<len) {
	131	numread=sock.read(creply+n, len-n);
	132	if(numread<0)
	133	throw AhuException("tcp read failed");
	134	n+=numread;
	135	}
	136
	137	reply=string(creply, len);
	138	delete[] creply;
	139
4106b16e	140	gettimeofday(&now, 0);
	141	q->tcpUsec = makeUsec(now - tv);
	142	q->answerSecond = now.tv_sec;
	143
a51b52c9	144	MOADNSParser mdp(reply);
f7896b52	145	// cout<<"Had correct TCP/IP response, "<<mdp.d_answers.size()<<" answers, aabit="<<mdp.d_header.aa<<endl;
36658df5	146	if(mdp.d_header.aa)
36658df5	147	g_authAnswers++;
f7896b52	148	g_OK++;
	149	}
	150	catch(NetworkError& ne)
	151	{
	152	cerr<<"Network error: "<<ne.what()<<endl;
	153	g_networkErrors++;
	154	}
	155	catch(...)
	156	{
	157	g_otherErrors++;
a51b52c9	158	}
	159
	160	/* read queries from stdin, put in vector
	161	launch n worker threads, each picks a query using AtomicCounter
	162	If a worker reaches the end of its queue, it stops */
	163
	164	AtomicCounter g_pos;
682549f2	165
22f4bd59	166	vector<BenchQuery> g_queries;
1d74012c	167
682549f2	168	static void* worker(void*)
a51b52c9	169	{
a51b52c9	170	for(;;) {
f7896b52	171	unsigned int pos = g_pos++;
f7896b52	172	if(pos >= g_queries.size())
a51b52c9	173	break;
1d74012c	174
1d74012c	175	doQuery(&g_queries[pos]); // this is safe as long as nobody inserts to g_queries
a51b52c9	176	}
	177	return 0;
	178	}
	179
a51b52c9	180	int main(int argc, char** argv)
	181	try
	182	{
36658df5	183	po::options_description desc("Allowed options"), hidden, alloptions;
	184	desc.add_options()
	185	("help,h", "produce help message")
	186	("verbose,v", "be verbose")
	187	("udp-first,u", "try UDP first")
74f463af	188	("file,f", po::value<string>(), "source file - if not specified, defaults to stdin")
a8f90a6a	189	("tcp-no-delay", po::value<bool>()->default_value(true), "use TCP_NODELAY socket option")
36658df5	190	("timeout-msec", po::value<int>()->default_value(10), "wait for this amount of milliseconds for an answer")
	191	("workers", po::value<int>()->default_value(100), "number of parallel workers");
	192
	193	hidden.add_options()
	194	("remote-host", po::value<string>(), "remote-host")
	195	("remote-port", po::value<int>()->default_value(53), "remote-port");
	196	alloptions.add(desc).add(hidden);
	197
	198	po::positional_options_description p;
	199	p.add("remote-host", 1);
	200	p.add("remote-port", 1);
	201
	202	po::store(po::command_line_parser(argc, argv).options(alloptions).positional(p).run(), g_vm);
	203	po::notify(g_vm);
	204
	205	if(g_vm.count("help")) {
	206	cout << desc<<endl;
	207	exit(EXIT_SUCCESS);
	208	}
a8f90a6a	209	g_tcpNoDelay = g_vm["tcp-no-delay"].as<bool>();
a8f90a6a	210
36658df5	211	g_onlyTCP = !g_vm.count("udp-first");
	212	g_verbose = g_vm.count("verbose");
	213	g_timeoutMsec = g_vm["timeout-msec"].as<int>();
	214
a51b52c9	215	reportAllTypes();
f7896b52	216
36658df5	217	if(g_vm["remote-host"].empty()) {
f7896b52	218	cerr<<"Syntax: tcpbench remote [port] < queries"<<endl;
f7896b52	219	cerr<<"Where queries is one query per line, format: qname qtype, just 1 space"<<endl;
36658df5	220	cerr<<desc<<endl;
f7896b52	221	exit(EXIT_FAILURE);
f7896b52	222	}
f7896b52	223
36658df5	224	g_dest = ComboAddress(g_vm["remote-host"].as<string>().c_str(), g_vm["remote-port"].as<int>());
	225
	226	unsigned int numworkers=g_vm["workers"].as<int>();
	227
	228	if(g_verbose) {
	229	cout<<"Sending queries to: "<<g_dest.toStringWithPort()<<endl;
	230	cout<<"Attempting UDP first: " << (g_onlyTCP ? "no" : "yes") <<endl;
	231	cout<<"Timeout: "<< g_timeoutMsec<<"msec"<<endl;
a8f90a6a	232	cout << "Using TCP_NODELAY: "<<g_tcpNoDelay<<endl;
36658df5	233	}
	234
	235
a51b52c9	236	pthread_t workers[numworkers];
a51b52c9	237
74f463af	238	FILE* fp;
	239	if(!g_vm.count("file"))
	240	fp=fdopen(0, "r");
	241	else {
	242	fp=fopen(g_vm["file"].as<string>().c_str(), "r");
	243	if(!fp)
	244	unixDie("Unable to open "+g_vm["file"].as<string>()+" for input");
	245	}
f7896b52	246	pair<string, string> q;
	247	string line;
	248	while(stringfgets(fp, line)) {
	249	trim_right(line);
	250	q=splitField(line, ' ');
22f4bd59	251	g_queries.push_back(BenchQuery(q.first, DNSRecordContent::TypeToNumber(q.second)));
a51b52c9	252	}
f7896b52	253	fclose(fp);
f7896b52	254
a51b52c9	255	for(unsigned int n = 0; n < numworkers; ++n) {
	256	pthread_create(&workers[n], 0, worker, 0);
	257	}
	258	for(unsigned int n = 0; n < numworkers; ++n) {
	259	void* status;
	260	pthread_join(workers[n], &status);
	261	}
4106b16e	262
	263	using namespace boost::accumulators;
	264	typedef accumulator_set<
	265	unsigned int
	266	, stats<boost::accumulators::tag::median(with_p_square_quantile),
	267	boost::accumulators::tag::mean(immediate)
	268	>
	269	> acc_t;
	270
	271	acc_t udpspeeds, tcpspeeds, qps;
	272
	273	typedef map<time_t, uint32_t> counts_t;
	274	counts_t counts;
	275
	276	BOOST_FOREACH(const BenchQuery& bq, g_queries) {
	277	counts[bq.answerSecond]++;
	278	udpspeeds(bq.udpUsec);
	279	tcpspeeds(bq.tcpUsec);
	280	}
	281
	282	BOOST_FOREACH(const counts_t::value_type& val, counts) {
	283	qps(val.second);
	284	}
	285
	286	cout<<"Average qps: "<<mean(qps)<<", median qps: "<<median(qps)<<endl;
	287	cout<<"Average UDP latency: "<<mean(udpspeeds)<<"usec, median: "<<median(udpspeeds)<<"usec"<<endl;
	288	cout<<"Average TCP latency: "<<mean(tcpspeeds)<<"usec, median: "<<median(tcpspeeds)<<"usec"<<endl;
	289
f7896b52	290	cout<<"OK: "<<g_OK<<", network errors: "<<g_networkErrors<<", other errors: "<<g_otherErrors<<endl;
36658df5	291	cout<<"Timeouts: "<<g_timeOuts<<endl;
36658df5	292	cout<<"Truncateds: "<<g_truncates<<", auth answers: "<<g_authAnswers<<endl;
a51b52c9	293	}
	294	catch(std::exception &e)
	295	{
	296	cerr<<"Fatal: "<<e.what()<<endl;
	297	}