some refactoring

Author: srprash

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

@@ -43,7 +43,6 @@ def __init__(self, delegate: SpanExporter, generator: MetricAttributeGenerator,
 
     @override
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
-        print("IN AWS EXPORTER:::")
         modified_spans: Sequence[ReadableSpan] = self._add_metric_attributes(spans)
         return self._delegate.export(modified_spans)
 

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

@@ -18,7 +18,7 @@
 from amazon.opentelemetry.distro.aws_span_metrics_processor_builder import AwsSpanMetricsProcessorBuilder
 from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 from amazon.opentelemetry.distro.otlp_udp_exporter import OtlpUdpMetricExporter
-from amazon.opentelemetry.distro.otlp_udp_exporter import OtlpUdpTraceExporter
+from amazon.opentelemetry.distro.otlp_udp_exporter import OtlpUdpSpanExporter
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
 from opentelemetry.sdk._configuration import (
     _get_exporter_names,
@@ -221,8 +221,9 @@ def _customize_sampler(sampler: Sampler) -> Sampler:
 def _customize_exporter(span_exporter: SpanExporter, resource: Resource) -> SpanExporter:
     if not _is_application_signals_enabled():
         return span_exporter
-    udp_span_exporter = OtlpUdpTraceExporter(endpoint="something")
-    return AwsMetricAttributesSpanExporterBuilder(udp_span_exporter, resource).build()
+    if _is_lambda_environment():
+        return AwsMetricAttributesSpanExporterBuilder(OtlpUdpSpanExporter(), resource).build()
+    return AwsMetricAttributesSpanExporterBuilder(span_exporter, resource).build()
 
 
 def _customize_span_processors(provider: TracerProvider, resource: Resource) -> None:
@@ -298,7 +299,8 @@ def create_exporter(self):
             temporality_dict[typ] = AggregationTemporality.DELTA
 
         if _is_lambda_environment():
-            return OtlpUdpMetricExporter(temporality_dict=temporality_dict)
+            # When running in Lambda, export Application Signals metrics over UDP
+            return OtlpUdpMetricExporter(preferred_temporality=temporality_dict)
 
         if protocol == "http/protobuf":
             application_signals_endpoint = os.environ.get(

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

@@ -1,6 +1,7 @@
-import typing
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Dict
 
+from opentelemetry.sdk.metrics._internal.aggregation import AggregationTemporality
+from opentelemetry.sdk.metrics.view import Aggregation
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from opentelemetry.sdk.metrics.export import MetricExporter
@@ -16,44 +17,57 @@
 
 import socket
 
-PROTOCOL_HEADER = "{\"format\":\"json\",\"version\":1}"
+DEFAULT_ENDPOINT = "127.0.0.1:2000"
+PROTOCOL_HEADER = '{"format":"json","version":1}\n'
 PROTOCOL_DELIMITER = '\n'
 
 
-class OtlpUdpExporterCommon:
+class UdpExporter:
     def __init__(
             self,
             endpoint: Optional[str] = None
     ):
-        self._endpoint = endpoint or "http://127.0.0.1:2000"
-        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)  #TODO: What does this mean?
-        self._socket.setblocking(False)  #TODO: Is this correct that we don't want to block?
+        self._endpoint = endpoint or DEFAULT_ENDPOINT
+        self._host, self._port = self._parse_endpoint(self._endpoint)
+        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self._socket.setblocking(False)
 
-    def send_data(self, data, format):
-        udp_data = {
-            "format": format,
-            "data": data
-        }
+    def send_data(self, data: str, format: str):
+        udp_data = f'{{"format":"{format}","data":{data}}}'
+        message = PROTOCOL_HEADER + udp_data
 
-        message = "%s%s%s" % (PROTOCOL_HEADER,
-                              PROTOCOL_DELIMITER,
-                              udp_data)
+        try:
+            print("Sending UDP data: ", message)  # TODO: remove
+            self._socket.sendto(message.encode('utf-8'), (self._host, int(self._port)))
+        except Exception as e:
+            print("Error sending UDP data: ", e)
 
-        print("Sending UDP data: ", message)
-        self._socket.sendto(message.encode('utf-8'), self._endpoint)
-        pass
+    def shutdown(self):
+        self._socket.close()
 
+    def _parse_endpoint(self, endpoint: str) -> tuple[str, int]:
+        try:
+            vals = endpoint.split(":")
+            host = vals[0]
+            port = int(vals[1])
+        except Exception as e:
+            raise ValueError(f"Invalid endpoint: {endpoint}") from e
 
-class OtlpUdpMetricExporter(OtlpUdpExporterCommon, MetricExporter):
-    def shutdown(self, timeout_millis: float = 30_000, **kwargs) -> None:
-        pass
+        return host, port
 
-    def force_flush(self, timeout_millis: float = 10_000) -> bool:
-        pass
 
-    def __init__(self, endpoint=None, temporality_dict=None):
-        OtlpUdpExporterCommon.__init__(self, endpoint)
-        MetricExporter.__init__(self, preferred_temporality=temporality_dict)
+class OtlpUdpMetricExporter(MetricExporter):
+    def __init__(
+            self,
+            endpoint: Optional[str] = None,
+            preferred_temporality: Dict[type, AggregationTemporality] = None,
+            preferred_aggregation: Dict[type, Aggregation] = None
+    ):
+        super().__init__(
+            preferred_temporality=preferred_temporality,
+            preferred_aggregation=preferred_aggregation,
+        )
+        self._udp_exporter = UdpExporter(endpoint=endpoint)
 
     @override
     def export(
@@ -62,26 +76,31 @@ def export(
             timeout_millis: float = 10_000,
             **kwargs,
     ) -> MetricExportResult:
-        # serialized_data = encode_metrics(metrics_data).SerializeToString()
-        # self.send_data(data=serialized_data, format="OTEL_V1_METRICS")
-        return MetricExportResult.SUCCESS
+        serialized_data = encode_metrics(metrics_data).SerializeToString()
+        self._udp_exporter.send_data(data=serialized_data, format="OTEL_V1_METRICS")  # TODO: Convert to constant
+        return MetricExportResult.SUCCESS  # TODO: send appropriate status back. Need to??
+
+    def force_flush(self, timeout_millis: float = 10_000) -> bool:
+        return True
+
+    def shutdown(self, timeout_millis: float = 30_000, **kwargs) -> None:
+        self._udp_exporter.shutdown()
 
 
-class OtlpUdpTraceExporter(SpanExporter):
-    def __init__(self, endpoint=None):
-        self._exp = OtlpUdpExporterCommon(endpoint=endpoint)
+class OtlpUdpSpanExporter(SpanExporter):
+    def __init__(self, endpoint: Optional[str] = None):
+        self._udp_exporter = UdpExporter(endpoint=endpoint)
 
     @override
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
-        print("SPAN EXPORTER CALLED")
         serialized_data = encode_spans(spans).SerializeToString()
-        self._exp.send_data(data=serialized_data, format="OTEL_V1_TRACES")
-        return SpanExportResult.SUCCESS
+        self._udp_exporter.send_data(data=serialized_data, format="OTEL_V1_TRACES")  # TODO: Convert to constant
+        return SpanExportResult.SUCCESS  # TODO: send appropriate status back. Need to??
 
     @override
     def force_flush(self, timeout_millis: int = 30000) -> bool:
         return True
 
     @override
     def shutdown(self) -> None:
-        pass
+        self._udp_exporter.shutdown()