From f7896b5258e1782df679f9a7f52c100b61c67bd8 Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@netherlabs.nl>
Date: Tue, 11 Jun 2013 13:39:50 +0200
Subject: [PATCH] make tcpbench production ready

---
 pdns/.gitignore  |  1 +
 pdns/misc.hh     |  6 ++++++
 pdns/tcpbench.cc | 56 ++++++++++++++++++++++++++++++++++++++++--------
 3 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/pdns/.gitignore b/pdns/.gitignore
index 414664751..974f72190 100644
--- a/pdns/.gitignore
+++ b/pdns/.gitignore
@@ -30,4 +30,5 @@
 /speedtest
 /toysdig
 /tsig-tests
+/tcpbench
 version_generated.h
diff --git a/pdns/misc.hh b/pdns/misc.hh
index 7b1489aaa..4da6bdadc 100644
--- a/pdns/misc.hh
+++ b/pdns/misc.hh
@@ -349,6 +349,12 @@ public:
       return atomic_exchange_and_add( &value_, +1 ) + 1;
     }
 
+    unsigned int operator++(int)
+    {
+      return atomic_exchange_and_add( &value_, +1 );
+    }
+
+
     unsigned int operator--()
     {
       return atomic_exchange_and_add( &value_, -1 ) - 1;
diff --git a/pdns/tcpbench.cc b/pdns/tcpbench.cc
index f30a6a80b..0f4c047da 100644
--- a/pdns/tcpbench.cc
+++ b/pdns/tcpbench.cc
@@ -8,8 +8,13 @@
 StatBag S;
 
 bool g_onlyTCP;
+AtomicCounter g_networkErrors, g_otherErrors, g_OK, g_truncates;
+
+// echo 1 > /proc/sys/net/ipv4/tcp_tw_recycle 
+
 
 void doQuery(const std::string& qname, uint16_t qtype, const ComboAddress& dest)
+try
 {
   vector<uint8_t> packet;
   DNSPacketWriter pw(packet, qname, qtype);
@@ -25,10 +30,14 @@ void doQuery(const std::string& qname, uint16_t qtype, const ComboAddress& dest)
     MOADNSParser mdp(reply);
     if(!mdp.d_header.tc)
       return;
+    g_truncates++;
   }
 
 
   Socket sock(InterNetwork, 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)));
 
   sock.connect(dest);
   uint16_t len;
@@ -56,8 +65,19 @@ void doQuery(const std::string& qname, uint16_t qtype, const ComboAddress& dest)
   delete[] creply;
   
   MOADNSParser mdp(reply);
- 
+  //  cout<<"Had correct TCP/IP response, "<<mdp.d_answers.size()<<" answers, aabit="<<mdp.d_header.aa<<endl;
+  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
@@ -74,12 +94,13 @@ struct Query
 
 vector<Query> g_queries;
 ComboAddress g_dest;
+
 void* worker(void*)
 {
   Query q;
   for(;;) {
-    unsigned int pos = ++g_pos; 
-    if(pos > g_queries.size())
+    unsigned int pos = g_pos++; 
+    if(pos >= g_queries.size())
       break;
     q=g_queries[pos];
     doQuery(q.qname, q.qtype, g_dest);
@@ -92,15 +113,31 @@ int main(int argc, char** argv)
 try
 {
   reportAllTypes();
-  g_onlyTCP=true;
-  g_dest = ComboAddress("127.0.0.1", 5300);
+  g_onlyTCP=false;
+
+  uint16_t port=53;
+  if(argc < 2) {
+    cerr<<"Syntax: tcpbench remote [port] < queries"<<endl;
+    cerr<<"Where queries is one query per line, format: qname qtype, just 1 space"<<endl;
+    exit(EXIT_FAILURE);
+  }
+  if(argc > 2)
+    port = atoi(argv[2]);
+
+  g_dest = ComboAddress(argv[1], port);
   unsigned int numworkers=100;
   pthread_t workers[numworkers];
 
-  for(unsigned int n = 0; n < 1000000; ++n) {
-    g_queries.push_back(Query("www.powerdns.com", QType::A));
+  FILE* fp=fdopen(0, "r");
+  pair<string, string> q;
+  string line;
+  while(stringfgets(fp, line)) {
+    trim_right(line);
+    q=splitField(line, ' ');
+    g_queries.push_back(Query(q.first, DNSRecordContent::TypeToNumber(q.second)));
   }
-
+  fclose(fp);
+    
   for(unsigned int n = 0; n < numworkers; ++n) {
     pthread_create(&workers[n], 0, worker, 0);
   }
@@ -108,7 +145,8 @@ try
     void* status;
     pthread_join(workers[n], &status);
   }
-
+  cout<<"OK: "<<g_OK<<", network errors: "<<g_networkErrors<<", other errors: "<<g_otherErrors<<endl;
+  cout<<"Truncateds: "<<g_truncates<<endl;
 }
 catch(std::exception &e)
 {
-- 
2.50.0