Add support for OTEL_TRACES_SAMPLER

Author: rads-1996

sdk/monitor/azure-monitor-opentelemetry/README.md

@@ -80,7 +80,7 @@ You can configure further with [OpenTelemetry environment variables][ot_env_vars
 | `OTEL_TRACES_EXPORTER` | If set to `None`, disables collection and export of distributed tracing telemetry. |
 | `OTEL_BLRP_SCHEDULE_DELAY` | Specifies the logging export interval in milliseconds. Defaults to 5000. |
 | `OTEL_BSP_SCHEDULE_DELAY` | Specifies the distributed tracing export interval in milliseconds. Defaults to 5000. |
-| `OTEL_TRACES_SAMPLER` | Specifies the sampler to be used for traces. Supports both [application_insights_sampling] and [rate_limited_sampling]. Use `microsoft.fixed.percentage` for the Application Insights sampler or `microsoft.rate_limited` for the Rate Limited sampler. |
+| `OTEL_TRACES_SAMPLER` | Specifies the sampler to be used for traces. Supports [`always_on`, `always_off`, `trace_id_ratio`, `parentbased_always_on`, `parentbased_always_off`, `parentbased_trace_id_ratio`], [application_insights_sampling] and [rate_limited_sampling]. Use `microsoft.fixed.percentage` for the Application Insights sampler or `microsoft.rate_limited` for the Rate Limited sampler. |
 | `OTEL_TRACES_SAMPLER_ARG` | Specifies the sampling parameter for the configured sampler. For the Application Insights sampler, this sets the ratio of distributed tracing telemetry to be [sampled][application_insights_sampling] with accepted values in the range [0,1]. Defaults to 1.0 (no sampling). For the Rate Limited sampler, this sets the maximum traces per second to be [sampled][rate_limited_sampler]. For example, 0.5 means one trace every two seconds, while 5.0 means five traces per second. |
 | `OTEL_PYTHON_DISABLED_INSTRUMENTATIONS` | Specifies which of the supported instrumentations to disable. Disabled instrumentations will not be instrumented as part of `configure_azure_monitor`. However, they can still be manually instrumented with `instrument()` directly. Accepts a comma-separated list of lowercase [Library Names](#officially-supported-instrumentations). For example, set to `"psycopg2,fastapi"` to disable the Psycopg2 and FastAPI instrumentations. Defaults to an empty list, enabling all supported instrumentations. |
 | `OTEL_EXPERIMENTAL_RESOURCE_DETECTORS` | An experimental OpenTelemetry environment variable used to specify Resource Detectors to be used to generate Resource Attributes. This is an experimental feature and the name of this variable and its behavior can change in a non-backwards compatible way. Defaults to "azure_app_service,azure_vm" to enable the [Azure Resource Detectors][ot_resource_detector_azure] for Azure App Service and Azure VM. To add or remove specific resource detectors, set the environment variable accordingly. See the [OpenTelemetry Python Resource Detector Documentation][ot_python_resource_detectors] for more. |
@@ -221,6 +221,7 @@ contact [[email protected]](mailto:[email protected]) with any additio
 [application_insights_namespace]: https://learn.microsoft.com/azure/azure-monitor/app/app-insights-overview
 [application_insights_sampling]: https://learn.microsoft.com/azure/azure-monitor/app/sampling
 [rate_limited_sampling]: https://learn.microsoft.com/azure/azure-monitor/app/opentelemetry-configuration
+[always_on]: https://opentelemetry.io/docs/languages/sdk-configuration/general/#otel_traces_sampler
 [connection_string_doc]: https://learn.microsoft.com/azure/azure-monitor/app/sdk-connection-string
 [distro_feature_request]: https://github.com/Azure/azure-sdk-for-python/issues/new
 [exporter_configuration_docs]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/monitor/azure-monitor-opentelemetry-exporter#configuration

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_configure.py

@@ -40,6 +40,8 @@
     SPAN_PROCESSORS_ARG,
     VIEWS_ARG,
     ENABLE_TRACE_BASED_SAMPLING_ARG,
+    SAMPLING_ARG,
+    SAMPLER_TYPE,
 )
 from azure.monitor.opentelemetry._types import ConfigurationValue
 from azure.monitor.opentelemetry.exporter._quickpulse import (  # pylint: disable=import-error,no-name-in-module
@@ -73,6 +75,7 @@
 from azure.monitor.opentelemetry._utils.configurations import (
     _get_configurations,
     _is_instrumentation_enabled,
+    _get_sampler_from_name,
 )
 from azure.monitor.opentelemetry._utils.instrumentation import (
     get_dist_dependency_conflicts,
@@ -149,7 +152,12 @@ def configure_azure_monitor(**kwargs) -> None:  # pylint: disable=C4758
 def _setup_tracing(configurations: Dict[str, ConfigurationValue]):
     resource: Resource = configurations[RESOURCE_ARG]  # type: ignore
     enable_performance_counters_config = configurations[ENABLE_PERFORMANCE_COUNTERS_ARG]
-    if SAMPLING_TRACES_PER_SECOND_ARG in configurations:
+    if SAMPLING_ARG in configurations:
+        sampler_arg = configurations[SAMPLING_ARG]
+        sampler_type = configurations[SAMPLER_TYPE]
+        sampler = _get_sampler_from_name(sampler_type, sampler_arg)
+        tracer_provider = TracerProvider(sampler=sampler, resource=resource)
+    elif SAMPLING_TRACES_PER_SECOND_ARG in configurations:
         traces_per_second = configurations[SAMPLING_TRACES_PER_SECOND_ARG]
         tracer_provider = TracerProvider(
             sampler=RateLimitedSampler(target_spans_per_second_limit=cast(float, traces_per_second)), resource=resource

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_constants.py

@@ -29,6 +29,24 @@
 FIXED_PERCENTAGE_SAMPLER = "microsoft.fixed.percentage"
 SAMPLING_TRACES_PER_SECOND_ARG = "traces_per_second"
 ENABLE_TRACE_BASED_SAMPLING_ARG = "enable_trace_based_sampling_for_logs"
+SAMPLER_TYPE = "sampler_type"
+SAMPLING_ARG = "sampling_arg"
+ALWAYS_ON_SAMPLER = "always_on"
+ALWAYS_OFF_SAMPLER = "always_off"
+TRACE_ID_RATIO_SAMPLER = "trace_id_ratio"
+PARENT_BASED_ALWAYS_ON_SAMPLER = "parentbased_always_on"
+PARENT_BASED_ALWAYS_OFF_SAMPLER = "parentbased_always_off"
+PARENT_BASED_TRACE_ID_RATIO_SAMPLER = "parentbased_trace_id_ratio"
+SUPPORTED_OTEL_SAMPLERS = (
+    RATE_LIMITED_SAMPLER,
+    FIXED_PERCENTAGE_SAMPLER,
+    ALWAYS_ON_SAMPLER,
+    ALWAYS_OFF_SAMPLER,
+    TRACE_ID_RATIO_SAMPLER,
+    PARENT_BASED_ALWAYS_ON_SAMPLER,
+    PARENT_BASED_ALWAYS_OFF_SAMPLER,
+    PARENT_BASED_TRACE_ID_RATIO_SAMPLER,
+)
 
 # --------------------Autoinstrumentation Configuration------------------------------------------
 

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_utils/configurations.py

@@ -23,6 +23,13 @@
 )
 from opentelemetry.sdk.resources import Resource
 
+from opentelemetry.sdk.trace.sampling import (
+    TraceIdRatioBased,
+    ALWAYS_OFF,
+    ALWAYS_ON,
+    ParentBased,
+)
+
 from azure.monitor.opentelemetry._constants import (
     _AZURE_APP_SERVICE_RESOURCE_DETECTOR_NAME,
     _AZURE_VM_RESOURCE_DETECTOR_NAME,
@@ -47,6 +54,15 @@
     RATE_LIMITED_SAMPLER,
     FIXED_PERCENTAGE_SAMPLER,
     ENABLE_TRACE_BASED_SAMPLING_ARG,
+    SUPPORTED_OTEL_SAMPLERS,
+    ALWAYS_OFF_SAMPLER,
+    ALWAYS_ON_SAMPLER,
+    TRACE_ID_RATIO_SAMPLER,
+    PARENT_BASED_ALWAYS_ON_SAMPLER,
+    PARENT_BASED_ALWAYS_OFF_SAMPLER,
+    PARENT_BASED_TRACE_ID_RATIO_SAMPLER,
+    SAMPLING_ARG,
+    SAMPLER_TYPE,
 )
 from azure.monitor.opentelemetry._types import ConfigurationValue
 from azure.monitor.opentelemetry._version import VERSION
@@ -145,6 +161,7 @@ def _default_resource(configurations):
         configurations[RESOURCE_ARG] = Resource.create(configurations[RESOURCE_ARG].attributes)
 
 
+# pylint: disable=too-many-statements, disable=too-many-branches
 def _default_sampling_ratio(configurations):
     default_value = 1.0
     sampler_type = environ.get(OTEL_TRACES_SAMPLER)
@@ -188,16 +205,75 @@ def _default_sampling_ratio(configurations):
             )
             configurations[SAMPLING_RATIO_ARG] = default_value
 
+    # Handle always_on sampler
+    elif sampler_type == ALWAYS_ON_SAMPLER:
+        configurations[SAMPLING_ARG] = 1.0
+        configurations[SAMPLER_TYPE] = ALWAYS_ON_SAMPLER
+
+    # Handle always_off sampler
+    elif sampler_type == ALWAYS_OFF_SAMPLER:
+        configurations[SAMPLING_ARG] = 0.0
+        configurations[SAMPLER_TYPE] = ALWAYS_OFF_SAMPLER
+
+    # Handle trace_id_ratio sampler
+    elif sampler_type == TRACE_ID_RATIO_SAMPLER:
+        try:
+            sampler_value = float(sampler_arg) if sampler_arg is not None else default_value
+            if sampler_value < 0 or sampler_value > 1:
+                _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a value between 0 and 1.")
+                sampler_value = default_value
+            else:
+                _logger.info("Using sampling value: %s", sampler_value)
+            configurations[SAMPLING_ARG] = sampler_value
+        except ValueError as e:
+            _logger.error(  # pylint: disable=C
+                _INVALID_FLOAT_MESSAGE,
+                OTEL_TRACES_SAMPLER_ARG,
+                default_value,
+                e,
+            )
+            configurations[SAMPLING_ARG] = sampler_value
+        configurations[SAMPLER_TYPE] = TRACE_ID_RATIO_SAMPLER
+
+    # Handle parentbased_always_on sampler
+    elif sampler_type == PARENT_BASED_ALWAYS_ON_SAMPLER:
+        configurations[SAMPLING_ARG] = 1.0
+        configurations[SAMPLER_TYPE] = PARENT_BASED_ALWAYS_ON_SAMPLER
+
+    # Handle parentbased_always_off sampler
+    elif sampler_type == PARENT_BASED_ALWAYS_OFF_SAMPLER:
+        configurations[SAMPLING_ARG] = 0.0
+        configurations[SAMPLER_TYPE] = PARENT_BASED_ALWAYS_OFF_SAMPLER
+
+    # Handle parentbased_traceidratio sampler
+    elif sampler_type == PARENT_BASED_TRACE_ID_RATIO_SAMPLER:
+        try:
+            sampler_value = float(sampler_arg) if sampler_arg is not None else default_value
+            if sampler_value < 0 or sampler_value > 1:
+                _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a value between 0 and 1.")
+                sampler_value = default_value
+            else:
+                _logger.info("Using sampling value: %s", sampler_value)
+            configurations[SAMPLING_ARG] = sampler_value
+        except ValueError as e:
+            _logger.error(  # pylint: disable=C
+                _INVALID_FLOAT_MESSAGE,
+                OTEL_TRACES_SAMPLER_ARG,
+                default_value,
+                e,
+            )
+            configurations[SAMPLING_ARG] = sampler_value
+        configurations[SAMPLER_TYPE] = PARENT_BASED_TRACE_ID_RATIO_SAMPLER
+
     # Handle all other cases (no sampler type specified or unsupported sampler type)
     else:
         if configurations.get(SAMPLING_RATIO_ARG) is None:
             configurations[SAMPLING_RATIO_ARG] = default_value
         if sampler_type is not None:
             _logger.error(  # pylint: disable=C
                 "Invalid argument for the sampler to be used for tracing. "
-                "Supported values are %s and %s. Defaulting to %s: %s",
-                RATE_LIMITED_SAMPLER,
-                FIXED_PERCENTAGE_SAMPLER,
+                "Supported values are %s. Defaulting to %s: %s",
+                SUPPORTED_OTEL_SAMPLERS,
                 FIXED_PERCENTAGE_SAMPLER,
                 configurations[SAMPLING_RATIO_ARG],
             )
@@ -259,3 +335,19 @@ def _is_instrumentation_enabled(configurations, lib_name):
 
 def _default_enable_trace_based_sampling(configurations):
     configurations.setdefault(ENABLE_TRACE_BASED_SAMPLING_ARG, False)
+
+
+def _get_sampler_from_name(sampler_type, sampler_arg):
+    if sampler_type == ALWAYS_ON_SAMPLER:
+        return ALWAYS_ON
+    if sampler_type == ALWAYS_OFF_SAMPLER:
+        return ALWAYS_OFF
+    if sampler_type == TRACE_ID_RATIO_SAMPLER:
+        ratio = float(sampler_arg) if sampler_arg is not None else 1.0
+        return TraceIdRatioBased(ratio)
+    if sampler_type == PARENT_BASED_ALWAYS_OFF_SAMPLER:
+        return ParentBased(ALWAYS_OFF)
+    if sampler_type == PARENT_BASED_TRACE_ID_RATIO_SAMPLER:
+        ratio = float(sampler_arg) if sampler_arg is not None else 1.0
+        return ParentBased(TraceIdRatioBased(ratio))
+    return ParentBased(ALWAYS_ON)