Fix inferred span cycles

inferred-spans/src/main/java/co/elastic/otel/profiler/CallTree.java

@@ -24,7 +24,6 @@
 import co.elastic.otel.common.ElasticAttributes;
 import co.elastic.otel.common.util.HexUtils;
 import co.elastic.otel.profiler.collections.LongHashSet;
-import co.elastic.otel.profiler.collections.LongList;
 import co.elastic.otel.profiler.pooling.ObjectPool;
 import co.elastic.otel.profiler.pooling.Recyclable;
 import io.opentelemetry.api.common.Attributes;
@@ -84,9 +83,9 @@ public class CallTree implements Recyclable {
   private boolean isSpan;
   private int depth;
 
-  @Nullable private LongList childIds;
+  @Nullable private ChildList childIds;
 
-  @Nullable private LongList maybeChildIds;
+  @Nullable private ChildList maybeChildIds;
 
   public CallTree() {}
 
@@ -504,7 +503,8 @@ protected Span asSpan(
             .setStartTimestamp(
                 clock.toEpochNanos(parentContext.getClockAnchor(), this.start),
                 TimeUnit.NANOSECONDS);
-    insertChildIdLinks(spanBuilder, Span.fromContext(parentOtelCtx).getSpanContext(), tempBuilder);
+    insertChildIdLinks(
+        spanBuilder, Span.fromContext(parentOtelCtx).getSpanContext(), parentContext, tempBuilder);
 
     // we're not interested in the very bottom of the stack which contains things like accepting and
     // handling connections
@@ -524,20 +524,27 @@ protected Span asSpan(
   }
 
   private void insertChildIdLinks(
-      SpanBuilder span, SpanContext parentContext, StringBuilder tempBuilder) {
+      SpanBuilder span,
+      SpanContext parentContext,
+      TraceContext nonInferredParent,
+      StringBuilder tempBuilder) {
     if (childIds == null || childIds.isEmpty()) {
       return;
     }
     for (int i = 0; i < childIds.getSize(); i++) {
-      tempBuilder.setLength(0);
-      HexUtils.appendLongAsHex(childIds.get(i), tempBuilder);
-      SpanContext spanContext =
-          SpanContext.create(
-              parentContext.getTraceId(),
-              tempBuilder.toString(),
-              parentContext.getTraceFlags(),
-              parentContext.getTraceState());
-      span.addLink(spanContext, CHILD_LINK_ATTRIBUTES);
+      // to avoid cycles, we only insert child-ids if the parent of the child is also
+      // the parent of the stack of inferred spans inserted
+      if (nonInferredParent.getSpanId() == childIds.getParentId(i)) {
+        tempBuilder.setLength(0);
+        HexUtils.appendLongAsHex(childIds.getId(i), tempBuilder);
+        SpanContext spanContext =
+            SpanContext.create(
+                parentContext.getTraceId(),
+                tempBuilder.toString(),
+                parentContext.getTraceFlags(),
+                parentContext.getTraceState());
+        span.addLink(spanContext, CHILD_LINK_ATTRIBUTES);
+      }
     }
   }
 
@@ -623,18 +630,18 @@ public void resetState() {
    *
    * @param id the child span id to add to this call tree element
    */
-  public void addMaybeChildId(long id) {
+  public void addMaybeChildId(long id, long parentId) {
     if (maybeChildIds == null) {
-      maybeChildIds = new LongList();
+      maybeChildIds = new ChildList();
     }
-    maybeChildIds.add(id);
+    maybeChildIds.add(id, parentId);
   }
 
-  public void addChildId(long id) {
+  public void addChildId(long id, long parentId) {
     if (childIds == null) {
-      childIds = new LongList();
+      childIds = new ChildList();
     }
-    childIds.add(id);
+    childIds.add(id, parentId);
   }
 
   public boolean hasChildIds() {
@@ -664,7 +671,11 @@ void giveChildIdsTo(CallTree giveTo) {
 
   void giveLastChildIdTo(CallTree giveTo) {
     if (childIds != null && !childIds.isEmpty()) {
-      giveTo.addChildId(childIds.remove(childIds.getSize() - 1));
+      int size = childIds.getSize();
+      long id = childIds.getId(size - 1);
+      long parentId = childIds.getParentId(size - 1);
+      giveTo.addChildId(id, parentId);
+      childIds.removeLast();
     }
   }
 
@@ -743,7 +754,7 @@ public void onActivation(byte[] active, long timestamp) {
         long spanId = TraceContext.getSpanId(active);
         activeSet.add(spanId);
         if (!isNestedActivation(topOfStack)) {
-          topOfStack.addMaybeChildId(spanId);
+          topOfStack.addMaybeChildId(spanId, TraceContext.getParentId(active));
         }
       }
     }
@@ -752,12 +763,12 @@ private boolean isNestedActivation(CallTree topOfStack) {
       return isAnyActive(topOfStack.childIds) || isAnyActive(topOfStack.maybeChildIds);
     }
 
-    private boolean isAnyActive(@Nullable LongList spanIds) {
+    private boolean isAnyActive(@Nullable ChildList spanIds) {
       if (spanIds == null) {
         return false;
       }
       for (int i = 0, size = spanIds.getSize(); i < size; i++) {
-        if (activeSet.contains(spanIds.get(i))) {
+        if (activeSet.contains(spanIds.getId(i))) {
           return true;
         }
       }

inferred-spans/src/main/java/co/elastic/otel/profiler/ChildList.java

@@ -0,0 +1,63 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package co.elastic.otel.profiler;
+
+import co.elastic.otel.profiler.collections.LongList;
+
+/** List for maintaining pairs of (spanId,parentIds) both represented as longs. */
+public class ChildList {
+
+  // this list contains the (spanId,parentIds) flattened
+  private LongList idsWithParentIds = new LongList();
+
+  public void add(long id, long parentId) {
+    idsWithParentIds.add(id);
+    idsWithParentIds.add(parentId);
+  }
+
+  public long getId(int index) {
+    return idsWithParentIds.get(index * 2);
+  }
+
+  public long getParentId(int index) {
+    return idsWithParentIds.get(index * 2 + 1);
+  }
+
+  public int getSize() {
+    return idsWithParentIds.getSize() / 2;
+  }
+
+  public void addAll(ChildList other) {
+    idsWithParentIds.addAll(other.idsWithParentIds);
+  }
+
+  public void clear() {
+    idsWithParentIds.clear();
+  }
+
+  public boolean isEmpty() {
+    return getSize() == 0;
+  }
+
+  public void removeLast() {
+    int size = idsWithParentIds.getSize();
+    idsWithParentIds.remove(size - 1);
+    idsWithParentIds.remove(size - 2);
+  }
+}

inferred-spans/src/main/java/co/elastic/otel/profiler/SamplingProfiler.java

@@ -862,15 +862,12 @@ private void set(
         @Nullable Span previousContext,
         long nanoTime,
         SpanAnchoredClock clock) {
-      TraceContext.serialize(
-          traceContext.getSpanContext(), clock.getAnchor(traceContext), traceContextBuffer);
+      TraceContext.serialize(traceContext, clock.getAnchor(traceContext), traceContextBuffer);
       this.threadId = threadId;
       this.activation = activation;
       if (previousContext != null) {
         TraceContext.serialize(
-            previousContext.getSpanContext(),
-            clock.getAnchor(previousContext),
-            previousContextBuffer);
+            previousContext, clock.getAnchor(previousContext), previousContextBuffer);
         rootContext = false;
       } else {
         rootContext = true;

inferred-spans/src/main/java/co/elastic/otel/profiler/TraceContext.java

@@ -25,6 +25,7 @@
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceFlags;
 import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.sdk.trace.ReadableSpan;
 import javax.annotation.Nullable;
 
 /**
@@ -34,28 +35,38 @@
  */
 public class TraceContext implements Recyclable {
 
-  public static final int SERIALIZED_LENGTH = 16 + 8 + 1 + 8;
+  public static final int SERIALIZED_LENGTH = 16 + 8 + 1 + 1 + 8 + 8;
   private long traceIdLow;
   private long traceIdHigh;
   private long id;
+
+  private boolean hasParentId;
+  private long parentId;
   private byte flags;
 
   private long clockAnchor;
 
   public TraceContext() {}
 
   // For testing only
-  static TraceContext fromSpanContextWithZeroClockAnchor(SpanContext ctx) {
+  static TraceContext fromSpanContextWithZeroClockAnchor(
+      SpanContext ctx, @Nullable String parentSpanId) {
     TraceContext result = new TraceContext();
-    result.filLFromSpanContext(ctx);
+    result.fillFromSpanContext(ctx, parentSpanId);
     result.clockAnchor = 0L;
     return result;
   }
 
-  private void filLFromSpanContext(SpanContext ctx) {
+  private void fillFromSpanContext(SpanContext ctx, @Nullable String parentSpanId) {
     id = HexUtils.hexToLong(ctx.getSpanId(), 0);
     traceIdHigh = HexUtils.hexToLong(ctx.getTraceId(), 0);
     traceIdLow = HexUtils.hexToLong(ctx.getTraceId(), 16);
+    if (parentSpanId != null) {
+      hasParentId = true;
+      parentId = HexUtils.hexToLong(parentSpanId, 0);
+    } else {
+      hasParentId = false;
+    }
     flags = ctx.getTraceFlags().asByte();
   }
 
@@ -73,6 +84,10 @@ public SpanContext toOtelSpanContext(StringBuilder temporaryBuilder) {
         traceIdStr, idStr, TraceFlags.fromByte(flags), TraceState.getDefault());
   }
 
+  public long getSpanId() {
+    return id;
+  }
+
   public boolean idEquals(@Nullable TraceContext o) {
     if (o == null) {
       return false;
@@ -89,7 +104,17 @@ public void deserialize(byte[] serialized) {
     traceIdHigh = ByteUtils.getLong(serialized, 8);
     id = ByteUtils.getLong(serialized, 16);
     flags = serialized[24];
-    clockAnchor = ByteUtils.getLong(serialized, 25);
+    hasParentId = serialized[25] != 0;
+    parentId = ByteUtils.getLong(serialized, 26);
+    clockAnchor = ByteUtils.getLong(serialized, 34);
+  }
+
+  public static long getParentId(byte[] serializedTraceContext) {
+    boolean hasParent = serializedTraceContext[25] != 0;
+    if (!hasParent) {
+      return 0L;
+    }
+    return ByteUtils.getLong(serializedTraceContext, 26);
   }
 
   public boolean traceIdAndIdEquals(byte[] otherSerialized) {
@@ -105,7 +130,13 @@ public boolean traceIdAndIdEquals(byte[] otherSerialized) {
     return id == otherId;
   }
 
-  public static void serialize(SpanContext ctx, long clockAnchor, byte[] buffer) {
+  public static void serialize(Span span, long clockAnchor, byte[] buffer) {
+    SpanContext ctx = span.getSpanContext();
+    SpanContext parentSpanCtx = SpanContext.getInvalid();
+    if (span instanceof ReadableSpan) {
+      parentSpanCtx = ((ReadableSpan) span).getParentSpanContext();
+    }
+
     long id = HexUtils.hexToLong(ctx.getSpanId(), 0);
     long traceIdHigh = HexUtils.hexToLong(ctx.getTraceId(), 0);
     long traceIdLow = HexUtils.hexToLong(ctx.getTraceId(), 16);
@@ -114,15 +145,29 @@ public static void serialize(SpanContext ctx, long clockAnchor, byte[] buffer) {
     ByteUtils.putLong(buffer, 8, traceIdHigh);
     ByteUtils.putLong(buffer, 16, id);
     buffer[24] = flags;
-    ByteUtils.putLong(buffer, 25, clockAnchor);
+    if (parentSpanCtx.isValid()) {
+      buffer[25] = 1;
+      ByteUtils.putLong(buffer, 26, HexUtils.hexToLong(parentSpanCtx.getSpanId(), 0));
+    } else {
+      buffer[25] = 0;
+      ByteUtils.putLong(buffer, 26, 0);
+    }
+    ByteUtils.putLong(buffer, 34, clockAnchor);
   }
 
   public void serialize(byte[] buffer) {
     ByteUtils.putLong(buffer, 0, traceIdLow);
     ByteUtils.putLong(buffer, 8, traceIdHigh);
     ByteUtils.putLong(buffer, 16, id);
     buffer[24] = flags;
-    ByteUtils.putLong(buffer, 25, clockAnchor);
+    if (hasParentId) {
+      buffer[25] = 1;
+      ByteUtils.putLong(buffer, 26, parentId);
+    } else {
+      buffer[25] = 0;
+      ByteUtils.putLong(buffer, 26, 0);
+    }
+    ByteUtils.putLong(buffer, 34, clockAnchor);
   }
 
   public byte[] serialize() {