Use a delegate sampler for each possible case in ParentBased Sampler (#1440)

Author: lzchen

opentelemetry-sdk/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 - Add meter reference to observers
   ([#1425](https://github.com/open-telemetry/opentelemetry-python/pull/1425))
+- Add local/remote samplers to parent based sampler
+  ([#1440](https://github.com/open-telemetry/opentelemetry-python/pull/1440))
 - Add `fields` to propagators
   ([#1374](https://github.com/open-telemetry/opentelemetry-python/pull/1374))
 

opentelemetry-sdk/src/opentelemetry/sdk/trace/sampling.py

@@ -27,7 +27,7 @@
 
 A `TraceIdRatioBased` sampler makes a random sampling result based on the sampling probability given.
 
-If the span being sampled has a parent, `ParentBased` will respect the parent span's sampling result. Otherwise, it returns the sampling result from the given delegate sampler.
+If the span being sampled has a parent, `ParentBased` will respect the parent delegate sampler. Otherwise, it returns the sampling result from the given root sampler.
 
 Currently, sampling results are always made during the creation of the span. However, this might not always be the case in the future (see `OTEP #115 <https://github.com/open-telemetry/oteps/pull/115>`_).
 
@@ -160,6 +160,13 @@ def get_description(self) -> str:
         return "AlwaysOnSampler"
 
 
+ALWAYS_OFF = StaticSampler(Decision.DROP)
+"""Sampler that never samples spans, regardless of the parent span's sampling decision."""
+
+ALWAYS_ON = StaticSampler(Decision.RECORD_AND_SAMPLE)
+"""Sampler that always samples spans, regardless of the parent span's sampling decision."""
+
+
 class TraceIdRatioBased(Sampler):
     """
     Sampler that makes sampling decisions probabalistically based on `rate`,
@@ -218,16 +225,33 @@ def get_description(self) -> str:
 
 class ParentBased(Sampler):
     """
-    If a parent is set, follows the same sampling decision as the parent.
-    Otherwise, uses the delegate provided at initialization to make a
+    If a parent is set, applies the respective delegate sampler.
+    Otherwise, uses the root provided at initialization to make a
     decision.
 
     Args:
-        delegate: The delegate sampler to use if parent is not set.
+        root: Sampler called for spans with no parent (root spans).
+        remote_parent_sampled: Sampler called for a remote sampled parent.
+        remote_parent_not_sampled: Sampler called for a remote parent that is
+            not sampled.
+        local_parent_sampled: Sampler called for a local sampled parent.
+        local_parent_not_sampled: Sampler called for a local parent that is
+            not sampled.
     """
 
-    def __init__(self, delegate: Sampler):
-        self._delegate = delegate
+    def __init__(
+        self,
+        root: Sampler,
+        remote_parent_sampled: Sampler = ALWAYS_ON,
+        remote_parent_not_sampled: Sampler = ALWAYS_OFF,
+        local_parent_sampled: Sampler = ALWAYS_ON,
+        local_parent_not_sampled: Sampler = ALWAYS_OFF,
+    ):
+        self._root = root
+        self._remote_parent_sampled = remote_parent_sampled
+        self._remote_parent_not_sampled = remote_parent_not_sampled
+        self._local_parent_sampled = local_parent_sampled
+        self._local_parent_not_sampled = local_parent_not_sampled
 
     def should_sample(
         self,
@@ -241,15 +265,22 @@ def should_sample(
         parent_span_context = get_current_span(
             parent_context
         ).get_span_context()
-        # respect the sampling flag of the parent if present
+        # default to the root sampler
+        sampler = self._root
+        # respect the sampling and remote flag of the parent if present
         if parent_span_context is not None and parent_span_context.is_valid:
-            decision = Decision.RECORD_AND_SAMPLE
-            if not parent_span_context.trace_flags.sampled:
-                decision = Decision.DROP
-                attributes = None
-            return SamplingResult(decision, attributes, trace_state)
-
-        return self._delegate.should_sample(
+            if parent_span_context.is_remote:
+                if parent_span_context.trace_flags.sampled:
+                    sampler = self._remote_parent_sampled
+                else:
+                    sampler = self._remote_parent_not_sampled
+            else:
+                if parent_span_context.trace_flags.sampled:
+                    sampler = self._local_parent_sampled
+                else:
+                    sampler = self._local_parent_not_sampled
+
+        return sampler.should_sample(
             parent_context=parent_context,
             trace_id=trace_id,
             name=name,
@@ -259,14 +290,17 @@ def should_sample(
         )
 
     def get_description(self):
-        return "ParentBased{{{}}}".format(self._delegate.get_description())
-
-
-ALWAYS_OFF = StaticSampler(Decision.DROP)
-"""Sampler that never samples spans, regardless of the parent span's sampling decision."""
+        return (
+            "ParentBased{{root:{},remoteParentSampled:{},remoteParentNotSampled:{},"
+            "localParentSampled:{},localParentNotSampled:{}}}".format(
+                self._root.get_description(),
+                self._remote_parent_sampled.get_description(),
+                self._remote_parent_not_sampled.get_description(),
+                self._local_parent_sampled.get_description(),
+                self._local_parent_not_sampled.get_description(),
+            )
+        )
 
-ALWAYS_ON = StaticSampler(Decision.RECORD_AND_SAMPLE)
-"""Sampler that always samples spans, regardless of the parent span's sampling decision."""
 
 DEFAULT_OFF = ParentBased(ALWAYS_OFF)
 """Sampler that respects its parent span's sampling decision, but otherwise never samples."""

opentelemetry-sdk/tests/trace/test_sampling.py

@@ -311,13 +311,15 @@ def test_probability_sampler_limits(self):
             almost_almost_always_on.bound, 0xFFFFFFFFFFFFFFFF,
         )
 
+    # pylint:disable=too-many-statements
     def exec_parent_based(self, parent_sampling_context):
         trace_state = trace.TraceState({"key": "value"})
         sampler = sampling.ParentBased(sampling.ALWAYS_ON)
+        # Check that the sampling decision matches the parent context if given
         with parent_sampling_context(
             self._create_parent_span(trace_flags=TO_DEFAULT)
         ) as context:
-            # Check that the sampling decision matches the parent context if given
+            # local, not sampled
             not_sampled_result = sampler.should_sample(
                 context,
                 0x7FFFFFFFFFFFFFFF,
@@ -329,11 +331,101 @@ def exec_parent_based(self, parent_sampling_context):
             self.assertEqual(not_sampled_result.attributes, {})
             self.assertEqual(not_sampled_result.trace_state, trace_state)
 
+        with parent_sampling_context(
+            self._create_parent_span(trace_flags=TO_DEFAULT)
+        ) as context:
+            sampler = sampling.ParentBased(
+                root=sampling.ALWAYS_OFF,
+                local_parent_not_sampled=sampling.ALWAYS_ON,
+            )
+            # local, not sampled -> opposite sampler
+            sampled_result = sampler.should_sample(
+                context,
+                0x7FFFFFFFFFFFFFFF,
+                "unsampled parent, sampling on",
+                attributes={"sampled": "false"},
+                trace_state=trace_state,
+            )
+            self.assertTrue(sampled_result.decision.is_sampled())
+            self.assertEqual(sampled_result.attributes, {"sampled": "false"})
+            self.assertEqual(sampled_result.trace_state, trace_state)
+
+        with parent_sampling_context(
+            self._create_parent_span(trace_flags=TO_SAMPLED)
+        ) as context:
+            sampler = sampling.ParentBased(sampling.ALWAYS_OFF)
+            # local, sampled
+            sampled_result = sampler.should_sample(
+                context,
+                0x8000000000000000,
+                "sampled parent, sampling off",
+                attributes={"sampled": "true"},
+                trace_state=trace_state,
+            )
+            self.assertTrue(sampled_result.decision.is_sampled())
+            self.assertEqual(sampled_result.attributes, {"sampled": "true"})
+            self.assertEqual(sampled_result.trace_state, trace_state)
+
         with parent_sampling_context(
             self._create_parent_span(trace_flags=TO_SAMPLED)
         ) as context:
-            sampler2 = sampling.ParentBased(sampling.ALWAYS_OFF)
-            sampled_result = sampler2.should_sample(
+            sampler = sampling.ParentBased(
+                root=sampling.ALWAYS_ON,
+                local_parent_sampled=sampling.ALWAYS_OFF,
+            )
+            # local, sampled -> opposite sampler
+            not_sampled_result = sampler.should_sample(
+                context,
+                0x7FFFFFFFFFFFFFFF,
+                "unsampled parent, sampling on",
+                attributes={"sampled": "false"},
+                trace_state=trace_state,
+            )
+            self.assertFalse(not_sampled_result.decision.is_sampled())
+            self.assertEqual(not_sampled_result.attributes, {})
+            self.assertEqual(not_sampled_result.trace_state, trace_state)
+
+        with parent_sampling_context(
+            self._create_parent_span(trace_flags=TO_DEFAULT, is_remote=True)
+        ) as context:
+            sampler = sampling.ParentBased(sampling.ALWAYS_ON)
+            # remote, not sampled
+            not_sampled_result = sampler.should_sample(
+                context,
+                0x7FFFFFFFFFFFFFFF,
+                "unsampled parent, sampling on",
+                attributes={"sampled": "false"},
+                trace_state=trace_state,
+            )
+            self.assertFalse(not_sampled_result.decision.is_sampled())
+            self.assertEqual(not_sampled_result.attributes, {})
+            self.assertEqual(not_sampled_result.trace_state, trace_state)
+
+        with parent_sampling_context(
+            self._create_parent_span(trace_flags=TO_DEFAULT, is_remote=True)
+        ) as context:
+            sampler = sampling.ParentBased(
+                root=sampling.ALWAYS_OFF,
+                remote_parent_not_sampled=sampling.ALWAYS_ON,
+            )
+            # remote, not sampled -> opposite sampler
+            sampled_result = sampler.should_sample(
+                context,
+                0x7FFFFFFFFFFFFFFF,
+                "unsampled parent, sampling on",
+                attributes={"sampled": "false"},
+                trace_state=trace_state,
+            )
+            self.assertTrue(sampled_result.decision.is_sampled())
+            self.assertEqual(sampled_result.attributes, {"sampled": "false"})
+            self.assertEqual(sampled_result.trace_state, trace_state)
+
+        with parent_sampling_context(
+            self._create_parent_span(trace_flags=TO_SAMPLED, is_remote=True)
+        ) as context:
+            sampler = sampling.ParentBased(sampling.ALWAYS_OFF)
+            # remote, sampled
+            sampled_result = sampler.should_sample(
                 context,
                 0x8000000000000000,
                 "sampled parent, sampling off",
@@ -344,10 +436,29 @@ def exec_parent_based(self, parent_sampling_context):
             self.assertEqual(sampled_result.attributes, {"sampled": "true"})
             self.assertEqual(sampled_result.trace_state, trace_state)
 
-        # for root span follow decision of delegate sampler
+        with parent_sampling_context(
+            self._create_parent_span(trace_flags=TO_SAMPLED, is_remote=True)
+        ) as context:
+            sampler = sampling.ParentBased(
+                root=sampling.ALWAYS_ON,
+                remote_parent_sampled=sampling.ALWAYS_OFF,
+            )
+            # remote, sampled -> opposite sampler
+            not_sampled_result = sampler.should_sample(
+                context,
+                0x7FFFFFFFFFFFFFFF,
+                "unsampled parent, sampling on",
+                attributes={"sampled": "false"},
+                trace_state=trace_state,
+            )
+            self.assertFalse(not_sampled_result.decision.is_sampled())
+            self.assertEqual(not_sampled_result.attributes, {})
+            self.assertEqual(not_sampled_result.trace_state, trace_state)
+
+        # for root span follow decision of root sampler
         with parent_sampling_context(trace.INVALID_SPAN) as context:
-            sampler3 = sampling.ParentBased(sampling.ALWAYS_OFF)
-            not_sampled_result = sampler3.should_sample(
+            sampler = sampling.ParentBased(sampling.ALWAYS_OFF)
+            not_sampled_result = sampler.should_sample(
                 context,
                 0x8000000000000000,
                 "parent, sampling off",
@@ -359,8 +470,8 @@ def exec_parent_based(self, parent_sampling_context):
             self.assertEqual(not_sampled_result.trace_state, trace_state)
 
         with parent_sampling_context(trace.INVALID_SPAN) as context:
-            sampler4 = sampling.ParentBased(sampling.ALWAYS_ON)
-            sampled_result = sampler4.should_sample(
+            sampler = sampling.ParentBased(sampling.ALWAYS_ON)
+            sampled_result = sampler.should_sample(
                 context,
                 0x8000000000000000,
                 "no parent, sampling on",