From: Oliver Chen <oliver.chen@nokia-sbell.com>
Date: Mon, 16 Jun 2025 00:52:24 +0000 (+0000)
Subject: dnsdist: add a count to track the number of query restarts
X-Git-Tag: dnsdist-2.0.0-beta1~13^2~3
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=487016c7419b4c38e7acace24672f983ad0c5f9d;p=thirdparty%2Fpdns.git

dnsdist: add a count to track the number of query restarts
---

diff --git a/pdns/dnsdistdist/dnsdist-idstate.hh b/pdns/dnsdistdist/dnsdist-idstate.hh
index a829d3589b..86db2dcbc2 100644
--- a/pdns/dnsdistdist/dnsdist-idstate.hh
+++ b/pdns/dnsdistdist/dnsdist-idstate.hh
@@ -168,6 +168,7 @@ struct InternalQueryState
   uint16_t cacheFlags{0}; // DNS flags as sent to the backend // 2
   uint16_t udpPayloadSize{0}; // Max UDP payload size from the query // 2
   dnsdist::Protocol protocol; // 1
+  uint8_t restartCount{0}; // 1
   bool ednsAdded{false};
   bool ecsAdded{false};
   bool skipCache{false};
diff --git a/pdns/dnsdistdist/dnsdist-lua-bindings-dnsquestion.cc b/pdns/dnsdistdist/dnsdist-lua-bindings-dnsquestion.cc
index 049b88a5c1..280fd82e40 100644
--- a/pdns/dnsdistdist/dnsdist-lua-bindings-dnsquestion.cc
+++ b/pdns/dnsdistdist/dnsdist-lua-bindings-dnsquestion.cc
@@ -656,6 +656,7 @@ void setupLuaBindingsDNSQuestion([[maybe_unused]] LuaContext& luaCtx)
     dnsResponse.asynchronous = true;
     dnsResponse.getMutableData() = *dnsResponse.ids.d_packet;
     dnsResponse.ids.d_proxyProtocolPayloadSize = 0;
+    dnsResponse.ids.restartCount++;
     auto query = dnsdist::getInternalQueryFromDQ(dnsResponse, false);
     return dnsdist::queueQueryResumptionEvent(std::move(query));
   });
@@ -667,5 +668,9 @@ void setupLuaBindingsDNSQuestion([[maybe_unused]] LuaContext& luaCtx)
   luaCtx.registerFunction<bool (DNSResponse::*)()>("getStaleCacheHit", [](DNSResponse& dnsResponse) {
     return dnsResponse.ids.staleCacheHit;
   });
+
+  luaCtx.registerFunction<uint8_t (DNSResponse::*)()>("getRestartCount", [](DNSResponse& dnsResponse) {
+    return dnsResponse.ids.restartCount;
+  });
 #endif /* DISABLE_NON_FFI_DQ_BINDINGS */
 }
diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
index 6feb6d9051..b818054b94 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
@@ -163,6 +163,7 @@ void dnsdist_ffi_dnsresponse_set_max_returned_ttl(dnsdist_ffi_dnsresponse_t* dr,
 void dnsdist_ffi_dnsresponse_clear_records_type(dnsdist_ffi_dnsresponse_t* dr, uint16_t qtype) __attribute__ ((visibility ("default")));
 bool dnsdist_ffi_dnsresponse_rebase(dnsdist_ffi_dnsresponse_t* dr, const char* initialName, size_t initialNameSize) __attribute__ ((visibility ("default")));
 bool dnsdist_ffi_dnsresponse_get_stale_cache_hit(const dnsdist_ffi_dnsresponse_t* dnsResponse) __attribute__ ((visibility ("default")));
+uint8_t dnsdist_ffi_dnsresponse_get_restart_count(const dnsdist_ffi_dnsresponse_t* dnsResponse) __attribute__ ((visibility ("default")));
 
 bool dnsdist_ffi_dnsquestion_set_async(dnsdist_ffi_dnsquestion_t* dq, uint16_t asyncID, uint16_t queryID, uint32_t timeoutMs) __attribute__ ((visibility ("default")));
 bool dnsdist_ffi_dnsresponse_set_async(dnsdist_ffi_dnsquestion_t* dq, uint16_t asyncID, uint16_t queryID, uint32_t timeoutMs) __attribute__ ((visibility ("default")));
diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.cc b/pdns/dnsdistdist/dnsdist-lua-ffi.cc
index d8b20f8c77..33ed5ea692 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi.cc
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi.cc
@@ -845,6 +845,11 @@ bool dnsdist_ffi_dnsresponse_get_stale_cache_hit(const dnsdist_ffi_dnsresponse_t
   return dnsResponse->dr->ids.staleCacheHit;
 }
 
+uint8_t dnsdist_ffi_dnsresponse_get_restart_count(const dnsdist_ffi_dnsresponse_t* dnsResponse)
+{
+  return dnsResponse->dr->ids.restartCount;
+}
+
 bool dnsdist_ffi_dnsquestion_set_async(dnsdist_ffi_dnsquestion_t* dq, uint16_t asyncID, uint16_t queryID, uint32_t timeoutMs)
 {
   try {
diff --git a/pdns/dnsdistdist/docs/reference/dq.rst b/pdns/dnsdistdist/docs/reference/dq.rst
index 08cb362330..9be06bbb59 100644
--- a/pdns/dnsdistdist/docs/reference/dq.rst
+++ b/pdns/dnsdistdist/docs/reference/dq.rst
@@ -450,6 +450,12 @@ DNSResponse object
 
     Get the indicator of whether the cache lookup hit a stale entry.
 
+  .. method:: DNSResponse:getRestartCount() -> uint8_t
+
+    .. versionadded:: 2.0.0
+
+    Get the current restart count, useful when the number of restart attempts is to be checked.
+
   .. method:: DNSResponse:editTTLs(func)
 
     The function ``func`` is invoked for every entry in the answer, authority and additional section.
diff --git a/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc b/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc
index 7504f0787c..b4b2333a3c 100644
--- a/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc
+++ b/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc
@@ -449,6 +449,10 @@ BOOST_AUTO_TEST_CASE(test_Response)
   {
     BOOST_CHECK_EQUAL(dnsdist_ffi_dnsresponse_get_stale_cache_hit(&lightDR), false);
   }
+
+  {
+    BOOST_CHECK_EQUAL(dnsdist_ffi_dnsresponse_get_restart_count(&lightDR), 0);
+  }
 }
 
 BOOST_AUTO_TEST_CASE(test_Server)