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

* RateLimitedSampler distro changes

* Fixed formatting errors

* Updated config and tests

* Updated CHANGELOG with URL

* Modified CHANGELOG

* Retrigger build

* Upgraded exporter version in setup.py

* Fixed argument and linting errors

* Addressed comments

* Bumped up the unreleased version number

* Updated CHANGELOG

* Updated CHANGELOG and README

* Updated url in README

* Fix link issue

* Fix link

* bandit

* capitalized rate limited sampler

Author: rads-1996

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

@@ -3,6 +3,8 @@
 ## 1.7.1 (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/README.md

@@ -67,6 +67,7 @@ You can use `configure_azure_monitor` to set up instrumentation for your app to
 | `resource` | Specifies the OpenTelemetry [Resource][ot_spec_resource] associated with your application. Passed in [Resource Attributes][ot_spec_resource_attributes] take priority over default attributes and those from [Resource Detectors][ot_python_resource_detectors]. | [OTEL_SERVICE_NAME][ot_spec_service_name], [OTEL_RESOURCE_ATTRIBUTES][ot_spec_resource_attributes], [OTEL_EXPERIMENTAL_RESOURCE_DETECTORS][ot_python_resource_detectors] |
 | `span_processors` | A list of [span processors][ot_span_processor] that will perform processing on each of your spans before they are exported. Useful for filtering/modifying telemetry. | `N/A` |
 | `views` | A list of [views][ot_view] that will be used to customize metrics exported by the SDK. | `N/A` |
+| `traces_per_second` | Configures the Rate Limited sampler by specifying the maximum number of traces to sample per second. When set, this automatically enables the rate-limited sampler. Alternatively, you can configure sampling using the `OTEL_TRACES_SAMPLER` and `OTEL_TRACES_SAMPLER_ARG` environment variables as described in the table below. Please note that the sampling configuration via environment variables will have precedence over the sampling exporter/distro options. | `N/A`
 
 You can configure further with [OpenTelemetry environment variables][ot_env_vars].
 
@@ -78,7 +79,8 @@ 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_ARG` | Specifies the ratio of distributed tracing telemetry to be [sampled][application_insights_sampling]. Accepted values are in the range [0,1]. Defaults to 1.0, meaning no telemetry is sampled out. |
+| `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_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. |
 
@@ -216,6 +218,7 @@ contact [[email protected]](mailto:[email protected]) with any additio
 [application_insights_live_metrics]: https://learn.microsoft.com/azure/azure-monitor/app/live-stream
 [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
 [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

@@ -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:
+        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
+        )
+    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 = "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,61 @@ 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)
+
+    # Handle rate-limited sampler
+    if sampler_type == RATE_LIMITED_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 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:
-            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 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}},
+            "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):