refactoring the builder and exporter classes

Author: srprash

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/OtlpUdpSpanExporter.java

@@ -41,24 +41,14 @@ class OtlpUdpSpanExporter implements SpanExporter {
 
   private static final Logger logger = Logger.getLogger(OtlpUdpSpanExporter.class.getName());
 
-  // The protocol header and delimiter is required for sending data to X-Ray Daemon or when running
-  // in Lambda.
-  // https://docs.aws.amazon.com/xray/latest/devguide/xray-api-sendingdata.html#xray-api-daemon
-  private static final String PROTOCOL_HEADER = "{\"format\": \"json\", \"version\": 1}";
-  private static final char PROTOCOL_DELIMITER = '\n';
-
-  // These prefixes help the backend identify if the spans payload is sampled or not.
-  private static final String FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX = "T1S";
-  private static final String FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX = "T1U";
-
   private final AtomicBoolean isShutdown = new AtomicBoolean();
 
   private final UdpSender sender;
-  private final boolean sampled;
+  private final String payloadPrefix;
 
-  OtlpUdpSpanExporter(UdpSender sender, boolean sampled) {
+  OtlpUdpSpanExporter(UdpSender sender, String payloadPrefix) {
     this.sender = sender;
-    this.sampled = sampled;
+    this.payloadPrefix = payloadPrefix;
   }
 
   @Override
@@ -71,13 +61,7 @@ public CompletableResultCode export(Collection<SpanData> spans) {
     ByteArrayOutputStream baos = new ByteArrayOutputStream();
     try {
       exportRequest.writeBinaryTo(baos);
-      String payload =
-          PROTOCOL_HEADER
-              + PROTOCOL_DELIMITER
-              + (sampled
-                  ? FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX
-                  : FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX)
-              + Base64.getEncoder().encodeToString(baos.toByteArray());
+      String payload = payloadPrefix + Base64.getEncoder().encodeToString(baos.toByteArray());
       sender.send(payload.getBytes(StandardCharsets.UTF_8));
       return CompletableResultCode.ofSuccess();
     } catch (Exception e) {
@@ -107,7 +91,7 @@ UdpSender getSender() {
   }
 
   // Visible for testing
-  boolean isSampled() {
-    return sampled;
+  String getPayloadPrefix() {
+    return payloadPrefix;
   }
 }

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/OtlpUdpSpanExporterBuilder.java

@@ -17,15 +17,23 @@
 
 import static java.util.Objects.requireNonNull;
 
-import java.net.SocketException;
-
 final class OtlpUdpSpanExporterBuilder {
 
   private static final String DEFAULT_HOST = "127.0.0.1";
   private static final int DEFAULT_PORT = 2000;
 
+  // The protocol header and delimiter is required for sending data to X-Ray Daemon or when running
+  // in Lambda.
+  // https://docs.aws.amazon.com/xray/latest/devguide/xray-api-sendingdata.html#xray-api-daemon
+  private static final String PROTOCOL_HEADER = "{\"format\": \"json\", \"version\": 1}";
+  private static final char PROTOCOL_DELIMITER = '\n';
+
+  // These prefixes help the backend identify if the spans payload is sampled or not.
+  private static final String FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX = "T1S";
+  private static final String FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX = "T1U";
+
   private UdpSender sender;
-  private boolean sampled = true;
+  private String tracePayloadPrefix = FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX;
 
   public OtlpUdpSpanExporterBuilder setEndpoint(String endpoint) {
     requireNonNull(endpoint, "endpoint must not be null");
@@ -40,20 +48,20 @@ public OtlpUdpSpanExporterBuilder setEndpoint(String endpoint) {
     return this;
   }
 
-  public OtlpUdpSpanExporterBuilder setSampled(boolean sampled) {
-    this.sampled = sampled;
+  public OtlpUdpSpanExporterBuilder setPayloadSampleDecision(TracePayloadSampleDecision decision) {
+    this.tracePayloadPrefix =
+        decision == TracePayloadSampleDecision.SAMPLED
+            ? FORMAT_OTEL_SAMPLED_TRACES_BINARY_PREFIX
+            : FORMAT_OTEL_UNSAMPLED_TRACES_BINARY_PREFIX;
     return this;
   }
 
   public OtlpUdpSpanExporter build() {
     if (sender == null) {
-      try {
-        this.sender = new UdpSender(DEFAULT_HOST, DEFAULT_PORT);
-      } catch (SocketException e) {
-        throw new IllegalStateException("Failed to create OtlpUdpSpanExporter", e);
-      }
+      this.sender = new UdpSender(DEFAULT_HOST, DEFAULT_PORT);
     }
-    return new OtlpUdpSpanExporter(this.sender, this.sampled);
+    return new OtlpUdpSpanExporter(
+        this.sender, PROTOCOL_HEADER + PROTOCOL_DELIMITER + tracePayloadPrefix);
   }
 
   // Only for testing
@@ -62,3 +70,8 @@ OtlpUdpSpanExporterBuilder setSender(UdpSender sender) {
     return this;
   }
 }
+
+enum TracePayloadSampleDecision {
+  SAMPLED,
+  UNSAMPLED
+}

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/UdpSender.java

@@ -34,7 +34,7 @@ class UdpSender {
   private DatagramSocket socket;
   private final InetSocketAddress endpoint;
 
-  public UdpSender(String host, int port) throws SocketException {
+  public UdpSender(String host, int port) {
     this.endpoint = new InetSocketAddress(host, port);
     try {
       this.socket = new DatagramSocket();
@@ -45,6 +45,9 @@ public UdpSender(String host, int port) throws SocketException {
 
   public CompletableResultCode shutdown() {
     try {
+      if (socket == null) {
+        return CompletableResultCode.ofSuccess();
+      }
       socket.close();
       return CompletableResultCode.ofSuccess();
     } catch (Exception e) {

awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/UdpExporterTest.java

@@ -39,17 +39,20 @@ public void testUdpExporterWithDefaults() {
     assertThat(sender.getEndpoint().getHostName())
         .isEqualTo("localhost"); // getHostName implicitly converts 127.0.0.1 to localhost
     assertThat(sender.getEndpoint().getPort()).isEqualTo(2000);
-    assertThat(exporter.isSampled()).isTrue();
+    assertThat(exporter.getPayloadPrefix()).endsWith("T1S");
   }
 
   @Test
   public void testUdpExporterWithCustomEndpointAndSample() {
     OtlpUdpSpanExporter exporter =
-        new OtlpUdpSpanExporterBuilder().setEndpoint("somehost:1000").setSampled(false).build();
+        new OtlpUdpSpanExporterBuilder()
+            .setEndpoint("somehost:1000")
+            .setPayloadSampleDecision(TracePayloadSampleDecision.UNSAMPLED)
+            .build();
     UdpSender sender = exporter.getSender();
     assertThat(sender.getEndpoint().getHostName()).isEqualTo("somehost");
     assertThat(sender.getEndpoint().getPort()).isEqualTo(1000);
-    assertThat(exporter.isSampled()).isFalse();
+    assertThat(exporter.getPayloadPrefix()).endsWith("T1U");
   }
 
   @Test
@@ -94,7 +97,10 @@ public void testExportWithSampledFalse() {
     SpanData spanData = buildSpanDataMock();
 
     OtlpUdpSpanExporter exporter =
-        new OtlpUdpSpanExporterBuilder().setSender(senderMock).setSampled(false).build();
+        new OtlpUdpSpanExporterBuilder()
+            .setSender(senderMock)
+            .setPayloadSampleDecision(TracePayloadSampleDecision.UNSAMPLED)
+            .build();
     exporter.export(Collections.singletonList(spanData));
 
     verify(senderMock, times(1)).send(any(byte[].class));