feat(dnsdist): Add infra to update Root SpanID and TraceID in Tracer

diff --git a/pdns/dnsdistdist/dnsdist-opentelemetry.cc b/pdns/dnsdistdist/dnsdist-opentelemetry.cc
index 81afbfa7227c53afbb0bd768ff880fc75da64530..f064f7c37f3b0d4b8e21615e59069a1cdd6966fd 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-opentelemetry.cc
+++ b/pdns/dnsdistdist/dnsdist-opentelemetry.cc
@@ -114,6 +114,30 @@ SpanID Tracer::addSpan([[maybe_unused]] const std::string& name, [[maybe_unused]
 #endif
 }
 
+void Tracer::setTraceID([[maybe_unused]] const TraceID& traceID)
+{
+#ifndef DISABLE_PROTOBUF
+  *d_traceid.lock() = traceID;
+#endif
+}
+
+void Tracer::setRootSpanID([[maybe_unused]] const SpanID& spanID)
+{
+#ifndef DISABLE_PROTOBUF
+  // XXX: We just assume that the first Span is the RootSpan
+  SpanID oldRootSpanID;
+  if (auto lockedSpans = d_spans.lock(); !lockedSpans->empty()) {
+    oldRootSpanID = lockedSpans->begin()->span_id;
+    lockedSpans->begin()->span_id = spanID;
+    for (auto& span : *lockedSpans) {
+      if (span.parent_span_id == oldRootSpanID) {
+        span.parent_span_id = spanID;
+      }
+    }
+  }
+#endif
+}
+
 // TODO: Figure out what to do with duplicate keys
 bool Tracer::setTraceAttribute([[maybe_unused]] const std::string& key, [[maybe_unused]] const AnyValue& value)
 {

diff --git a/pdns/dnsdistdist/dnsdist-opentelemetry.hh b/pdns/dnsdistdist/dnsdist-opentelemetry.hh
index cd7436d703b933050569c2cc78f49b7f9fbad473..95344c0c0aecb63e29c12ba1957ed9b9c732bda0 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-opentelemetry.hh
+++ b/pdns/dnsdistdist/dnsdist-opentelemetry.hh
@@ -79,6 +79,20 @@ public:
     return std::shared_ptr<Tracer>(new Tracer);
   }
 
+  /**
+   * @brief Set the TraceID
+   *
+   * @param traceID
+   */
+  void setTraceID(const TraceID& traceID);
+
+  /**
+   * @brief Set the SpanID for the root and re-parent
+   *
+   * @param spanID
+   */
+  void setRootSpanID(const SpanID& spanID);
+
   /**
    * @brief Add an attribute to the Trace
    *

diff --git a/pdns/dnsdistdist/test-dnsdist-opentelemetry_cc.cc b/pdns/dnsdistdist/test-dnsdist-opentelemetry_cc.cc
index 5449f61f22951e0c3e070990b4e19dbadc75ccc2..c6afa31729e04bc52df4334ae00681e6f3ca23d4 100644 (file)
--- a/pdns/dnsdistdist/test-dnsdist-opentelemetry_cc.cc
+++ b/pdns/dnsdistdist/test-dnsdist-opentelemetry_cc.cc
@@ -190,5 +190,65 @@ BOOST_AUTO_TEST_CASE(getOTProtobuf)
   BOOST_TEST(data.size() >= 110U);
 }
 
+BOOST_AUTO_TEST_CASE(setTraceID)
+{
+  auto tracer = pdns::trace::dnsdist::Tracer::getTracer();
+  auto oldTraceID = tracer->getTraceID();
+  auto newTraceID = pdns::trace::TraceID::getRandomTraceID();
+  while (oldTraceID == newTraceID) {
+    newTraceID.makeRandom();
+  }
+  tracer->setTraceID(newTraceID);
+  BOOST_CHECK_EQUAL(tracer->getTraceID(), newTraceID);
+}
+
+BOOST_AUTO_TEST_CASE(setRootSpanID)
+{
+  auto tracer = pdns::trace::dnsdist::Tracer::getTracer();
+  // Setting the root span id without any spans should not error
+  auto newRootSpanID = pdns::trace::SpanID::getRandomSpanID();
+  tracer->setRootSpanID(newRootSpanID);
+  BOOST_CHECK_EQUAL(pdns::trace::s_emptySpanID, tracer->getRootSpanID());
+
+  // Add a Span, this'll be the root span
+  auto originalRootSpanID = tracer->openSpan("rootspan").getSpanID();
+  BOOST_CHECK_EQUAL(originalRootSpanID, tracer->getRootSpanID());
+
+  // Add 4 spans, all children of the root span
+  for (auto i = 0; i < 4; i++) {
+    tracer->openSpan("span " + std::to_string(i), originalRootSpanID);
+  }
+
+  // Now set the new root span ID
+  while (newRootSpanID == originalRootSpanID) {
+    newRootSpanID.makeRandom();
+  }
+  tracer->setRootSpanID(newRootSpanID);
+  BOOST_CHECK_EQUAL(tracer->getRootSpanID(), newRootSpanID);
+
+  // Verify the parent_span_id is updated
+  auto data = tracer->getTracesData();
+  for (auto i = 0; i < 4; i++) {
+    BOOST_CHECK_EQUAL(data.resource_spans.at(0).scope_spans.at(0).spans.at(i + 1).parent_span_id, newRootSpanID);
+  }
+
+  // New tracer, so we can easily test if non-root parent span IDs are not updated
+  tracer = pdns::trace::dnsdist::Tracer::getTracer();
+  // Add a Span, this'll be the root span
+  originalRootSpanID = tracer->openSpan("rootspan").getSpanID();
+  BOOST_CHECK_EQUAL(originalRootSpanID, tracer->getRootSpanID());
+
+  tracer->openSpan("span one", pdns::trace::SpanID::getRandomSpanID());
+  tracer->openSpan("span two", pdns::trace::SpanID::getRandomSpanID());
+  while (newRootSpanID == originalRootSpanID) {
+    newRootSpanID.makeRandom();
+  }
+  tracer->setRootSpanID(newRootSpanID);
+  data = tracer->getTracesData();
+  for (auto i = 0; i < 2; i++) {
+    BOOST_CHECK_NE(data.resource_spans.at(0).scope_spans.at(0).spans.at(i + 1).parent_span_id, newRootSpanID);
+  }
+}
+
 BOOST_AUTO_TEST_SUITE_END()
 #endif // DISABLE_PROTOBUF