Avoid ArrayList copying from TraceInterceptors

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -1211,15 +1211,15 @@ public AgentDataStreamsMonitoring getDataStreamsMonitoring() {
     return dataStreamsMonitoring;
   }
 
-  private final RatelimitedLogger rlLog = new RatelimitedLogger(log, 1, MINUTES);
+  private static final RatelimitedLogger rlLog = new RatelimitedLogger(log, 1, MINUTES);
 
   /**
    * We use the sampler to know if the trace has to be reported/written. The sampler is called on
    * the first span (root span) of the trace. If the trace is marked as a sample, we report it.
    *
    * @param trace a list of the spans related to the same trace
    */
-  void write(final List<DDSpan> trace) {
+  void write(final SpanList trace) {
     if (trace.isEmpty() || !trace.get(0).traceConfig().isTraceEnabled()) {
       return;
     }
@@ -1257,30 +1257,59 @@ void write(final List<DDSpan> trace) {
     }
   }
 
-  private List<DDSpan> interceptCompleteTrace(List<DDSpan> trace) {
-    if (!interceptors.isEmpty() && !trace.isEmpty()) {
-      Collection<? extends MutableSpan> interceptedTrace = new ArrayList<>(trace);
-      for (final TraceInterceptor interceptor : interceptors) {
-        try {
-          // If one TraceInterceptor throws an exception, then continue with the next one
-          interceptedTrace = interceptor.onTraceComplete(interceptedTrace);
-        } catch (Throwable e) {
-          String interceptorName = interceptor.getClass().getName();
-          rlLog.warn("Throwable raised in TraceInterceptor {}", interceptorName, e);
-        }
-        if (interceptedTrace == null) {
-          interceptedTrace = emptyList();
-        }
+  private List<DDSpan> interceptCompleteTrace(SpanList originalTrace) {
+    return interceptCompleteTrace(interceptors, originalTrace);
+  }
+
+  static final List<DDSpan> interceptCompleteTrace(
+      SortedSet<TraceInterceptor> interceptors, SpanList originalTrace) {
+    if (interceptors.isEmpty()) {
+      return originalTrace;
+    }
+    if (originalTrace.isEmpty()) {
+      return SpanList.EMPTY;
+    }
+
+    // Using TraceList to optimize the common case where the interceptors,
+    // don't alter the list.  If the interceptors just return the provided
+    // List, then no need to copy to another List.
+
+    // As an extra precaution, also check the modCount before and after on
+    // the TraceList, since TraceInterceptor could put some other type of
+    // object into the List.
+
+    // There is still a risk that a TraceInterceptor holds onto the provided
+    // List and modifies it later on, but we cannot safeguard against
+    // every possible misuse.
+    Collection<? extends MutableSpan> interceptedTrace = originalTrace;
+    int originalModCount = originalTrace.modCount();
+
+    for (final TraceInterceptor interceptor : interceptors) {
+      try {
+        // If one TraceInterceptor throws an exception, then continue with the next one
+        interceptedTrace = interceptor.onTraceComplete(interceptedTrace);
+      } catch (Throwable e) {
+        String interceptorName = interceptor.getClass().getName();
+        rlLog.warn("Throwable raised in TraceInterceptor {}", interceptorName, e);
       }
+      if (interceptedTrace == null) {
+        interceptedTrace = emptyList();
+      }
+    }
 
-      trace = new ArrayList<>(interceptedTrace.size());
+    if (interceptedTrace == null || interceptedTrace.isEmpty()) {
+      return SpanList.EMPTY;
+    } else if (interceptedTrace == originalTrace && originalTrace.modCount() == originalModCount) {
+      return originalTrace;
+    } else {
+      SpanList trace = new SpanList(interceptedTrace.size());
       for (final MutableSpan span : interceptedTrace) {
         if (span instanceof DDSpan) {
           trace.add((DDSpan) span);
         }
       }
+      return trace;
     }
-    return trace;
   }
 
   @Override

dd-trace-core/src/main/java/datadog/trace/core/PendingTrace.java

@@ -327,7 +327,7 @@ private int write(boolean isPartial) {
     if (!spans.isEmpty()) {
       try (Recording recording = tracer.writeTimer()) {
         // Only one writer at a time
-        final List<DDSpan> trace;
+        final SpanList trace;
         int completedSpans = 0;
         synchronized (this) {
           if (!isPartial) {
@@ -346,10 +346,10 @@ private int write(boolean isPartial) {
           // count(s) will be incremented, and any new spans added during the period that the count
           // was negative will be written by someone even if we don't write them right now.
           if (size > 0 && (!isPartial || size >= tracer.getPartialFlushMinSpans())) {
-            trace = new ArrayList<>(size);
+            trace = new SpanList(size);
             completedSpans = enqueueSpansToWrite(trace, writeRunningSpans);
           } else {
-            trace = EMPTY;
+            trace = SpanList.EMPTY;
           }
         }
         if (!trace.isEmpty()) {

dd-trace-core/src/main/java/datadog/trace/core/SpanList.java

@@ -0,0 +1,35 @@
+package datadog.trace.core;
+
+import java.util.ArrayList;
+
+/** ArrayList that exposes modCount to allow for an optimization in TraceInterceptor handling */
+final class SpanList extends ArrayList<DDSpan> {
+  static final SpanList EMPTY = new SpanList(0);
+
+  /**
+   * Convenience function for creating a SpanList containing a single DDSpan. Meant as replacement
+   * for Collections.singletonList when creating a SpanList.
+   *
+   * @param span != null
+   * @return a SpanList
+   */
+  static final SpanList of(DDSpan span) {
+    SpanList list = new SpanList(1);
+    list.add(span);
+    return list;
+  }
+
+  /**
+   * Constructs a SpanList with the specified capacity
+   *
+   * @param capacity
+   */
+  SpanList(int capacity) {
+    super(capacity);
+  }
+
+  /** The modifcation count of the List - can be used to check if the List has been altered */
+  int modCount() {
+    return this.modCount;
+  }
+}

dd-trace-core/src/main/java/datadog/trace/core/StreamingTraceCollector.java

@@ -4,7 +4,6 @@
 import datadog.trace.api.time.TimeSource;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.core.monitor.HealthMetrics;
-import java.util.Collections;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import javax.annotation.Nonnull;
 
@@ -72,7 +71,7 @@ PublishState onPublish(DDSpan span) {
       tracer.onRootSpanPublished(rootSpan);
     }
     healthMetrics.onFinishSpan();
-    tracer.write(Collections.singletonList(span));
+    tracer.write(SpanList.of(span));
     return PublishState.WRITTEN;
   }
 

dd-trace-core/src/test/groovy/datadog/trace/core/TraceInterceptorTest.groovy

@@ -188,7 +188,7 @@ class TraceInterceptorTest extends DDCoreSpecification {
     when:
     DDSpan span = (DDSpan) tracer.startSpan("test", "test")
     span.phasedFinish()
-    tracer.write([span])
+    tracer.write(SpanList.of(span))
 
     then:
     notThrown(Throwable)

dd-trace-core/src/test/java/datadog/trace/core/CoreTracerTest2.java

@@ -8,10 +8,17 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import datadog.trace.api.Config;
+import datadog.trace.api.interceptor.MutableSpan;
+import datadog.trace.api.interceptor.TraceInterceptor;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.core.CoreTracer.CoreSpanBuilder;
 import datadog.trace.core.CoreTracer.ReusableSingleSpanBuilder;
 import datadog.trace.core.CoreTracer.ReusableSingleSpanBuilderThreadLocalCache;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.SortedSet;
+import java.util.TreeSet;
 import org.junit.jupiter.api.Test;
 
 // named CoreTracerTest2 to avoid collision with Groovy which appears to have messed up test
@@ -179,4 +186,114 @@ public void start_not_inUse() {
     ReusableSingleSpanBuilder builder = new ReusableSingleSpanBuilder(TRACER);
     assertThrows(AssertionError.class, () -> builder.start());
   }
+
+  @Test
+  public void noInterceptorsTest() {
+    // No interceptors should return the original list to avoid allocation
+    DDSpan span = (DDSpan) TRACER.startSpan("foo", "foo");
+
+    SpanList list = SpanList.of(span);
+    List<DDSpan> interceptedList =
+        CoreTracer.interceptCompleteTrace(Collections.emptySortedSet(), list);
+    assertSame(list, interceptedList);
+  }
+
+  @Test
+  public void interceptNoInterceptors() {
+    // No interceptors should return the original list to avoid allocation
+    DDSpan span = (DDSpan) TRACER.startSpan("foo", "foo");
+
+    SpanList list = SpanList.of(span);
+    List<DDSpan> interceptedList =
+        CoreTracer.interceptCompleteTrace(Collections.emptySortedSet(), list);
+    assertSame(list, interceptedList);
+  }
+
+  @Test
+  public void interceptEmptyList() {
+    DDSpan span = (DDSpan) TRACER.startSpan("foo", "foo");
+    SortedSet<TraceInterceptor> interceptors = interceptors((list) -> SpanList.of(span));
+
+    SpanList list = new SpanList(0); // not using EMPTY deliberately
+    List<DDSpan> interceptedList = CoreTracer.interceptCompleteTrace(interceptors, list);
+    assertTrue(interceptedList.isEmpty());
+  }
+
+  @Test
+  public void interceptUnchanged() {
+    SortedSet<TraceInterceptor> interceptors = interceptors((list) -> list, (list) -> list);
+
+    DDSpan span = (DDSpan) TRACER.startSpan("foo", "foo");
+    SpanList list = SpanList.of(span);
+    List<DDSpan> interceptedList = CoreTracer.interceptCompleteTrace(interceptors, list);
+    assertSame(list, interceptedList);
+  }
+
+  @Test
+  public void interceptNewList() {
+    DDSpan substituteSpan = (DDSpan) TRACER.startSpan("sub", "sub");
+    SpanList substituteList = SpanList.of(substituteSpan);
+    SortedSet<TraceInterceptor> interceptors =
+        interceptors((list) -> list, (list) -> substituteList);
+
+    DDSpan span = (DDSpan) TRACER.startSpan("foo", "foo");
+    SpanList list = SpanList.of(span);
+    List<DDSpan> interceptedList = CoreTracer.interceptCompleteTrace(interceptors, list);
+    assertEquals(1, interceptedList.size());
+    assertEquals("sub", interceptedList.get(0).getOperationName());
+  }
+
+  @Test
+  public void interceptAlteredList() {
+    // This is an unlikely case and arguably not something we need to support
+
+    DDSpan substituteSpan = (DDSpan) TRACER.startSpan("sub", "sub");
+    SortedSet<TraceInterceptor> interceptors =
+        interceptors(
+            (list) -> list,
+            (list) -> {
+              List erasedList = (List) list;
+              erasedList.clear();
+              erasedList.add(substituteSpan);
+              return erasedList;
+            });
+
+    DDSpan span = (DDSpan) TRACER.startSpan("foo", "foo");
+    SpanList list = SpanList.of(span);
+    List<DDSpan> interceptedList = CoreTracer.interceptCompleteTrace(interceptors, list);
+    assertNotSame(interceptedList, list);
+    assertEquals(1, interceptedList.size());
+    assertEquals("sub", interceptedList.get(0).getOperationName());
+  }
+
+  static final SortedSet<TraceInterceptor> interceptors(TestInterceptor... interceptors) {
+    TreeSet<TraceInterceptor> interceptorSet =
+        new TreeSet<>((lhs, rhs) -> Integer.compare(lhs.priority(), rhs.priority()));
+    for (int i = 0; i < interceptors.length; ++i) {
+      int priority = i;
+      TestInterceptor interceptor = interceptors[i];
+
+      interceptorSet.add(
+          new TraceInterceptor() {
+            @Override
+            public int priority() {
+              return priority;
+            }
+
+            @Override
+            public Collection<? extends MutableSpan> onTraceComplete(
+                Collection<? extends MutableSpan> trace) {
+              return interceptor.onTraceComplete(trace);
+            }
+          });
+    }
+    return interceptorSet;
+  }
+
+  // Matches TraceInterceptor but priority is implied in interceptors
+  // Only having onTraceComplete allows this to @FunctionalInterface and a little nicer for a test
+  @FunctionalInterface
+  interface TestInterceptor {
+    Collection<? extends MutableSpan> onTraceComplete(Collection<? extends MutableSpan> trace);
+  }
 }