Add AlwaysRecordSampler

Author: majanjua-amzn

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 __all__ = [
+    "AlwaysRecordSampler",
     "ComposableSampler",
     "SamplingIntent",
     "composable_always_off",
@@ -25,6 +26,7 @@
 
 from ._always_off import composable_always_off
 from ._always_on import composable_always_on
+from ._always_record import AlwaysRecordSampler
 from ._composable import ComposableSampler, SamplingIntent
 from ._parent_threshold import composable_parent_threshold
 from ._sampler import composite_sampler

opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/_always_record.py

@@ -0,0 +1,81 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Sequence
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult
+from opentelemetry.trace import Link, SpanKind
+from opentelemetry.trace.span import TraceState
+from opentelemetry.util.types import Attributes
+
+
+class AlwaysRecordSampler(Sampler):
+    """
+    This sampler will return the sampling result of the provided `_root_sampler`, unless the
+    sampling result contains the sampling decision `Decision.DROP`, in which case, a
+    new sampling result will be returned that is functionally equivalent to the original, except that
+    it contains the sampling decision `SamplingDecision.RECORD_ONLY`. This ensures that all
+    spans are recorded, with no change to sampling.
+
+    The intended use case of this sampler is to provide a means of sending all spans to a
+    processor without having an impact on the sampling rate. This may be desirable if a user wishes
+    to count or otherwise measure all spans produced in a service, without incurring the cost of 100%
+    sampling.
+    """
+
+    _root_sampler: Sampler
+
+    def __init__(self, root_sampler: Sampler):
+        if not root_sampler:
+            raise ValueError("root_sampler must not be None")
+        self._root_sampler = root_sampler
+
+    def should_sample(
+        self,
+        parent_context: Optional["Context"],
+        trace_id: int,
+        name: str,
+        kind: Optional[SpanKind] = None,
+        attributes: Attributes = None,
+        links: Optional[Sequence["Link"]] = None,
+        trace_state: Optional["TraceState"] = None,
+    ) -> "SamplingResult":
+        result: SamplingResult = self._root_sampler.should_sample(
+            parent_context,
+            trace_id,
+            name,
+            kind,
+            attributes,
+            links,
+            trace_state,
+        )
+        if result.decision is Decision.DROP:
+            result = _wrap_result_with_record_only_result(result, attributes)
+        return result
+
+    def get_description(self):
+        return (
+            "AlwaysRecordSampler{" + self._root_sampler.get_description() + "}"
+        )
+
+
+def _wrap_result_with_record_only_result(
+    result: SamplingResult, attributes: Attributes
+) -> SamplingResult:
+    return SamplingResult(
+        Decision.RECORD_ONLY,
+        attributes,
+        result.trace_state,
+    )

opentelemetry-sdk/tests/trace/composite_sampler/__init__.py renamed to opentelemetry-sdk/tests/trace/sampling_experimental/__init__.py

@@ -0,0 +1,77 @@
+from unittest import TestCase
+from unittest.mock import MagicMock
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace._sampling_experimental import AlwaysRecordSampler
+from opentelemetry.sdk.trace.sampling import (
+    Decision,
+    Sampler,
+    SamplingResult,
+    StaticSampler,
+)
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.span import TraceState
+from opentelemetry.util.types import Attributes
+
+
+class TestAlwaysRecordSampler(TestCase):
+    def setUp(self):
+        self.mock_sampler: Sampler = MagicMock()
+        self.sampler: Sampler = AlwaysRecordSampler(self.mock_sampler)
+
+    def test_get_description(self):
+        static_sampler: Sampler = StaticSampler(Decision.DROP)
+        test_sampler: Sampler = AlwaysRecordSampler(static_sampler)
+        self.assertEqual(
+            "AlwaysRecordSampler{AlwaysOffSampler}",
+            test_sampler.get_description(),
+        )
+
+    def test_record_and_sample_sampling_decision(self):
+        self.validate_should_sample(
+            Decision.RECORD_AND_SAMPLE, Decision.RECORD_AND_SAMPLE
+        )
+
+    def test_record_only_sampling_decision(self):
+        self.validate_should_sample(Decision.RECORD_ONLY, Decision.RECORD_ONLY)
+
+    def test_drop_sampling_decision(self):
+        self.validate_should_sample(Decision.DROP, Decision.RECORD_ONLY)
+
+    def validate_should_sample(
+        self, root_decision: Decision, expected_decision: Decision
+    ):
+        root_result: SamplingResult = _build_root_sampling_result(
+            root_decision
+        )
+        self.mock_sampler.should_sample.return_value = root_result
+        actual_result: SamplingResult = self.sampler.should_sample(
+            parent_context=Context(),
+            trace_id=0,
+            name="name",
+            kind=SpanKind.CLIENT,
+            attributes={"key": root_decision.name},
+            trace_state=TraceState(),
+        )
+
+        if root_decision == expected_decision:
+            self.assertEqual(actual_result, root_result)
+            self.assertEqual(actual_result.decision, root_decision)
+        else:
+            self.assertNotEqual(actual_result, root_result)
+            self.assertEqual(actual_result.decision, expected_decision)
+
+        self.assertEqual(actual_result.attributes, root_result.attributes)
+        self.assertEqual(actual_result.trace_state, root_result.trace_state)
+
+
+def _build_root_sampling_result(sampling_decision: Decision):
+    sampling_attr: Attributes = {"key": sampling_decision.name}
+    sampling_trace_state: TraceState = TraceState()
+    sampling_trace_state.add("key", sampling_decision.name)
+    sampling_result: SamplingResult = SamplingResult(
+        decision=sampling_decision,
+        attributes=sampling_attr,
+        trace_state=sampling_trace_state,
+    )
+    return sampling_result