From d3674927978c36afdcfba4acd008956ea8d387ea Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Tue, 16 Jul 2024 15:08:45 +0200
Subject: [PATCH] dnsdist: Only include the expensive LuaContext header when
 required

---
 pdns/dnsdistdist/dnsdist-backend.cc      |  2 ++
 pdns/dnsdistdist/dnsdist-console.cc      |  1 +
 pdns/dnsdistdist/dnsdist-dynblocks.hh    |  2 ++
 pdns/dnsdistdist/dnsdist-healthchecks.cc |  1 +
 pdns/dnsdistdist/dnsdist-lbpolicies.cc   | 26 +++++++++++++++---------
 pdns/dnsdistdist/dnsdist-lbpolicies.hh   |  9 ++------
 pdns/dnsdistdist/dnsdist-lua-ffi.hh      |  2 ++
 pdns/dnsdistdist/dnsdist-lua.hh          |  8 +++++++-
 pdns/dnsdistdist/dnsdist-rules.hh        |  1 +
 pdns/dnsdistdist/dnsdist-tcp.cc          |  1 +
 pdns/dnsdistdist/dnsdist-web.cc          |  1 +
 pdns/dnsdistdist/dnsdist.hh              |  3 ---
 pdns/dnsdistdist/test-dnsdistrules_cc.cc |  1 -
 13 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/pdns/dnsdistdist/dnsdist-backend.cc b/pdns/dnsdistdist/dnsdist-backend.cc
index eb5e3c6c60..11b00667b9 100644
--- a/pdns/dnsdistdist/dnsdist-backend.cc
+++ b/pdns/dnsdistdist/dnsdist-backend.cc
@@ -19,6 +19,8 @@
  * along with this program; if not, write to the Free Software
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  */
+#include <boost/format.hpp>
+
 #include "config.h"
 #include "dnsdist.hh"
 #include "dnsdist-backend.hh"
diff --git a/pdns/dnsdistdist/dnsdist-console.cc b/pdns/dnsdistdist/dnsdist-console.cc
index f5ed5c705f..2b382977fa 100644
--- a/pdns/dnsdistdist/dnsdist-console.cc
+++ b/pdns/dnsdistdist/dnsdist-console.cc
@@ -44,6 +44,7 @@
 #include "dnsdist.hh"
 #include "dnsdist-console.hh"
 #include "dnsdist-crypto.hh"
+#include "dnsdist-lua.hh"
 #include "threadname.hh"
 
 static LockGuarded<std::vector<pair<timeval, string>>> s_confDelta;
diff --git a/pdns/dnsdistdist/dnsdist-dynblocks.hh b/pdns/dnsdistdist/dnsdist-dynblocks.hh
index f700e16416..7c240f5fd2 100644
--- a/pdns/dnsdistdist/dnsdist-dynblocks.hh
+++ b/pdns/dnsdistdist/dnsdist-dynblocks.hh
@@ -33,6 +33,8 @@ extern "C"
 #include "dnsdist-lua-inspection-ffi.h"
 }
 
+#include "ext/luawrapper/include/LuaContext.hpp"
+
 // dnsdist_ffi_stat_node_t is a lightuserdata
 template <>
 struct LuaContext::Pusher<dnsdist_ffi_stat_node_t*>
diff --git a/pdns/dnsdistdist/dnsdist-healthchecks.cc b/pdns/dnsdistdist/dnsdist-healthchecks.cc
index 02c0803746..7f02b24f76 100644
--- a/pdns/dnsdistdist/dnsdist-healthchecks.cc
+++ b/pdns/dnsdistdist/dnsdist-healthchecks.cc
@@ -24,6 +24,7 @@
 #include "tcpiohandler-mplexer.hh"
 #include "dnswriter.hh"
 #include "dolog.hh"
+#include "dnsdist-lua.hh"
 #include "dnsdist-random.hh"
 #include "dnsdist-tcp.hh"
 #include "dnsdist-nghttp2.hh"
diff --git a/pdns/dnsdistdist/dnsdist-lbpolicies.cc b/pdns/dnsdistdist/dnsdist-lbpolicies.cc
index 6fb2c24b43..337a062669 100644
--- a/pdns/dnsdistdist/dnsdist-lbpolicies.cc
+++ b/pdns/dnsdistdist/dnsdist-lbpolicies.cc
@@ -338,24 +338,30 @@ ServerPolicy::ServerPolicy(const std::string& name_, const std::string& code): d
   auto ret = tmpContext.executeCode<ServerPolicy::ffipolicyfunc_t>(code);
 }
 
-thread_local ServerPolicy::PerThreadState ServerPolicy::t_perThreadState;
+struct ServerPolicy::PerThreadState
+{
+  LuaContext d_luaContext;
+  std::unordered_map<std::string, ffipolicyfunc_t> d_policies;
+};
+
+thread_local std::unique_ptr<ServerPolicy::PerThreadState> ServerPolicy::t_perThreadState;
 
 const ServerPolicy::ffipolicyfunc_t& ServerPolicy::getPerThreadPolicy() const
 {
   auto& state = t_perThreadState;
-  if (!state.d_initialized) {
-    setupLuaLoadBalancingContext(state.d_luaContext);
-    state.d_initialized = true;
+  if (!state) {
+    state = std::make_unique<ServerPolicy::PerThreadState>();
+    setupLuaLoadBalancingContext(state->d_luaContext);
   }
 
-  const auto& it = state.d_policies.find(d_name);
-  if (it != state.d_policies.end()) {
-    return it->second;
+  const auto& policyIt = state->d_policies.find(d_name);
+  if (policyIt != state->d_policies.end()) {
+    return policyIt->second;
   }
 
-  auto newPolicy = state.d_luaContext.executeCode<ServerPolicy::ffipolicyfunc_t>(d_perThreadPolicyCode);
-  state.d_policies[d_name] = std::move(newPolicy);
-  return state.d_policies.at(d_name);
+  auto newPolicy = state->d_luaContext.executeCode<ServerPolicy::ffipolicyfunc_t>(d_perThreadPolicyCode);
+  state->d_policies[d_name] = std::move(newPolicy);
+  return state->d_policies.at(d_name);
 }
 
 std::shared_ptr<DownstreamState> ServerPolicy::getSelectedBackend(const ServerPolicy::NumberedServerVector& servers, DNSQuestion& dq) const
diff --git a/pdns/dnsdistdist/dnsdist-lbpolicies.hh b/pdns/dnsdistdist/dnsdist-lbpolicies.hh
index ce95addab9..3ddc4ec592 100644
--- a/pdns/dnsdistdist/dnsdist-lbpolicies.hh
+++ b/pdns/dnsdistdist/dnsdist-lbpolicies.hh
@@ -69,15 +69,10 @@ public:
   }
 
 private:
-  struct PerThreadState
-  {
-    LuaContext d_luaContext;
-    std::unordered_map<std::string, ffipolicyfunc_t> d_policies;
-    bool d_initialized{false};
-  };
+  struct PerThreadState;
 
   const ffipolicyfunc_t& getPerThreadPolicy() const;
-  static thread_local PerThreadState t_perThreadState;
+  static thread_local std::unique_ptr<PerThreadState> t_perThreadState;
 
 public:
   std::string d_name;
diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.hh b/pdns/dnsdistdist/dnsdist-lua-ffi.hh
index a91003b971..d620c24554 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi.hh
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi.hh
@@ -27,6 +27,8 @@ extern "C" {
 #include "dnsdist-lua-ffi-interface.h"
 }
 
+#include "ext/luawrapper/include/LuaContext.hpp"
+
 // dnsdist_ffi_dnsquestion_t is a lightuserdata
 template<>
 struct LuaContext::Pusher<dnsdist_ffi_dnsquestion_t*> {
diff --git a/pdns/dnsdistdist/dnsdist-lua.hh b/pdns/dnsdistdist/dnsdist-lua.hh
index e8123a387c..e2be6dca1e 100644
--- a/pdns/dnsdistdist/dnsdist-lua.hh
+++ b/pdns/dnsdistdist/dnsdist-lua.hh
@@ -21,11 +21,17 @@
  */
 #pragma once
 
+#include <random>
+
 #include "dolog.hh"
 #include "dnsdist.hh"
 #include "dnsdist-dnsparser.hh"
 #include "dnsparser.hh"
-#include <random>
+
+#include "ext/luawrapper/include/LuaContext.hpp"
+
+extern RecursiveLockGuarded<LuaContext> g_lua;
+extern std::string g_outputBuffer; // locking for this is ok, as locked by g_luamutex
 
 class SpoofAction : public DNSAction
 {
diff --git a/pdns/dnsdistdist/dnsdist-rules.hh b/pdns/dnsdistdist/dnsdist-rules.hh
index ac136968a4..83ea08dbd7 100644
--- a/pdns/dnsdistdist/dnsdist-rules.hh
+++ b/pdns/dnsdistdist/dnsdist-rules.hh
@@ -30,6 +30,7 @@
 #include "dnsdist.hh"
 #include "dnsdist-ecs.hh"
 #include "dnsdist-kvs.hh"
+#include "dnsdist-lua.hh"
 #include "dnsdist-lua-ffi.hh"
 #include "dolog.hh"
 #include "dnsparser.hh"
diff --git a/pdns/dnsdistdist/dnsdist-tcp.cc b/pdns/dnsdistdist/dnsdist-tcp.cc
index e5f9fcdb0a..1e02000fe4 100644
--- a/pdns/dnsdistdist/dnsdist-tcp.cc
+++ b/pdns/dnsdistdist/dnsdist-tcp.cc
@@ -23,6 +23,7 @@
 #include <thread>
 #include <netinet/tcp.h>
 #include <queue>
+#include <boost/format.hpp>
 
 #include "dnsdist.hh"
 #include "dnsdist-concurrent-connections.hh"
diff --git a/pdns/dnsdistdist/dnsdist-web.cc b/pdns/dnsdistdist/dnsdist-web.cc
index 7790e755f0..d4f392920f 100644
--- a/pdns/dnsdistdist/dnsdist-web.cc
+++ b/pdns/dnsdistdist/dnsdist-web.cc
@@ -38,6 +38,7 @@
 #include "dnsdist-dynbpf.hh"
 #include "dnsdist-frontend.hh"
 #include "dnsdist-healthchecks.hh"
+#include "dnsdist-lua.hh"
 #include "dnsdist-metrics.hh"
 #include "dnsdist-prometheus.hh"
 #include "dnsdist-rings.hh"
diff --git a/pdns/dnsdistdist/dnsdist.hh b/pdns/dnsdistdist/dnsdist.hh
index 5754af6a70..e682172780 100644
--- a/pdns/dnsdistdist/dnsdist.hh
+++ b/pdns/dnsdistdist/dnsdist.hh
@@ -22,7 +22,6 @@
 #pragma once
 
 #include "config.h"
-#include "ext/luawrapper/include/LuaContext.hpp"
 
 #include <condition_variable>
 #include <memory>
@@ -890,8 +889,6 @@ public:
 };
 
 void responderThread(std::shared_ptr<DownstreamState> dss);
-extern RecursiveLockGuarded<LuaContext> g_lua;
-extern std::string g_outputBuffer; // locking for this is ok, as locked by g_luamutex
 
 class DNSDistPacketCache;
 
diff --git a/pdns/dnsdistdist/test-dnsdistrules_cc.cc b/pdns/dnsdistdist/test-dnsdistrules_cc.cc
index 0cda248317..7457fb3bb1 100644
--- a/pdns/dnsdistdist/test-dnsdistrules_cc.cc
+++ b/pdns/dnsdistdist/test-dnsdistrules_cc.cc
@@ -10,7 +10,6 @@
 
 #include "dnsdist-rules.hh"
 
-void checkParameterBound(const std::string& parameter, uint64_t value, size_t max);
 void checkParameterBound(const std::string& parameter, uint64_t value, size_t max)
 {
   if (value > max) {
-- 
2.47.2