WIP

Author: PerfectSlayer

dd-smoke-tests/spring-boot-rabbit/src/test/java/datadog/smoketest/SpringBootRabbitSmokeTest.java

@@ -1,7 +1,7 @@
 package datadog.smoketest;
 
 import static datadog.smoketest.trace.SpanMatcher.span;
-import static datadog.smoketest.trace.TraceMatcher.SORT_BY_PARENT_CHAIN;
+import static datadog.smoketest.trace.TraceMatcher.SORT_BY_ANCESTRY;
 import static datadog.smoketest.trace.TraceMatcher.trace;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
@@ -58,19 +58,19 @@
  * {@linkplain SmokeTraceAssertions order-independent subset} mode ({@code
  * unordered().ignoreAdditionalTraces()}): each matcher matches a distinct received trace and extras
  * are ignored. This is stronger than the Groovy set-membership check — each round-trip trace is
- * matched count-exact (all 12 spans, in {@link TraceMatcher#SORT_BY_PARENT_CHAIN parent-chain
- * order}), verifying every AMQP operation (publish/deliver/consume, both directions) <em>and</em>
- * its cross-service linkage — while staying robust to the timing-dependent extras (the broker emits
- * its connection-setup commands and per-ack traces as their own single-span traces, in
+ * matched count-exact (all 12 spans, in {@link TraceMatcher#SORT_BY_ANCESTRY ancestry order}),
+ * verifying every AMQP operation (publish/deliver/consume, both directions) <em>and</em> its
+ * cross-service linkage — while staying robust to the timing-dependent extras (the broker emits its
+ * connection-setup commands and per-ack traces as their own single-span traces, in
  * non-deterministic count and order).
  *
  * <p>Two design points, informed by inspecting the live trace collection:
  *
  * <ul>
- *   <li><b>Parent-chain order, not start time</b> — the 12-span round-trip is a strict linear
- *       chain, but its spans start within the same tick and race, so {@code SORT_BY_START_TIME} is
- *       unstable across runs. {@code SORT_BY_PARENT_CHAIN} orders spans by parent links
- *       (timestamp-independent) and asserts the trace is exactly that chain.
+ *   <li><b>Ancestry order, not start time</b> — the 12-span round-trip is a strict linear chain,
+ *       but its spans start within the same tick and race, so {@code SORT_BY_START_TIME} is
+ *       unstable across runs. {@code SORT_BY_ANCESTRY} orders each parent before its child
+ *       (timestamp-independent along the chain), giving a stable positional order.
  *   <li><b>Accumulate, don't isolate</b> — {@code .retainAcrossTests()} keeps traces from app
  *       startup onward, because {@code basic.qos}/{@code basic.consume}/{@code queue.declare} are
  *       emitted when the consumers start (before any test method), which a per-method session
@@ -145,12 +145,12 @@ void roundtripsProduceFullAmqpTraceStructure() throws IOException {
   }
 
   // The full distributed round-trip as one strict parent->child chain across both services and the
-  // broker. SORT_BY_PARENT_CHAIN orders the spans root->leaf and asserts the trace IS exactly this
-  // linear chain: HTTP entrypoint -> publish -> receiver consumes+forwards -> sender consumes
-  // reply.
+  // broker. SORT_BY_ANCESTRY orders the spans parent->child (stable for a linear chain), and the 12
+  // count-exact positional matchers pin the shape: HTTP entrypoint -> publish -> receiver
+  // consumes+forwards -> sender consumes reply.
   private static TraceMatcher roundTrip() {
     return trace(
-        SORT_BY_PARENT_CHAIN,
+        SORT_BY_ANCESTRY,
         sp("spring-rabbit-0", "servlet.request", "GET /roundtrip/{message}").root(),
         sp("spring-rabbit-0", "spring.handler", "WebController.roundtrip"),
         sp("spring-rabbit-0", "amqp.command", "basic.publish <default> -> otherqueue"),

dd-smoke-tests/src/main/java/datadog/smoketest/trace/SmokeTraceAssertions.java

@@ -28,11 +28,11 @@
  * <p>A single matching algorithm is tuned by three combinable {@link Options}:
  *
  * <ul>
- *   <li>{@link #SORT_BY_START_TIME} (the {@code sorter}) — sort the received traces before matching,
- *       making positional matching deterministic.
+ *   <li>{@link #SORT_BY_START_TIME} (the {@code sorter}) — sort the received traces before
+ *       matching, making positional matching deterministic.
  *   <li>{@link #UNORDERED} — a matcher may match <em>any</em> distinct (unconsumed) trace rather
- *       than the one at its position; the received order stops mattering (so a {@code sorter} has no
- *       effect in this mode).
+ *       than the one at its position; the received order stops mattering (so a {@code sorter} has
+ *       no effect in this mode).
  *   <li>{@link #IGNORE_ADDITIONAL_TRACES} — don't require every received trace to be matched; extra
  *       traces are ignored (a subset assertion). Without it, the counts must be equal.
  * </ul>
@@ -54,18 +54,24 @@
  * test-agent), both decoded to {@link DecodedTrace}.
  */
 public final class SmokeTraceAssertions {
-  /** Sorts traces by the earliest span start time. */
+  /*
+   * Trace comparators.
+   */
+  /** Trace comparator to sort by start time. */
   public static final Comparator<DecodedTrace> TRACE_START_TIME_COMPARATOR =
       Comparator.comparingLong(SmokeTraceAssertions::earliestStart);
 
-  /** Don't require every received trace to be matched — ignore the extras (a subset assertion). */
+  /*
+   * Trace assertion options.
+   */
+  /** Ignores additional traces. If there are more traces than expected, do not fail. */
   public static final UnaryOperator<Options> IGNORE_ADDITIONAL_TRACES =
       Options::ignoreAdditionalTraces;
 
-  /** Let each matcher match any distinct trace rather than the one at its position. */
+  /** Allows matchers to match any distinct trace rather than the one at its position. */
   public static final UnaryOperator<Options> UNORDERED = Options::unorder;
 
-  /** Sorts traces by earliest span start time before matching. */
+  /** Sorts traces by start time. */
   public static final UnaryOperator<Options> SORT_BY_START_TIME =
       options -> options.sort(TRACE_START_TIME_COMPARATOR);
 
@@ -129,18 +135,21 @@ public static void assertTraces(
         if (skippedTraces.get(traceIndex)) {
           continue;
         }
-        if (matcher.matches(traces.get(traceIndex).getSpans(), matcherIndex)) {
+        try {
+          matcher.assertTrace(traces.get(traceIndex).getSpans(), matcherIndex);
           matched = true;
           if (opts.unordered) {
             skippedTraces.set(traceIndex);
           } else {
             skippedTraces.set(0, traceIndex);
           }
           break;
+        } catch (AssertionError ignored) {
+          // Swallow assertion errors, keep looking for a match
         }
       }
       if (!matched) {
-        assertionFailure().message("No trace matches matcher # "+ matcherIndex).buildAndThrow();
+        assertionFailure().message("No trace matches matcher # " + matcherIndex).buildAndThrow();
       }
     }
   }

dd-smoke-tests/src/main/java/datadog/smoketest/trace/TraceMatcher.java

@@ -1,173 +1,144 @@
 package datadog.smoketest.trace;
 
 import static java.util.Comparator.comparingLong;
+import static java.util.stream.Collectors.toSet;
 
 import datadog.trace.test.agent.decoder.DecodedSpan;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.UnaryOperator;
+import org.opentest4j.AssertionFailedError;
 
 /**
- * Thin trace matcher for smoke tests: one {@link SpanMatcher} per expected span in a trace.
- * Companion of {@link SpanMatcher}; see its class doc for the serialized-model rationale.
+ * This class is a helper class to verify a trace structure.
  *
- * <p>Spans are matched positionally. By default they are matched in the order received; sort them
- * first for a deterministic position:
+ * <p>To get a {@link TraceMatcher}, use the static factory methods: {@link #trace(SpanMatcher...)}
+ * with the expected {@link SpanMatcher}s (one per expected span), or {@link #trace(UnaryOperator,
+ * SpanMatcher...)} to configure the checks with a {@link Options} object.
  *
- * <ul>
- *   <li>{@link #SORT_BY_START_TIME} — by span start time. Beware: spans that start within the same
- *       clock tick (e.g. a publish and its broker-side deliver) can race, so this order is <em>not
- *       stable</em> for tightly-coupled concurrent spans.
- *   <li>{@link #SORT_BY_PARENT_CHAIN} — root→leaf following parent links. Stable regardless of
- *       timestamps, but only defined for a <em>single linear chain</em> (each span has at most one
- *       child); it throws otherwise. Ideal for a straight request→...→response distributed trace.
- * </ul>
+ * <p>{@link #SORT_BY_START_TIME} can be used as predefined configuration to sort spans by start
+ * time, and {@link #SORT_BY_ANCESTRY} to sort spans by ancestry, root spans (or which parents are
+ * not present in the trace chunk) first, followed by their children by start time, depth-first.
  *
- * <p>A sorted order also makes {@link SpanMatcher#childOfIndex(int)} / {@link
- * SpanMatcher#childOfPrevious()} deterministic.
+ * @see SmokeTraceAssertions
+ * @see SpanMatcher
  */
 public final class TraceMatcher {
-  /** Sorts a trace's spans by start time. */
+  /*
+   * Span comparators.
+   */
+  /** Span comparator to sort by start time. */
   public static final Comparator<DecodedSpan> START_TIME_COMPARATOR =
-      comparingLong(DecodedSpan::getStart);
+      comparingLong(DecodedSpan::getStart).thenComparingLong(DecodedSpan::getSpanId);
 
-  /** Configures {@link #trace(UnaryOperator, SpanMatcher...)} to sort spans by start time. */
+  /*
+   * Span assertion options.
+   */
+  /** Sorts spans by start time. */
   public static final UnaryOperator<Options> SORT_BY_START_TIME =
       options -> options.sort(START_TIME_COMPARATOR);
 
   /**
-   * Configures {@link #trace(UnaryOperator, SpanMatcher...)} to order spans root→leaf by following
-   * parent links (see {@link Options#linearizeByParentChain()}). Timestamp-independent, so unlike
-   * {@link #SORT_BY_START_TIME} it is stable for concurrent spans — but only valid for a single
-   * linear chain.
+   * Sorts spans by ancestry, root spans (or which parents are absent from the trace chunk) first,
+   * followed by their children by start time, depth-first.
    */
-  public static final UnaryOperator<Options> SORT_BY_PARENT_CHAIN = Options::linearizeByParentChain;
+  public static final UnaryOperator<Options> SORT_BY_ANCESTRY = Options::sortByAncestry;
 
   private final Options options;
-  private final SpanMatcher[] spanMatchers;
+  private final SpanMatcher[] matchers;
 
   private TraceMatcher(Options options, SpanMatcher[] spanMatchers) {
     if (spanMatchers.length == 0) {
       throw new IllegalArgumentException("No span matchers provided");
     }
     this.options = options;
-    this.spanMatchers = spanMatchers;
+    this.matchers = spanMatchers;
   }
 
-  /** Checks a trace structure, one matcher per expected span, in the order received. */
+  /**
+   * Checks a trace structure.
+   *
+   * @param matchers The matchers to verify the trace structure.
+   */
   public static TraceMatcher trace(SpanMatcher... matchers) {
     return new TraceMatcher(new Options(), matchers);
   }
 
-  /** Checks a trace structure with the given options (e.g. {@link #SORT_BY_START_TIME}). */
+  /**
+   * Checks a trace structure.
+   *
+   * @param options The {@link TraceMatcher.Options} to configure the checks.
+   * @param matchers The matchers to verify the trace structure.
+   */
   public static TraceMatcher trace(UnaryOperator<Options> options, SpanMatcher... matchers) {
     return new TraceMatcher(options.apply(new Options()), matchers);
   }
 
   void assertTrace(List<DecodedSpan> spans, int traceIndex) {
-    if (spans.size() != spanMatchers.length) {
-      throw new AssertionError(
-          "Expected "
-              + spanMatchers.length
-              + " spans in trace "
-              + traceIndex
-              + " but got "
-              + spans.size()
-              + ": "
-              + spans);
+    if (spans.size() != this.matchers.length) {
+      throw new AssertionFailedError(
+          "Invalid number of spans for trace " + traceIndex + " : " + spans,
+          this.matchers.length,
+          spans.size());
     }
-    List<DecodedSpan> ordered;
-    if (options.linearizeByParentChain) {
-      ordered = chainOrder(spans);
-    } else {
-      ordered = new ArrayList<>(spans);
-      if (options.comparator != null) {
-        ordered.sort(options.comparator);
-      }
-    }
-    for (int i = 0; i < spanMatchers.length; i++) {
-      spanMatchers[i].assertSpan(ordered, i);
+    if (this.options.sortByAncestry) {
+      spans = sortByAncestry(spans);
+    } else if (this.options.comparator != null) {
+      spans = new ArrayList<>(spans);
+      spans.sort(this.options.comparator);
     }
-  }
-
-  /**
-   * Whether this matcher accepts {@code spans} as a whole trace, without throwing — the boolean
-   * counterpart of {@link #assertTrace} used by order-independent matching (see {@link
-   * SmokeTraceAssertions#assertTraces} with {@link SmokeTraceAssertions#IGNORE_ADDITIONAL_TRACES}),
-   * where each matcher must find some received trace it fits. A non-linear trace under {@link
-   * #SORT_BY_PARENT_CHAIN} counts as "does not match" here (rather than propagating) so probing can
-   * move on.
-   */
-  boolean matches(List<DecodedSpan> spans, int traceIndex) {
-    try {
-      assertTrace(spans, traceIndex);
-      return true;
-    } catch (AssertionError | IllegalStateException ignored) {
-      return false;
+    for (int i = 0; i < this.matchers.length; i++) {
+      this.matchers[i].assertSpan(spans, i);
     }
   }
 
-  /**
-   * Orders spans root→leaf by following parent links. Requires a single linear chain: exactly one
-   * root (parent id 0) and every span at most one child, with all spans reachable from the root.
-   * Throws {@link IllegalStateException} otherwise, so a non-chain trace fails loudly rather than
-   * being silently mis-ordered.
-   */
-  private static List<DecodedSpan> chainOrder(List<DecodedSpan> spans) {
-    DecodedSpan root = null;
-    Map<Long, DecodedSpan> childByParent = new HashMap<>();
+  private static List<DecodedSpan> sortByAncestry(List<DecodedSpan> spans) {
+    Set<Long> spanIds = spans.stream().map(DecodedSpan::getSpanId).collect(toSet());
+    Map<Long, List<DecodedSpan>> spansByParentId = new HashMap<>();
     for (DecodedSpan span : spans) {
-      if (span.getParentId() == 0L) {
-        if (root != null) {
-          throw new IllegalStateException(
-              "SORT_BY_PARENT_CHAIN requires a single root span (parent id 0); found more than one");
-        }
-        root = span;
-      } else if (childByParent.put(span.getParentId(), span) != null) {
-        throw new IllegalStateException(
-            "SORT_BY_PARENT_CHAIN requires a linear chain, but span "
-                + span.getParentId()
-                + " has more than one child");
+      long parentId = span.getParentId();
+      if (parentId != 0 && !spanIds.contains(parentId)) {
+        parentId = 0;
       }
+      spansByParentId.computeIfAbsent(parentId, k -> new ArrayList<>()).add(span);
     }
-    if (root == null) {
-      throw new IllegalStateException("SORT_BY_PARENT_CHAIN requires a root span (parent id 0)");
-    }
+    spansByParentId.forEach((k, v) -> v.sort(START_TIME_COMPARATOR));
+
     List<DecodedSpan> ordered = new ArrayList<>(spans.size());
-    for (DecodedSpan current = root;
-        current != null;
-        current = childByParent.get(current.getSpanId())) {
-      ordered.add(current);
-    }
-    if (ordered.size() != spans.size()) {
-      throw new IllegalStateException(
-          "SORT_BY_PARENT_CHAIN: trace is not a single chain ("
-              + ordered.size()
-              + " of "
-              + spans.size()
-              + " spans reachable from the root)");
-    }
+    appendChildren(ordered, spansByParentId.get(0L), spansByParentId);
     return ordered;
   }
 
+  private static void appendChildren(
+      List<DecodedSpan> orderedSpan,
+      List<DecodedSpan> children,
+      Map<Long, List<DecodedSpan>> spansByParentId) {
+    for (DecodedSpan child : children) {
+      orderedSpan.add(child);
+      List<DecodedSpan> grandChildren = spansByParentId.get(child.getSpanId());
+      if (grandChildren != null) {
+        appendChildren(orderedSpan, grandChildren, spansByParentId);
+      }
+    }
+  }
+
   public static final class Options {
-    Comparator<DecodedSpan> comparator = null; // null => keep received order
-    boolean linearizeByParentChain = false;
+    private Comparator<DecodedSpan> comparator = null;
+    private boolean sortByAncestry = false;
 
     public Options sort(Comparator<DecodedSpan> comparator) {
       this.comparator = comparator;
+      this.sortByAncestry = false;
       return this;
     }
 
-    /**
-     * Order spans root→leaf by parent links instead of by a comparator (see {@link
-     * #SORT_BY_PARENT_CHAIN}).
-     */
-    public Options linearizeByParentChain() {
-      this.linearizeByParentChain = true;
+    Options sortByAncestry() {
+      this.comparator = null;
+      this.sortByAncestry = true;
       return this;
     }
   }