Add experimental composite sampler (#4714)

* Add experimental consistent sampler

* Fix import

* Fix lint

* CHANGELOG

* Fix test name

* Cleanup and match spec better

* Format

* Cleanup

* Formatting

---------

Co-authored-by: Riccardo Magliocchetti <[email protected]>

Author: anuraaga

CHANGELOG.md

@@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Add experimental composite samplers
+  ([#4714](https://github.com/open-telemetry/opentelemetry-python/pull/4714))
 - Filter duplicate logs out of some internal `logger`'s logs on the export logs path that might otherwise endlessly log or cause a recursion depth exceeded issue in cases where logging itself results in an exception.
   ([#4695](https://github.com/open-telemetry/opentelemetry-python/pull/4695)).
 - docs: linked the examples with their github source code location and added Prometheus example

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

@@ -0,0 +1,31 @@
+# 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.
+
+__all__ = [
+    "ComposableSampler",
+    "SamplingIntent",
+    "composable_always_off",
+    "composable_always_on",
+    "composable_parent_threshold",
+    "composable_traceid_ratio_based",
+    "composite_sampler",
+]
+
+
+from ._always_off import composable_always_off
+from ._always_on import composable_always_on
+from ._composable import ComposableSampler, SamplingIntent
+from ._parent_threshold import composable_parent_threshold
+from ._sampler import composite_sampler
+from ._traceid_ratio import composable_traceid_ratio_based

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

@@ -0,0 +1,55 @@
+# 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 __future__ import annotations
+
+from typing import Sequence
+
+from opentelemetry.context import Context
+from opentelemetry.trace import Link, SpanKind, TraceState
+from opentelemetry.util.types import Attributes
+
+from ._composable import ComposableSampler, SamplingIntent
+from ._util import INVALID_THRESHOLD
+
+_intent = SamplingIntent(threshold=INVALID_THRESHOLD, threshold_reliable=False)
+
+
+class _ComposableAlwaysOffSampler(ComposableSampler):
+    def sampling_intent(
+        self,
+        parent_ctx: Context | None,
+        name: str,
+        span_kind: SpanKind | None,
+        attributes: Attributes,
+        links: Sequence[Link] | None,
+        trace_state: TraceState | None = None,
+    ) -> SamplingIntent:
+        return _intent
+
+    def get_description(self) -> str:
+        return "ComposableAlwaysOff"
+
+
+_always_off = _ComposableAlwaysOffSampler()
+
+
+def composable_always_off() -> ComposableSampler:
+    """Returns a composable sampler that does not sample any span.
+
+    - Always returns a SamplingIntent with no threshold, indicating all spans should be dropped
+    - Sets threshold_reliable to false
+    - Does not add any attributes
+    """
+    return _always_off

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

@@ -0,0 +1,55 @@
+# 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 __future__ import annotations
+
+from typing import Sequence
+
+from opentelemetry.context import Context
+from opentelemetry.trace import Link, SpanKind, TraceState
+from opentelemetry.util.types import Attributes
+
+from ._composable import ComposableSampler, SamplingIntent
+from ._util import MIN_THRESHOLD
+
+_intent = SamplingIntent(threshold=MIN_THRESHOLD)
+
+
+class _ComposableAlwaysOnSampler(ComposableSampler):
+    def sampling_intent(
+        self,
+        parent_ctx: Context | None,
+        name: str,
+        span_kind: SpanKind | None,
+        attributes: Attributes,
+        links: Sequence[Link] | None,
+        trace_state: TraceState | None = None,
+    ) -> SamplingIntent:
+        return _intent
+
+    def get_description(self) -> str:
+        return "ComposableAlwaysOn"
+
+
+_always_on = _ComposableAlwaysOnSampler()
+
+
+def composable_always_on() -> ComposableSampler:
+    """Returns a composable sampler that samples all spans.
+
+    - Always returns a SamplingIntent with threshold set to sample all spans (threshold = 0)
+    - Sets threshold_reliable to true
+    - Does not add any attributes
+    """
+    return _always_on

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

@@ -0,0 +1,61 @@
+# 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 __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Callable, Protocol, Sequence
+
+from opentelemetry.context import Context
+from opentelemetry.trace import Link, SpanKind, TraceState
+from opentelemetry.util.types import Attributes
+
+
+@dataclass(frozen=True)
+class SamplingIntent:
+    """Information to make a consistent sampling decision."""
+
+    threshold: int
+    """The sampling threshold value. A lower threshold increases the likelihood of sampling."""
+
+    threshold_reliable: bool = field(default=True)
+    """Indicates whether the threshold is reliable for Span-to-Metrics estimation."""
+
+    attributes: Attributes = field(default=None)
+    """Any attributes to be added to a sampled span."""
+
+    update_trace_state: Callable[[TraceState], TraceState] = field(
+        default=lambda ts: ts
+    )
+    """Any updates to be made to trace state."""
+
+
+class ComposableSampler(Protocol):
+    """A sampler that can be composed to make a final sampling decision."""
+
+    def sampling_intent(
+        self,
+        parent_ctx: Context | None,
+        name: str,
+        span_kind: SpanKind | None,
+        attributes: Attributes,
+        links: Sequence[Link] | None,
+        trace_state: TraceState | None,
+    ) -> SamplingIntent:
+        """Returns information to make a sampling decision."""
+        ...  # pylint: disable=unnecessary-ellipsis
+
+    def get_description(self) -> str:
+        """Returns a description of the sampler."""
+        ...  # pylint: disable=unnecessary-ellipsis

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

@@ -0,0 +1,89 @@
+# 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 __future__ import annotations
+
+from typing import Sequence
+
+from opentelemetry.context import Context
+from opentelemetry.trace import Link, SpanKind, TraceState, get_current_span
+from opentelemetry.util.types import Attributes
+
+from ._composable import ComposableSampler, SamplingIntent
+from ._trace_state import OtelTraceState
+from ._util import (
+    INVALID_THRESHOLD,
+    MIN_THRESHOLD,
+    is_valid_threshold,
+)
+
+
+class _ComposableParentThreshold(ComposableSampler):
+    def __init__(self, root_sampler: ComposableSampler):
+        self._root_sampler = root_sampler
+        self._description = f"ComposableParentThreshold{{root={root_sampler.get_description()}}}"
+
+    def sampling_intent(
+        self,
+        parent_ctx: Context | None,
+        name: str,
+        span_kind: SpanKind | None,
+        attributes: Attributes,
+        links: Sequence[Link] | None,
+        trace_state: TraceState | None = None,
+    ) -> SamplingIntent:
+        parent_span = get_current_span(parent_ctx)
+        parent_span_ctx = parent_span.get_span_context()
+        is_root = not parent_span_ctx.is_valid
+        if is_root:
+            return self._root_sampler.sampling_intent(
+                parent_ctx, name, span_kind, attributes, links, trace_state
+            )
+
+        ot_trace_state = OtelTraceState.parse(trace_state)
+
+        if is_valid_threshold(ot_trace_state.threshold):
+            return SamplingIntent(
+                threshold=ot_trace_state.threshold,
+                threshold_reliable=True,
+            )
+
+        threshold = (
+            MIN_THRESHOLD
+            if parent_span_ctx.trace_flags.sampled
+            else INVALID_THRESHOLD
+        )
+        return SamplingIntent(threshold=threshold, threshold_reliable=False)
+
+    def get_description(self) -> str:
+        return self._description
+
+
+def composable_parent_threshold(
+    root_sampler: ComposableSampler,
+) -> ComposableSampler:
+    """Returns a consistent sampler that respects the sampling decision of
+    the parent span or falls-back to the given sampler if it is a root span.
+
+    - For spans without a parent context, delegate to the root sampler
+    - For spans with a parent context, returns a SamplingIntent that propagates the parent's sampling decision
+    - Returns the parent's threshold if available; otherwise, if the parent's sampled flag is set,
+    returns threshold=0; otherwise, if the parent's sampled flag is not set, no threshold is returned.
+    - Sets threshold_reliable to match the parent’s reliability, which is true if the parent had a threshold.
+    - Does not add any attributes
+
+    Args:
+        root_sampler: The root sampler to use for spans without a parent context.
+    """
+    return _ComposableParentThreshold(root_sampler)