dnsdist: Avoid allocating memory in LB policies for small number of servers

The gist of this commit is that we do not need to dynamically allocate
memory in our load-balancing policies when we have a few number of
servers to consider: we can very well use a stack-based array instead.

diff --git a/pdns/dnsdistdist/dnsdist-lbpolicies.cc b/pdns/dnsdistdist/dnsdist-lbpolicies.cc
index eb5a3aa3ae53dcddcc074c937beacb01dca0e80e..8d424036546935c4c1815fbe83aa9139dabe6881 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-lbpolicies.cc
+++ b/pdns/dnsdistdist/dnsdist-lbpolicies.cc
@@ -29,36 +29,46 @@
 GlobalStateHolder<ServerPolicy> g_policy;
 bool g_roundrobinFailOnNoServer{false};
 
-// get server with least outstanding queries, and within those, with the lowest order, and within those: the fastest
-shared_ptr<DownstreamState> leastOutstanding(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
+template <class T> static std::shared_ptr<DownstreamState> getLeastOutstanding(const ServerPolicy::NumberedServerVector& servers, T& poss)
 {
-  if (servers.size() == 1 && servers[0].second->isUp()) {
-    return servers[0].second;
-  }
-
-  vector<pair<std::tuple<int,int,double>, size_t>> poss;
   /* so you might wonder, why do we go through this trouble? The data on which we sort could change during the sort,
      which would suck royally and could even lead to crashes. So first we snapshot on what we sort, and then we sort */
-  poss.reserve(servers.size());
   size_t position = 0;
-  for(const auto& d : servers) {
-    if(d.second->isUp()) {
-      poss.emplace_back(std::make_tuple(d.second->outstanding.load(), d.second->d_config.order, d.second->latencyUsec), position);
+  for (const auto& d : servers) {
+    if (d.second->isUp()) {
+      poss[position] = std::make_pair(std::make_tuple(d.second->outstanding.load(), d.second->d_config.order, d.second->latencyUsec), position);
     }
     ++position;
   }
 
-  if (poss.empty()) {
+  if (position == 0) {
     return shared_ptr<DownstreamState>();
   }
 
-  nth_element(poss.begin(), poss.begin(), poss.end(), [](const decltype(poss)::value_type& a, const decltype(poss)::value_type& b) { return a.first < b.first; });
+  nth_element(poss.begin(), poss.begin(), poss.begin() + position, [](const typename T::value_type& a, const typename T::value_type& b) { return a.first < b.first; });
   return servers.at(poss.begin()->second).second;
 }
 
+// get server with least outstanding queries, and within those, with the lowest order, and within those: the fastest
+shared_ptr<DownstreamState> leastOutstanding(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
+{
+  if (servers.size() == 1 && servers[0].second->isUp()) {
+    return servers[0].second;
+  }
+
+  if (servers.size() <= 16) {
+    std::array<pair<std::tuple<int,int,double>, size_t>, 16> poss;
+    return getLeastOutstanding(servers, poss);
+  }
+
+  vector<pair<std::tuple<int,int,double>, size_t>> poss;
+  poss.resize(servers.size());
+  return getLeastOutstanding(servers, poss);
+}
+
 shared_ptr<DownstreamState> firstAvailable(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
 {
-  for(auto& d : servers) {
+  for (auto& d : servers) {
     if (d.second->isUp() && d.second->qps.checkOnly()) {
       return d.second;
     }
@@ -68,30 +78,12 @@ shared_ptr<DownstreamState> firstAvailable(const ServerPolicy::NumberedServerVec
 
 double g_weightedBalancingFactor = 0;
 
-static shared_ptr<DownstreamState> valrandom(unsigned int val, const ServerPolicy::NumberedServerVector& servers)
+template <class T> static std::shared_ptr<DownstreamState> getValRandom(const ServerPolicy::NumberedServerVector& servers, T& poss, const unsigned int val, const double targetLoad)
 {
-  vector<pair<int, size_t>> poss;
-  poss.reserve(servers.size());
   int sum = 0;
   int max = std::numeric_limits<int>::max();
-  double targetLoad = std::numeric_limits<double>::max();
-
-  if (g_weightedBalancingFactor > 0) {
-    /* we start with one, representing the query we are currently handling */
-    double currentLoad = 1;
-    size_t totalWeight = 0;
-    for (const auto& pair : servers) {
-      if (pair.second->isUp()) {
-        currentLoad += pair.second->outstanding;
-        totalWeight += pair.second->d_config.d_weight;
-      }
-    }
-
-    if (totalWeight > 0) {
-      targetLoad = (currentLoad / totalWeight) * g_weightedBalancingFactor;
-    }
-  }
 
+  size_t position = 0;
   for (const auto& d : servers) {      // w=1, w=10 -> 1, 11
     if (d.second->isUp() && (g_weightedBalancingFactor == 0 || (d.second->outstanding <= (targetLoad * d.second->d_config.d_weight)))) {
       // Don't overflow sum when adding high weights
@@ -101,17 +93,18 @@ static shared_ptr<DownstreamState> valrandom(unsigned int val, const ServerPolic
         sum += d.second->d_config.d_weight;
       }
 
-      poss.emplace_back(sum, d.first);
+      poss[position]  = std::make_pair(sum, d.first);
+      position++;
     }
   }
 
   // Catch poss & sum are empty to avoid SIGFPE
-  if (poss.empty() || sum == 0) {
+  if (position == 0 || sum == 0) {
     return shared_ptr<DownstreamState>();
   }
 
   int r = val % sum;
-  auto p = upper_bound(poss.begin(), poss.end(),r, [](int r_, const decltype(poss)::value_type& a) { return  r_ < a.first;});
+  auto p = upper_bound(poss.begin(), poss.begin() + position, r, [](int r_, const typename T::value_type& a) { return  r_ < a.first;});
   if (p == poss.end()) {
     return shared_ptr<DownstreamState>();
   }
@@ -119,6 +112,36 @@ static shared_ptr<DownstreamState> valrandom(unsigned int val, const ServerPolic
   return servers.at(p->second - 1).second;
 }
 
+static shared_ptr<DownstreamState> valrandom(const unsigned int val, const ServerPolicy::NumberedServerVector& servers)
+{
+  double targetLoad = std::numeric_limits<double>::max();
+
+  if (g_weightedBalancingFactor > 0) {
+    /* we start with one, representing the query we are currently handling */
+    double currentLoad = 1;
+    size_t totalWeight = 0;
+    for (const auto& pair : servers) {
+      if (pair.second->isUp()) {
+        currentLoad += pair.second->outstanding;
+        totalWeight += pair.second->d_config.d_weight;
+      }
+    }
+
+    if (totalWeight > 0) {
+      targetLoad = (currentLoad / totalWeight) * g_weightedBalancingFactor;
+    }
+  }
+
+  if (servers.size() <= 16) {
+    std::array<pair<int, size_t>, 16> poss;
+    return getValRandom(servers, poss, val, targetLoad);
+  }
+
+  vector<pair<int, size_t>> poss;
+  poss.resize(servers.size());
+  return getValRandom(servers, poss, val, targetLoad);
+}
+
 shared_ptr<DownstreamState> wrandom(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
 {
   return valrandom(random(), servers);