feat: add manual_keep and manual_drop endpoints for parametric test (#271)

* feat: add manual_keep and manual_drop endpoints for parametric test

* chore: fmt

* chore: fmt 2

* fix: missing includes

Author: genesor

test/system-tests/main.cpp

@@ -120,6 +120,14 @@ int main(int argc, char* argv[]) {
            [&handler](const httplib::Request& req, httplib::Response& res) {
              handler.on_stats_flush(req, res);
            });
+  svr.Post("/trace/span/manual_drop",
+           [&handler](const httplib::Request& req, httplib::Response& res) {
+             handler.on_manual_drop(req, res);
+           });
+  svr.Post("/trace/span/manual_keep",
+           [&handler](const httplib::Request& req, httplib::Response& res) {
+             handler.on_manual_keep(req, res);
+           });
 
   // Not implemented
   svr.Post("/trace/span/set_metric",

test/system-tests/request_handler.cpp

@@ -1,7 +1,9 @@
 #include "request_handler.h"
 
 #include <datadog/optional.h>
+#include <datadog/sampling_priority.h>
 #include <datadog/span_config.h>
+#include <datadog/trace_segment.h>
 #include <datadog/tracer.h>
 #include <datadog/tracer_config.h>
 
@@ -216,19 +218,19 @@ void RequestHandler::on_set_meta(const httplib::Request& req,
   res.status = 200;
 }
 
-void RequestHandler::on_set_metric(const httplib::Request& /* req */,
+void RequestHandler::on_set_metric(const httplib::Request& req,
                                    httplib::Response& res) {
-  const auto request_json = nlohmann::json::parse(res.body);
+  const auto request_json = nlohmann::json::parse(req.body);
 
   auto span_id = utils::get_if_exists<uint64_t>(request_json, "span_id");
   if (!span_id) {
-    VALIDATION_ERROR(res, "on_set_meta: missing `span_id` field.");
+    VALIDATION_ERROR(res, "on_set_metric: missing `span_id` field.");
   }
 
   auto span_it = active_spans_.find(*span_id);
   if (span_it == active_spans_.cend()) {
     const auto msg =
-        "on_set_meta: span not found for id " + std::to_string(*span_id);
+        "on_set_metric: span not found for id " + std::to_string(*span_id);
     VALIDATION_ERROR(res, msg);
   }
 
@@ -239,6 +241,52 @@ void RequestHandler::on_set_metric(const httplib::Request& /* req */,
   res.status = 200;
 }
 
+void RequestHandler::on_manual_keep(const httplib::Request& req,
+                                    httplib::Response& res) {
+  const auto request_json = nlohmann::json::parse(req.body);
+
+  auto span_id = utils::get_if_exists<uint64_t>(request_json, "span_id");
+  if (!span_id) {
+    VALIDATION_ERROR(res, "on_manual_keep: missing `span_id` field.");
+  }
+
+  auto span_it = active_spans_.find(*span_id);
+  if (span_it == active_spans_.cend()) {
+    const auto msg =
+        "on_manual_keep: span not found for id " + std::to_string(*span_id);
+    VALIDATION_ERROR(res, msg);
+  }
+
+  auto& span = span_it->second;
+  span.trace_segment().override_sampling_priority(
+      datadog::tracing::SamplingPriority::USER_KEEP);
+
+  res.status = 200;
+}
+
+void RequestHandler::on_manual_drop(const httplib::Request& req,
+                                    httplib::Response& res) {
+  const auto request_json = nlohmann::json::parse(req.body);
+
+  auto span_id = utils::get_if_exists<uint64_t>(request_json, "span_id");
+  if (!span_id) {
+    VALIDATION_ERROR(res, "on_manual_drop: missing `span_id` field.");
+  }
+
+  auto span_it = active_spans_.find(*span_id);
+  if (span_it == active_spans_.cend()) {
+    const auto msg =
+        "on_manual_drop: span not found for id " + std::to_string(*span_id);
+    VALIDATION_ERROR(res, msg);
+  }
+
+  auto& span = span_it->second;
+  span.trace_segment().override_sampling_priority(
+      datadog::tracing::SamplingPriority::USER_DROP);
+
+  res.status = 200;
+}
+
 void RequestHandler::on_inject_headers(const httplib::Request& req,
                                        httplib::Response& res) {
   const auto request_json = nlohmann::json::parse(req.body);

test/system-tests/request_handler.h

@@ -25,6 +25,8 @@ class RequestHandler final {
   void on_span_end(const httplib::Request& req, httplib::Response& res);
   void on_set_meta(const httplib::Request& req, httplib::Response& res);
   void on_set_metric(const httplib::Request& /* req */, httplib::Response& res);
+  void on_manual_keep(const httplib::Request& req, httplib::Response& res);
+  void on_manual_drop(const httplib::Request& req, httplib::Response& res);
   void on_inject_headers(const httplib::Request& req, httplib::Response& res);
   void on_extract_headers(const httplib::Request& req, httplib::Response& res);
   void on_span_flush(const httplib::Request& /* req */, httplib::Response& res);