implement parent-based & ratio-based samplers

Signed-off-by: Ahmad Karimi <[email protected]>

api/envoy/extensions/tracers/opentelemetry/samplers/v3/parent_based_sampler.proto

@@ -0,0 +1,31 @@
+syntax = "proto3";
+
+package envoy.extensions.tracers.opentelemetry.samplers.v3;
+
+import "envoy/config/core/v3/extension.proto";
+
+import "udpa/annotations/status.proto";
+
+option java_package = "io.envoyproxy.envoy.extensions.tracers.opentelemetry.samplers.v3";
+option java_outer_classname = "ParentBasedSamplerProto";
+option java_multiple_files = true;
+option go_package = "github.com/envoyproxy/go-control-plane/envoy/extensions/tracers/opentelemetry/samplers/v3;samplersv3";
+option (udpa.annotations.file_status).package_version_status = ACTIVE;
+
+// [#protodoc-title: Parent Based Sampler config]
+// Configuration for the "ParentBased" Sampler extension.
+// The sampler follows the "ParentBased" implementation from the OpenTelemetry
+// SDK specification.
+//
+// See:
+// `ParentBased sampler specification <https://opentelemetry.io/docs/specs/otel/trace/sdk/#parentbased>`_
+// [#extension: envoy.tracers.opentelemetry.samplers.parent_based]
+
+message ParentBasedSamplerConfig {
+  // Specifies the sampler to be used by this sampler.
+  // The configured sampler will be used if the parent trace ID is not passed to Envoy
+  //
+  // required
+  // [#extension-category: envoy.tracers.opentelemetry.samplers]
+  config.core.v3.TypedExtensionConfig wrapped_sampler = 1;
+}

api/envoy/extensions/tracers/opentelemetry/samplers/v3/trace_id_ratio_based_sampler.proto

@@ -0,0 +1,29 @@
+syntax = "proto3";
+
+package envoy.extensions.tracers.opentelemetry.samplers.v3;
+
+import "udpa/annotations/status.proto";
+
+option java_package = "io.envoyproxy.envoy.extensions.tracers.opentelemetry.samplers.v3";
+option java_outer_classname = "TraceIdRatioBasedSamplerProto";
+option java_multiple_files = true;
+option go_package = "github.com/envoyproxy/go-control-plane/envoy/extensions/tracers/opentelemetry/samplers/v3;samplersv3";
+option (udpa.annotations.file_status).package_version_status = ACTIVE;
+
+// [#protodoc-title: Trace Id Ratio Based Sampler config]
+// Configuration for the "TraceIdRatioBased" Sampler extension.
+// The sampler follows the "TraceIdRatioBased" implementation from the OpenTelemetry
+// SDK specification.
+//
+// See:
+// `TraceIdRatioBased sampler specification <https://opentelemetry.io/docs/specs/otel/trace/sdk/#traceidratiobased>`_
+// [#extension: envoy.tracers.opentelemetry.samplers.trace_id_ratio_based]
+
+message TraceIdRatioBasedSamplerConfig {
+  // This is a required value in the [0, 1] interval, If the given trace_id
+  // falls into a given ratio of all possible trace_id values, ShouldSample will
+  // return RECORD_AND_SAMPLE.
+  // required
+  // [#extension-category: envoy.tracers.opentelemetry.samplers]
+  double ratio = 1;
+}

bazel/repository_locations.bzl

@@ -600,6 +600,8 @@ REPOSITORY_LOCATIONS_SPEC = dict(
             "envoy.tracers.opentelemetry",
             "envoy.tracers.opentelemetry.samplers.always_on",
             "envoy.tracers.opentelemetry.samplers.dynatrace",
+            "envoy.tracers.opentelemetry.samplers.parent_based",
+            "envoy.tracers.opentelemetry.samplers.trace_id_ratio_based",
         ],
         release_date = "2024-10-07",
         cpe = "N/A",

source/extensions/extensions_build_config.bzl

@@ -294,8 +294,10 @@ EXTENSIONS = {
     # OpenTelemetry tracer samplers
     #
 
-    "envoy.tracers.opentelemetry.samplers.always_on":         "//source/extensions/tracers/opentelemetry/samplers/always_on:config",
-    "envoy.tracers.opentelemetry.samplers.dynatrace":         "//source/extensions/tracers/opentelemetry/samplers/dynatrace:config",
+    "envoy.tracers.opentelemetry.samplers.always_on":                     "//source/extensions/tracers/opentelemetry/samplers/always_on:config",
+    "envoy.tracers.opentelemetry.samplers.dynatrace":                     "//source/extensions/tracers/opentelemetry/samplers/dynatrace:config",
+    "envoy.tracers.opentelemetry.samplers.parent_based":                  "//source/extensions/tracers/opentelemetry/samplers/parent_based:config",
+    "envoy.tracers.opentelemetry.samplers.trace_id_ratio_based":          "//source/extensions/tracers/opentelemetry/samplers/trace_id_ratio_based:config",
 
     #
     # Transport sockets

source/extensions/extensions_metadata.yaml

@@ -1236,6 +1236,20 @@ envoy.tracers.opentelemetry.samplers.always_on:
   status: wip
   type_urls:
   - envoy.extensions.tracers.opentelemetry.samplers.v3.AlwaysOnSamplerConfig
+envoy.tracers.opentelemetry.samplers.parent_based:
+  categories:
+  - envoy.tracers.opentelemetry.samplers
+  security_posture: unknown
+  status: wip
+  type_urls:
+  - envoy.extensions.tracers.opentelemetry.samplers.v3.ParentBasedSamplerConfig
+envoy.tracers.opentelemetry.samplers.trace_id_ratio_based:
+  categories:
+  - envoy.tracers.opentelemetry.samplers
+  security_posture: unknown
+  status: wip
+  type_urls:
+  - envoy.extensions.tracers.opentelemetry.samplers.v3.TraceIdRatioBasedSamplerConfig
 envoy.tracers.opentelemetry.samplers.dynatrace:
   categories:
   - envoy.tracers.opentelemetry.samplers

source/extensions/tracers/opentelemetry/samplers/parent_based/BUILD

@@ -0,0 +1,34 @@
+load(
+    "//bazel:envoy_build_system.bzl",
+    "envoy_cc_extension",
+    "envoy_cc_library",
+    "envoy_extension_package",
+)
+
+licenses(["notice"])  # Apache 2
+
+envoy_extension_package()
+
+envoy_cc_extension(
+    name = "config",
+    srcs = ["config.cc"],
+    hdrs = ["config.h"],
+    deps = [
+        ":parent_based_sampler_lib",
+        "//envoy/registry",
+        "//source/common/config:utility_lib",
+        "//source/extensions/tracers/opentelemetry/samplers/trace_id_ratio_based:config",
+        "@envoy_api//envoy/extensions/tracers/opentelemetry/samplers/v3:pkg_cc_proto",
+    ],
+)
+
+envoy_cc_library(
+    name = "parent_based_sampler_lib",
+    srcs = ["parent_based_sampler.cc"],
+    hdrs = ["parent_based_sampler.h"],
+    deps = [
+        "//source/extensions/tracers/opentelemetry:opentelemetry_tracer_lib",
+        "//source/extensions/tracers/opentelemetry/samplers:sampler_lib",
+        "@envoy_api//envoy/extensions/tracers/opentelemetry/samplers/v3:pkg_cc_proto",
+    ],
+)

source/extensions/tracers/opentelemetry/samplers/parent_based/config.cc

@@ -0,0 +1,38 @@
+#include "source/extensions/tracers/opentelemetry/samplers/parent_based/config.h"
+
+#include "envoy/extensions/tracers/opentelemetry/samplers/v3/parent_based_sampler.pb.validate.h"
+#include "envoy/server/tracer_config.h"
+
+#include "source/common/config/utility.h"
+#include "source/extensions/tracers/opentelemetry/samplers/parent_based/parent_based_sampler.h"
+#include "source/extensions/tracers/opentelemetry/samplers/trace_id_ratio_based/config.h"
+
+namespace Envoy {
+namespace Extensions {
+namespace Tracers {
+namespace OpenTelemetry {
+
+SamplerSharedPtr
+ParentBasedSamplerFactory::createSampler(const Protobuf::Message& config,
+                                         Server::Configuration::TracerFactoryContext& context) {
+  auto mptr = Envoy::Config::Utility::translateAnyToFactoryConfig(
+      dynamic_cast<const ProtobufWkt::Any&>(config), context.messageValidationVisitor(), *this);
+  const auto& proto_config = MessageUtil::downcastAndValidate<
+      const envoy::extensions::tracers::opentelemetry::samplers::v3::ParentBasedSamplerConfig&>(
+      *mptr, context.messageValidationVisitor());
+  auto* factory =
+      Envoy::Config::Utility::getFactory<SamplerFactory>(proto_config.wrapped_sampler());
+  SamplerSharedPtr wrapped_sampler =
+      factory->createSampler(proto_config.wrapped_sampler().typed_config(), context);
+  return std::make_shared<ParentBasedSampler>(config, context, wrapped_sampler);
+}
+
+/**
+ * Static registration for the Env sampler factory. @see RegisterFactory.
+ */
+REGISTER_FACTORY(ParentBasedSamplerFactory, SamplerFactory);
+
+} // namespace OpenTelemetry
+} // namespace Tracers
+} // namespace Extensions
+} // namespace Envoy

source/extensions/tracers/opentelemetry/samplers/parent_based/config.h

@@ -0,0 +1,42 @@
+#pragma once
+
+#include <string>
+
+#include "envoy/extensions/tracers/opentelemetry/samplers/v3/parent_based_sampler.pb.h"
+#include "envoy/registry/registry.h"
+
+#include "source/extensions/tracers/opentelemetry/samplers/sampler.h"
+
+namespace Envoy {
+namespace Extensions {
+namespace Tracers {
+namespace OpenTelemetry {
+
+/**
+ * Config registration for the ParentBasedSampler. @see SamplerFactory.
+ */
+class ParentBasedSamplerFactory : public SamplerFactory {
+public:
+  /**
+   * @brief Create a Sampler. @see ParentBasedSampler
+   *
+   * @param config Protobuf config for the sampler.
+   * @param context A reference to the TracerFactoryContext.
+   * @return SamplerSharedPtr
+   */
+  SamplerSharedPtr createSampler(const Protobuf::Message& config,
+                                 Server::Configuration::TracerFactoryContext& context) override;
+
+  ProtobufTypes::MessagePtr createEmptyConfigProto() override {
+    return std::make_unique<
+        envoy::extensions::tracers::opentelemetry::samplers::v3::ParentBasedSamplerConfig>();
+  }
+  std::string name() const override { return "envoy.tracers.opentelemetry.samplers.parent_based"; }
+};
+
+DECLARE_FACTORY(ParentBasedSamplerFactory);
+
+} // namespace OpenTelemetry
+} // namespace Tracers
+} // namespace Extensions
+} // namespace Envoy

source/extensions/tracers/opentelemetry/samplers/parent_based/parent_based_sampler.cc

@@ -0,0 +1,39 @@
+#include "source/extensions/tracers/opentelemetry/samplers/parent_based/parent_based_sampler.h"
+
+#include <memory>
+#include <sstream>
+#include <string>
+
+#include "source/extensions/tracers/opentelemetry/span_context.h"
+
+namespace Envoy {
+namespace Extensions {
+namespace Tracers {
+namespace OpenTelemetry {
+
+SamplingResult ParentBasedSampler::shouldSample(const absl::optional<SpanContext> parent_context,
+                                                const std::string& trace_id,
+                                                const std::string& name, OTelSpanKind kind,
+                                                OptRef<const Tracing::TraceContext> trace_context,
+                                                const std::vector<SpanContext>& links) {
+  if (!parent_context.has_value() || parent_context->traceId().empty()) {
+    return wrapped_sampler_->shouldSample(parent_context, trace_id, name, kind, trace_context,
+                                          links);
+  }
+
+  SamplingResult result;
+  result.tracestate = parent_context.value().tracestate();
+  if (parent_context->sampled()) {
+    result.decision = Decision::RecordAndSample;
+  } else {
+    result.decision = Decision::Drop;
+  }
+  return result;
+}
+
+std::string ParentBasedSampler::getDescription() const { return "ParentBasedSampler"; }
+
+} // namespace OpenTelemetry
+} // namespace Tracers
+} // namespace Extensions
+} // namespace Envoy

source/extensions/tracers/opentelemetry/samplers/parent_based/parent_based_sampler.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include <memory>
+
+#include "envoy/extensions/tracers/opentelemetry/samplers/v3/parent_based_sampler.pb.h"
+#include "envoy/server/factory_context.h"
+
+#include "source/common/common/logger.h"
+#include "source/extensions/tracers/opentelemetry/samplers/sampler.h"
+
+namespace Envoy {
+namespace Extensions {
+namespace Tracers {
+namespace OpenTelemetry {
+
+/**
+ * This is a sampler decorator. If the parent context is empty or doesn't have a valid traceId,
+ * the ParentBasedSampler will delegate the decision to the wrapped sampler.
+ * Otherwise, it will decide based on the sampled flag of the parent context:
+ * if parent_context->sampled -> return RecordAndSample
+ * else -> return Decision::Drop
+ * Check https://opentelemetry.io/docs/specs/otel/trace/sdk/#parentbased for the official docs and
+ * https://github.com/open-telemetry/opentelemetry-cpp/blob/eb2b9753ea2df64079e07d40489388ea1b323108/sdk/src/trace/samplers/parent.cc#L30
+ * for an official implementation
+ */
+class ParentBasedSampler : public Sampler, Logger::Loggable<Logger::Id::tracing> {
+public:
+  explicit ParentBasedSampler(const Protobuf::Message& /*config*/,
+                              Server::Configuration::TracerFactoryContext& /*context*/,
+                              SamplerSharedPtr wrapped_sampler)
+      : wrapped_sampler_(wrapped_sampler) {}
+  SamplingResult shouldSample(const absl::optional<SpanContext> parent_context,
+                              const std::string& trace_id, const std::string& name,
+                              OTelSpanKind spankind,
+                              OptRef<const Tracing::TraceContext> trace_context,
+                              const std::vector<SpanContext>& links) override;
+  std::string getDescription() const override;
+
+private:
+  SamplerSharedPtr wrapped_sampler_;
+};
+
+} // namespace OpenTelemetry
+} // namespace Tracers
+} // namespace Extensions
+} // namespace Envoy

source/extensions/tracers/opentelemetry/samplers/trace_id_ratio_based/BUILD

@@ -0,0 +1,34 @@
+load(
+    "//bazel:envoy_build_system.bzl",
+    "envoy_cc_extension",
+    "envoy_cc_library",
+    "envoy_extension_package",
+)
+
+licenses(["notice"])  # Apache 2
+
+envoy_extension_package()
+
+envoy_cc_extension(
+    name = "config",
+    srcs = ["config.cc"],
+    hdrs = ["config.h"],
+    deps = [
+        ":trace_id_ratio_based_sampler_lib",
+        "//envoy/registry",
+        "//source/common/config:utility_lib",
+        "@envoy_api//envoy/extensions/tracers/opentelemetry/samplers/v3:pkg_cc_proto",
+    ],
+)
+
+envoy_cc_library(
+    name = "trace_id_ratio_based_sampler_lib",
+    srcs = ["trace_id_ratio_based_sampler.cc"],
+    hdrs = ["trace_id_ratio_based_sampler.h"],
+    deps = [
+        "//source/common/common:safe_memcpy_lib",
+        "//source/extensions/tracers/opentelemetry:opentelemetry_tracer_lib",
+        "//source/extensions/tracers/opentelemetry/samplers:sampler_lib",
+        "@envoy_api//envoy/extensions/tracers/opentelemetry/samplers/v3:pkg_cc_proto",
+    ],
+)

source/extensions/tracers/opentelemetry/samplers/trace_id_ratio_based/config.cc

@@ -0,0 +1,36 @@
+#include "source/extensions/tracers/opentelemetry/samplers/trace_id_ratio_based/config.h"
+
+#include "envoy/extensions/tracers/opentelemetry/samplers/v3/trace_id_ratio_based_sampler.pb.h"
+#include "envoy/extensions/tracers/opentelemetry/samplers/v3/trace_id_ratio_based_sampler.pb.validate.h"
+#include "envoy/server/tracer_config.h"
+
+#include "source/common/config/utility.h"
+#include "source/extensions/tracers/opentelemetry/samplers/trace_id_ratio_based/trace_id_ratio_based_sampler.h"
+
+namespace Envoy {
+namespace Extensions {
+namespace Tracers {
+namespace OpenTelemetry {
+
+SamplerSharedPtr TraceIdRatioBasedSamplerFactory::createSampler(
+    const Protobuf::Message& config, Server::Configuration::TracerFactoryContext& context) {
+  auto mptr = Envoy::Config::Utility::translateAnyToFactoryConfig(
+      dynamic_cast<const ProtobufWkt::Any&>(config), context.messageValidationVisitor(), *this);
+
+  const auto& proto_config =
+      MessageUtil::downcastAndValidate<const envoy::extensions::tracers::opentelemetry::samplers::
+                                           v3::TraceIdRatioBasedSamplerConfig&>(
+          *mptr, context.messageValidationVisitor());
+
+  return std::make_shared<TraceIdRatioBasedSampler>(proto_config, context);
+}
+
+/**
+ * Static registration for the Env sampler factory. @see RegisterFactory.
+ */
+REGISTER_FACTORY(TraceIdRatioBasedSamplerFactory, SamplerFactory);
+
+} // namespace OpenTelemetry
+} // namespace Tracers
+} // namespace Extensions
+} // namespace Envoy