diff --git a/.github/actions/spell-check/expect.txt b/.github/actions/spell-check/expect.txt index 51742aa7028d..289f683d0fa1 100644 --- a/.github/actions/spell-check/expect.txt +++ b/.github/actions/spell-check/expect.txt @@ -1015,6 +1015,7 @@ pickchashed pickclosest pickhashed picknamehashed +pickselfweighted pickrandom pickrandomsample pickwhashed diff --git a/docs/lua-records/functions.rst b/docs/lua-records/functions.rst index f6bde0bdf759..2bd42a3d5bdd 100644 --- a/docs/lua-records/functions.rst +++ b/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. diff --git a/pdns/lua-record.cc b/pdns/lua-record.cc index 71f55b47aa13..b64482464f1e 100644 --- a/pdns/lua-record.cc +++ b/pdns/lua-record.cc @@ -71,6 +71,8 @@ class IsUpOracle CheckState(time_t _lastAccess): lastAccess(_lastAccess) {} /* current status */ std::atomic status{false}; + /* current weight */ + std::atomic weight{0}; /* first check ? */ std::atomic 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<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([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(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(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(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 options) { + vector< pair > items; + opts_t opts; + if(options) { + opts = *options; + } + + items.reserve(ips.capacity()); + bool available = false; + + vector 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(s_lua_record_ctx->bestwho, items).toString(); + } + + // All units down, apply backupSelector on all candidates + return pickWeightedRandom(items).toString(); + }); + lua.writeFunction("pickrandomsample", [](int n, const iplist_t& ips) { vector items = convStringList(ips); return pickRandomSample(n, items); diff --git a/regression-tests.auth-py/test_LuaRecords.py b/regression-tests.auth-py/test_LuaRecords.py index dc4664584534..04a105ce93de 100644 --- a/regression-tests.auth-py/test_LuaRecords.py +++ b/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("

hi!

Programming in Lua !

", "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