From e605e1615001f125f0ef85208bd964c36ba06272 Mon Sep 17 00:00:00 2001
From: ph1 <ph1@slurpee2.phoned.ca>
Date: Sat, 9 Oct 2021 20:35:23 -0600
Subject: [PATCH] Rule for outstanding limits

---
 pdns/dnsdist-console.cc                 |  1 +
 pdns/dnsdist.hh                         |  1 +
 pdns/dnsdistdist/dnsdist-backend.cc     | 11 +++++++++++
 pdns/dnsdistdist/dnsdist-rules.hh       | 22 ++++++++++++++++++++++
 pdns/dnsdistdist/docs/rules-actions.rst | 13 +++++++++++++
 5 files changed, 48 insertions(+)

diff --git a/pdns/dnsdist-console.cc b/pdns/dnsdist-console.cc
index 72e9718913..e20018dc09 100644
--- a/pdns/dnsdist-console.cc
+++ b/pdns/dnsdist-console.cc
@@ -541,6 +541,7 @@ const std::vector<ConsoleKeyword> g_consoleKeywords{
   { "OrRule", true, "selectors", "Matches the traffic if one or more of the the selectors rules does match" },
   { "PoolAction", true, "poolname", "set the packet into the specified pool" },
   { "PoolAvailableRule", true, "poolname", "Check whether a pool has any servers available to handle queries" },
+  { "PoolOutstandingRule", true, "poolname, limit", "Check weather a pool has outstanding queries above limit" },
   { "printDNSCryptProviderFingerprint", true, "\"/path/to/providerPublic.key\"", "display the fingerprint of the provided resolver public key" },
   { "ProbaRule", true, "probability", "Matches queries with a given probability. 1.0 means always" },
   { "ProxyProtocolValueRule", true, "type [, value]", "matches queries with a specified Proxy Protocol TLV value of that type, optionally matching the content of the option as well" },
diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index 2f4d6480fc..d823c600de 100644
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -913,6 +913,7 @@ struct ServerPool
   std::shared_ptr<DNSDistPacketCache> packetCache{nullptr};
   std::shared_ptr<ServerPolicy> policy{nullptr};
 
+  size_t poolLoad();
   size_t countServers(bool upOnly);
   const std::shared_ptr<ServerPolicy::NumberedServerVector> getServers();
   void addServer(shared_ptr<DownstreamState>& server);
diff --git a/pdns/dnsdistdist/dnsdist-backend.cc b/pdns/dnsdistdist/dnsdist-backend.cc
index ceef5574b4..d72605794b 100644
--- a/pdns/dnsdistdist/dnsdist-backend.cc
+++ b/pdns/dnsdistdist/dnsdist-backend.cc
@@ -218,6 +218,17 @@ size_t ServerPool::countServers(bool upOnly)
   return count;
 }
 
+size_t ServerPool::poolLoad()
+{
+  size_t load = 0;
+  auto servers = d_servers.read_lock();
+  for (const auto& server : **servers) {
+    size_t serverOutstanding = std::get<1>(server)->outstanding.load();
+    load += serverOutstanding;
+  }
+  return load;
+}
+
 const std::shared_ptr<ServerPolicy::NumberedServerVector> ServerPool::getServers()
 {
   std::shared_ptr<ServerPolicy::NumberedServerVector> result;
diff --git a/pdns/dnsdistdist/dnsdist-rules.hh b/pdns/dnsdistdist/dnsdist-rules.hh
index 6e0e5d9232..dda702e435 100644
--- a/pdns/dnsdistdist/dnsdist-rules.hh
+++ b/pdns/dnsdistdist/dnsdist-rules.hh
@@ -1082,6 +1082,28 @@ private:
   std::string d_poolname;
 };
 
+class PoolOutstandingRule : public DNSRule
+{
+public:
+  PoolOutstandingRule(const std::string& poolname, const size_t limit) : d_pools(&g_pools), d_poolname(poolname), d_limit(limit)
+  {
+  }
+
+  bool matches(const DNSQuestion* dq) const override
+  {
+    return (getPool(*d_pools, d_poolname)->poolLoad()) > d_limit;
+  }
+
+  string toString() const override
+  {
+    return "pool '" + d_poolname + "' outstanding > " + std::to_string(d_limit);
+  }
+private:
+  mutable LocalStateHolder<pools_t> d_pools;
+  std::string d_poolname;
+  size_t d_limit;
+};
+
 class KeyValueStoreLookupRule: public DNSRule
 {
 public:
diff --git a/pdns/dnsdistdist/docs/rules-actions.rst b/pdns/dnsdistdist/docs/rules-actions.rst
index 3493d7f67c..e408e56ee4 100644
--- a/pdns/dnsdistdist/docs/rules-actions.rst
+++ b/pdns/dnsdistdist/docs/rules-actions.rst
@@ -761,6 +761,19 @@ These ``DNSRule``\ s be one of the following items:
 
   :param string poolname: Pool to check
 
+.. function:: PoolOutstandingRule(poolname, limit)
+
+  Check whether a pool has total outstanding queries above limit
+
+  .. code-block:: Lua
+
+    --- Send queries to spill over pool if default pool is under pressure
+    addAction(PoolOutstandingRule("", 5000), PoolAction("spillover"))
+
+  :param string poolname: Pool to check
+  :param int limit: Total outstanding limit
+
+
 Combining Rules
 ~~~~~~~~~~~~~~~
 
-- 
2.47.2