datadog: set sampling rate (#740)

Set the sampling rate in the Datadog exported span if span was sampled with the ProbabilitySampler.

Co-authored-by: Diego Hurtado <[email protected]>

ext/opentelemetry-ext-datadog/src/opentelemetry/ext/datadog/constants.py

@@ -2,3 +2,5 @@
 AUTO_REJECT = 0
 AUTO_KEEP = 1
 USER_KEEP = 2
+SAMPLE_RATE_METRIC_KEY = "_sample_rate"
+SAMPLING_PRIORITY_KEY = "_sampling_priority_v1"

ext/opentelemetry-ext-datadog/src/opentelemetry/ext/datadog/exporter.py

@@ -25,7 +25,7 @@
 from opentelemetry.trace.status import StatusCanonicalCode
 
 # pylint:disable=relative-beyond-top-level
-from .constants import DD_ORIGIN
+from .constants import DD_ORIGIN, SAMPLE_RATE_METRIC_KEY
 
 logger = logging.getLogger(__name__)
 
@@ -136,6 +136,10 @@ def _translate_to_datadog(self, spans):
             if origin and parent_id == 0:
                 datadog_span.set_tag(DD_ORIGIN, origin)
 
+            sampling_rate = _get_sampling_rate(span)
+            if sampling_rate is not None:
+                datadog_span.set_metric(SAMPLE_RATE_METRIC_KEY, sampling_rate)
+
             # span events and span links are not supported
 
             datadog_spans.append(datadog_span)
@@ -216,3 +220,13 @@ def _get_origin(span):
     ctx = span.get_context()
     origin = ctx.trace_state.get(DD_ORIGIN)
     return origin
+
+
+def _get_sampling_rate(span):
+    ctx = span.get_context()
+    return (
+        span.sampler.rate
+        if ctx.trace_flags.sampled
+        and isinstance(span.sampler, trace_api.sampling.ProbabilitySampler)
+        else None
+    )

ext/opentelemetry-ext-datadog/tests/test_datadog_exporter.py

@@ -440,3 +440,35 @@ def test_origin(self):
         ]
         expected = ["origin-service", None]
         self.assertListEqual(actual, expected)
+
+    def test_sampling_rate(self):
+        context = trace_api.SpanContext(
+            trace_id=0x000000000000000000000000DEADBEEF,
+            span_id=0x34BF92DEEFC58C92,
+            is_remote=False,
+            trace_flags=trace_api.TraceFlags(trace_api.TraceFlags.SAMPLED),
+        )
+        sampler = trace_api.sampling.ProbabilitySampler(0.5)
+
+        span = trace.Span(
+            name="sampled", context=context, parent=None, sampler=sampler
+        )
+        span.start()
+        span.end()
+
+        # pylint: disable=protected-access
+        exporter = datadog.DatadogSpanExporter()
+        datadog_spans = [
+            span.to_dict() for span in exporter._translate_to_datadog([span])
+        ]
+
+        self.assertEqual(len(datadog_spans), 1)
+
+        actual = [
+            span["metrics"].get(datadog.constants.SAMPLE_RATE_METRIC_KEY)
+            if "metrics" in span
+            else None
+            for span in datadog_spans
+        ]
+        expected = [0.5]
+        self.assertListEqual(actual, expected)