diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py index d8d24ac3e..3f3f6df63 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py @@ -288,6 +288,28 @@ def _export_unsampled_span_for_lambda(trace_provider: TracerProvider, resource: ) +def _export_unsampled_span_for_agent_observability(trace_provider: TracerProvider, resource: Resource = None): + if not is_agent_observability_enabled(): + return + + # Get the traces endpoint from environment + traces_endpoint = os.environ.get(OTEL_EXPORTER_OTLP_TRACES_ENDPOINT) + + if not traces_endpoint: + # No traces endpoint configured, skip unsampled span export + _logger.warning("No traces endpoint configured for agent observability unsampled spans") + return + + # Create the appropriate span exporter based on the endpoint + if _is_aws_otlp_endpoint(traces_endpoint, "xray"): + span_exporter = OTLPAwsSpanExporter(endpoint=traces_endpoint, logger_provider=get_logger_provider()) + else: + span_exporter = OTLPSpanExporter(endpoint=traces_endpoint) + + # Add the unsampled span processor + trace_provider.add_span_processor(BatchUnsampledSpanProcessor(span_exporter=span_exporter)) + + def _is_defer_to_workers_enabled(): return os.environ.get(OTEL_AWS_PYTHON_DEFER_TO_WORKERS_ENABLED_CONFIG, "false").strip().lower() == "true" @@ -426,6 +448,9 @@ def _customize_span_processors(provider: TracerProvider, resource: Resource) -> if _is_lambda_environment(): provider.add_span_processor(AwsLambdaSpanProcessor()) + # Add unsampled span export for agent observability + _export_unsampled_span_for_agent_observability(provider, resource) + if not _is_application_signals_enabled(): return diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py index 9df1b81ff..370070e71 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py @@ -26,6 +26,7 @@ _customize_span_exporter, _customize_span_processors, _export_unsampled_span_for_lambda, + _export_unsampled_span_for_agent_observability, _init_logging, _is_application_signals_enabled, _is_application_signals_runtime_enabled, @@ -576,6 +577,24 @@ def test_customize_span_processors_lambda(self): os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None) os.environ.pop("AWS_LAMBDA_FUNCTION_NAME", None) + def test_customize_span_processors_with_agent_observability(self): + """Test that _customize_span_processors calls _export_unsampled_span_for_agent_observability""" + mock_tracer_provider: TracerProvider = MagicMock() + + with patch( + "amazon.opentelemetry.distro.aws_opentelemetry_configurator._export_unsampled_span_for_agent_observability" + ) as mock_agent_observability: + # Test that agent observability function is called regardless of application signals + _customize_span_processors(mock_tracer_provider, Resource.get_empty()) + mock_agent_observability.assert_called_once_with(mock_tracer_provider, Resource.get_empty()) + + # Reset and test with application signals enabled + mock_agent_observability.reset_mock() + os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", "True") + _customize_span_processors(mock_tracer_provider, Resource.get_empty()) + mock_agent_observability.assert_called_once_with(mock_tracer_provider, Resource.get_empty()) + os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None) + def test_application_signals_exporter_provider(self): # Check default protocol - HTTP, as specified by AwsOpenTelemetryDistro. exporter: OTLPMetricExporterMixin = ApplicationSignalsExporterProvider().create_exporter() @@ -653,6 +672,57 @@ def test_export_unsampled_span_for_lambda(self): os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None) os.environ.pop("AWS_LAMBDA_FUNCTION_NAME", None) + def test_export_unsampled_span_for_agent_observability(self): + mock_tracer_provider: TracerProvider = MagicMock() + + # Test when agent observability is disabled (default) + _export_unsampled_span_for_agent_observability(mock_tracer_provider, Resource.get_empty()) + self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0) + + # Test when agent observability is enabled with AWS endpoint (the default case) + os.environ["AGENT_OBSERVABILITY_ENABLED"] = "true" + os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = "https://xray.us-east-1.amazonaws.com/v1/traces" + _export_unsampled_span_for_agent_observability(mock_tracer_provider, Resource.get_empty()) + self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 1) + processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0] + self.assertIsInstance(processor, BatchUnsampledSpanProcessor) + + # Clean up + os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None) + os.environ.pop("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", None) + + def test_export_unsampled_span_for_agent_observability_uses_aws_exporter(self): + """Test that OTLPAwsSpanExporter is used for AWS endpoints""" + mock_tracer_provider: TracerProvider = MagicMock() + + with patch( + "amazon.opentelemetry.distro.aws_opentelemetry_configurator.OTLPAwsSpanExporter" + ) as mock_aws_exporter: + with patch( + "amazon.opentelemetry.distro.aws_opentelemetry_configurator.BatchUnsampledSpanProcessor" + ) as mock_processor: + with patch( + "amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_logger_provider" + ) as mock_logger_provider: + os.environ["AGENT_OBSERVABILITY_ENABLED"] = "true" + os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = "https://xray.us-east-1.amazonaws.com/v1/traces" + + _export_unsampled_span_for_agent_observability(mock_tracer_provider, Resource.get_empty()) + + # Verify OTLPAwsSpanExporter is created with correct parameters + mock_aws_exporter.assert_called_once_with( + endpoint="https://xray.us-east-1.amazonaws.com/v1/traces", + logger_provider=mock_logger_provider.return_value, + ) + # Verify BatchUnsampledSpanProcessor wraps the exporter + mock_processor.assert_called_once_with(span_exporter=mock_aws_exporter.return_value) + # Verify processor is added to tracer provider + mock_tracer_provider.add_span_processor.assert_called_once_with(mock_processor.return_value) + + # Clean up + os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None) + os.environ.pop("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", None) + def test_customize_metric_exporter(self): metric_readers = [] views = []