Pass in context instead of parentcontext into samplers

opentelemetry-sdk/CHANGELOG.md

@@ -8,6 +8,8 @@
   ([#1265](https://github.com/open-telemetry/opentelemetry-python/pull/1265))
 - Allow None in sequence attributes values
   ([#998](https://github.com/open-telemetry/opentelemetry-python/pull/998))
+- Samplers to accept parent Context
+  ([#1267](https://github.com/open-telemetry/opentelemetry-python/pull/1267))
 
 ## Version 0.14b0
 

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

@@ -757,6 +757,7 @@ def start_span(  # pylint: disable=too-many-locals
                 "parent_span_context must be a SpanContext or None."
             )
 
+        # is_valid determines root span
         if parent_span_context is None or not parent_span_context.is_valid:
             parent_span_context = None
             trace_id = self.source.ids_generator.generate_trace_id()
@@ -773,7 +774,7 @@ def start_span(  # pylint: disable=too-many-locals
         # The sampler may also add attributes to the newly-created span, e.g.
         # to include information about the sampling result.
         sampling_result = self.source.sampler.should_sample(
-            parent_span_context, trace_id, name, attributes, links,
+            context, trace_id, name, attributes, links,
         )
 
         trace_flags = (

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

@@ -66,7 +66,8 @@
 from typing import Optional, Sequence
 
 # pylint: disable=unused-import
-from opentelemetry.trace import Link, SpanContext
+from opentelemetry.context import Context
+from opentelemetry.trace import Link, get_current_span
 from opentelemetry.util.types import Attributes
 
 
@@ -113,11 +114,11 @@ class Sampler(abc.ABC):
     @abc.abstractmethod
     def should_sample(
         self,
-        parent_span_context: Optional["SpanContext"],
+        parent_context: Optional["Context"],
         trace_id: int,
         name: str,
         attributes: Attributes = None,
-        links: Sequence["Link"] = (),
+        links: Sequence["Link"] = None,
     ) -> "SamplingResult":
         pass
 
@@ -134,11 +135,11 @@ def __init__(self, decision: "Decision"):
 
     def should_sample(
         self,
-        parent_span_context: Optional["SpanContext"],
+        parent_context: Optional["Context"],
         trace_id: int,
         name: str,
         attributes: Attributes = None,
-        links: Sequence["Link"] = (),
+        links: Sequence["Link"] = None,
     ) -> "SamplingResult":
         if self._decision is Decision.DROP:
             return SamplingResult(self._decision)
@@ -188,11 +189,11 @@ def bound(self) -> int:
 
     def should_sample(
         self,
-        parent_span_context: Optional["SpanContext"],
+        parent_context: Optional["Context"],
         trace_id: int,
         name: str,
-        attributes: Attributes = None,  # TODO
-        links: Sequence["Link"] = (),
+        attributes: Attributes = None,
+        links: Sequence["Link"] = None,
     ) -> "SamplingResult":
         decision = Decision.DROP
         if trace_id & self.TRACE_ID_LIMIT < self.bound:
@@ -220,22 +221,27 @@ def __init__(self, delegate: Sampler):
 
     def should_sample(
         self,
-        parent_span_context: Optional["SpanContext"],
+        parent_context: Optional["Context"],
         trace_id: int,
         name: str,
-        attributes: Attributes = None,  # TODO
-        links: Sequence["Link"] = (),
+        attributes: Attributes = None,
+        links: Sequence["Link"] = None,
     ) -> "SamplingResult":
-        if parent_span_context is not None:
+        if parent_context is not None:
+            parent_span_context = get_current_span(
+                parent_context
+            ).get_span_context()
+            # only drop if parent exists and is not a root span
             if (
-                not parent_span_context.is_valid
-                or not parent_span_context.trace_flags.sampled
+                parent_span_context is not None
+                and parent_span_context.is_valid
+                and not parent_span_context.trace_flags.sampled
             ):
                 return SamplingResult(Decision.DROP)
             return SamplingResult(Decision.RECORD_AND_SAMPLE, attributes)
 
         return self._delegate.should_sample(
-            parent_span_context=parent_span_context,
+            parent_context=parent_context,
             trace_id=trace_id,
             name=name,
             attributes=attributes,

opentelemetry-sdk/tests/trace/test_sampling.py

@@ -109,24 +109,34 @@ def test_always_off(self):
         self.assertEqual(sampled_always_on.attributes, {})
 
     def test_default_on(self):
+        context = trace.set_span_in_context(
+            trace.DefaultSpan(
+                trace.SpanContext(
+                    0xDEADBEEF,
+                    0xDEADBEF0,
+                    is_remote=False,
+                    trace_flags=TO_DEFAULT,
+                )
+            )
+        )
         no_record_default_on = sampling.DEFAULT_ON.should_sample(
-            trace.SpanContext(
-                0xDEADBEEF, 0xDEADBEF0, is_remote=False, trace_flags=TO_DEFAULT
-            ),
-            0xDEADBEF1,
-            0xDEADBEF2,
-            "unsampled parent, sampling on",
+            context, 0xDEADBEF1, 0xDEADBEF2, "unsampled parent, sampling on",
         )
         self.assertFalse(no_record_default_on.decision.is_sampled())
         self.assertEqual(no_record_default_on.attributes, {})
 
+        context = trace.set_span_in_context(
+            trace.DefaultSpan(
+                trace.SpanContext(
+                    0xDEADBEEF,
+                    0xDEADBEF0,
+                    is_remote=False,
+                    trace_flags=TO_SAMPLED,
+                )
+            )
+        )
         sampled_default_on = sampling.DEFAULT_ON.should_sample(
-            trace.SpanContext(
-                0xDEADBEEF, 0xDEADBEF0, is_remote=False, trace_flags=TO_SAMPLED
-            ),
-            0xDEADBEF1,
-            0xDEADBEF2,
-            {"sampled parent": "sampling on"},
+            context, 0xDEADBEF1, 0xDEADBEF2, {"sampled parent": "sampling on"},
         )
         self.assertTrue(sampled_default_on.decision.is_sampled())
         self.assertEqual(
@@ -142,24 +152,34 @@ def test_default_on(self):
         )
 
     def test_default_off(self):
+        context = trace.set_span_in_context(
+            trace.DefaultSpan(
+                trace.SpanContext(
+                    0xDEADBEEF,
+                    0xDEADBEF0,
+                    is_remote=False,
+                    trace_flags=TO_DEFAULT,
+                )
+            )
+        )
         no_record_default_off = sampling.DEFAULT_OFF.should_sample(
-            trace.SpanContext(
-                0xDEADBEEF, 0xDEADBEF0, is_remote=False, trace_flags=TO_DEFAULT
-            ),
-            0xDEADBEF1,
-            0xDEADBEF2,
-            "unsampled parent, sampling off",
+            context, 0xDEADBEF1, 0xDEADBEF2, "unsampled parent, sampling off",
         )
         self.assertFalse(no_record_default_off.decision.is_sampled())
         self.assertEqual(no_record_default_off.attributes, {})
 
+        context = trace.set_span_in_context(
+            trace.DefaultSpan(
+                trace.SpanContext(
+                    0xDEADBEEF,
+                    0xDEADBEF0,
+                    is_remote=False,
+                    trace_flags=TO_SAMPLED,
+                )
+            )
+        )
         sampled_default_off = sampling.DEFAULT_OFF.should_sample(
-            trace.SpanContext(
-                0xDEADBEEF, 0xDEADBEF0, is_remote=False, trace_flags=TO_SAMPLED
-            ),
-            0xDEADBEF1,
-            0xDEADBEF2,
-            {"sampled parent": "sampling on"},
+            context, 0xDEADBEF1, 0xDEADBEF2, {"sampled parent": "sampling on"},
         )
         self.assertTrue(sampled_default_off.decision.is_sampled())
         self.assertEqual(
@@ -275,32 +295,45 @@ def test_probability_sampler_limits(self):
 
     def test_parent_based(self):
         sampler = sampling.ParentBased(sampling.ALWAYS_ON)
-        # Check that the sampling decision matches the parent context if given
-        self.assertFalse(
-            sampler.should_sample(
+        context = trace.set_span_in_context(
+            trace.DefaultSpan(
                 trace.SpanContext(
                     0xDEADBEF0,
                     0xDEADBEF1,
                     is_remote=False,
                     trace_flags=TO_DEFAULT,
-                ),
-                0x7FFFFFFFFFFFFFFF,
-                0xDEADBEEF,
-                "span name",
+                )
+            )
+        )
+        # Check that the sampling decision matches the parent context if given
+        self.assertFalse(
+            sampler.should_sample(
+                context, 0x7FFFFFFFFFFFFFFF, 0xDEADBEEF, "span name",
             ).decision.is_sampled()
         )
 
-        sampler2 = sampling.ParentBased(sampling.ALWAYS_OFF)
-        self.assertTrue(
-            sampler2.should_sample(
+        context = trace.set_span_in_context(
+            trace.DefaultSpan(
                 trace.SpanContext(
                     0xDEADBEF0,
                     0xDEADBEF1,
                     is_remote=False,
                     trace_flags=TO_SAMPLED,
-                ),
-                0x8000000000000000,
-                0xDEADBEEF,
-                "span name",
+                )
+            )
+        )
+        sampler2 = sampling.ParentBased(sampling.ALWAYS_OFF)
+        self.assertTrue(
+            sampler2.should_sample(
+                context, 0x8000000000000000, 0xDEADBEEF, "span name",
+            ).decision.is_sampled()
+        )
+
+        # root span always sampled for parentbased
+        context = trace.set_span_in_context(trace.INVALID_SPAN)
+        sampler3 = sampling.ParentBased(sampling.ALWAYS_OFF)
+        self.assertTrue(
+            sampler3.should_sample(
+                context, 0x8000000000000000, 0xDEADBEEF, "span name",
             ).decision.is_sampled()
         )