From a1a787dca089f2531abf7246bdb99a26ffb95791 Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@netherlabs.nl>
Date: Sat, 31 Oct 2015 14:38:56 +0100
Subject: [PATCH] implement new metric: real-memory-usage, plus hook up &
 enable function based metrics in dnsdist

---
 pdns/dnsdist-carbon.cc |  6 ++++--
 pdns/dnsdist-tcp.cc    |  8 ++------
 pdns/dnsdist.hh        |  6 ++++--
 pdns/misc.cc           | 20 ++++++++++++++++++++
 pdns/misc.hh           |  2 ++
 5 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/pdns/dnsdist-carbon.cc b/pdns/dnsdist-carbon.cc
index 3bdecc125..e9fc295db 100644
--- a/pdns/dnsdist-carbon.cc
+++ b/pdns/dnsdist-carbon.cc
@@ -45,12 +45,14 @@ try
 	str<<"dnsdist."<<hostname<<".main."<<e.first<<' ';
 	if(const auto& val = boost::get<DNSDistStats::stat_t*>(&e.second))
 	  str<<(*val)->load();
+	else if (const auto& val = boost::get<double*>(&e.second))
+	  str<<**val;
 	else
-	  str<<*boost::get<double*>(e.second);
+	  str<<(*boost::get<DNSDistStats::statfunction_t>(&e.second))(e.first);
 	str<<' '<<now<<"\r\n";
       }
       const string msg = str.str();
-      
+
       int ret = waitForRWData(s.getHandle(), false, 1 , 0); 
       if(ret <= 0 ) {
 	infolog("Unable to write data to carbon server on %s: %s", localCarbon->server.toStringWithPort(), (ret<0 ? strerror(errno) : "Timeout"));
diff --git a/pdns/dnsdist-tcp.cc b/pdns/dnsdist-tcp.cc
index 52538bae7..2ebbbf009 100644
--- a/pdns/dnsdist-tcp.cc
+++ b/pdns/dnsdist-tcp.cc
@@ -45,14 +45,12 @@ using std::atomic;
 
 static int setupTCPDownstream(const ComboAddress& remote)
 {  
-  
   vinfolog("TCP connecting to downstream %s", remote.toStringWithPort());
   int sock = SSocket(remote.sin4.sin_family, SOCK_STREAM, 0);
   SConnect(sock, remote);
   return sock;
 }
 
-
 struct ConnectionInfo
 {
   int fd;
@@ -61,8 +59,7 @@ struct ConnectionInfo
 
 void* tcpClientThread(int pipefd);
 
-
-  // Should not be called simultaneously!
+// Should not be called simultaneously!
 void TCPClientCollection::addTCPClientThread()
 {  
   vinfolog("Adding TCP Client thread");
@@ -86,7 +83,6 @@ void* tcpClientThread(int pipefd)
      
   typedef std::function<bool(ComboAddress, DNSName, uint16_t, dnsheader*)> blockfilter_t;
   blockfilter_t blockFilter = 0;
-
   
   {
     std::lock_guard<std::mutex> lock(g_luamutex);
@@ -199,7 +195,7 @@ void* tcpClientThread(int pipefd)
         ds->queries++;
         ds->outstanding++;
 
-	if(qtype == QType::AXFR)  // XXX fixme we really need to do better
+	if(qtype == QType::AXFR || qtype == QType::IXFR)  // XXX fixme we really need to do better
 	  break;
 
       retry:; 
diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index a1c318bcb..09f8aefd9 100644
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -29,7 +29,8 @@ struct DNSDistStats
   stat_t latency0_1{0}, latency1_10{0}, latency10_50{0}, latency50_100{0}, latency100_1000{0}, latencySlow{0};
   
   double latencyAvg100{0}, latencyAvg1000{0}, latencyAvg10000{0}, latencyAvg1000000{0};
-  typedef boost::variant<stat_t*, double*> entry_t;
+  typedef std::function<uint64_t(const std::string&)> statfunction_t;
+  typedef boost::variant<stat_t*, double*, statfunction_t> entry_t;
   std::vector<std::pair<std::string, entry_t>> entries{
     {"responses", &responses}, {"servfail-responses", &servfailResponses},
     {"queries", &queries}, {"acl-drops", &aclDrops},
@@ -41,7 +42,8 @@ struct DNSDistStats
     {"latency10-50", &latency10_50}, {"latency50-100", &latency50_100}, 
     {"latency100-1000", &latency100_1000}, {"latency-slow", &latencySlow},
     {"latency-avg100", &latencyAvg100}, {"latency-avg1000", &latencyAvg1000}, 
-    {"latency-avg10000", &latencyAvg10000}, {"latency-avg1000000", &latencyAvg1000000}
+    {"latency-avg10000", &latencyAvg10000}, {"latency-avg1000000", &latencyAvg1000000},
+    {"real-memory-usage", getRealMemoryUsage}
   };
 };
 
diff --git a/pdns/misc.cc b/pdns/misc.cc
index 299010401..ea4b3f6cd 100644
--- a/pdns/misc.cc
+++ b/pdns/misc.cc
@@ -1113,3 +1113,23 @@ DNSName getTSIGAlgoName(TSIGHashEnum& algoEnum)
   }
   throw PDNSException("getTSIGAlgoName does not understand given algorithm, please fix!");
 }
+
+uint64_t getRealMemoryUsage(const std::string&)
+{
+#ifdef __linux__
+  ifstream ifs("/proc/"+std::to_string(getpid())+"/smaps");
+  if(!ifs)
+    return 0;
+  string line;
+  uint64_t bytes=0;
+  string header("Private_Dirty:");
+  while(getline(ifs, line)) {
+    if(boost::starts_with(line, header)) {
+      bytes += atoi(line.c_str() + header.length() +1)*1024;
+    }
+  }
+  return bytes;
+#else
+  return 0;
+#endif
+}
diff --git a/pdns/misc.hh b/pdns/misc.hh
index b967ebe5f..b72016a0e 100644
--- a/pdns/misc.hh
+++ b/pdns/misc.hh
@@ -622,6 +622,8 @@ int closesocket(int fd);
 bool setCloseOnExec(int sock);
 uint64_t udpErrorStats(const std::string& str);
 
+uint64_t getRealMemoryUsage(const std::string&);
+
 template<typename T, typename... Args>
 std::unique_ptr<T> make_unique(Args&&... args)
 {
-- 
2.40.0