From 596bf48219cb82379bf872e746dd7f626b582342 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Wed, 27 Mar 2019 12:10:37 +0100
Subject: [PATCH] rec: Add a 'rebalanced-queries' metric

---
 pdns/pdns_recursor.cc              | 10 +++++++---
 pdns/rec-snmp.cc                   |  2 ++
 pdns/rec_channel_rec.cc            |  2 ++
 pdns/recursordist/RECURSOR-MIB.txt | 10 ++++++++++
 pdns/recursordist/docs/metrics.rst |  6 ++++++
 pdns/syncres.hh                    |  1 +
 6 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/pdns/pdns_recursor.cc b/pdns/pdns_recursor.cc
index 7437e7185..065833eb5 100644
--- a/pdns/pdns_recursor.cc
+++ b/pdns/pdns_recursor.cc
@@ -2889,9 +2889,13 @@ static unsigned int selectWorker(unsigned int hash)
 
   unsigned int worker = hash % g_numWorkerThreads;
   /* at least one server has to be below the average load */
-  while(load[worker] > targetLoad) {
-    // cerr<<"worker "<<worker<<" is above the target load, selecting another one"<<endl;
-    worker = (worker + 1) % g_numWorkerThreads;
+  if (load[worker] > targetLoad) {
+    ++g_stats.rebalancedQueries;
+    do {
+      // cerr<<"worker "<<worker<<" is above the target load, selecting another one"<<endl;
+      worker = (worker + 1) % g_numWorkerThreads;
+    }
+    while(load[worker] > targetLoad);
   }
 
   return /* skip handler */ 1 + g_numDistributorThreads + worker;
diff --git a/pdns/rec-snmp.cc b/pdns/rec-snmp.cc
index 415a36ff3..45b48cea8 100644
--- a/pdns/rec-snmp.cc
+++ b/pdns/rec-snmp.cc
@@ -114,6 +114,7 @@ static const oid dnssecAuthenticDataQueriesOID[] = { RECURSOR_STATS_OID, 95 };
 static const oid dnssecCheckDisabledQueriesOID[] = { RECURSOR_STATS_OID, 96 };
 static const oid variableResponsesOID[] = { RECURSOR_STATS_OID, 97 };
 static const oid specialMemoryUsageOID[] = { RECURSOR_STATS_OID, 98 };
+static const oid rebalancedQueriesOID[] = { RECURSOR_STATS_OID, 99 };
 
 static std::unordered_map<oid, std::string> s_statsMap;
 
@@ -320,5 +321,6 @@ RecursorSNMPAgent::RecursorSNMPAgent(const std::string& name, const std::string&
   registerCounter64Stat("policy-result-truncate", policyResultTruncateOID, OID_LENGTH(policyResultTruncateOID));
   registerCounter64Stat("policy-result-custom", policyResultCustomOID, OID_LENGTH(policyResultCustomOID));
   registerCounter64Stat("special-memory-usage", specialMemoryUsageOID, OID_LENGTH(specialMemoryUsageOID));
+  registerCounter64Stat("rebalanced-queries", rebalancedQueriesOID, OID_LENGTH(rebalancedQueriesOID));
 #endif /* HAVE_NET_SNMP */
 }
diff --git a/pdns/rec_channel_rec.cc b/pdns/rec_channel_rec.cc
index bb5bfba56..453885bad 100644
--- a/pdns/rec_channel_rec.cc
+++ b/pdns/rec_channel_rec.cc
@@ -1102,6 +1102,8 @@ void registerAllStats()
   addGetStat("policy-result-truncate", &g_stats.policyResults[DNSFilterEngine::PolicyKind::Truncate]);
   addGetStat("policy-result-custom", &g_stats.policyResults[DNSFilterEngine::PolicyKind::Custom]);
 
+  addGetStat("rebalanced-queries", &g_stats.rebalancedQueries);
+
   /* make sure that the ECS stats are properly initialized */
   SyncRes::clearECSStats();
   for (size_t idx = 0; idx < SyncRes::s_ecsResponsesBySubnetSize4.size(); idx++) {
diff --git a/pdns/recursordist/RECURSOR-MIB.txt b/pdns/recursordist/RECURSOR-MIB.txt
index 2af4621d0..6b9fac22e 100644
--- a/pdns/recursordist/RECURSOR-MIB.txt
+++ b/pdns/recursordist/RECURSOR-MIB.txt
@@ -817,6 +817,14 @@ specialMemoryUsage OBJECT-TYPE
         "Memory usage (more precise bbut expensive to retrieve)"
     ::= { stats 98 }
 
+rebalancedQueries OBJECT-TYPE
+    SYNTAX Counter64
+    MAX-ACCESS read-only
+    STATUS current
+    DESCRIPTION
+        "Number of queries re-distributed because the first selected worker thread was above the target load"
+    ::= { stats 99 }
+
 ---
 --- Traps / Notifications
 ---
@@ -957,6 +965,8 @@ recGroup OBJECT-GROUP
         dnssecAuthenticDataQueries,
         dnssecCheckDisabledQueries,
         variableResponses,
+        specialMemoryUsage,
+        rebalancedQueries,
         trapReason
     }
     STATUS current
diff --git a/pdns/recursordist/docs/metrics.rst b/pdns/recursordist/docs/metrics.rst
index 3fb0ee71e..1fbf80576 100644
--- a/pdns/recursordist/docs/metrics.rst
+++ b/pdns/recursordist/docs/metrics.rst
@@ -399,6 +399,12 @@ questions
 ^^^^^^^^^
 counts all end-user initiated queries with the RD bit   set
 
+rebalanced-queries
+^^^^^^^^^^^^^^^^^^
+.. versionadded:: 4.1.12
+
+number of queries balanced to a different worker thread because the first selected one was above the target load configured with 'distribution-load-factor'
+
 resource-limits
 ^^^^^^^^^^^^^^^
 counts number of queries that could not be   performed because of resource limits
diff --git a/pdns/syncres.hh b/pdns/syncres.hh
index d808668b4..bba453e2c 100644
--- a/pdns/syncres.hh
+++ b/pdns/syncres.hh
@@ -957,6 +957,7 @@ struct RecursorStats
   std::atomic<uint64_t> dnssecValidations; // should be the sum of all dnssecResult* stats
   std::map<vState, std::atomic<uint64_t> > dnssecResults;
   std::map<DNSFilterEngine::PolicyKind, std::atomic<uint64_t> > policyResults;
+  std::atomic<uint64_t> rebalancedQueries{0};
 };
 
 //! represents a running TCP/IP client session
-- 
2.40.0