Merge pull request #10692 from n0tlu5/sulton/self-weighted

auth: added self weighted lua function

.github/actions/spell-check/expect.txt

@@ -1015,6 +1015,7 @@ pickchashed
 pickclosest
 pickhashed
 picknamehashed
+pickselfweighted
 pickrandom
 pickrandomsample
 pickwhashed

docs/lua-records/functions.rst

@@ -128,6 +128,32 @@ Record creation functions
 
   This function also works for CNAME or TXT records.
 
+.. function:: pickselfweighted(url, addresses[, options])
+
+  Selects an IP address from the supplied list, weighted according to the results of `isUp` checks. Each address is evaluated, and if its associated weight (from `isUp`) is greater than 0, it is considered for selection using a weighted hash based on `bestwho`. If no address is "up" the function defaults to a random selection.
+
+  :param string url: The health check url to retrieve.
+  :param addresses: A list of IP addresses to evaluate.
+  :param options: Table of options for this specific check, see below.
+
+  Various options can be set in the ``options`` parameter:
+
+  - ``selector``: used to pick the address(es) from the subset of available addresses of the selected set. Choices include 'pickclosest', 'random', 'hashed', 'all' (default 'random').
+  - ``backupSelector``: used to pick the address from all addresses if all addresses are down. Choices include 'pickclosest', 'random', 'hashed', 'all' (default 'random').
+  - ``source``: Source address to check from
+  - ``timeout``: Maximum time in seconds that you allow the check to take (default 2)
+  - ``stringmatch``: check ``url`` for this string, only declare 'up' if found
+  - ``useragent``: Set the HTTP "User-Agent" header in the requests. By default it is set to "PowerDNS Authoritative Server"
+  - ``byteslimit``: Limit the maximum download size to ``byteslimit`` bytes (default 0 meaning no limit).
+
+  An example of a list of address sets:
+
+  .. code-block:: lua
+
+    pickselfweighted("http://example.com/weight", { "192.0.2.20", "203.0.113.4", "203.0.113.2" })
+
+  This function is ideal for scenarios where candidates can self-determine their weights, while also providing fallback behavior when all addresses are down.
+
 .. function:: pickrandomsample(number, values)
 
   Returns N random values from the list supplied.

pdns/lua-record.cc

@@ -71,6 +71,8 @@ class IsUpOracle
     CheckState(time_t _lastAccess): lastAccess(_lastAccess) {}
     /* current status */
     std::atomic<bool> status{false};
+    /* current weight */
+    std::atomic<int> weight{0};
     /* first check ? */
     std::atomic<bool> first{true};
     /* last time the status was accessed */
@@ -83,9 +85,10 @@ class IsUpOracle
     d_checkerThreadStarted.clear();
   }
   ~IsUpOracle() = default;
-  bool isUp(const ComboAddress& remote, const opts_t& opts);
-  bool isUp(const ComboAddress& remote, const std::string& url, const opts_t& opts);
-  bool isUp(const CheckDesc& cd);
+  int isUp(const ComboAddress& remote, const opts_t& opts);
+  int isUp(const ComboAddress& remote, const std::string& url, const opts_t& opts);
+  //NOLINTNEXTLINE(readability-identifier-length)
+  int isUp(const CheckDesc& cd);
 
 private:
   void checkURL(const CheckDesc& cd, const bool status, const bool first = false)
@@ -128,14 +131,26 @@ class IsUpOracle
         throw std::runtime_error(boost::str(boost::format("unable to match content with `%s`") % cd.opts.at("stringmatch")));
       }
 
-      if(!status) {
-        g_log<<Logger::Info<<"LUA record monitoring declaring "<<remstring<<" UP for URL "<<cd.url<<"!"<<endl;
+      int weight = 0;
+      try {
+        weight = stoi(content);
+        if(!status) {
+          g_log<<Logger::Info<<"LUA record monitoring declaring "<<remstring<<" UP for URL "<<cd.url<<"!"<<" with WEIGHT "<<content<<"!"<<endl;
+        }
+      }
+      catch (const std::exception&) {
+        if(!status) {
+          g_log<<Logger::Info<<"LUA record monitoring declaring "<<remstring<<" UP for URL "<<cd.url<<"!"<<endl;
+        }
       }
+
+      setWeight(cd, weight);
       setUp(cd);
     }
     catch(std::exception& ne) {
       if(status || first)
         g_log<<Logger::Info<<"LUA record monitoring declaring "<<remstring<<" DOWN for URL "<<cd.url<<", error: "<<ne.what()<<endl;
+      setWeight(cd, 0);
       setDown(cd);
     }
   }
@@ -228,8 +243,16 @@ class IsUpOracle
     }
   }
 
+  //NOLINTNEXTLINE(readability-identifier-length)
+  void setWeight(const CheckDesc& cd, int weight){
+    auto statuses = d_statuses.write_lock();
+    auto& state = (*statuses)[cd];
+    state->weight = weight;
+  }
+
   void setDown(const ComboAddress& rem, const std::string& url=std::string(), const opts_t& opts = opts_t())
   {
+    //NOLINTNEXTLINE(readability-identifier-length)
     CheckDesc cd{rem, url, opts};
     setStatus(cd, false);
   }
@@ -252,7 +275,8 @@ class IsUpOracle
   }
 };
 
-bool IsUpOracle::isUp(const CheckDesc& cd)
+//NOLINTNEXTLINE(readability-identifier-length)
+int IsUpOracle::isUp(const CheckDesc& cd)
 {
   if (!d_checkerThreadStarted.test_and_set()) {
     d_checkerThread = std::make_unique<std::thread>([this] { return checkThread(); });
@@ -263,7 +287,10 @@ bool IsUpOracle::isUp(const CheckDesc& cd)
     auto iter = statuses->find(cd);
     if (iter != statuses->end()) {
       iter->second->lastAccess = now;
-      return iter->second->status;
+      if (iter->second->weight > 0) {
+        return iter->second->weight;
+      }
+      return static_cast<int>(iter->second->status);
     }
   }
   // try to parse options so we don't insert any malformed content
@@ -277,16 +304,16 @@ bool IsUpOracle::isUp(const CheckDesc& cd)
       (*statuses)[cd] = std::make_unique<CheckState>(now);
     }
   }
-  return false;
+  return 0;
 }
 
-bool IsUpOracle::isUp(const ComboAddress& remote, const opts_t& opts)
+int IsUpOracle::isUp(const ComboAddress& remote, const opts_t& opts)
 {
   CheckDesc cd{remote, "", opts};
   return isUp(cd);
 }
 
-bool IsUpOracle::isUp(const ComboAddress& remote, const std::string& url, const opts_t& opts)
+int IsUpOracle::isUp(const ComboAddress& remote, const std::string& url, const opts_t& opts)
 {
   CheckDesc cd{remote, url, opts};
   return isUp(cd);
@@ -1173,6 +1200,40 @@ static void setupLuaRecords(LuaContext& lua) // NOLINT(readability-function-cogn
       return pickRandom<string>(items);
     });
 
+  /*
+   * Based on the hash of `bestwho`, returns an IP address from the list
+   * supplied, weighted according to the results of isUp calls.
+   * @example pickselfweighted('http://example.com/weight', { "192.0.2.20", "203.0.113.4", "203.0.113.2" })
+   */
+  lua.writeFunction("pickselfweighted", [](const std::string& url,
+                                             const iplist_t& ips,
+                                             boost::optional<opts_t> options) {
+      vector< pair<int, ComboAddress> > items;
+      opts_t opts;
+      if(options) {
+        opts = *options;
+      }
+
+      items.reserve(ips.capacity());
+      bool available = false;
+
+      vector<ComboAddress> conv = convComboAddressList(ips);
+      for (auto& entry : conv) {
+        int weight = 0;
+        weight = g_up.isUp(entry, url, opts);
+        if(weight>0) {
+          available = true;
+        }
+        items.emplace_back(weight, entry);
+      }
+      if(available) {
+        return pickWeightedHashed<ComboAddress>(s_lua_record_ctx->bestwho, items).toString();
+      }
+
+      // All units down, apply backupSelector on all candidates
+      return pickWeightedRandom<ComboAddress>(items).toString();
+    });
+
   lua.writeFunction("pickrandomsample", [](int n, const iplist_t& ips) {
       vector<string> items = convStringList(ips);
 	  return pickRandomSample<string>(n, items);

regression-tests.auth-py/test_LuaRecords.py

@@ -27,6 +27,8 @@ def do_GET(self):
         self._set_headers()
         if self.path == '/ping.json':
             self.wfile.write(bytes('{"ping":"pong"}', 'utf-8'))
+        if self.path == '/weight.txt':
+            self.wfile.write(bytes('12', 'utf-8'))
         else:
             self.wfile.write(bytes("<html><body><h1>hi!</h1><h2>Programming in Lua !</h2></body></html>", "utf-8"))
 
@@ -78,6 +80,7 @@ class TestLuaRecords(AuthTest):
 randn-txt.example.org.       3600 IN LUA  TXT   "pickrandomsample( 2, {{ 'bob', 'alice', 'john' }} )"
 v6-bogus.rand.example.org.   3600 IN LUA  AAAA  "pickrandom({{'{prefix}.101', '{prefix}.102'}})"
 v6.rand.example.org.         3600 IN LUA  AAAA  "pickrandom({{ '2001:db8:a0b:12f0::1', 'fe80::2a1:9bff:fe9b:f268' }})"
+selfweighted.example.org.    3600 IN LUA  A     "pickselfweighted('http://selfweighted.example.org:8080/weight.txt',{{'{prefix}.101', '{prefix}.102'}})"
 closest.geo                  3600 IN LUA  A     "pickclosest({{ '1.1.1.2', '1.2.3.4' }})"
 empty.rand.example.org.      3600 IN LUA  A     "pickrandom()"
 timeout.example.org.         3600 IN LUA  A     "; local i = 0 ;  while i < 1000 do pickrandom() ; i = i + 1 end return '1.2.3.4'"
@@ -254,6 +257,24 @@ def testEmptyRandom(self):
         res = self.sendUDPQuery(query)
         self.assertRcodeEqual(res, dns.rcode.SERVFAIL)
 
+    def testSelfWeighted(self):
+        """
+        Test the selfweighted() function with a set of A records
+        """
+        expected = [dns.rrset.from_text('selfweighted.example.org.', 0, dns.rdataclass.IN, 'A',
+                                        '{prefix}.101'.format(prefix=self._PREFIX)),
+                    dns.rrset.from_text('selfweighted.example.org.', 0, dns.rdataclass.IN, 'A',
+                                        '{prefix}.102'.format(prefix=self._PREFIX))]
+        query = dns.message.make_query('selfweighted.example.org', 'A')
+        res = self.sendUDPQuery(query)
+
+        # wait for health checks to happen
+        time.sleep(3)
+
+        res = self.sendUDPQuery(query)
+        self.assertRcodeEqual(res, dns.rcode.NOERROR)
+        self.assertAnyRRsetInAnswer(res, expected)
+
     def testPickRandomSampleTxt(self):
         """
         Basic pickrandomsample() test with a set of TXT records