Add support for setting the server selection policy on a per pool basis

diff --git a/pdns/README-dnsdist.md b/pdns/README-dnsdist.md
index e3d615be80ae1dd356e1d62efd2eaf05c8141884..958638d4de7c4a35805fc13827e78c98055ba3c4 100644 (file)
--- a/pdns/README-dnsdist.md
+++ b/pdns/README-dnsdist.md
@@ -1491,6 +1491,9 @@ instantiate a server with additional parameters
     * `showServerPolicy()`: show name of currently operational server selection policy
     * `newServerPolicy(name, function)`: create a policy object from a Lua function
     * `setServFailWhenNoServer(bool)`: if set, return a ServFail when no servers are available, instead of the default behaviour of dropping the query
+    * `setPoolServerPolicy(policy, pool)`: set the server selection policy for this pool to that policy
+    * `setPoolServerPolicyLua(name, function, poool)`: set the server selection policy for this pool to one named 'name' and provided by 'function'
+    * `showPoolServerPolicy()`: show server selection policy for this pool
  * Available policies:
     * `firstAvailable`: Pick first server that has not exceeded its QPS limit, ordered by the server 'order' parameter
     * `whashed`: Weighted hashed ('sticky') distribution over available servers, based on the server 'weight' parameter

diff --git a/pdns/dnsdist-console.cc b/pdns/dnsdist-console.cc
index 40fc8d7bd55e5ec02cfeddac1d88b0cc7a518e62..e73b9541c665aed7b06eca015329caeab680e636 100644 (file)
--- a/pdns/dnsdist-console.cc
+++ b/pdns/dnsdist-console.cc
@@ -366,6 +366,8 @@ const std::vector<ConsoleKeyword> g_consoleKeywords{
   { "setMaxTCPQueriesPerConnection", true, "n", "set the maximum number of queries in an incoming TCP connection. 0 means unlimited" },
   { "setMaxTCPQueuedConnections", true, "n", "set the maximum number of TCP connections queued (waiting to be picked up by a client thread)" },
   { "setMaxUDPOutstanding", true, "n", "set the maximum number of outstanding UDP queries to a given backend server. This can only be set at configuration time and defaults to 10240" },
+  { "setPoolServerPolicy", true, "policy, pool", "set the server selection policy for this pool to that policy" },
+  { "setPoolServerPolicy", true, "name, func, pool", "set the server selection policy for this pool to one named 'name' and provided by 'function'" },
   { "setQueryCount", true, "bool", "set whether queries should be counted" },
   { "setQueryCountFilter", true, "func", "filter queries that would be counted, where `func` is a function with parameter `dq` which decides whether a query should and how it should be counted" },
   { "setRingBuffersSize", true, "n", "set the capacity of the ringbuffers used for live traffic inspection to `n`" },
@@ -383,6 +385,7 @@ const std::vector<ConsoleKeyword> g_consoleKeywords{
   { "showCacheHitResponseRules", true, "", "show all defined cache hit response rules" },
   { "showDNSCryptBinds", true, "", "display the currently configured DNSCrypt binds" },
   { "showDynBlocks", true, "", "show dynamic blocks in force" },
+  { "showPoolServerPolicy", true, "pool", "show server selection policy for this pool" },
   { "showResponseLatency", true, "", "show a plot of the response time latency distribution" },
   { "showResponseRules", true, "", "show all defined response rules" },
   { "showRules", true, "", "show all defined rules" },

diff --git a/pdns/dnsdist-lua2.cc b/pdns/dnsdist-lua2.cc
index 7fab8b942f93c729238ed85b8c94de5f14e335e1..e749325973b0c1b9e107698b9e9fbf389ed4df5f 100644 (file)
--- a/pdns/dnsdist-lua2.cc
+++ b/pdns/dnsdist-lua2.cc
@@ -1291,4 +1291,27 @@ void moreLua(bool client)
         }
 #endif /* HAVE_NET_SNMP */
       });
+
+    g_lua.writeFunction("setPoolServerPolicy", [](ServerPolicy policy, string pool) {
+        setLuaSideEffect();
+        auto localPools = g_pools.getCopy();
+        setPoolPolicy(localPools, pool, std::make_shared<ServerPolicy>(policy));
+      });
+
+    g_lua.writeFunction("setPoolServerPolicyLua", [](string name, policyfunc_t policy, string pool) {
+        setLuaSideEffect();
+        auto localPools = g_pools.getCopy();
+        setPoolPolicy(localPools, pool, std::make_shared<ServerPolicy>(ServerPolicy{name, policy}));
+      });
+
+    g_lua.writeFunction("showPoolServerPolicy", [](string pool) {
+        setLuaSideEffect();
+        auto localPools = g_pools.getCopy();
+        auto poolObj = getPool(localPools, pool);
+        if (poolObj->policy == NULL) {
+          g_outputBuffer=g_policy.getLocal()->name+"\n";
+        } else {
+          g_outputBuffer=poolObj->policy->name+"\n";
+        }
+      });
 }

diff --git a/pdns/dnsdist-tcp.cc b/pdns/dnsdist-tcp.cc
index 2e2eac64f5e1fbaec4872adf7b62400349187a3d..eadbd8cf965d002b9b030d4709d4ac06a099d836 100644 (file)
--- a/pdns/dnsdist-tcp.cc
+++ b/pdns/dnsdist-tcp.cc
@@ -364,11 +364,15 @@ void* tcpClientThread(int pipefd)
 
         std::shared_ptr<ServerPool> serverPool = getPool(*localPools, poolname);
         std::shared_ptr<DNSDistPacketCache> packetCache = nullptr;
-       {
-         std::lock_guard<std::mutex> lock(g_luamutex);
-         ds = localPolicy->policy(serverPool->servers, &dq);
-         packetCache = serverPool->packetCache;
-       }
+        auto policy = localPolicy->policy;
+        if (serverPool->policy != NULL) {
+          policy = serverPool->policy->policy;
+        }
+        {
+          std::lock_guard<std::mutex> lock(g_luamutex);
+          ds = policy(serverPool->servers, &dq);
+          packetCache = serverPool->packetCache;
+        }
 
         if (dq.useECS && ds && ds->useECS) {
           uint16_t newLen = dq.len;

diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index 32585f6a77713c7ae68a09d6b6e598751393849d..d8e58b7b04a424f493a90a9e9990d6df46cce914 100644 (file)
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -692,11 +692,23 @@ std::shared_ptr<ServerPool> createPoolIfNotExists(pools_t& pools, const string&
     if (!poolName.empty())
       vinfolog("Creating pool %s", poolName);
     pool = std::make_shared<ServerPool>();
+    pool->policy = NULL;
     pools.insert(std::pair<std::string,std::shared_ptr<ServerPool> >(poolName, pool));
   }
   return pool;
 }
 
+void setPoolPolicy(pools_t& pools, const string& poolName, std::shared_ptr<ServerPolicy> policy)
+{
+  std::shared_ptr<ServerPool> pool = createPoolIfNotExists(pools, poolName);
+  if (!poolName.empty()) {
+    vinfolog("Setting pool %s server selection policy to %s", poolName, policy->name);
+  } else {
+    vinfolog("Setting default pool server selection policy to %s", policy->name);
+  }
+  pool->policy = policy;
+}
+
 void addServerToPool(pools_t& pools, const string& poolName, std::shared_ptr<DownstreamState> server)
 {
   std::shared_ptr<ServerPool> pool = createPoolIfNotExists(pools, poolName);
@@ -1146,11 +1158,14 @@ try
       DownstreamState* ss = nullptr;
       std::shared_ptr<ServerPool> serverPool = getPool(*localPools, poolname);
       std::shared_ptr<DNSDistPacketCache> packetCache = nullptr;
-      auto policy=localPolicy->policy;
+      auto policy = localPolicy->policy;
+      if (serverPool->policy != NULL) {
+        policy = serverPool->policy->policy;
+      }
       {
-       std::lock_guard<std::mutex> lock(g_luamutex);
-       ss = policy(serverPool->servers, &dq).get();
-       packetCache = serverPool->packetCache;
+        std::lock_guard<std::mutex> lock(g_luamutex);
+        ss = policy(serverPool->servers, &dq).get();
+        packetCache = serverPool->packetCache;
       }
 
       bool ednsAdded = false;

diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index 9b222c09f07f3c8205b458392a32e52f053c5e11..180da4a71e268089aeeec0299b121c2766175dce 100644 (file)
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -616,8 +616,10 @@ struct ServerPool
 
   NumberedVector<shared_ptr<DownstreamState>> servers;
   std::shared_ptr<DNSDistPacketCache> packetCache{nullptr};
+  std::shared_ptr<ServerPolicy> policy;
 };
 using pools_t=map<std::string,std::shared_ptr<ServerPool>>;
+void setPoolPolicy(pools_t& pools, const string& poolName, std::shared_ptr<ServerPolicy> policy);
 void addServerToPool(pools_t& pools, const string& poolName, std::shared_ptr<DownstreamState> server);
 void removeServerFromPool(pools_t& pools, const string& poolName, std::shared_ptr<DownstreamState> server);