Temporarily Default to XRay Remote Sampler is Sampler is not Specified

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py

@@ -14,6 +14,7 @@
     AwsMetricAttributesSpanExporterBuilder,
 )
 from amazon.opentelemetry.distro.aws_span_metrics_processor_builder import AwsSpanMetricsProcessorBuilder
+from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
 from opentelemetry.sdk._configuration import (
     _get_exporter_names,
@@ -26,7 +27,10 @@
     _init_metrics,
     _OTelSDKConfigurator,
 )
-from opentelemetry.sdk.environment_variables import _OTEL_PYTHON_LOGGING_AUTO_INSTRUMENTATION_ENABLED
+from opentelemetry.sdk.environment_variables import (
+    _OTEL_PYTHON_LOGGING_AUTO_INSTRUMENTATION_ENABLED,
+    OTEL_TRACES_SAMPLER_ARG,
+)
 from opentelemetry.sdk.extension.aws.resource.ec2 import AwsEc2ResourceDetector
 from opentelemetry.sdk.extension.aws.resource.ecs import AwsEcsResourceDetector
 from opentelemetry.sdk.extension.aws.resource.eks import AwsEksResourceDetector
@@ -81,8 +85,7 @@ def _initialize_components(auto_instrumentation_version):
         _get_exporter_names("metrics"),
         _get_exporter_names("logs"),
     )
-    sampler_name = _get_sampler()
-    sampler = _import_sampler(sampler_name)
+
     id_generator_name = _get_id_generator()
     id_generator = _import_id_generator(id_generator_name)
     # if env var OTEL_RESOURCE_ATTRIBUTES is given, it will read the service_name
@@ -100,6 +103,9 @@ def _initialize_components(auto_instrumentation_version):
         ]
     ).merge(Resource.create(auto_resource))
 
+    sampler_name = _get_sampler()
+    sampler = _custom_import_sampler(sampler_name, resource)
+
     _init_tracing(
         exporters=trace_exporters,
         id_generator=id_generator,
@@ -137,6 +143,36 @@ def _init_tracing(
     set_tracer_provider(trace_provider)
 
 
+def _custom_import_sampler(sampler_name: str, resource: Resource) -> Sampler:
+    # TODO: Remove `sampler_name is None` condition when xray sampler is configured here:
+    # https://github.com/aws/amazon-cloudwatch-agent-operator/blob/main/pkg/instrumentation/defaultinstrumentation.go#L90
+    if sampler_name is None or sampler_name == "xray":
+        # Example env var value
+        # OTEL_TRACES_SAMPLER_ARG=endpoint=http://localhost:2000,polling_interval=360
+        sampler_argument_env: str = os.getenv(OTEL_TRACES_SAMPLER_ARG, None)
+        endpoint: str = None
+        polling_interval: int = None
+
+        if sampler_argument_env is not None:
+            args = sampler_argument_env.split(",")
+            for arg in args:
+                key_value = arg.split("=", 1)
+                if len(key_value) != 2:
+                    continue
+                if key_value[0] == "endpoint":
+                    endpoint = key_value[1]
+                elif key_value[0] == "polling_interval":
+                    try:
+                        polling_interval = int(key_value[1])
+                    except ValueError as error:
+                        _logger.error("polling_interval in OTEL_TRACES_SAMPLER_ARG must be a number: %s", error)
+
+        _logger.debug("XRay Sampler Endpoint: %s", str(endpoint))
+        _logger.debug("XRay Sampler Polling Interval: %s", str(polling_interval))
+        return AwsXRayRemoteSampler(resource=resource, endpoint=endpoint, polling_interval=polling_interval)
+    return _import_sampler(sampler_name)
+
+
 def _customize_sampler(sampler: Sampler) -> Sampler:
     if not is_smp_enabled():
         return sampler

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py

@@ -3,12 +3,19 @@
 import os
 import time
 from unittest import TestCase
+from unittest.mock import patch
 
-from amazon.opentelemetry.distro.aws_opentelemetry_configurator import AwsOpenTelemetryConfigurator
+from amazon.opentelemetry.distro.aws_opentelemetry_configurator import (
+    AwsOpenTelemetryConfigurator,
+    _custom_import_sampler,
+)
 from amazon.opentelemetry.distro.aws_opentelemetry_distro import AwsOpenTelemetryDistro
+from amazon.opentelemetry.distro.sampler._aws_xray_sampling_client import _AwsXRaySamplingClient
+from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 from opentelemetry.environment_variables import OTEL_LOGS_EXPORTER, OTEL_METRICS_EXPORTER, OTEL_TRACES_EXPORTER
 from opentelemetry.sdk.environment_variables import OTEL_TRACES_SAMPLER, OTEL_TRACES_SAMPLER_ARG
 from opentelemetry.sdk.trace import Span, Tracer, TracerProvider
+from opentelemetry.sdk.trace.sampling import Sampler
 from opentelemetry.trace import get_tracer_provider
 
 
@@ -52,3 +59,93 @@ def test_trace_id_ratio_sampler(self):
                 span.end()
             # Configured for 1%, confirm there are at most 5% to account for randomness and reduce test flakiness.
             self.assertGreater(0.05, num_sampled / num_spans)
+
+    # Test method for importing xray sampler
+    # Cannot test this logic via `aws_otel_configurator.configure()` because that will
+    # attempt to setup tracer provider again, which can be only be done once (already done)
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_rule_poller", lambda x: None)
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_target_poller", lambda x: None)
+    def test_import_xray_sampler_without_environment_arguments(self):
+        os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
+
+        # May log http request error as xray sampler will attempt to fetch rules
+        xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
+        xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
+        self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertEqual(
+            xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://127.0.0.1:2000/GetSamplingRules"
+        )
+
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_rule_poller", lambda x: None)
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_target_poller", lambda x: None)
+    def test_import_xray_sampler_with_valid_environment_arguments(self):
+        os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
+        os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint=http://localhost:2000,polling_interval=600")
+
+        # May log http request error as xray sampler will attempt to fetch rules
+        xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
+        xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
+        self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 600)
+        self.assertEqual(
+            xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://localhost:2000/GetSamplingRules"
+        )
+
+        os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
+        os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "polling_interval=123")
+
+        # May log http request error as xray sampler will attempt to fetch rules
+        xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
+        xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
+        self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 123)
+        self.assertEqual(
+            xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://127.0.0.1:2000/GetSamplingRules"
+        )
+
+        os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
+        os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint=http://cloudwatch-agent.amazon-cloudwatch:2000")
+
+        # May log http request error as xray sampler will attempt to fetch rules
+        xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
+        xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
+        self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertEqual(
+            xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint,
+            "http://cloudwatch-agent.amazon-cloudwatch:2000/GetSamplingRules",
+        )
+
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_rule_poller", lambda x: None)
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_target_poller", lambda x: None)
+    def test_import_xray_sampler_with_invalid_environment_arguments(self):
+        os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
+        os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint=h=tt=p://=loca=lho=st:2000,polling_interval=FOOBAR")
+
+        # May log http request error as xray sampler will attempt to fetch rules
+        xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
+        xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
+        self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertEqual(
+            xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint,
+            "h=tt=p://=loca=lho=st:2000/GetSamplingRules",
+        )
+
+        os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
+        os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, ",,=,==,,===,")
+
+        # May log http request error as xray sampler will attempt to fetch rules
+        xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
+        xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
+        self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertEqual(
+            xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://127.0.0.1:2000/GetSamplingRules"
+        )
+
+        os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
+        os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint,polling_interval")
+
+        # May log http request error as xray sampler will attempt to fetch rules
+        xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
+        xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
+        self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertEqual(
+            xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://127.0.0.1:2000/GetSamplingRules"
+        )