Avoid ArrayList copying from TraceInterceptors (#10828)

Skip copying to another array if interceptedTrace == trace

Reducing allocation in interceptCompleteTrace

Introduced TraceList which is just an ArrayList that exposes the internal modCount

Being able to access the modCount allows interceptCompleteTrace to check if the list post-interception has been modified

If the list is the same & is unmodified, then interceptCompleteTrace is able to bypass making a defensive copy

As a further optimization, use of TraceList is pushed up to callers of interceptCompleteTrace and TraceCollector.write

That does mean StreamingTraceCollector can no longer used a singletonList, but the savings throughput the pipeline outweigh the cost of the extra Object[] in that case

Merge branch 'master' into dougqh/trace-intercept-optimization

spotless

Merge branch 'dougqh/trace-intercept-optimization' of github.com:DataDog/dd-trace-java into dougqh/trace-intercept-optimization

Merge branch 'master' into dougqh/trace-intercept-optimization

Clean-up

- Renamed TraceList -> SpanList
- Fix broken test - use SpanList in Groovy test

Refactor & tests

Refactored the code to clean-up the edge case handling, quick return on interceptor empty and trace empty.

That reduced the nesting level and made the code easier to read.  Also flipped the code that handles unaltered lists, that reads a bit better.

Added tests for all these cases to CoreTracerTest2.  To make this easier to test in isolation, created a static interceptCompleteTrace that can be called directly by the test methods.

Merge branch 'master' into dougqh/trace-intercept-optimization

Adding some Javadoc

Merge branch 'master' into dougqh/trace-intercept-optimization

codeNarc - removing unnecessary import

Merge branch 'master' into dougqh/trace-intercept-optimization

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: dougqh

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);
+  }
 }