Add support for OTEL_TRACE_SAMPLER and OTEL_TRACE_SAMPLER_ARG env variables (open-telemetry#1496)

srikanthccv · web-flow · commit 9a1f594b7512 · 2021-01-18T21:08:45.000-08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -34,6 +34,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#1285](https://github.com/open-telemetry/opentelemetry-python/pull/1285))
 - Added `__repr__` for `DefaultSpan`, added `trace_flags` to `__repr__` of
   `SpanContext` ([#1485](https://github.com/open-telemetry/opentelemetry-python/pull/1485)])
+- `opentelemetry-sdk` Add support for OTEL_TRACE_SAMPLER and OTEL_TRACE_SAMPLER_ARG env variables
+  ([#1496](https://github.com/open-telemetry/opentelemetry-python/pull/1496))
 
 ### Changed
 - `opentelemetry-exporter-zipkin` Updated zipkin exporter status code and error tag
diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py
@@ -60,6 +60,8 @@
 SPAN_EVENT_COUNT_LIMIT = Configuration().get("SPAN_EVENT_COUNT_LIMIT", 1000)
 SPAN_LINK_COUNT_LIMIT = Configuration().get("SPAN_LINK_COUNT_LIMIT", 1000)
 VALID_ATTR_VALUE_TYPES = (bool, str, int, float)
+# pylint: disable=protected-access
+TRACE_SAMPLER = sampling._get_from_env_or_default()
 
 
 class SpanProcessor:
@@ -891,7 +893,7 @@ def use_span(
 class TracerProvider(trace_api.TracerProvider):
     def __init__(
         self,
-        sampler: sampling.Sampler = sampling.DEFAULT_ON,
+        sampler: sampling.Sampler = TRACE_SAMPLER,
         resource: Resource = Resource.create({}),
         shutdown_on_exit: bool = True,
         active_span_processor: Union[
diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/sampling.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/sampling.py
@@ -59,9 +59,46 @@
     # created spans will now be sampled by the TraceIdRatioBased sampler
     with trace.get_tracer(__name__).start_as_current_span("Test Span"):
         ...
+
+The tracer sampler can also be configured via environment variables ``OTEL_TRACE_SAMPLER`` and ``OTEL_TRACE_SAMPLER_ARG`` (only if applicable).
+The list of known values for ``OTEL_TRACE_SAMPLER`` are:
+
+    * always_on - Sampler that always samples spans, regardless of the parent span's sampling decision.
+    * always_off - Sampler that never samples spans, regardless of the parent span's sampling decision.
+    * traceidratio - Sampler that samples probabalistically based on rate.
+    * parentbased_always_on - (default) Sampler that respects its parent span's sampling decision, but otherwise always samples.
+    * parentbased_always_off - Sampler that respects its parent span's sampling decision, but otherwise never samples.
+    * parentbased_traceidratio - Sampler that respects its parent span's sampling decision, but otherwise samples probabalistically based on rate.
+
+Sampling probability can be set with ``OTEL_TRACE_SAMPLER_ARG`` if the sampler is traceidratio or parentbased_traceidratio, when not provided rate will be set to 1.0 (maximum rate possible).
+
+
+Prev example but with environment vairables. Please make sure to set the env ``OTEL_TRACE_SAMPLER=traceidratio`` and ``OTEL_TRACE_SAMPLER_ARG=0.001``.
+
+.. code:: python
+
+    from opentelemetry import trace
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import (
+        ConsoleSpanExporter,
+        SimpleExportSpanProcessor,
+    )
+
+    trace.set_tracer_provider(TracerProvider())
+
+    # set up an exporter for sampled spans
+    trace.get_tracer_provider().add_span_processor(
+        SimpleExportSpanProcessor(ConsoleSpanExporter())
+    )
+
+    # created spans will now be sampled by the TraceIdRatioBased sampler with rate 1/1000.
+    with trace.get_tracer(__name__).start_as_current_span("Test Span"):
+        ...
 """
 import abc
 import enum
+import os
+from logging import getLogger
 from types import MappingProxyType
 from typing import Optional, Sequence
 
@@ -71,6 +108,8 @@
 from opentelemetry.trace.span import TraceState
 from opentelemetry.util.types import Attributes
 
+_logger = getLogger(__name__)
+
 
 class Decision(enum.Enum):
     # IsRecording() == false, span will not be recorded and all events and attributes will be dropped.
@@ -307,3 +346,43 @@ def get_description(self):
 
 DEFAULT_ON = ParentBased(ALWAYS_ON)
 """Sampler that respects its parent span's sampling decision, but otherwise always samples."""
+
+
+class ParentBasedTraceIdRatio(ParentBased):
+    """
+    Sampler that respects its parent span's sampling decision, but otherwise
+    samples probabalistically based on `rate`.
+    """
+
+    def __init__(self, rate: float):
+        root = TraceIdRatioBased(rate=rate)
+        super().__init__(root=root)
+
+
+_KNOWN_SAMPLERS = {
+    "always_on": ALWAYS_ON,
+    "always_off": ALWAYS_OFF,
+    "parentbased_always_on": DEFAULT_ON,
+    "parentbased_always_off": DEFAULT_OFF,
+    "traceidratio": TraceIdRatioBased,
+    "parentbased_traceidratio": ParentBasedTraceIdRatio,
+}
+
+
+def _get_from_env_or_default() -> Sampler:
+    trace_sampler = os.getenv(
+        "OTEL_TRACE_SAMPLER", "parentbased_always_on"
+    ).lower()
+    if trace_sampler not in _KNOWN_SAMPLERS:
+        _logger.warning("Couldn't recognize sampler %s.", trace_sampler)
+        trace_sampler = "parentbased_always_on"
+
+    if trace_sampler in ("traceidratio", "parentbased_traceidratio"):
+        try:
+            rate = float(os.getenv("OTEL_TRACE_SAMPLER_ARG"))
+        except ValueError:
+            _logger.warning("Could not convert TRACE_SAMPLER_ARG to float.")
+            rate = 1.0
+        return _KNOWN_SAMPLERS[trace_sampler](rate)
+
+    return _KNOWN_SAMPLERS[trace_sampler]
diff --git a/opentelemetry-sdk/tests/trace/test_trace.py b/opentelemetry-sdk/tests/trace/test_trace.py
@@ -139,6 +139,9 @@ def test_tracer_provider_accepts_concurrent_multi_span_processor(self):
 
 
 class TestTracerSampling(unittest.TestCase):
+    def tearDown(self):
+        reload(trace)
+
     def test_default_sampler(self):
         tracer = new_tracer()
 
@@ -159,6 +162,12 @@ def test_default_sampler(self):
             trace_api.TraceFlags.SAMPLED,
         )
 
+    def test_default_sampler_type(self):
+        tracer_provider = trace.TracerProvider()
+        self.assertIsInstance(tracer_provider.sampler, sampling.ParentBased)
+        # pylint: disable=protected-access
+        self.assertEqual(tracer_provider.sampler._root, sampling.ALWAYS_ON)
+
     def test_sampler_no_sampling(self):
         tracer_provider = trace.TracerProvider(sampling.ALWAYS_OFF)
         tracer = tracer_provider.get_tracer(__name__)
@@ -179,6 +188,36 @@ def test_sampler_no_sampling(self):
             trace_api.TraceFlags.DEFAULT,
         )
 
+    @mock.patch.dict("os.environ", {"OTEL_TRACE_SAMPLER": "always_off"})
+    def test_sampler_with_env(self):
+        # pylint: disable=protected-access
+        reload(trace)
+        tracer_provider = trace.TracerProvider()
+        self.assertIsInstance(tracer_provider.sampler, sampling.StaticSampler)
+        self.assertEqual(
+            tracer_provider.sampler._decision, sampling.Decision.DROP
+        )
+
+        tracer = tracer_provider.get_tracer(__name__)
+
+        root_span = tracer.start_span(name="root span", context=None)
+        # Should be no-op
+        self.assertIsInstance(root_span, trace_api.DefaultSpan)
+
+    @mock.patch.dict(
+        "os.environ",
+        {
+            "OTEL_TRACE_SAMPLER": "parentbased_traceidratio",
+            "OTEL_TRACE_SAMPLER_ARG": "0.25",
+        },
+    )
+    def test_ratio_sampler_with_env(self):
+        # pylint: disable=protected-access
+        reload(trace)
+        tracer_provider = trace.TracerProvider()
+        self.assertIsInstance(tracer_provider.sampler, sampling.ParentBased)
+        self.assertEqual(tracer_provider.sampler._root.rate, 0.25)
+
 
 class TestSpanCreation(unittest.TestCase):
     def test_start_span_invalid_spancontext(self):
