set BSP batch size when on Lambda

Author: srprash

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

@@ -86,9 +86,13 @@ public class AwsApplicationSignalsCustomizerProvider
   private static final String OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_CONFIG =
       "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT";
   private static final String AWS_XRAY_DAEMON_ADDRESS_CONFIG = "AWS_XRAY_DAEMON_ADDRESS";
-
   private static final String DEFAULT_UDP_ENDPOINT = "127.0.0.1:2000";
 
+  // UDP packet can be upto 64KB. To limit the packet size, we limit the exported batch size.
+  // This is a bit of a magic number, as there is no simple way to tell how many spans can make a
+  // 64KB batch since spans can vary in size.
+  private static final int LAMBDA_SPAN_EXPORT_BATCH_SIZE = 10;
+
   public void customize(AutoConfigurationCustomizer autoConfiguration) {
     autoConfiguration.addPropertiesCustomizer(this::customizeProperties);
     autoConfiguration.addResourceCustomizer(this::customizeResource);
@@ -168,7 +172,9 @@ private SdkTracerProviderBuilder customizeTracerProviderBuilder(
       // Signals metrics.
       if (isLambdaEnvironment()) {
         tracerProviderBuilder.addSpanProcessor(
-            AwsUnsampledOnlySpanProcessorBuilder.create().build());
+            AwsUnsampledOnlySpanProcessorBuilder.create()
+                .setMaxExportBatchSize(LAMBDA_SPAN_EXPORT_BATCH_SIZE)
+                .build());
         return tracerProviderBuilder;
       }
 

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

@@ -31,18 +31,30 @@ public static AwsUnsampledOnlySpanProcessorBuilder create() {
           .setPayloadSampleDecision(TracePayloadSampleDecision.UNSAMPLED)
           .build();
 
+  // Default batch size to be same as Otel BSP default
+  private int maxExportBatchSize = 512;
+
   public AwsUnsampledOnlySpanProcessorBuilder setSpanExporter(SpanExporter exporter) {
     requireNonNull(exporter, "exporter cannot be null");
     this.exporter = exporter;
     return this;
   }
 
+  public AwsUnsampledOnlySpanProcessorBuilder setMaxExportBatchSize(int maxExportBatchSize) {
+    this.maxExportBatchSize = maxExportBatchSize;
+    return this;
+  }
+
   public AwsUnsampledOnlySpanProcessor build() {
     BatchSpanProcessor bsp =
-        BatchSpanProcessor.builder(exporter).setExportUnsampledSpans(true).build();
+        BatchSpanProcessor.builder(exporter)
+            .setExportUnsampledSpans(true)
+            .setMaxExportBatchSize(maxExportBatchSize)
+            .build();
     return new AwsUnsampledOnlySpanProcessor(bsp);
   }
 
+  // Visible for testing
   SpanExporter getSpanExporter() {
     return exporter;
   }

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

@@ -59,6 +59,7 @@ public void testDefaultSpanProcessor() {
         .contains(
             "spanExporter=software.amazon.opentelemetry.javaagent.providers.OtlpUdpSpanExporter");
     assertThat(delegateBspString).contains("exportUnsampledSpans=true");
+    assertThat(delegateBspString).contains("maxExportBatchSize=512");
   }
 
   @Test
@@ -78,6 +79,19 @@ public void testSpanProcessorWithExporter() {
     assertThat(delegateBspString).contains("exportUnsampledSpans=true");
   }
 
+  @Test
+  public void testSpanProcessorWithBatchSize() {
+    AwsUnsampledOnlySpanProcessorBuilder builder =
+        AwsUnsampledOnlySpanProcessorBuilder.create().setMaxExportBatchSize(100);
+    AwsUnsampledOnlySpanProcessor unsampledSP = builder.build();
+
+    SpanProcessor delegate = unsampledSP.getDelegate();
+    assertThat(delegate).isInstanceOf(BatchSpanProcessor.class);
+    BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate;
+    String delegateBspString = delegateBsp.toString();
+    assertThat(delegateBspString).contains("maxExportBatchSize=100");
+  }
+
   @Test
   public void testStartAddsAttributeToSampledSpan() {
     SpanContext mockSpanContext = mock(SpanContext.class);