Allow customization of parent-override behaviour for inferred-spans

inferred-spans/src/main/java/io/opentelemetry/contrib/inferredspans/InferredSpansAutoConfig.java

@@ -6,12 +6,15 @@
 package io.opentelemetry.contrib.inferredspans;
 
 import com.google.auto.service.AutoService;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -34,6 +37,8 @@ public class InferredSpansAutoConfig implements AutoConfigurationCustomizerProvi
   static final String INTERVAL_OPTION = "otel.inferred.spans.interval";
   static final String DURATION_OPTION = "otel.inferred.spans.duration";
   static final String LIB_DIRECTORY_OPTION = "otel.inferred.spans.lib.directory";
+  static final String PARENT_OVERRIDE_HANDLER_OPTION =
+      "otel.inferred.spans.parent.override.handler";
 
   @Override
   public void customize(AutoConfigurationCustomizer config) {
@@ -56,6 +61,12 @@ public void customize(AutoConfigurationCustomizer config) {
             applier.applyDuration(DURATION_OPTION, builder::profilingDuration);
             applier.applyString(LIB_DIRECTORY_OPTION, builder::profilerLibDirectory);
 
+            String parentOverrideHandlerName = properties.getString(PARENT_OVERRIDE_HANDLER_OPTION);
+            if (parentOverrideHandlerName != null && !parentOverrideHandlerName.isEmpty()) {
+              builder.parentOverrideHandler(
+                  constructParentOverrideHandler(parentOverrideHandlerName));
+            }
+
             providerBuilder.addSpanProcessor(builder.build());
           } else {
             log.finest(
@@ -65,6 +76,16 @@ public void customize(AutoConfigurationCustomizer config) {
         });
   }
 
+  @SuppressWarnings("unchecked")
+  private static BiConsumer<SpanBuilder, SpanContext> constructParentOverrideHandler(String name) {
+    try {
+      Class<?> clazz = Class.forName(name);
+      return (BiConsumer<SpanBuilder, SpanContext>) clazz.getConstructor().newInstance();
+    } catch (Exception e) {
+      throw new IllegalArgumentException("Could not construct parent override handler", e);
+    }
+  }
+
   private static class PropertiesApplier {
 
     private final ConfigProperties properties;

inferred-spans/src/main/java/io/opentelemetry/contrib/inferredspans/InferredSpansProcessorBuilder.java

@@ -5,12 +5,16 @@
 
 package io.opentelemetry.contrib.inferredspans;
 
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.contrib.inferredspans.internal.CallTree;
 import io.opentelemetry.contrib.inferredspans.internal.InferredSpansConfiguration;
 import io.opentelemetry.contrib.inferredspans.internal.SpanAnchoredClock;
 import java.io.File;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.BiConsumer;
 import javax.annotation.Nullable;
 
 @SuppressWarnings("CanIgnoreReturnValueSuggester")
@@ -48,6 +52,8 @@ public class InferredSpansProcessorBuilder {
 
   @Nullable private File activationEventsFile = null;
   @Nullable private File jfrFile = null;
+  private BiConsumer<SpanBuilder, SpanContext> parentOverrideHandler =
+      CallTree.DEFAULT_PARENT_OVERRIDE;
 
   InferredSpansProcessorBuilder() {}
 
@@ -64,7 +70,8 @@ public InferredSpansProcessor build() {
             excludedClasses,
             profilerInterval,
             profilingDuration,
-            profilerLibDirectory);
+            profilerLibDirectory,
+            parentOverrideHandler);
     return new InferredSpansProcessor(
         config, clock, startScheduledProfiling, activationEventsFile, jfrFile);
   }
@@ -188,4 +195,17 @@ InferredSpansProcessorBuilder jfrFile(@Nullable File jfrFile) {
     this.jfrFile = jfrFile;
     return this;
   }
+
+  /**
+   * Defines the action to perform when a inferred span is discovered to actually be the parent of a
+   * normal span. The first argument of the handler is the modifiable inferred span, the second
+   * argument the span context of the normal span which should be somehow marked as child of the
+   * inferred one. By default, a span link is added to the inferred span to represent this
+   * relationship.
+   */
+  InferredSpansProcessorBuilder parentOverrideHandler(
+      BiConsumer<SpanBuilder, SpanContext> handler) {
+    this.parentOverrideHandler = handler;
+    return this;
+  }
 }

inferred-spans/src/main/java/io/opentelemetry/contrib/inferredspans/internal/CallTree.java

@@ -25,6 +25,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
 import org.agrona.collections.LongHashSet;
@@ -50,9 +51,12 @@ public class CallTree implements Recyclable {
 
   private static final int INITIAL_CHILD_SIZE = 2;
 
-  private static final Attributes CHILD_LINK_ATTRIBUTES =
+  public static final Attributes CHILD_LINK_ATTRIBUTES =
       Attributes.builder().put(LINK_IS_CHILD, true).build();
 
+  public static final BiConsumer<SpanBuilder, SpanContext> DEFAULT_PARENT_OVERRIDE =
+      (inferredSpan, child) -> inferredSpan.addLink(child, CHILD_LINK_ATTRIBUTES);
+
   @Nullable private CallTree parent;
   protected int count;
   private List<CallTree> children = new ArrayList<>(INITIAL_CHILD_SIZE);
@@ -427,6 +431,7 @@ int spanify(
       @Nullable Span parentSpan,
       TraceContext parentContext,
       SpanAnchoredClock clock,
+      BiConsumer<SpanBuilder, SpanContext> normalSpanParentOverride,
       StringBuilder tempBuilder,
       Tracer tracer) {
     int createdSpans = 0;
@@ -437,7 +442,15 @@ int spanify(
     Span span = null;
     if (!isPillar() || isLeaf()) {
       createdSpans++;
-      span = asSpan(root, parentSpan, parentContext, tracer, clock, tempBuilder);
+      span =
+          asSpan(
+              root,
+              parentSpan,
+              parentContext,
+              tracer,
+              clock,
+              normalSpanParentOverride,
+              tempBuilder);
       this.isSpan = true;
     }
     List<CallTree> children = getChildren();
@@ -450,6 +463,7 @@ int spanify(
                   span != null ? span : parentSpan,
                   parentContext,
                   clock,
+                  normalSpanParentOverride,
                   tempBuilder,
                   tracer);
     }
@@ -462,6 +476,7 @@ protected Span asSpan(
       TraceContext parentContext,
       Tracer tracer,
       SpanAnchoredClock clock,
+      BiConsumer<SpanBuilder, SpanContext> normalSpanParentOverride,
       StringBuilder tempBuilder) {
 
     Context parentOtelCtx;
@@ -494,7 +509,11 @@ protected Span asSpan(
                 clock.toEpochNanos(parentContext.getClockAnchor(), this.start),
                 TimeUnit.NANOSECONDS);
     insertChildIdLinks(
-        spanBuilder, Span.fromContext(parentOtelCtx).getSpanContext(), parentContext, tempBuilder);
+        spanBuilder,
+        Span.fromContext(parentOtelCtx).getSpanContext(),
+        parentContext,
+        normalSpanParentOverride,
+        tempBuilder);
 
     // we're not interested in the very bottom of the stack which contains things like accepting and
     // handling connections
@@ -517,6 +536,7 @@ private void insertChildIdLinks(
       SpanBuilder span,
       SpanContext parentContext,
       TraceContext nonInferredParent,
+      BiConsumer<SpanBuilder, SpanContext> normalSpanParentOverride,
       StringBuilder tempBuilder) {
     if (childIds == null || childIds.isEmpty()) {
       return;
@@ -527,13 +547,13 @@ private void insertChildIdLinks(
       if (nonInferredParent.getSpanId() == childIds.getParentId(i)) {
         tempBuilder.setLength(0);
         HexUtils.appendLongAsHex(childIds.getId(i), tempBuilder);
-        SpanContext spanContext =
+        SpanContext childSpanContext =
             SpanContext.create(
                 parentContext.getTraceId(),
                 tempBuilder.toString(),
                 parentContext.getTraceFlags(),
                 parentContext.getTraceState());
-        span.addLink(spanContext, CHILD_LINK_ATTRIBUTES);
+        normalSpanParentOverride.accept(span, childSpanContext);
       }
     }
   }
@@ -863,13 +883,18 @@ private static CallTree findCommonAncestor(CallTree previousTopOfStack, CallTree
      * possible to update the parent ID of a regular span so that it correctly reflects being a
      * child of an inferred span.
      */
-    public int spanify(SpanAnchoredClock clock, Tracer tracer) {
+    public int spanify(
+        SpanAnchoredClock clock,
+        Tracer tracer,
+        BiConsumer<SpanBuilder, SpanContext> normalSpanOverride) {
       StringBuilder tempBuilder = new StringBuilder();
       int createdSpans = 0;
       List<CallTree> callTrees = getChildren();
       for (int i = 0, size = callTrees.size(); i < size; i++) {
         createdSpans +=
-            callTrees.get(i).spanify(this, null, rootContext, clock, tempBuilder, tracer);
+            callTrees
+                .get(i)
+                .spanify(this, null, rootContext, clock, normalSpanOverride, tempBuilder, tracer);
       }
       return createdSpans;
     }

inferred-spans/src/main/java/io/opentelemetry/contrib/inferredspans/internal/InferredSpansConfiguration.java

@@ -5,9 +5,12 @@
 
 package io.opentelemetry.contrib.inferredspans.internal;
 
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.contrib.inferredspans.WildcardMatcher;
 import java.time.Duration;
 import java.util.List;
+import java.util.function.BiConsumer;
 import javax.annotation.Nullable;
 
 public class InferredSpansConfiguration {
@@ -22,8 +25,8 @@ public class InferredSpansConfiguration {
   private final List<WildcardMatcher> excludedClasses;
   private final Duration profilerInterval;
   private final Duration profilingDuration;
-
   @Nullable private final String profilerLibDirectory;
+  private final BiConsumer<SpanBuilder, SpanContext> parentOverrideHandler;
 
   @SuppressWarnings("TooManyParameters")
   public InferredSpansConfiguration(
@@ -37,7 +40,8 @@ public InferredSpansConfiguration(
       List<WildcardMatcher> excludedClasses,
       Duration profilerInterval,
       Duration profilingDuration,
-      @Nullable String profilerLibDirectory) {
+      @Nullable String profilerLibDirectory,
+      BiConsumer<SpanBuilder, SpanContext> parentOverrideHandler) {
     this.profilerLoggingEnabled = profilerLoggingEnabled;
     this.backupDiagnosticFiles = backupDiagnosticFiles;
     this.asyncProfilerSafeMode = asyncProfilerSafeMode;
@@ -49,6 +53,7 @@ public InferredSpansConfiguration(
     this.profilerInterval = profilerInterval;
     this.profilingDuration = profilingDuration;
     this.profilerLibDirectory = profilerLibDirectory;
+    this.parentOverrideHandler = parentOverrideHandler;
   }
 
   public boolean isProfilingLoggingEnabled() {
@@ -100,4 +105,8 @@ public String getProfilerLibDirectory() {
   public boolean isPostProcessingEnabled() {
     return postProcessingEnabled;
   }
+
+  public BiConsumer<SpanBuilder, SpanContext> getParentOverrideHandler() {
+    return parentOverrideHandler;
+  }
 }

inferred-spans/src/main/java/io/opentelemetry/contrib/inferredspans/internal/SamplingProfiler.java

@@ -952,7 +952,10 @@ private void stopProfiling(SamplingProfiler samplingProfiler) {
           callTree.end(
               samplingProfiler.callTreePool, samplingProfiler.getInferredSpansMinDurationNs());
           int createdSpans =
-              callTree.spanify(samplingProfiler.getClock(), samplingProfiler.tracerProvider.get());
+              callTree.spanify(
+                  samplingProfiler.getClock(),
+                  samplingProfiler.tracerProvider.get(),
+                  samplingProfiler.config.getParentOverrideHandler());
           if (logger.isLoggable(Level.FINE)) {
             if (createdSpans > 0) {
               logger.log(

inferred-spans/src/test/java/io/opentelemetry/contrib/inferredspans/InferredSpansAutoConfigTest.java

@@ -11,6 +11,8 @@
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.contrib.inferredspans.internal.InferredSpansConfiguration;
@@ -23,6 +25,7 @@
 import java.nio.file.Path;
 import java.time.Duration;
 import java.util.List;
+import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -40,6 +43,12 @@ public void resetGlobalOtel() {
     OtelReflectionUtils.shutdownAndResetGlobalOtel();
   }
 
+  public static class NoOpParentOverrideHandler implements BiConsumer<SpanBuilder, SpanContext> {
+
+    @Override
+    public void accept(SpanBuilder spanBuilder, SpanContext spanContext) {}
+  }
+
   @Test
   @DisabledOnOpenJ9
   public void checkAllOptions(@TempDir Path tmpDir) {
@@ -57,7 +66,10 @@ public void checkAllOptions(@TempDir Path tmpDir) {
             .put(InferredSpansAutoConfig.EXCLUDED_CLASSES_OPTION, "blub,test*.test2")
             .put(InferredSpansAutoConfig.INTERVAL_OPTION, "2s")
             .put(InferredSpansAutoConfig.DURATION_OPTION, "3s")
-            .put(InferredSpansAutoConfig.LIB_DIRECTORY_OPTION, libDir)) {
+            .put(InferredSpansAutoConfig.LIB_DIRECTORY_OPTION, libDir)
+            .put(
+                InferredSpansAutoConfig.PARENT_OVERRIDE_HANDLER_OPTION,
+                NoOpParentOverrideHandler.class.getName())) {
 
       OpenTelemetry otel = GlobalOpenTelemetry.get();
       List<SpanProcessor> processors = OtelReflectionUtils.getSpanProcessors(otel);
@@ -81,6 +93,7 @@ public void checkAllOptions(@TempDir Path tmpDir) {
       assertThat(config.getProfilingInterval()).isEqualTo(Duration.ofSeconds(2));
       assertThat(config.getProfilingDuration()).isEqualTo(Duration.ofSeconds(3));
       assertThat(config.getProfilerLibDirectory()).isEqualTo(libDir);
+      assertThat(config.getParentOverrideHandler()).isInstanceOf(NoOpParentOverrideHandler.class);
     }
   }
 

inferred-spans/src/test/java/io/opentelemetry/contrib/inferredspans/internal/CallTreeSpanifyTest.java

@@ -48,7 +48,10 @@ void testSpanification() throws Exception {
       setup.profiler.setProfilingSessionOngoing(true);
       CallTree.Root callTree =
           CallTreeTest.getCallTree(setup, new String[] {" dd   ", " cc   ", " bbb  ", "aaaaee"});
-      assertThat(callTree.spanify(nanoClock, setup.sdk.getTracer("dummy-tracer"))).isEqualTo(4);
+      assertThat(
+              callTree.spanify(
+                  nanoClock, setup.sdk.getTracer("dummy-tracer"), CallTree.DEFAULT_PARENT_OVERRIDE))
+          .isEqualTo(4);
       assertThat(setup.getSpans()).hasSize(5);
       assertThat(setup.getSpans().stream().map(SpanData::getName))
           .containsExactly(
@@ -158,7 +161,8 @@ void testCallTreeWithActiveSpan() {
                     .build());
 
     try (OpenTelemetrySdk outputSdk = sdkBuilder.build()) {
-      root.spanify(nanoClock, outputSdk.getTracer("dummy-tracer"));
+      root.spanify(
+          nanoClock, outputSdk.getTracer("dummy-tracer"), CallTree.DEFAULT_PARENT_OVERRIDE);
 
       List<SpanData> spans = exporter.getFinishedSpanItems();
       assertThat(spans).hasSize(2);
@@ -206,7 +210,8 @@ void testSpanWithInvertedActivation() {
                     .addSpanProcessor(SimpleSpanProcessor.create(exporter))
                     .build());
     try (OpenTelemetrySdk outputSdk = sdkBuilder.build()) {
-      root.spanify(nanoClock, outputSdk.getTracer("dummy-tracer"));
+      root.spanify(
+          nanoClock, outputSdk.getTracer("dummy-tracer"), CallTree.DEFAULT_PARENT_OVERRIDE);
 
       List<SpanData> spans = exporter.getFinishedSpanItems();
       assertThat(spans).hasSize(1);
@@ -249,7 +254,8 @@ void testSpanWithNestedActivation() {
                     .addSpanProcessor(SimpleSpanProcessor.create(exporter))
                     .build());
     try (OpenTelemetrySdk outputSdk = sdkBuilder.build()) {
-      root.spanify(nanoClock, outputSdk.getTracer("dummy-tracer"));
+      root.spanify(
+          nanoClock, outputSdk.getTracer("dummy-tracer"), CallTree.DEFAULT_PARENT_OVERRIDE);
 
       List<SpanData> spans = exporter.getFinishedSpanItems();
       assertThat(spans).hasSize(1);

inferred-spans/src/test/java/io/opentelemetry/contrib/inferredspans/internal/CallTreeTest.java

@@ -869,7 +869,10 @@ private Map<String, SpanData> assertCallTree(
     assertThat(actualResult).isEqualTo(expectedResult.toString());
 
     if (expectedSpans != null) {
-      root.spanify(nanoClock, profilerSetup.sdk.getTracer("dummy-inferred-spans-tracer"));
+      root.spanify(
+          nanoClock,
+          profilerSetup.sdk.getTracer("dummy-inferred-spans-tracer"),
+          CallTree.DEFAULT_PARENT_OVERRIDE);
       Map<String, SpanData> spans =
           profilerSetup.getSpans().stream()
               .collect(toMap(s -> s.getName().replaceAll(".*#", ""), Function.identity()));