From 6bbf74571895f4844945e9590e41c4d6634e3190 Mon Sep 17 00:00:00 2001 From: majanjua-amzn Date: Thu, 27 Nov 2025 15:56:34 -0800 Subject: [PATCH] feat: Add AlwaysRecordSampler --- CHANGELOG.md | 2 + .../trace/_sampling_experimental/__init__.py | 2 + .../_sampling_experimental/_always_record.py | 81 +++++++++++++++++++ .../__init__.py | 0 .../test_always_off.py | 0 .../test_always_on.py | 0 .../test_always_record.py | 77 ++++++++++++++++++ .../test_sampler.py | 0 .../test_traceid_ratio.py | 0 .../test_tracestate.py | 0 10 files changed, 162 insertions(+) create mode 100644 opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/_always_record.py rename opentelemetry-sdk/tests/trace/{composite_sampler => sampling_experimental}/__init__.py (100%) rename opentelemetry-sdk/tests/trace/{composite_sampler => sampling_experimental}/test_always_off.py (100%) rename opentelemetry-sdk/tests/trace/{composite_sampler => sampling_experimental}/test_always_on.py (100%) create mode 100644 opentelemetry-sdk/tests/trace/sampling_experimental/test_always_record.py rename opentelemetry-sdk/tests/trace/{composite_sampler => sampling_experimental}/test_sampler.py (100%) rename opentelemetry-sdk/tests/trace/{composite_sampler => sampling_experimental}/test_traceid_ratio.py (100%) rename opentelemetry-sdk/tests/trace/{composite_sampler => sampling_experimental}/test_tracestate.py (100%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 15001c9e9ad..9e16221be48 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +- feat: Add AlwaysRecordSampler + ([#4823](https://github.com/open-telemetry/opentelemetry-python/pull/4823)) - `opentelemetry-api`: Convert objects of any type other than AnyValue in attributes to string to be exportable ([#4808](https://github.com/open-telemetry/opentelemetry-python/pull/4808)) - docs: Added sqlcommenter example diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/__init__.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/__init__.py index 1a8c372276d..4a4e64bfc6b 100644 --- a/opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/__init__.py +++ b/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 diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/_always_record.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/_always_record.py new file mode 100644 index 00000000000..84637bc1d00 --- /dev/null +++ b/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, + ) diff --git a/opentelemetry-sdk/tests/trace/composite_sampler/__init__.py b/opentelemetry-sdk/tests/trace/sampling_experimental/__init__.py similarity index 100% rename from opentelemetry-sdk/tests/trace/composite_sampler/__init__.py rename to opentelemetry-sdk/tests/trace/sampling_experimental/__init__.py diff --git a/opentelemetry-sdk/tests/trace/composite_sampler/test_always_off.py b/opentelemetry-sdk/tests/trace/sampling_experimental/test_always_off.py similarity index 100% rename from opentelemetry-sdk/tests/trace/composite_sampler/test_always_off.py rename to opentelemetry-sdk/tests/trace/sampling_experimental/test_always_off.py diff --git a/opentelemetry-sdk/tests/trace/composite_sampler/test_always_on.py b/opentelemetry-sdk/tests/trace/sampling_experimental/test_always_on.py similarity index 100% rename from opentelemetry-sdk/tests/trace/composite_sampler/test_always_on.py rename to opentelemetry-sdk/tests/trace/sampling_experimental/test_always_on.py diff --git a/opentelemetry-sdk/tests/trace/sampling_experimental/test_always_record.py b/opentelemetry-sdk/tests/trace/sampling_experimental/test_always_record.py new file mode 100644 index 00000000000..c556eb9f4bc --- /dev/null +++ b/opentelemetry-sdk/tests/trace/sampling_experimental/test_always_record.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 diff --git a/opentelemetry-sdk/tests/trace/composite_sampler/test_sampler.py b/opentelemetry-sdk/tests/trace/sampling_experimental/test_sampler.py similarity index 100% rename from opentelemetry-sdk/tests/trace/composite_sampler/test_sampler.py rename to opentelemetry-sdk/tests/trace/sampling_experimental/test_sampler.py diff --git a/opentelemetry-sdk/tests/trace/composite_sampler/test_traceid_ratio.py b/opentelemetry-sdk/tests/trace/sampling_experimental/test_traceid_ratio.py similarity index 100% rename from opentelemetry-sdk/tests/trace/composite_sampler/test_traceid_ratio.py rename to opentelemetry-sdk/tests/trace/sampling_experimental/test_traceid_ratio.py diff --git a/opentelemetry-sdk/tests/trace/composite_sampler/test_tracestate.py b/opentelemetry-sdk/tests/trace/sampling_experimental/test_tracestate.py similarity index 100% rename from opentelemetry-sdk/tests/trace/composite_sampler/test_tracestate.py rename to opentelemetry-sdk/tests/trace/sampling_experimental/test_tracestate.py