Export unsampled span in AWS Lambda environment

Author: wangzlei

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.

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()

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"
 

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)

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])

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")