From 47237d5993d967dd43c307e6b710b0e833b23e25 Mon Sep 17 00:00:00 2001
From: Pieter Lexis <pieter.lexis@powerdns.com>
Date: Mon, 27 Oct 2025 20:15:30 +0100
Subject: [PATCH] feat(dnsdist): Add OT trace for each rule

---
 pdns/dnsdistdist/dnsdist.cc                   | 24 +++++++++++++++++--
 .../test_OpenTelemetryTracing.py              | 10 ++++----
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/pdns/dnsdistdist/dnsdist.cc b/pdns/dnsdistdist/dnsdist.cc
index d64c5d6fe0..787bd72d9e 100644
--- a/pdns/dnsdistdist/dnsdist.cc
+++ b/pdns/dnsdistdist/dnsdist.cc
@@ -466,7 +466,16 @@ bool applyRulesToResponse(const std::vector<dnsdist::rules::ResponseRuleAction>&
 
   DNSResponseAction::Action action = DNSResponseAction::Action::None;
   std::string ruleresult;
+  SpanID parentSpanID;
+  if (auto tracer = dnsResponse.ids.getTracer(); tracer != nullptr && dnsResponse.ids.tracingEnabled) {
+    parentSpanID = tracer->getLastSpanID();
+  }
+
   for (const auto& rrule : respRuleActions) {
+    pdns::trace::dnsdist::Tracer::Closer ruleCloser;
+    if (auto tracer = dnsResponse.ids.getTracer(); tracer != nullptr && dnsResponse.ids.tracingEnabled) {
+      ruleCloser = tracer->openSpan("Rule: " + rrule.d_name, parentSpanID);
+    }
     if (rrule.d_rule->matches(&dnsResponse)) {
       ++rrule.d_rule->d_matches;
       action = (*rrule.d_action)(&dnsResponse, &ruleresult);
@@ -522,7 +531,7 @@ bool processResponseAfterRules(PacketBuffer& response, DNSResponse& dnsResponse,
 {
   pdns::trace::dnsdist::Tracer::Closer closer;
   if (auto tracer = dnsResponse.ids.getTracer(); tracer != nullptr && dnsResponse.ids.tracingEnabled) {
-    closer = tracer->openSpan("processResponseAfterRules", tracer->getLastSpanID());
+    closer = tracer->openSpan("processResponseAfterRules", tracer->getLastSpanIDForName("processResponse"));
   }
   bool zeroScope = false;
   if (!fixUpResponse(response, dnsResponse.ids.qname, dnsResponse.ids.origFlags, dnsResponse.ids.ednsAdded, dnsResponse.ids.ecsAdded, dnsResponse.ids.useZeroScope ? &zeroScope : nullptr)) {
@@ -1031,11 +1040,21 @@ static bool applyRulesChainToQuery(const std::vector<dnsdist::rules::RuleAction>
   string ruleresult;
   bool drop = false;
 
+  SpanID parentSpanID;
+  if (auto tracer = dnsQuestion.ids.getTracer(); tracer != nullptr) {
+    parentSpanID = tracer->getLastSpanID();
+  }
+
   for (const auto& rule : rules) {
     if (!rule.d_rule->matches(&dnsQuestion)) {
       continue;
     }
 
+    pdns::trace::dnsdist::Tracer::Closer ruleCloser;
+    if (auto tracer = dnsQuestion.ids.getTracer(); tracer != nullptr) {
+      ruleCloser = tracer->openSpan("Rule: " + rule.d_name, parentSpanID);
+    }
+
     rule.d_rule->d_matches++;
     action = (*rule.d_action)(&dnsQuestion, &ruleresult);
     if (processRulesResult(action, dnsQuestion, ruleresult, drop)) {
@@ -1461,7 +1480,8 @@ static ServerPolicy::SelectedBackend selectBackendForOutgoingQuery(DNSQuestion&
 {
   pdns::trace::dnsdist::Tracer::Closer closer;
   if (auto tracer = dnsQuestion.ids.getTracer(); tracer != nullptr && dnsQuestion.ids.tracingEnabled) {
-    closer = tracer->openSpan("selectBackendForOutgoingQuery", tracer->getLastSpanID());
+    // Not exactly processQuery, but it works for now
+    closer = tracer->openSpan("selectBackendForOutgoingQuery", tracer->getLastSpanIDForName("processQuery"));
   }
 
   const auto& policy = serverPool.policy != nullptr ? *serverPool.policy : *dnsdist::configuration::getCurrentRuntimeConfiguration().d_lbPolicy;
diff --git a/regression-tests.dnsdist/test_OpenTelemetryTracing.py b/regression-tests.dnsdist/test_OpenTelemetryTracing.py
index a4176867f0..fc7a34bfa2 100644
--- a/regression-tests.dnsdist/test_OpenTelemetryTracing.py
+++ b/regression-tests.dnsdist/test_OpenTelemetryTracing.py
@@ -117,9 +117,11 @@ class DNSDistOpenTelemetryProtobufBaseTest(DNSDistOpenTelemetryProtobufTest):
         funcs = {
             "processQuery",
             "applyRulesToQuery",
+            "Rule: Enable tracing",
             "selectBackendForOutgoingQuery",
             "processResponse",
             "applyRulesToResponse",
+            "Rule: Do PB logging",
         }
 
         if useTCP:
@@ -195,8 +197,8 @@ newServer{address="127.0.0.1:%d"}
 rl = newRemoteLogger('127.0.0.1:%d')
 setOpenTelemetryTracing(true)
 
-addAction(AllRule(), SetTraceAction(true))
-addResponseAction(AllRule(), RemoteLogResponseAction(rl))
+addAction(AllRule(), SetTraceAction(true), {name="Enable tracing"})
+addResponseAction(AllRule(), RemoteLogResponseAction(rl), {name="Do PB logging"})
 """
 
     def testBasic(self):
@@ -260,8 +262,8 @@ newServer{address="127.0.0.1:%d"}
 rl = newRemoteLogger('127.0.0.1:%d')
 setOpenTelemetryTracing(true)
 
-addAction(AllRule(), SetTraceAction(true))
-addResponseAction(AllRule(), RemoteLogResponseAction(rl, nil, false, {}, {}, true))
+addAction(AllRule(), SetTraceAction(true), {name="Enable tracing"})
+addResponseAction(AllRule(), RemoteLogResponseAction(rl, nil, false, {}, {}, true), {name="Do PB logging"})
 """
 
     def testBasic(self):
-- 
2.47.3