From 288df57220f9ca99606c31341da465541aeb3fa9 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 30 Aug 2024 14:14:36 -0700 Subject: [PATCH 1/9] Export unsampled span in AWS Lambda environment --- .../distro/_aws_attribute_keys.py | 1 + .../aws_batch_unsampled_span_processor.py | 37 ++++++++++++++++++ .../distro/aws_opentelemetry_configurator.py | 17 ++++++++ .../opentelemetry/distro/otlp_udp_exporter.py | 15 +++++-- ...test_aws_batch_unsampled_span_processor.py | 39 +++++++++++++++++++ .../distro/test_otlp_udp_exporter.py | 9 ++++- 6 files changed, 113 insertions(+), 5 deletions(-) create mode 100644 aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py create mode 100644 aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py index e093daf96..923025894 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py @@ -11,6 +11,7 @@ AWS_REMOTE_RESOURCE_IDENTIFIER: str = "aws.remote.resource.identifier" AWS_SDK_DESCENDANT: str = "aws.sdk.descendant" AWS_CONSUMER_PARENT_SPAN_KIND: str = "aws.consumer.parent.span.kind" +AWS_TRACE_FLAG_UNSAMPLED: str = "aws.trace.flag.unsampled" # AWS_#_NAME attributes are not supported in python as they are not part of the Semantic Conventions. # TODOļ¼šMove to Semantic Conventions when these attributes are added. diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py new file mode 100644 index 000000000..0934fc9a0 --- /dev/null +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py @@ -0,0 +1,37 @@ +import logging +from typing import Optional + +from amazon.opentelemetry.distro._aws_attribute_keys import AWS_TRACE_FLAG_UNSAMPLED +from opentelemetry.context import Context +from opentelemetry.sdk.trace import ReadableSpan, Span +from opentelemetry.sdk.trace.export import BatchSpanProcessor as BaseBatchSpanProcessor + +logger = logging.getLogger(__name__) + +SPANUNSAMPLED_FLAG = "OTEL_AWS_APP_SIGNALS_ENABLED" + + +class BatchUnsampledSpanProcessor(BaseBatchSpanProcessor): + + def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None: + if not span.context.trace_flags.sampled: + span.set_attribute(AWS_TRACE_FLAG_UNSAMPLED, "True") + + def on_end(self, span: ReadableSpan) -> None: + if span.context.trace_flags.sampled: + return + + if self.done: + logger.warning("Already shutdown, dropping span.") + return + + if len(self.queue) == self.max_queue_size: + if not self._spans_dropped: + logger.warning("Queue is full, likely spans will be dropped.") + self._spans_dropped = True + + self.queue.appendleft(span) + + if len(self.queue) >= self.max_export_batch_size: + with self.condition: + self.condition.notify() 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 99829e5a0..935e74caa 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 @@ -12,6 +12,7 @@ from amazon.opentelemetry.distro.attribute_propagating_span_processor_builder import ( AttributePropagatingSpanProcessorBuilder, ) +from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor from amazon.opentelemetry.distro.aws_metric_attributes_span_exporter_builder import ( AwsMetricAttributesSpanExporterBuilder, ) @@ -153,6 +154,7 @@ def _init_tracing( span_exporter: SpanExporter = exporter_class(**exporter_args) span_exporter = _customize_exporter(span_exporter, resource) trace_provider.add_span_processor(BatchSpanProcessor(span_exporter)) + _export_unsampled_span_for_lambda(trace_provider, resource) _customize_span_processors(trace_provider, resource) @@ -162,6 +164,21 @@ def _init_tracing( # END The OpenTelemetry Authors code +def _export_unsampled_span_for_lambda(trace_provider: TracerProvider, resource: Resource = None): + if not _is_application_signals_enabled(): + return + if not _is_lambda_environment(): + return + + traces_endpoint = os.environ.get(AWS_XRAY_DAEMON_ADDRESS_CONFIG, "127.0.0.1:2000") + + span_exporter = AwsMetricAttributesSpanExporterBuilder( + OTLPUdpSpanExporter(endpoint=traces_endpoint, sampled=False), resource + ).build() + + trace_provider.add_span_processor(BatchUnsampledSpanProcessor(span_exporter)) + + def _is_defer_to_workers_enabled(): return os.environ.get(OTEL_AWS_PYTHON_DEFER_TO_WORKERS_ENABLED_CONFIG, "false").strip().lower() == "true" diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_udp_exporter.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_udp_exporter.py index 0b7f8f8c0..1849384fc 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_udp_exporter.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_udp_exporter.py @@ -20,7 +20,10 @@ DEFAULT_ENDPOINT = "127.0.0.1:2000" PROTOCOL_HEADER = '{"format":"json","version":1}\n' FORMAT_OTEL_METRICS_BINARY_PREFIX = "M1" -FORMAT_OTEL_TRACES_BINARY_PREFIX = "T1" + +# TODO: update sampled and unsampled prefix later +FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX = "T1" +FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX = "T1" _logger: Logger = getLogger(__name__) @@ -98,15 +101,21 @@ def shutdown(self, timeout_millis: float = 30_000, **kwargs) -> None: class OTLPUdpSpanExporter(SpanExporter): - def __init__(self, endpoint: Optional[str] = None): + def __init__(self, endpoint: Optional[str] = None, sampled: bool = True): self._udp_exporter = UdpExporter(endpoint=endpoint) + self._sampled = sampled @override def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: serialized_data = encode_spans(spans).SerializeToString() try: - self._udp_exporter.send_data(data=serialized_data, signal_format_prefix=FORMAT_OTEL_TRACES_BINARY_PREFIX) + prefix = ( + FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX + if self._sampled + else FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX + ) + self._udp_exporter.send_data(data=serialized_data, signal_format_prefix=prefix) return SpanExportResult.SUCCESS except Exception as exc: # pylint: disable=broad-except _logger.error("Error exporting spans: %s", exc) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py new file mode 100644 index 000000000..0af2f8b5f --- /dev/null +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py @@ -0,0 +1,39 @@ +from unittest import TestCase +from unittest.mock import MagicMock, patch + +from amazon.opentelemetry.distro._aws_attribute_keys import AWS_TRACE_FLAG_UNSAMPLED +from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor +from opentelemetry.trace import TraceFlags + + +class TestBatchUnsampledSpanProcessor(TestCase): + + def setUp(self): + self.mock_exporter = MagicMock() + self.processor = BatchUnsampledSpanProcessor(self.mock_exporter) + + @patch("opentelemetry.sdk.trace.Span") + def test_on_end_sampled(self, mock_span_class): + trace_flags = TraceFlags(TraceFlags.SAMPLED) + + mock_span = mock_span_class.return_value + mock_span.context.trace_flags = trace_flags + + self.processor.on_start(mock_span) + self.processor.on_end(mock_span) + + self.assertEqual(len(self.processor.queue), 0) + mock_span.set_attribute.assert_not_called() + + @patch("opentelemetry.sdk.trace.Span") + def test_on_end_not_sampled(self, mock_span_class): + + trace_flags = TraceFlags(0) + mock_span = mock_span_class.return_value + mock_span.context.trace_flags = trace_flags + + self.processor.on_start(mock_span) + self.processor.on_end(mock_span) + + self.assertEqual(len(self.processor.queue), 1) + self.assertIn(AWS_TRACE_FLAG_UNSAMPLED, mock_span.set_attribute.call_args_list[0][0][0]) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py index a97e02321..85eb84e5a 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py @@ -12,10 +12,13 @@ OTLPUdpMetricExporter, OTLPUdpSpanExporter, UdpExporter, + FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX, ) from opentelemetry.sdk.metrics._internal.export import MetricExportResult from opentelemetry.sdk.trace.export import SpanExportResult +from build.python.amazon.opentelemetry.distro.otlp_udp_exporter import FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX + class TestUdpExporter(TestCase): @@ -102,9 +105,11 @@ def test_export(self, mock_udp_exporter, mock_encode_spans): mock_udp_exporter_instance = mock_udp_exporter.return_value mock_encoded_data = MagicMock() mock_encode_spans.return_value.SerializeToString.return_value = mock_encoded_data - exporter = OTLPUdpSpanExporter() + exporter = OTLPUdpSpanExporter(sampled=False) result = exporter.export(MagicMock()) - mock_udp_exporter_instance.send_data.assert_called_once_with(data=mock_encoded_data, signal_format_prefix="T1") + mock_udp_exporter_instance.send_data.assert_called_once_with( + data=mock_encoded_data, signal_format_prefix=FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX + ) self.assertEqual(result, SpanExportResult.SUCCESS) @patch("amazon.opentelemetry.distro.otlp_udp_exporter.encode_spans") From f91d294f1149c774e1b5bc239c246e387fb8e9ee Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 30 Aug 2024 15:43:44 -0700 Subject: [PATCH 2/9] Fix unittest typo --- .../tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py index 85eb84e5a..436194fb4 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py @@ -17,8 +17,6 @@ from opentelemetry.sdk.metrics._internal.export import MetricExportResult from opentelemetry.sdk.trace.export import SpanExportResult -from build.python.amazon.opentelemetry.distro.otlp_udp_exporter import FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX - class TestUdpExporter(TestCase): From ec982938c1e3f6657106bfe7e443675ae2294f7b Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 30 Aug 2024 15:55:16 -0700 Subject: [PATCH 3/9] fix unit test coverage --- .../test_aws_batch_unsampled_span_processor.py | 2 +- .../distro/test_otlp_udp_exporter.py | 16 +++++++++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py index 0af2f8b5f..6b91ee8ec 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py @@ -10,7 +10,7 @@ class TestBatchUnsampledSpanProcessor(TestCase): def setUp(self): self.mock_exporter = MagicMock() - self.processor = BatchUnsampledSpanProcessor(self.mock_exporter) + self.processor = BatchUnsampledSpanProcessor(self.mock_exporter, max_queue_size=1, max_export_batch_size=1) @patch("opentelemetry.sdk.trace.Span") def test_on_end_sampled(self, mock_span_class): diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py index 436194fb4..0061567c7 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py @@ -12,6 +12,7 @@ OTLPUdpMetricExporter, OTLPUdpSpanExporter, UdpExporter, + FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX, FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX, ) from opentelemetry.sdk.metrics._internal.export import MetricExportResult @@ -99,7 +100,7 @@ class TestOTLPUdpSpanExporter(unittest.TestCase): @patch("amazon.opentelemetry.distro.otlp_udp_exporter.encode_spans") @patch("amazon.opentelemetry.distro.otlp_udp_exporter.UdpExporter") - def test_export(self, mock_udp_exporter, mock_encode_spans): + def test_export_unsampled_span(self, mock_udp_exporter, mock_encode_spans): mock_udp_exporter_instance = mock_udp_exporter.return_value mock_encoded_data = MagicMock() mock_encode_spans.return_value.SerializeToString.return_value = mock_encoded_data @@ -110,6 +111,19 @@ def test_export(self, mock_udp_exporter, mock_encode_spans): ) self.assertEqual(result, SpanExportResult.SUCCESS) + @patch("amazon.opentelemetry.distro.otlp_udp_exporter.encode_spans") + @patch("amazon.opentelemetry.distro.otlp_udp_exporter.UdpExporter") + def test_export_sampled_span(self, mock_udp_exporter, mock_encode_spans): + mock_udp_exporter_instance = mock_udp_exporter.return_value + mock_encoded_data = MagicMock() + mock_encode_spans.return_value.SerializeToString.return_value = mock_encoded_data + exporter = OTLPUdpSpanExporter() + result = exporter.export(MagicMock()) + mock_udp_exporter_instance.send_data.assert_called_once_with( + data=mock_encoded_data, signal_format_prefix=FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX + ) + self.assertEqual(result, SpanExportResult.SUCCESS) + @patch("amazon.opentelemetry.distro.otlp_udp_exporter.encode_spans") @patch("amazon.opentelemetry.distro.otlp_udp_exporter.UdpExporter") def test_export_with_exception(self, mock_udp_exporter, mock_encode_spans): From c881fd7334a2a132010cf013525231df6802d230 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 30 Aug 2024 16:04:43 -0700 Subject: [PATCH 4/9] fix unit test coverage --- ...test_aws_batch_unsampled_span_processor.py | 24 +++++++++++++++---- .../distro/test_otlp_udp_exporter.py | 4 ++-- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py index 6b91ee8ec..22bdb14d2 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py @@ -21,6 +21,7 @@ def test_on_end_sampled(self, mock_span_class): self.processor.on_start(mock_span) self.processor.on_end(mock_span) + self.processor.shutdown() self.assertEqual(len(self.processor.queue), 0) mock_span.set_attribute.assert_not_called() @@ -29,11 +30,24 @@ def test_on_end_sampled(self, mock_span_class): def test_on_end_not_sampled(self, mock_span_class): trace_flags = TraceFlags(0) - mock_span = mock_span_class.return_value - mock_span.context.trace_flags = trace_flags + mock_span1 = mock_span_class.return_value + mock_span1.context.trace_flags = trace_flags - self.processor.on_start(mock_span) - self.processor.on_end(mock_span) + self.processor.on_start(mock_span1) + self.processor.on_end(mock_span1) + + mock_span2 = mock_span_class.return_value + mock_span2.context.trace_flags = trace_flags + self.processor.on_start(mock_span2) + self.processor.on_end(mock_span2) + + self.assertEqual(len(self.processor.queue), 1) + self.assertIn(AWS_TRACE_FLAG_UNSAMPLED, mock_span1.set_attribute.call_args_list[0][0][0]) + + mock_span2 = mock_span_class.return_value + mock_span2.context.trace_flags = trace_flags + self.processor.on_start(mock_span2) + self.processor.on_end(mock_span2) self.assertEqual(len(self.processor.queue), 1) - self.assertIn(AWS_TRACE_FLAG_UNSAMPLED, mock_span.set_attribute.call_args_list[0][0][0]) + self.assertIn(AWS_TRACE_FLAG_UNSAMPLED, mock_span2.set_attribute.call_args_list[0][0][0]) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py index 0061567c7..aa9d19330 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py @@ -8,12 +8,12 @@ from amazon.opentelemetry.distro.otlp_udp_exporter import ( DEFAULT_ENDPOINT, + FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX, + FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX, PROTOCOL_HEADER, OTLPUdpMetricExporter, OTLPUdpSpanExporter, UdpExporter, - FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX, - FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX, ) from opentelemetry.sdk.metrics._internal.export import MetricExportResult from opentelemetry.sdk.trace.export import SpanExportResult From d2bbcb6121a65e438bd218786a4e87b2f25fa878 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 30 Aug 2024 16:13:51 -0700 Subject: [PATCH 5/9] fix unit test coverage --- .../distro/test_aws_batch_unsampled_span_processor.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py index 22bdb14d2..e28f71e1e 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py @@ -21,7 +21,6 @@ def test_on_end_sampled(self, mock_span_class): self.processor.on_start(mock_span) self.processor.on_end(mock_span) - self.processor.shutdown() self.assertEqual(len(self.processor.queue), 0) mock_span.set_attribute.assert_not_called() @@ -44,10 +43,10 @@ def test_on_end_not_sampled(self, mock_span_class): self.assertEqual(len(self.processor.queue), 1) self.assertIn(AWS_TRACE_FLAG_UNSAMPLED, mock_span1.set_attribute.call_args_list[0][0][0]) + self.processor.shutdown() mock_span2 = mock_span_class.return_value mock_span2.context.trace_flags = trace_flags self.processor.on_start(mock_span2) self.processor.on_end(mock_span2) - self.assertEqual(len(self.processor.queue), 1) - self.assertIn(AWS_TRACE_FLAG_UNSAMPLED, mock_span2.set_attribute.call_args_list[0][0][0]) + self.assertEqual(len(self.processor.queue), 0) From a1154dfe0b02e25673464b8495fb778bf338eec2 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 30 Aug 2024 16:25:28 -0700 Subject: [PATCH 6/9] fix unit test coverage --- .../tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py index aa9d19330..6d3e4bdbb 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py @@ -92,6 +92,7 @@ def test_export_with_exception(self, mock_udp_exporter, mock_encode_metrics): def test_shutdown(self, mock_udp_exporter): mock_udp_exporter_instance = mock_udp_exporter.return_value exporter = OTLPUdpMetricExporter() + exporter.force_flush() exporter.shutdown() mock_udp_exporter_instance.shutdown.assert_called_once() @@ -141,4 +142,5 @@ def test_shutdown(self, mock_udp_exporter): mock_udp_exporter_instance = mock_udp_exporter.return_value exporter = OTLPUdpSpanExporter() exporter.shutdown() + exporter.force_flush() mock_udp_exporter_instance.shutdown.assert_called_once() From a1d322d1cd7461fad6c03302a66966e69cbe93f4 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 30 Aug 2024 20:16:10 -0700 Subject: [PATCH 7/9] fix unit test coverage --- .../distro/aws_batch_unsampled_span_processor.py | 4 ++++ .../distro/aws_opentelemetry_configurator.py | 2 +- .../test_aws_batch_unsampled_span_processor.py | 2 ++ .../test_aws_opentelementry_configurator.py | 15 +++++++++++++++ 4 files changed, 22 insertions(+), 1 deletion(-) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py index 0934fc9a0..d27f2c58b 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py @@ -1,3 +1,5 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 import logging from typing import Optional @@ -13,6 +15,7 @@ class BatchUnsampledSpanProcessor(BaseBatchSpanProcessor): + # pylint: disable=no-self-use def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None: if not span.context.trace_flags.sampled: span.set_attribute(AWS_TRACE_FLAG_UNSAMPLED, "True") @@ -26,6 +29,7 @@ def on_end(self, span: ReadableSpan) -> None: return if len(self.queue) == self.max_queue_size: + # pylint: disable=access-member-before-definition if not self._spans_dropped: logger.warning("Queue is full, likely spans will be dropped.") self._spans_dropped = True 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 935e74caa..db0892d8c 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 @@ -154,8 +154,8 @@ def _init_tracing( span_exporter: SpanExporter = exporter_class(**exporter_args) span_exporter = _customize_exporter(span_exporter, resource) trace_provider.add_span_processor(BatchSpanProcessor(span_exporter)) - _export_unsampled_span_for_lambda(trace_provider, resource) + _export_unsampled_span_for_lambda(trace_provider, resource) _customize_span_processors(trace_provider, resource) set_tracer_provider(trace_provider) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py index e28f71e1e..004d2f1d1 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_batch_unsampled_span_processor.py @@ -1,3 +1,5 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 from unittest import TestCase from unittest.mock import MagicMock, patch 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 d8bfc5756..06247fd27 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 @@ -7,6 +7,7 @@ from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler from amazon.opentelemetry.distro.attribute_propagating_span_processor import AttributePropagatingSpanProcessor +from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor from amazon.opentelemetry.distro.aws_metric_attributes_span_exporter import AwsMetricAttributesSpanExporter from amazon.opentelemetry.distro.aws_opentelemetry_configurator import ( ApplicationSignalsExporterProvider, @@ -15,6 +16,7 @@ _customize_exporter, _customize_sampler, _customize_span_processors, + _export_unsampled_span_for_lambda, _is_application_signals_enabled, _is_defer_to_workers_enabled, _is_wsgi_master_process, @@ -351,6 +353,19 @@ def test_initialize_components_called_when_deferred_disabled(self, mock_initiali mock_initialize_components.assert_called_once() os.environ.pop("IS_WSGI_MASTER_PROCESS_ALREADY_SEEN", None) + def test_export_unsampled_span_for_lambda(self): + mock_tracer_provider: TracerProvider = MagicMock() + _export_unsampled_span_for_lambda(mock_tracer_provider, Resource.get_empty()) + self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0) + + os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", "True") + os.environ.setdefault("AWS_LAMBDA_FUNCTION_NAME", "myfunction") + _export_unsampled_span_for_lambda(mock_tracer_provider, Resource.get_empty()) + self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 1) + first_processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0] + self.assertIsInstance(first_processor, BatchUnsampledSpanProcessor) + os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None) + def validate_distro_environ(): tc: TestCase = TestCase() From 0fd43f8d006e3531b2171090e3c1ce690360b7f6 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 30 Aug 2024 20:28:29 -0700 Subject: [PATCH 8/9] fix lint --- .../opentelemetry/distro/aws_batch_unsampled_span_processor.py | 1 + .../opentelemetry/distro/test_aws_opentelementry_configurator.py | 1 + 2 files changed, 2 insertions(+) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py index d27f2c58b..785a04941 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py @@ -32,6 +32,7 @@ def on_end(self, span: ReadableSpan) -> None: # pylint: disable=access-member-before-definition if not self._spans_dropped: logger.warning("Queue is full, likely spans will be dropped.") + # pylint: disable=attribute-defined-outside-init self._spans_dropped = True self.queue.appendleft(span) 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 06247fd27..eab86d08e 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 @@ -39,6 +39,7 @@ from opentelemetry.trace import get_tracer_provider +# pylint: disable=too-many-public-methods class TestAwsOpenTelemetryConfigurator(TestCase): """Tests AwsOpenTelemetryConfigurator and AwsOpenTelemetryDistro From 8108f4d86ac9cd49ff2cc9f986ef0c8bc117d56c Mon Sep 17 00:00:00 2001 From: Wang Date: Wed, 4 Sep 2024 10:39:26 -0700 Subject: [PATCH 9/9] address comments --- .../opentelemetry/distro/aws_batch_unsampled_span_processor.py | 2 +- .../opentelemetry/distro/aws_opentelemetry_configurator.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py index 785a04941..36e239188 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_batch_unsampled_span_processor.py @@ -18,7 +18,7 @@ class BatchUnsampledSpanProcessor(BaseBatchSpanProcessor): # pylint: disable=no-self-use def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None: if not span.context.trace_flags.sampled: - span.set_attribute(AWS_TRACE_FLAG_UNSAMPLED, "True") + span.set_attribute(AWS_TRACE_FLAG_UNSAMPLED, True) def on_end(self, span: ReadableSpan) -> None: if span.context.trace_flags.sampled: 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 db0892d8c..6211186a4 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 @@ -155,7 +155,6 @@ def _init_tracing( span_exporter = _customize_exporter(span_exporter, resource) trace_provider.add_span_processor(BatchSpanProcessor(span_exporter)) - _export_unsampled_span_for_lambda(trace_provider, resource) _customize_span_processors(trace_provider, resource) set_tracer_provider(trace_provider) @@ -276,6 +275,8 @@ def _customize_span_processors(provider: TracerProvider, resource: Resource) -> if not _is_application_signals_enabled(): return + _export_unsampled_span_for_lambda(provider, resource) + # Construct and set local and remote attributes span processor provider.add_span_processor(AttributePropagatingSpanProcessorBuilder().build())