pdns/dnstcpbench.cc

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