Use a delegate sampler for each possible case in ParentBased Sampler

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))
 
 ## Version 0.16b0
 

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

@@ -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`,
@@ -219,15 +226,32 @@ 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
+    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

@@ -314,10 +314,11 @@ def test_probability_sampler_limits(self):
     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,
@@ -332,8 +333,9 @@ def exec_parent_based(self, parent_sampling_context):
         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(sampling.ALWAYS_OFF)
+            # local, sampled
+            sampled_result = sampler.should_sample(
                 context,
                 0x8000000000000000,
                 "sampled parent, sampling off",
@@ -344,10 +346,42 @@ 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_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_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",
+                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)
+
+        # 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 +393,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",