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.

Author: dougqh

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

@@ -1211,7 +1211,7 @@ 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
@@ -1257,49 +1257,57 @@ void write(final SpanList trace) {
   }
 
   private List<DDSpan> interceptCompleteTrace(SpanList originalTrace) {
-    if (!interceptors.isEmpty() && !originalTrace.isEmpty()) {
-      // 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();
-        }
+    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();
       }
+    }
 
-      if (interceptedTrace == null || interceptedTrace.isEmpty()) {
-    	return SpanList.EMPTY;
-      } else if (interceptedTrace != originalTrace || originalTrace.modCount() != originalModCount) {
-        SpanList trace = new SpanList(interceptedTrace.size());
-        for (final MutableSpan span : interceptedTrace) {
-          if (span instanceof DDSpan) {
-            trace.add((DDSpan) span);
-          }
+    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;
-      } else {
-        return originalTrace;
       }
-    } else {
-      return originalTrace;
+      return trace;
     }
   }
 

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