17795654203

Committed 17 Sep 2025 11:10AM UTC coverage: 66.037% (+0.03%) from 66.012%

Build # 17795654203

Build Type

Pull #15808

github

Committed by

web-flow

Commit Message

Merge 103733592 into a7f305b95

Pull Request Pull Request #15808: rec: allow overriding remote address from preoutquery()

Run Details

42442 of 92980 branches covered (45.65%)

Branch coverage included in aggregate %.

7 of 7 new or added lines in 2 files covered. (100.0%)

18 existing lines in 6 files now uncovered.

128910 of 166497 relevant lines covered (77.42%)

5194927.21 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

72.36

/pdns/dnsdistdist/dnsdist-lbpolicies.cc

/*
 * This file is part of PowerDNS or dnsdist.
 * Copyright -- PowerDNS.COM B.V. and its contributors
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of version 2 of the GNU General Public License as
 * published by the Free Software Foundation.
 *
 * In addition, for the avoidance of any doubt, permission is granted to
 * link this program with OpenSSL and to (re)distribute the binaries
 * produced as the result of such linking.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include "dnsdist.hh"
#include "dnsdist-lbpolicies.hh"
#include "dnsdist-lua.hh"
#include "dnsdist-lua-ffi.hh"
#include "dolog.hh"
#include "dns_random.hh"

static constexpr size_t s_staticArrayCutOff = 16;
template <typename T> using DynamicIndexArray = std::vector<std::pair<T, size_t>>;
template <typename T> using StaticIndexArray = std::array<std::pair<T, size_t>, s_staticArrayCutOff>;

template <class T> static std::shared_ptr<DownstreamState> getLeastOutstanding(const ServerPolicy::NumberedServerVector& servers, 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 */
  size_t usableServers = 0;
  for (const auto& d : servers) {
    if (d.second->isUp()) {
      poss.at(usableServers) = std::pair(std::tuple(d.second->outstanding.load(), d.second->d_config.order, d.second->getRelevantLatencyUsec()), d.first);
      usableServers++;
    }
  }

  if (usableServers == 0) {
    return shared_ptr<DownstreamState>();
  }

  std::nth_element(poss.begin(), poss.begin(), poss.begin() + usableServers, [](const typename T::value_type& a, const typename T::value_type& b) { return a.first < b.first; });
  // minus 1 because the NumberedServerVector starts at 1 for Lua
  return servers.at(poss.begin()->second - 1).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)
{
  (void)dq;
  using LeastOutstandingType = std::tuple<int,int,double>;

  if (servers.size() == 1 && servers[0].second->isUp()) {
    return servers[0].second;
  }

  if (servers.size() <= s_staticArrayCutOff) {
    StaticIndexArray<LeastOutstandingType> poss;
    return getLeastOutstanding(servers, poss);
  }

  DynamicIndexArray<LeastOutstandingType> poss;
  poss.resize(servers.size());
  return getLeastOutstanding(servers, poss);
}

shared_ptr<DownstreamState> firstAvailable(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
{
  for (auto& d : servers) {
    if (d.second->isUp() && (!d.second->d_qpsLimiter || d.second->d_qpsLimiter->checkOnly())) {
      return d.second;
    }
  }
  return leastOutstanding(servers, dq);
}

template <class T> static std::shared_ptr<DownstreamState> getValRandom(const ServerPolicy::NumberedServerVector& servers, T& poss, const unsigned int val, const double targetLoad)
{
  constexpr int max = std::numeric_limits<int>::max();
  int sum = 0;

  size_t usableServers = 0;
  const auto weightedBalancingFactor = dnsdist::configuration::getImmutableConfiguration().d_weightedBalancingFactor;
  for (const auto& d : servers) {      // w=1, w=10 -> 1, 11
    if (d.second->isUp() && (weightedBalancingFactor == 0 || (static_cast<double>(d.second->outstanding.load()) <= (targetLoad * d.second->d_config.d_weight)))) {
      // Don't overflow sum when adding high weights
      if (d.second->d_config.d_weight > max - sum) {
        sum = max;
      } else {
        sum += d.second->d_config.d_weight;
      }

      poss.at(usableServers) = std::pair(sum, d.first);
      usableServers++;
    }
  }

  // Catch the case where usableServers or sum are equal to 0 to avoid a SIGFPE
  if (usableServers == 0 || sum == 0) {
    return shared_ptr<DownstreamState>();
  }

  int r = val % sum;
  auto p = std::upper_bound(poss.begin(), poss.begin() + usableServers, r, [](int r_, const typename T::value_type& a) { return  r_ < a.first;});
  if (p == poss.begin() + usableServers) {
    return shared_ptr<DownstreamState>();
  }

  // minus 1 because the NumberedServerVector starts at 1 for Lua
  return servers.at(p->second - 1).second;
}

static shared_ptr<DownstreamState> valrandom(const unsigned int val, const ServerPolicy::NumberedServerVector& servers)
{
  using ValRandomType = int;
  double targetLoad = std::numeric_limits<double>::max();
  const auto weightedBalancingFactor = dnsdist::configuration::getImmutableConfiguration().d_weightedBalancingFactor;
  if (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 / static_cast<double>(totalWeight)) * weightedBalancingFactor;
    }
  }

  if (servers.size() <= s_staticArrayCutOff) {
    StaticIndexArray<ValRandomType> poss;
    return getValRandom(servers, poss, val, targetLoad);
  }

  DynamicIndexArray<ValRandomType> poss;
  poss.resize(servers.size());
  return getValRandom(servers, poss, val, targetLoad);
}

shared_ptr<DownstreamState> wrandom(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
{
  (void)dq;
  return valrandom(dns_random_uint32(), servers);
}

shared_ptr<DownstreamState> whashedFromHash(const ServerPolicy::NumberedServerVector& servers, size_t hash)
{
  return valrandom(hash, servers);
}

shared_ptr<DownstreamState> whashed(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
{
  const auto hashPerturbation = dnsdist::configuration::getImmutableConfiguration().d_hashPerturbation;
  return whashedFromHash(servers, dq->ids.qname.hash(hashPerturbation));
}

shared_ptr<DownstreamState> chashedFromHash(const ServerPolicy::NumberedServerVector& servers, size_t qhash)
{
  unsigned int sel = std::numeric_limits<unsigned int>::max();
  unsigned int min = std::numeric_limits<unsigned int>::max();
  shared_ptr<DownstreamState> ret = nullptr, first = nullptr;

  double targetLoad = std::numeric_limits<double>::max();
  const auto consistentHashBalancingFactor = dnsdist::configuration::getImmutableConfiguration().d_consistentHashBalancingFactor;
  if (consistentHashBalancingFactor > 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 / static_cast<double>(totalWeight)) * consistentHashBalancingFactor;
    }
  }

  for (const auto& d: servers) {
    if (d.second->isUp() && (consistentHashBalancingFactor == 0 || static_cast<double>(d.second->outstanding.load()) <= (targetLoad * d.second->d_config.d_weight))) {
      // make sure hashes have been computed
      if (!d.second->hashesComputed) {
        d.second->hash();
      }
      {
        const auto& server = d.second;
        auto hashes = server->hashes.read_lock();
        // we want to keep track of the last hash
        if (min > *(hashes->begin())) {
          min = *(hashes->begin());
          first = server;
        }

        auto hash_it = std::lower_bound(hashes->begin(), hashes->end(), qhash);
        if (hash_it != hashes->end()) {
          if (*hash_it < sel) {
            sel = *hash_it;
            ret = server;
          }
        }
      }
    }
  }
  if (ret != nullptr) {
    return ret;
  }
  if (first != nullptr) {
    return first;
  }
  return shared_ptr<DownstreamState>();
}

shared_ptr<DownstreamState> chashed(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
{
  const auto hashPerturbation = dnsdist::configuration::getImmutableConfiguration().d_hashPerturbation;
  return chashedFromHash(servers, dq->ids.qname.hash(hashPerturbation));
}

shared_ptr<DownstreamState> roundrobin(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
{
  (void)dq;
  if (servers.empty()) {
    return shared_ptr<DownstreamState>();
  }

  vector<size_t> candidates;
  candidates.reserve(servers.size());

  for (auto& d : servers) {
    if (d.second->isUp()) {
      candidates.push_back(d.first);
    }
  }

  if (candidates.empty()) {
    if (dnsdist::configuration::getCurrentRuntimeConfiguration().d_roundrobinFailOnNoServer) {
      return shared_ptr<DownstreamState>();
    }
    for (auto& d : servers) {
      candidates.push_back(d.first);
    }
  }

  static unsigned int counter;
  return servers.at(candidates.at((counter++) % candidates.size()) - 1).second;
}

shared_ptr<DownstreamState> orderedWrandUntag(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dnsq)
{
  if (servers.empty()) {
    return {};
  }

  ServerPolicy::NumberedServerVector candidates;
  candidates.reserve(servers.size());

  int curOrder = std::numeric_limits<int>::max();
  unsigned int curNumber = 1;

  for (const auto& svr : servers) {
    if (svr.second->isUp() && (!dnsq->ids.qTag || dnsq->ids.qTag->count(svr.second->getNameWithAddr()) == 0)) {
      // the servers in a pool are already sorted in ascending order by its 'order', see ``ServerPool::addServer()``
      if (svr.second->d_config.order > curOrder) {
        break;
      }
      curOrder = svr.second->d_config.order;
      candidates.push_back(ServerPolicy::NumberedServer(curNumber++, svr.second));
    }
  }

  if (candidates.empty()) {
    return {};
  }

  return wrandom(candidates, dnsq);
}

std::shared_ptr<const ServerPolicy::NumberedServerVector> getDownstreamCandidates(const std::string& poolName)
{
  std::shared_ptr<ServerPool> pool = getPool(poolName);
  return pool->getServers();
}

std::shared_ptr<ServerPool> createPoolIfNotExists(const string& poolName)
{
  {
    const auto& pools = dnsdist::configuration::getCurrentRuntimeConfiguration().d_pools;
    const auto poolIt = pools.find(poolName);
    if (poolIt != pools.end()) {
      return poolIt->second;
    }
  }

  if (!poolName.empty()) {
    vinfolog("Creating pool %s", poolName);
  }

  auto pool = std::make_shared<ServerPool>();
  dnsdist::configuration::updateRuntimeConfiguration([&poolName,&pool](dnsdist::configuration::RuntimeConfiguration& config) {
    config.d_pools.emplace(poolName, pool);
  });

  return pool;
}

void setPoolPolicy(const string& poolName, std::shared_ptr<ServerPolicy> policy)
{
  std::shared_ptr<ServerPool> pool = createPoolIfNotExists(poolName);
  if (!poolName.empty()) {
    vinfolog("Setting pool %s server selection policy to %s", poolName, policy->getName());
  } else {
    vinfolog("Setting default pool server selection policy to %s", policy->getName());
  }
  pool->policy = std::move(policy);
}

void addServerToPool(const string& poolName, std::shared_ptr<DownstreamState> server)
{
  std::shared_ptr<ServerPool> pool = createPoolIfNotExists(poolName);
  if (!poolName.empty()) {
    vinfolog("Adding server to pool %s", poolName);
  } else {
    vinfolog("Adding server to default pool");
  }
  pool->addServer(server);
}

void removeServerFromPool(const string& poolName, std::shared_ptr<DownstreamState> server)
{
  std::shared_ptr<ServerPool> pool = getPool(poolName);

  if (!poolName.empty()) {
    vinfolog("Removing server from pool %s", poolName);
  }
  else {
    vinfolog("Removing server from default pool");
  }

  pool->removeServer(server);
}

std::shared_ptr<ServerPool> getPool(const std::string& poolName)
{
  const auto& pools = dnsdist::configuration::getCurrentRuntimeConfiguration().d_pools;
  auto poolIt = pools.find(poolName);
  if (poolIt == pools.end()) {
    throw std::out_of_range("No pool named " + poolName);
  }

  return poolIt->second;
}

ServerPolicy::ServerPolicy(const std::string& name_, const std::string& code): d_name(name_), d_perThreadPolicyCode(code), d_isLua(true), d_isFFI(true), d_isPerThread(true)
{
  LuaContext tmpContext;
  setupLuaLoadBalancingContext(tmpContext);
  auto ret = tmpContext.executeCode<ServerPolicy::ffipolicyfunc_t>(code);
}

struct ServerPolicy::PerThreadState
{
  LuaContext d_luaContext;
  std::unordered_map<std::string, ffipolicyfunc_t> d_policies;
};

thread_local std::unique_ptr<ServerPolicy::PerThreadState> ServerPolicy::t_perThreadState;

const ServerPolicy::ffipolicyfunc_t& ServerPolicy::getPerThreadPolicy() const
{
  auto& state = t_perThreadState;
  if (!state) {
    state = std::make_unique<ServerPolicy::PerThreadState>();
    setupLuaLoadBalancingContext(state->d_luaContext);
  }

  const auto& policyIt = state->d_policies.find(d_name);
  if (policyIt != state->d_policies.end()) {
    return policyIt->second;
  }

  auto newPolicy = state->d_luaContext.executeCode<ServerPolicy::ffipolicyfunc_t>(d_perThreadPolicyCode);
  state->d_policies[d_name] = std::move(newPolicy);
  return state->d_policies.at(d_name);
}

std::shared_ptr<DownstreamState> ServerPolicy::getSelectedBackend(const ServerPolicy::NumberedServerVector& servers, DNSQuestion& dq) const
{
  std::shared_ptr<DownstreamState> selectedBackend{nullptr};

  if (d_isLua) {
    if (!d_isFFI) {
      auto lock = g_lua.lock();
      selectedBackend = d_policy(servers, &dq);
    }
    else {
      dnsdist_ffi_dnsquestion_t dnsq(&dq);
      dnsdist_ffi_servers_list_t serversList(servers);
      unsigned int selected = 0;

      if (!d_isPerThread) {
        auto lock = g_lua.lock();
        selected = d_ffipolicy(&serversList, &dnsq);
      }
      else {
        const auto& policy = getPerThreadPolicy();
        selected = policy(&serversList, &dnsq);
      }

      if (selected >= servers.size()) {
        /* invalid offset, meaning that there is no server available */
        return {};
      }

      selectedBackend = servers.at(selected).second;
    }
  }
  else {
    selectedBackend = d_policy(servers, &dq);
  }

  return selectedBackend;
}

namespace dnsdist::lbpolicies
{
const std::vector<std::shared_ptr<ServerPolicy>>& getBuiltInPolicies()
{
  static const std::vector<std::shared_ptr<ServerPolicy>> s_policies{
    std::make_shared<ServerPolicy>("firstAvailable", firstAvailable, false),
    std::make_shared<ServerPolicy>("roundrobin", roundrobin, false),
    std::make_shared<ServerPolicy>("wrandom", wrandom, false),
    std::make_shared<ServerPolicy>("whashed", whashed, false),
    std::make_shared<ServerPolicy>("chashed", chashed, false),
    std::make_shared<ServerPolicy>("orderedWrandUntag", orderedWrandUntag, false),
    std::make_shared<ServerPolicy>("leastOutstanding", leastOutstanding, false)};
  return s_policies;
}
}

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
23	#include "dnsdist.hh"
24	#include "dnsdist-lbpolicies.hh"
25	#include "dnsdist-lua.hh"
26	#include "dnsdist-lua-ffi.hh"
27	#include "dolog.hh"
28	#include "dns_random.hh"
29
30	static constexpr size_t s_staticArrayCutOff = 16;
31	template <typename T> using DynamicIndexArray = std::vector<std::pair<T, size_t>>;
32	template <typename T> using StaticIndexArray = std::array<std::pair<T, size_t>, s_staticArrayCutOff>;
33
34	template <class T> static std::shared_ptr<DownstreamState> getLeastOutstanding(const ServerPolicy::NumberedServerVector& servers, T& poss)
35	{	32✔
36	/* so you might wonder, why do we go through this trouble? The data on which we sort could change during the sort,
37	which would suck royally and could even lead to crashes. So first we snapshot on what we sort, and then we sort */
38	size_t usableServers = 0;	32✔
39	for (const auto& d : servers) {	32!
40	if (d.second->isUp()) {	17!
41	poss.at(usableServers) = std::pair(std::tuple(d.second->outstanding.load(), d.second->d_config.order, d.second->getRelevantLatencyUsec()), d.first);	×
42	usableServers++;	×
43	}	×
44	}	17✔
45
46	if (usableServers == 0) {	32!
47	return shared_ptr<DownstreamState>();	32✔
48	}	32✔
49
50	std::nth_element(poss.begin(), poss.begin(), poss.begin() + usableServers, [](const typename T::value_type& a, const typename T::value_type& b) { return a.first < b.first; });	×
51	// minus 1 because the NumberedServerVector starts at 1 for Lua
52	return servers.at(poss.begin()->second - 1).second;	×
53	}	32✔
54
55	// get server with least outstanding queries, and within those, with the lowest order, and within those: the fastest
56	shared_ptr<DownstreamState> leastOutstanding(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
57	{	4,229✔
58	(void)dq;	4,229✔
59	using LeastOutstandingType = std::tuple<int,int,double>;	4,229✔
60
61	if (servers.size() == 1 && servers[0].second->isUp()) {	4,229✔
62	return servers[0].second;	4,197✔
63	}	4,197✔
64
65	if (servers.size() <= s_staticArrayCutOff) {	32!
66	StaticIndexArray<LeastOutstandingType> poss;	32✔
67	return getLeastOutstanding(servers, poss);	32✔
68	}	32✔
69
70	DynamicIndexArray<LeastOutstandingType> poss;	×
71	poss.resize(servers.size());	×
72	return getLeastOutstanding(servers, poss);	×
73	}	32✔
74
75	shared_ptr<DownstreamState> firstAvailable(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
76	{	224✔
77	for (auto& d : servers) {	224!
78	if (d.second->isUp() && (!d.second->d_qpsLimiter \|\| d.second->d_qpsLimiter->checkOnly())) {	224!
79	return d.second;	224✔
80	}	224✔
81	}	224✔
82	return leastOutstanding(servers, dq);	×
83	}	224✔
84
85	template <class T> static std::shared_ptr<DownstreamState> getValRandom(const ServerPolicy::NumberedServerVector& servers, T& poss, const unsigned int val, const double targetLoad)
86	{	1,100✔
87	constexpr int max = std::numeric_limits<int>::max();	1,100✔
88	int sum = 0;	1,100✔
89
90	size_t usableServers = 0;	1,100✔
91	const auto weightedBalancingFactor = dnsdist::configuration::getImmutableConfiguration().d_weightedBalancingFactor;	1,100✔
92	for (const auto& d : servers) { // w=1, w=10 -> 1, 11	2,000!
93	if (d.second->isUp() && (weightedBalancingFactor == 0 \|\| (static_cast<double>(d.second->outstanding.load()) <= (targetLoad * d.second->d_config.d_weight)))) {	2,000!
94	// Don't overflow sum when adding high weights
95	if (d.second->d_config.d_weight > max - sum) {	2,000!
96	sum = max;	200✔
97	} else {	1,800✔
98	sum += d.second->d_config.d_weight;	1,800✔
99	}	1,800✔
100
101	poss.at(usableServers) = std::pair(sum, d.first);	2,000✔
102	usableServers++;	2,000✔
103	}	2,000✔
104	}	2,000✔
105
106	// Catch the case where usableServers or sum are equal to 0 to avoid a SIGFPE
107	if (usableServers == 0 \|\| sum == 0) {	1,100!
108	return shared_ptr<DownstreamState>();	×
109	}	×
110
111	int r = val % sum;	1,100✔
112	auto p = std::upper_bound(poss.begin(), poss.begin() + usableServers, r, [](int r_, const typename T::value_type& a) { return r_ < a.first;});	2,000✔
113	if (p == poss.begin() + usableServers) {	1,100!
114	return shared_ptr<DownstreamState>();	×
115	}	×
116
117	// minus 1 because the NumberedServerVector starts at 1 for Lua
118	return servers.at(p->second - 1).second;	1,100✔
119	}	1,100✔
120
121	static shared_ptr<DownstreamState> valrandom(const unsigned int val, const ServerPolicy::NumberedServerVector& servers)
122	{	1,100✔
123	using ValRandomType = int;	1,100✔
124	double targetLoad = std::numeric_limits<double>::max();	1,100✔
125	const auto weightedBalancingFactor = dnsdist::configuration::getImmutableConfiguration().d_weightedBalancingFactor;	1,100✔
126	if (weightedBalancingFactor > 0) {	1,100!
127	/* we start with one, representing the query we are currently handling */
128	double currentLoad = 1;	×
129	size_t totalWeight = 0;	×
130	for (const auto& pair : servers) {	×
131	if (pair.second->isUp()) {	×
132	currentLoad += pair.second->outstanding;	×
133	totalWeight += pair.second->d_config.d_weight;	×
134	}	×
135	}	×
136
137	if (totalWeight > 0) {	×
138	targetLoad = (currentLoad / static_cast<double>(totalWeight)) * weightedBalancingFactor;	×
139	}	×
140	}	×
141
142	if (servers.size() <= s_staticArrayCutOff) {	1,100!
143	StaticIndexArray<ValRandomType> poss;	1,100✔
144	return getValRandom(servers, poss, val, targetLoad);	1,100✔
145	}	1,100✔
146
147	DynamicIndexArray<ValRandomType> poss;	×
148	poss.resize(servers.size());	×
149	return getValRandom(servers, poss, val, targetLoad);	×
150	}	1,100✔
151
152	shared_ptr<DownstreamState> wrandom(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
153	{	900✔
154	(void)dq;	900✔
155	return valrandom(dns_random_uint32(), servers);	900✔
156	}	900✔
157
158	shared_ptr<DownstreamState> whashedFromHash(const ServerPolicy::NumberedServerVector& servers, size_t hash)
159	{	200✔
160	return valrandom(hash, servers);	200✔
161	}	200✔
162
163	shared_ptr<DownstreamState> whashed(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
164	{	200✔
165	const auto hashPerturbation = dnsdist::configuration::getImmutableConfiguration().d_hashPerturbation;	200✔
166	return whashedFromHash(servers, dq->ids.qname.hash(hashPerturbation));	200✔
167	}	200✔
168
169	shared_ptr<DownstreamState> chashedFromHash(const ServerPolicy::NumberedServerVector& servers, size_t qhash)
170	{	200✔
171	unsigned int sel = std::numeric_limits<unsigned int>::max();	200✔
172	unsigned int min = std::numeric_limits<unsigned int>::max();	200✔
173	shared_ptr<DownstreamState> ret = nullptr, first = nullptr;	200✔
174
175	double targetLoad = std::numeric_limits<double>::max();	200✔
176	const auto consistentHashBalancingFactor = dnsdist::configuration::getImmutableConfiguration().d_consistentHashBalancingFactor;	200✔
177	if (consistentHashBalancingFactor > 0) {	200!
178	/* we start with one, representing the query we are currently handling */
179	double currentLoad = 1;	×
180	size_t totalWeight = 0;	×
181	for (const auto& pair : servers) {	×
182	if (pair.second->isUp()) {	×
183	currentLoad += pair.second->outstanding;	×
184	totalWeight += pair.second->d_config.d_weight;	×
185	}	×
186	}	×
187
188	if (totalWeight > 0) {	×
189	targetLoad = (currentLoad / static_cast<double>(totalWeight)) * consistentHashBalancingFactor;	×
190	}	×
191	}	×
192
193	for (const auto& d: servers) {	400✔
194	if (d.second->isUp() && (consistentHashBalancingFactor == 0 \|\| static_cast<double>(d.second->outstanding.load()) <= (targetLoad * d.second->d_config.d_weight))) {	400!
195	// make sure hashes have been computed
196	if (!d.second->hashesComputed) {	400!
197	d.second->hash();	×
198	}	×
199	{	400✔
200	const auto& server = d.second;	400✔
201	auto hashes = server->hashes.read_lock();	400✔
202	// we want to keep track of the last hash
203	if (min > *(hashes->begin())) {	400✔
204	min = *(hashes->begin());	200✔
205	first = server;	200✔
206	}	200✔
207
208	auto hash_it = std::lower_bound(hashes->begin(), hashes->end(), qhash);	400✔
209	if (hash_it != hashes->end()) {	400!
210	if (*hash_it < sel) {	400✔
211	sel = *hash_it;	297✔
212	ret = server;	297✔
213	}	297✔
214	}	400✔
215	}	400✔
216	}	400✔
217	}	400✔
218	if (ret != nullptr) {	200!
219	return ret;	200✔
220	}	200✔
UNCOV 221	if (first != nullptr) {	×
UNCOV 222	return first;	×
UNCOV 223	}	×
224	return shared_ptr<DownstreamState>();	×
UNCOV 225	}	×
226
227	shared_ptr<DownstreamState> chashed(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
228	{	200✔
229	const auto hashPerturbation = dnsdist::configuration::getImmutableConfiguration().d_hashPerturbation;	200✔
230	return chashedFromHash(servers, dq->ids.qname.hash(hashPerturbation));	200✔
231	}	200✔
232
233	shared_ptr<DownstreamState> roundrobin(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dq)
234	{	62✔
235	(void)dq;	62✔
236	if (servers.empty()) {	62!
237	return shared_ptr<DownstreamState>();	×
238	}	×
239
240	vector<size_t> candidates;	62✔
241	candidates.reserve(servers.size());	62✔
242
243	for (auto& d : servers) {	124✔
244	if (d.second->isUp()) {	124✔
245	candidates.push_back(d.first);	100✔
246	}	100✔
247	}	124✔
248
249	if (candidates.empty()) {	62✔
250	if (dnsdist::configuration::getCurrentRuntimeConfiguration().d_roundrobinFailOnNoServer) {	2!
251	return shared_ptr<DownstreamState>();	2✔
252	}	2✔
253	for (auto& d : servers) {	×
254	candidates.push_back(d.first);	×
255	}	×
256	}	×
257
258	static unsigned int counter;	60✔
259	return servers.at(candidates.at((counter++) % candidates.size()) - 1).second;	60✔
260	}	62✔
261
262	shared_ptr<DownstreamState> orderedWrandUntag(const ServerPolicy::NumberedServerVector& servers, const DNSQuestion* dnsq)
263	{	500✔
264	if (servers.empty()) {	500!
265	return {};	×
266	}	×
267
268	ServerPolicy::NumberedServerVector candidates;	500✔
269	candidates.reserve(servers.size());	500✔
270
271	int curOrder = std::numeric_limits<int>::max();	500✔
272	unsigned int curNumber = 1;	500✔
273
274	for (const auto& svr : servers) {	1,700✔
275	if (svr.second->isUp() && (!dnsq->ids.qTag \|\| dnsq->ids.qTag->count(svr.second->getNameWithAddr()) == 0)) {	1,700✔
276	// the servers in a pool are already sorted in ascending order by its 'order', see ``ServerPool::addServer()``
277	if (svr.second->d_config.order > curOrder) {	1,100✔
278	break;	300✔
279	}	300✔
280	curOrder = svr.second->d_config.order;	800✔
281	candidates.push_back(ServerPolicy::NumberedServer(curNumber++, svr.second));	800✔
282	}	800✔
283	}	1,700✔
284
285	if (candidates.empty()) {	500!
286	return {};	×
287	}	×
288
289	return wrandom(candidates, dnsq);	500✔
290	}	500✔
291
292	std::shared_ptr<const ServerPolicy::NumberedServerVector> getDownstreamCandidates(const std::string& poolName)
293	{	×
294	std::shared_ptr<ServerPool> pool = getPool(poolName);	×
295	return pool->getServers();	×
296	}	×
297
298	std::shared_ptr<ServerPool> createPoolIfNotExists(const string& poolName)
299	{	1,435✔
300	{	1,435✔
301	const auto& pools = dnsdist::configuration::getCurrentRuntimeConfiguration().d_pools;	1,435✔
302	const auto poolIt = pools.find(poolName);	1,435✔
303	if (poolIt != pools.end()) {	1,435✔
304	return poolIt->second;	592✔
305	}	592✔
306	}	1,435✔
307
308	if (!poolName.empty()) {	843✔
309	vinfolog("Creating pool %s", poolName);	95✔
310	}	95✔
311
312	auto pool = std::make_shared<ServerPool>();	843✔
313	dnsdist::configuration::updateRuntimeConfiguration([&poolName,&pool](dnsdist::configuration::RuntimeConfiguration& config) {	843✔
314	config.d_pools.emplace(poolName, pool);	843✔
315	});	843✔
316
317	return pool;	843✔
318	}	1,435✔
319
320	void setPoolPolicy(const string& poolName, std::shared_ptr<ServerPolicy> policy)
321	{	2✔
322	std::shared_ptr<ServerPool> pool = createPoolIfNotExists(poolName);	2✔
323	if (!poolName.empty()) {	2!
324	vinfolog("Setting pool %s server selection policy to %s", poolName, policy->getName());	×
325	} else {	2✔
326	vinfolog("Setting default pool server selection policy to %s", policy->getName());	2!
327	}	2✔
328	pool->policy = std::move(policy);	2✔
329	}	2✔
330
331	void addServerToPool(const string& poolName, std::shared_ptr<DownstreamState> server)
332	{	911✔
333	std::shared_ptr<ServerPool> pool = createPoolIfNotExists(poolName);	911✔
334	if (!poolName.empty()) {	911✔
335	vinfolog("Adding server to pool %s", poolName);	98✔
336	} else {	817✔
337	vinfolog("Adding server to default pool");	813✔
338	}	813✔
339	pool->addServer(server);	911✔
340	}	911✔
341
342	void removeServerFromPool(const string& poolName, std::shared_ptr<DownstreamState> server)
343	{	15✔
344	std::shared_ptr<ServerPool> pool = getPool(poolName);	15✔
345
346	if (!poolName.empty()) {	15✔
347	vinfolog("Removing server from pool %s", poolName);	9✔
348	}	9✔
349	else {	6✔
350	vinfolog("Removing server from default pool");	6✔
351	}	6✔
352
353	pool->removeServer(server);	15✔
354	}	15✔
355
356	std::shared_ptr<ServerPool> getPool(const std::string& poolName)
357	{	5,888✔
358	const auto& pools = dnsdist::configuration::getCurrentRuntimeConfiguration().d_pools;	5,888✔
359	auto poolIt = pools.find(poolName);	5,888✔
360	if (poolIt == pools.end()) {	5,888✔
361	throw std::out_of_range("No pool named " + poolName);	6✔
362	}	6✔
363
364	return poolIt->second;	5,882✔
365	}	5,888✔
366
367	ServerPolicy::ServerPolicy(const std::string& name_, const std::string& code): d_name(name_), d_perThreadPolicyCode(code), d_isLua(true), d_isFFI(true), d_isPerThread(true)	2✔
368	{	2✔
369	LuaContext tmpContext;	2✔
370	setupLuaLoadBalancingContext(tmpContext);	2✔
371	auto ret = tmpContext.executeCode<ServerPolicy::ffipolicyfunc_t>(code);	2✔
372	}	2✔
373
374	struct ServerPolicy::PerThreadState
375	{
376	LuaContext d_luaContext;
377	std::unordered_map<std::string, ffipolicyfunc_t> d_policies;
378	};
379
380	thread_local std::unique_ptr<ServerPolicy::PerThreadState> ServerPolicy::t_perThreadState;
381
382	const ServerPolicy::ffipolicyfunc_t& ServerPolicy::getPerThreadPolicy() const
383	{	20✔
384	auto& state = t_perThreadState;	20✔
385	if (!state) {	20✔
386	state = std::make_unique<ServerPolicy::PerThreadState>();	2✔
387	setupLuaLoadBalancingContext(state->d_luaContext);	2✔
388	}	2✔
389
390	const auto& policyIt = state->d_policies.find(d_name);	20✔
391	if (policyIt != state->d_policies.end()) {	20✔
392	return policyIt->second;	18✔
393	}	18✔
394
395	auto newPolicy = state->d_luaContext.executeCode<ServerPolicy::ffipolicyfunc_t>(d_perThreadPolicyCode);	2✔
396	state->d_policies[d_name] = std::move(newPolicy);	2✔
397	return state->d_policies.at(d_name);	2✔
398	}	20✔
399
400	std::shared_ptr<DownstreamState> ServerPolicy::getSelectedBackend(const ServerPolicy::NumberedServerVector& servers, DNSQuestion& dq) const
401	{	7,857✔
402	std::shared_ptr<DownstreamState> selectedBackend{nullptr};	7,857✔
403
404	if (d_isLua) {	7,857✔
405	if (!d_isFFI) {	2,042✔
406	auto lock = g_lua.lock();	2,020✔
407	selectedBackend = d_policy(servers, &dq);	2,020✔
408	}	2,020✔
409	else {	22✔
410	dnsdist_ffi_dnsquestion_t dnsq(&dq);	22✔
411	dnsdist_ffi_servers_list_t serversList(servers);	22✔
412	unsigned int selected = 0;	22✔
413
414	if (!d_isPerThread) {	22✔
415	auto lock = g_lua.lock();	2✔
416	selected = d_ffipolicy(&serversList, &dnsq);	2✔
417	}	2✔
418	else {	20✔
419	const auto& policy = getPerThreadPolicy();	20✔
420	selected = policy(&serversList, &dnsq);	20✔
421	}	20✔
422
423	if (selected >= servers.size()) {	22✔
424	/* invalid offset, meaning that there is no server available */
425	return {};	2✔
426	}	2✔
427
428	selectedBackend = servers.at(selected).second;	20✔
429	}	20✔
430	}	2,042✔
431	else {	5,815✔
432	selectedBackend = d_policy(servers, &dq);	5,815✔
433	}	5,815✔
434
435	return selectedBackend;	7,855✔
436	}	7,857✔
437
438	namespace dnsdist::lbpolicies
439	{
440	const std::vector<std::shared_ptr<ServerPolicy>>& getBuiltInPolicies()
441	{	816✔
442	static const std::vector<std::shared_ptr<ServerPolicy>> s_policies{	816✔
443	std::make_shared<ServerPolicy>("firstAvailable", firstAvailable, false),	816✔
444	std::make_shared<ServerPolicy>("roundrobin", roundrobin, false),	816✔
445	std::make_shared<ServerPolicy>("wrandom", wrandom, false),	816✔
446	std::make_shared<ServerPolicy>("whashed", whashed, false),	816✔
447	std::make_shared<ServerPolicy>("chashed", chashed, false),	816✔
448	std::make_shared<ServerPolicy>("orderedWrandUntag", orderedWrandUntag, false),	816✔
449	std::make_shared<ServerPolicy>("leastOutstanding", leastOutstanding, false)};	816✔
450	return s_policies;	816✔
451	}	816✔
452	}

PowerDNS / pdns / 17795654203

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous