From 04f6e9692a0dd25455e8e52fa21e6c79ee6dd46c Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Fri, 6 Nov 2020 21:33:30 -0800
Subject: [PATCH] ircd::net::dns::resolver: Improve timeout cycle; avoid false
 positives under load.

---
 include/ircd/net/dns_resolver.h |  2 +
 ircd/net_dns_resolver.cc        | 65 +++++++++++++++++++++++++--------
 2 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/include/ircd/net/dns_resolver.h b/include/ircd/net/dns_resolver.h
index 36c6d4221..76a7ae170 100644
--- a/include/ircd/net/dns_resolver.h
+++ b/include/ircd/net/dns_resolver.h
@@ -46,6 +46,7 @@ struct ircd::net::dns::resolver
 	steady_point send_last;                      // Time of last send
 	std::deque<uint16_t> sendq;                  // Queue of frames for rate-limiting
 	ip::udp::socket ns;                          // A pollable activity object
+	bool recv_idle {false};                      // Timeout worker won't run if false
 
 	// util
 	void add_server(const ipport &);
@@ -67,6 +68,7 @@ struct ircd::net::dns::resolver
 	void handle_reply(const header &, const const_buffer &body, tag &);
 	void handle_reply(const ipport &, const header &, const const_buffer &body);
 	void handle(const ipport &, const mutable_buffer &);
+	void handle_interrupt(ctx::ctx *const &) noexcept;
 	std::tuple<net::ipport, mutable_buffer> recv_recv(const mutable_buffer &);
 	void recv_worker();
 	ctx::context recv_context;
diff --git a/ircd/net_dns_resolver.cc b/ircd/net_dns_resolver.cc
index 37bb00be2..3196cdebb 100644
--- a/ircd/net_dns_resolver.cc
+++ b/ircd/net_dns_resolver.cc
@@ -286,9 +286,12 @@ ircd::net::dns::resolver::timeout_worker()
 {
 	while(1)
 	{
+		// Dock here until somebody submits a request into the tag map. Also
+		// wait until recv_idle is asserted which indicates the UDP queue has
+		// been exhausted.
 		dock.wait([this]
 		{
-			return !tags.empty();
+			return !tags.empty() && recv_idle;
 		});
 
 		check_timeouts(milliseconds(timeout));
@@ -526,37 +529,69 @@ catch(const std::exception &e)
 std::tuple<ircd::net::ipport, ircd::mutable_buffer>
 ircd::net::dns::resolver::recv_recv(const mutable_buffer &buf)
 {
+	static const ip::udp::socket::message_flags flags
+	{
+		0
+	};
+
 	const asio::mutable_buffers_1 bufs
 	{
 		buf
 	};
 
-	const auto interruption{[this](ctx::ctx *const &)
-	{
-		if(this->ns.is_open())
-			this->ns.cancel();
-	}};
-
+	// First try a non-blocking receive to find and return anything in the
+	// queue. If this comes back as -EAGAIN we'll assert recv_idle and then
+	// conduct the normal blocking receive.
+	boost::system::error_code ec;
 	ip::udp::endpoint ep;
-	size_t recv; continuation
+	size_t recv
 	{
-		continuation::asio_predicate, interruption, [this, &bufs, &recv, &ep]
-		(auto &yield)
-		{
-			recv = ns.async_receive_from(bufs, ep, yield);
-		}
+		ns.receive_from(bufs, ep, flags, ec)
 	};
 
+	assert(!ec || recv == 0);
+	assert(!ec || ec == boost::system::errc::resource_unavailable_try_again);
+
+	// branch on any ec, not just -EAGAIN; this time it can throw...
+	if(likely(ec))
+	{
+		const scope_restore recv_idle
+		{
+			this->recv_idle, true
+		};
+
+		const auto interruption
+		{
+			std::bind(&resolver::handle_interrupt, this, ph::_1)
+		};
+
+		continuation
+		{
+			continuation::asio_predicate, interruption, [this, &bufs, &recv, &ep]
+			(auto &yield)
+			{
+				recv = ns.async_receive_from(bufs, ep, yield);
+			}
+		};
+	}
+
 	return
 	{
-		make_ipport(ep),
-		mutable_buffer
+		make_ipport(ep), mutable_buffer
 		{
 			data(buf), recv
 		}
 	};
 }
 
+void
+ircd::net::dns::resolver::handle_interrupt(ctx::ctx *const &interruptor)
+noexcept
+{
+	if(!ns.is_open())
+		ns.cancel();
+}
+
 void
 ircd::net::dns::resolver::handle(const ipport &from,
                                  const mutable_buffer &buf)