Added changes for rate limited sampler (azure-distro changes)

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

@@ -3,6 +3,8 @@
 ## 1.6.14 (Unreleased)
 
 ### Features Added
+- Added configuration changes for RateLimited Sampler.
+  ([#41976](https://github.com/Azure/azure-sdk-for-python/pull/41976))
 
 ### Breaking Changes
 

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

@@ -35,6 +35,7 @@
     LOGGING_FORMATTER_ARG,
     RESOURCE_ARG,
     SAMPLING_RATIO_ARG,
+    SAMPLING_TRACES_PER_SECOND_ARG,
     SPAN_PROCESSORS_ARG,
     VIEWS_ARG,
 )
@@ -50,6 +51,7 @@
     ApplicationInsightsSampler,
     AzureMonitorMetricExporter,
     AzureMonitorTraceExporter,
+    RateLimitedSampler,
 )
 from azure.monitor.opentelemetry.exporter._utils import (  # pylint: disable=import-error,no-name-in-module
     _is_attach_enabled,
@@ -133,10 +135,21 @@ def configure_azure_monitor(**kwargs) -> None:  # pylint: disable=C4758
 
 def _setup_tracing(configurations: Dict[str, ConfigurationValue]):
     resource: Resource = configurations[RESOURCE_ARG]  # type: ignore
-    sampling_ratio = configurations[SAMPLING_RATIO_ARG]
-    tracer_provider = TracerProvider(
-        sampler=ApplicationInsightsSampler(sampling_ratio=cast(float, sampling_ratio)), resource=resource
-    )
+    if SAMPLING_TRACES_PER_SECOND_ARG in configurations:
+        sampling_traces_per_second = configurations[SAMPLING_TRACES_PER_SECOND_ARG]
+        tracer_provider = TracerProvider(
+            sampler=RateLimitedSampler(
+                target_spans_per_second_limit=cast(float, sampling_traces_per_second)
+            ),
+            resource=resource
+        )
+    else:
+        sampling_ratio = configurations[SAMPLING_RATIO_ARG]
+        tracer_provider = TracerProvider(
+            sampler=ApplicationInsightsSampler(sampling_ratio=cast(float, sampling_ratio)), resource=resource
+        )
+
+
     for span_processor in configurations[SPAN_PROCESSORS_ARG]:  # type: ignore
         tracer_provider.add_span_processor(span_processor)  # type: ignore
     if configurations.get(ENABLE_LIVE_METRICS_ARG):

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

@@ -24,7 +24,9 @@
 SAMPLING_RATIO_ARG = "sampling_ratio"
 SPAN_PROCESSORS_ARG = "span_processors"
 VIEWS_ARG = "views"
-
+RATE_LIMITED_SAMPLER = "microsoft.rate_limited"
+FIXED_PERCENTAGE_SAMPLER = "microsoft.fixed.percentage"
+SAMPLING_TRACES_PER_SECOND_ARG = "sampling_traces_per_second"
 
 # --------------------Autoinstrumentation Configuration------------------------------------------
 

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

@@ -19,6 +19,7 @@
 from opentelemetry.sdk.environment_variables import (
     OTEL_EXPERIMENTAL_RESOURCE_DETECTORS,
     OTEL_TRACES_SAMPLER_ARG,
+    OTEL_TRACES_SAMPLER
 )
 from opentelemetry.sdk.resources import Resource
 
@@ -39,25 +40,25 @@
     LOGGING_FORMAT_ENV_ARG,
     RESOURCE_ARG,
     SAMPLING_RATIO_ARG,
+    SAMPLING_TRACES_PER_SECOND_ARG,
     SPAN_PROCESSORS_ARG,
     VIEWS_ARG,
+    RATE_LIMITED_SAMPLER,
+    FIXED_PERCENTAGE_SAMPLER,
 )
 from azure.monitor.opentelemetry._types import ConfigurationValue
 from azure.monitor.opentelemetry._version import VERSION
 
 
 _INVALID_FLOAT_MESSAGE = "Value of %s must be a float. Defaulting to %s: %s"
+_INVALID_TRACES_PER_SECOND_MESSAGE = "Value of %s must be a positive number for traces per second. Defaulting to %s: %s"
 _SUPPORTED_RESOURCE_DETECTORS = (
     _AZURE_APP_SERVICE_RESOURCE_DETECTOR_NAME,
     _AZURE_VM_RESOURCE_DETECTOR_NAME,
 )
-# TODO: remove when sampler uses env var instead
-SAMPLING_RATIO_ENV_VAR = OTEL_TRACES_SAMPLER_ARG
-
 
 _logger = getLogger(__name__)
 
-
 def _get_configurations(**kwargs) -> Dict[str, ConfigurationValue]:
     configurations = {}
 
@@ -137,21 +138,67 @@ def _default_resource(configurations):
         configurations[RESOURCE_ARG] = Resource.create(configurations[RESOURCE_ARG].attributes)
 
 
-# TODO: remove when sampler uses env var instead
 def _default_sampling_ratio(configurations):
-    default = 1.0
-    if SAMPLING_RATIO_ENV_VAR in environ:
+    default_value = 1.0
+    sampler_type = environ.get(OTEL_TRACES_SAMPLER)
+    sampler_arg = environ.get(OTEL_TRACES_SAMPLER_ARG)
+
+    # Early exit if no sampler argument is provided - always set default sampling_ratio
+    if sampler_arg is None:
+        _logger.error("OTEL_TRACES_SAMPLER_ARG is not set.")
+        configurations[SAMPLING_RATIO_ARG] = default_value
+        return
+
+    # Handle rate-limited sampler
+    if sampler_type == RATE_LIMITED_SAMPLER:
         try:
-            default = float(environ[SAMPLING_RATIO_ENV_VAR])
+            sampler_value = float(sampler_arg)
+            if sampler_value < 0:
+                _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a non-negative number.")
+                sampler_value = default_value
+            else:
+                _logger.info("Using rate limited sampler: %s traces per second", sampler_value)
+            configurations[SAMPLING_TRACES_PER_SECOND_ARG] = sampler_value
+        except ValueError as e:
+            _logger.error(  # pylint: disable=C
+                _INVALID_TRACES_PER_SECOND_MESSAGE,
+                OTEL_TRACES_SAMPLER_ARG,
+                default_value,
+                e,
+            )
+            configurations[SAMPLING_TRACES_PER_SECOND_ARG] = default_value
+
+    # Handle fixed percentage sampler
+    elif sampler_type == FIXED_PERCENTAGE_SAMPLER:
+        try:
+            sampler_value = float(sampler_arg)
+            if sampler_value < 0:
+                _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a non-negative number.")
+                sampler_value = default_value
+            else:
+                _logger.info("Using sampling ratio: %s", sampler_value)
+            configurations[SAMPLING_RATIO_ARG] = sampler_value
         except ValueError as e:
             _logger.error(  # pylint: disable=C
                 _INVALID_FLOAT_MESSAGE,
-                SAMPLING_RATIO_ENV_VAR,
-                default,
+                OTEL_TRACES_SAMPLER_ARG,
+                default_value,
                 e,
             )
-    configurations[SAMPLING_RATIO_ARG] = default
+            configurations[SAMPLING_RATIO_ARG] = default_value
 
+    # Handle all other cases (no sampler type specified or unsupported sampler type)
+    else:
+        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,
+                OTEL_TRACES_SAMPLER,
+                OTEL_TRACES_SAMPLER_ARG,
+            )
 
 def _default_instrumentation_options(configurations):
     otel_disabled_instrumentations = _get_otel_disabled_instrumentations()

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

@@ -82,7 +82,7 @@
     install_requires=[
         "azure-core<2.0.0,>=1.28.0",
         "azure-core-tracing-opentelemetry~=1.0.0b11",
-        "azure-monitor-opentelemetry-exporter~=1.0.0b40",
+        "azure-monitor-opentelemetry-exporter~=1.0.0b41",
         "opentelemetry-instrumentation-django~=0.57b0",
         "opentelemetry-instrumentation-fastapi~=0.57b0",
         "opentelemetry-instrumentation-flask~=0.57b0",

sdk/monitor/azure-monitor-opentelemetry/tests/test_configure.py

@@ -309,6 +309,66 @@ def test_setup_tracing(
                 tp_init_mock.add_span_processor.assert_has_calls([call(custom_sp), call(bsp_init_mock)])
                 self.assertEqual(settings_mock.tracing_implementation, opentelemetry_span_mock)
 
+    @patch(
+        "azure.monitor.opentelemetry._configure.BatchSpanProcessor",
+    )
+    @patch(
+        "azure.monitor.opentelemetry._configure.AzureMonitorTraceExporter",
+    )
+    @patch(
+        "azure.monitor.opentelemetry._configure.set_tracer_provider",
+    )
+    @patch(
+        "azure.monitor.opentelemetry._configure.TracerProvider",
+        autospec=True,
+    )
+    @patch(
+        "azure.monitor.opentelemetry._configure.RateLimitedSampler",
+    )
+    def test_setup_tracing_rate_limited_sampler(
+        self,
+        sampler_mock,
+        tp_mock,
+        set_tracer_provider_mock,
+        trace_exporter_mock,
+        bsp_mock,
+    ):
+        sampler_init_mock = Mock()
+        sampler_mock.return_value = sampler_init_mock
+        tp_init_mock = Mock()
+        tp_mock.return_value = tp_init_mock
+        trace_exp_init_mock = Mock()
+        trace_exporter_mock.return_value = trace_exp_init_mock
+        bsp_init_mock = Mock()
+        bsp_mock.return_value = bsp_init_mock
+        custom_sp = Mock()
+
+        settings_mock = Mock()
+        opentelemetry_span_mock = Mock()
+
+        configurations = {
+            "connection_string": "test_cs",
+            "instrumentation_options": {"azure_sdk": {"enabled": True}},
+            "sampling_traces_per_second": 2.0,
+            "span_processors": [custom_sp],
+            "resource": TEST_RESOURCE,
+        }
+        with patch("azure.monitor.opentelemetry._configure._is_instrumentation_enabled") as instr_mock:
+            instr_mock.return_value = True
+            with patch.dict('sys.modules', {
+                'azure.core.settings': Mock(settings=settings_mock),
+                'azure.core.tracing.ext.opentelemetry_span': Mock(OpenTelemetrySpan=opentelemetry_span_mock)
+            }):
+                _setup_tracing(configurations)
+                sampler_mock.assert_called_once_with(target_spans_per_second_limit=2.0)
+                tp_mock.assert_called_once_with(sampler=sampler_init_mock, resource=TEST_RESOURCE)
+                set_tracer_provider_mock.assert_called_once_with(tp_init_mock)
+                trace_exporter_mock.assert_called_once_with(**configurations)
+                bsp_mock.assert_called_once_with(trace_exp_init_mock)
+                self.assertEqual(tp_init_mock.add_span_processor.call_count, 2)
+                tp_init_mock.add_span_processor.assert_has_calls([call(custom_sp), call(bsp_init_mock)])
+                self.assertEqual(settings_mock.tracing_implementation, opentelemetry_span_mock)
+
 
     @patch("azure.monitor.opentelemetry._configure.getLogger")
     def test_setup_logging(self, get_logger_mock):