Add span processing function from application code

instrumentation-api-semconv/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/ContextSpanProcessor.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.instrumenter;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ImplicitContextKeyed;
+import io.opentelemetry.instrumentation.api.internal.ContextSpanProcessorImpl;
+import java.util.function.BiConsumer;
+
+/**
+ * A span processing function that can be stored in context. When stored in context this function
+ * will be applied to all spans create while the context is active. Processing function is called
+ * synchronously on the execution thread, it should not throw or block the execution thread.
+ *
+ * <p><strong>NOTE:</strong> This API is <strong>EXPERIMENTAL</strong>, it may be removed or
+ * changed.
+ */
+public interface ContextSpanProcessor extends ImplicitContextKeyed {
+
+  /**
+   * Wrap a {@link BiConsumer} so that it can be stored in context as a span processing function.
+   */
+  static ContextSpanProcessor wrap(BiConsumer<Context, Span> processor) {
+    return new ContextSpanProcessorImpl(processor);
+  }
+}

instrumentation-api-semconv/src/main/java/io/opentelemetry/instrumentation/api/internal/ContextSpanProcessorImpl.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.internal;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.instrumentation.api.instrumenter.ContextSpanProcessor;
+import java.util.function.BiConsumer;
+import javax.annotation.Nullable;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public final class ContextSpanProcessorImpl implements ContextSpanProcessor {
+  private static final ContextKey<BiConsumer<Context, Span>> KEY =
+      ContextKey.named("opentelemetry-context-span-processor");
+
+  private final BiConsumer<Context, Span> processor;
+
+  public ContextSpanProcessorImpl(BiConsumer<Context, Span> processor) {
+    this.processor = processor;
+  }
+
+  public static void onStart(Context context, Span span) {
+    BiConsumer<Context, Span> processor = fromContextOrNull(context);
+    if (processor != null) {
+      processor.accept(context, span);
+    }
+  }
+
+  @Override
+  public Context storeInContext(Context context) {
+    return storeInContext(context, processor);
+  }
+
+  // instrumented by ContextSpanProcessorUtilInstrumentation
+  public static Context storeInContext(Context context, BiConsumer<Context, Span> processor) {
+    return context.with(KEY, processor);
+  }
+
+  @Nullable
+  // instrumented by ContextSpanProcessorImplInstrumentation
+  public static BiConsumer<Context, Span> fromContextOrNull(Context context) {
+    return context.get(KEY);
+  }
+}

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/Instrumenter.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.internal.ContextSpanProcessorInvoker;
 import io.opentelemetry.instrumentation.api.internal.SupportabilityMetrics;
 import java.time.Instant;
 import java.util.ArrayList;
@@ -183,6 +184,7 @@ private Context doStart(Context parentContext, REQUEST request, @Nullable Instan
 
     spanBuilder.setAllAttributes(attributes);
     Span span = spanBuilder.startSpan();
+    ContextSpanProcessorInvoker.onStart(context, span);
     context = context.with(span);
 
     for (ContextCustomizer<? super REQUEST> contextCustomizer : contextCustomizers) {

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/internal/ContextSpanProcessorInvoker.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.internal;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import java.lang.reflect.Method;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public final class ContextSpanProcessorInvoker {
+  private static final Method onStartMethod = getMethod();
+
+  private static Method getMethod() {
+    try {
+      Class<?> clazz =
+          Class.forName("io.opentelemetry.instrumentation.api.internal.ContextSpanProcessorImpl");
+      return clazz.getMethod("onStart", Context.class, Span.class);
+    } catch (ClassNotFoundException | NoSuchMethodException exception) {
+      return null;
+    }
+  }
+
+  public static void onStart(Context context, Span span) {
+    if (onStartMethod == null) {
+      return;
+    }
+    try {
+      onStartMethod.invoke(null, context, span);
+    } catch (Exception exception) {
+      throw new IllegalStateException(exception);
+    }
+  }
+
+  private ContextSpanProcessorInvoker() {}
+}

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/internal/SpanKey.java

@@ -76,11 +76,13 @@ private SpanKey(ContextKey<Span> key) {
     this.key = key;
   }
 
+  // instumented by SpanKeyInstrumentation
   public Context storeInContext(Context context, Span span) {
     return context.with(key, span);
   }
 
   @Nullable
+  // instumented by SpanKeyInstrumentation
   public Span fromContextOrNull(Context context) {
     return context.get(key);
   }

instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/ContextSpanProcessorImplInstrumentation.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.instrumentationapi;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import application.io.opentelemetry.api.trace.Span;
+import application.io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.internal.ContextSpanProcessorImpl;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.context.AgentContextStorage;
+import java.util.function.BiConsumer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+final class ContextSpanProcessorImplInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named(
+        "application.io.opentelemetry.instrumentation.api.internal.ContextSpanProcessorImpl");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("storeInContext")
+            .and(takesArgument(0, named("application.io.opentelemetry.context.Context")))
+            .and(takesArgument(1, BiConsumer.class)),
+        this.getClass().getName() + "$StoreInContextAdvice");
+    transformer.applyAdviceToMethod(
+        named("fromContextOrNull")
+            .and(takesArgument(0, named("application.io.opentelemetry.context.Context"))),
+        this.getClass().getName() + "$FromContextOrNullAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class StoreInContextAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class)
+    public static Object onEnter() {
+      return null;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Argument(0) Context applicationContext,
+        @Advice.Argument(1) BiConsumer<Context, Span> processor,
+        @Advice.Return(readOnly = false) Context newApplicationContext) {
+
+      io.opentelemetry.context.Context agentContext =
+          AgentContextStorage.getAgentContext(applicationContext);
+
+      BiConsumer<io.opentelemetry.context.Context, io.opentelemetry.api.trace.Span> agentProcessor =
+          ContextSpanProcessorWrapper.wrap(processor);
+      io.opentelemetry.context.Context newAgentContext =
+          ContextSpanProcessorImpl.storeInContext(agentContext, agentProcessor);
+
+      newApplicationContext = AgentContextStorage.toApplicationContext(newAgentContext);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class FromContextOrNullAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class)
+    public static Object onEnter() {
+      return null;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Argument(0) Context applicationContext,
+        @Advice.Return(readOnly = false) BiConsumer<Context, Span> processor) {
+
+      io.opentelemetry.context.Context agentContext =
+          AgentContextStorage.getAgentContext(applicationContext);
+      BiConsumer<io.opentelemetry.context.Context, io.opentelemetry.api.trace.Span> agentProcessor =
+          ContextSpanProcessorImpl.fromContextOrNull(agentContext);
+      processor = ContextSpanProcessorWrapper.unwrap(agentProcessor);
+    }
+  }
+}

instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/ContextSpanProcessorWrapper.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.instrumentationapi;
+
+import application.io.opentelemetry.api.trace.Span;
+import application.io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.context.AgentContextStorage;
+import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.trace.Bridging;
+import java.util.function.BiConsumer;
+
+/**
+ * Wrapper that translates agent context and span to application context and span before invoking
+ * the delegate.
+ */
+public final class ContextSpanProcessorWrapper
+    implements BiConsumer<io.opentelemetry.context.Context, io.opentelemetry.api.trace.Span> {
+
+  private final BiConsumer<Context, Span> delegate;
+
+  private ContextSpanProcessorWrapper(BiConsumer<Context, Span> delegate) {
+    this.delegate = delegate;
+  }
+
+  public static BiConsumer<io.opentelemetry.context.Context, io.opentelemetry.api.trace.Span> wrap(
+      BiConsumer<Context, Span> processor) {
+    return new ContextSpanProcessorWrapper(processor);
+  }
+
+  public static BiConsumer<Context, Span> unwrap(
+      BiConsumer<io.opentelemetry.context.Context, io.opentelemetry.api.trace.Span> processor) {
+    if (processor instanceof ContextSpanProcessorWrapper) {
+      return ((ContextSpanProcessorWrapper) processor).delegate;
+    }
+    return null;
+  }
+
+  @Override
+  public void accept(
+      io.opentelemetry.context.Context context, io.opentelemetry.api.trace.Span span) {
+    delegate.accept(
+        AgentContextStorage.toApplicationContext(context), Bridging.toApplication(span));
+  }
+}

instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/InstrumentationApiInstrumentationModule.java

@@ -28,6 +28,9 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return asList(new HttpRouteStateInstrumentation(), new SpanKeyInstrumentation());
+    return asList(
+        new HttpRouteStateInstrumentation(),
+        new SpanKeyInstrumentation(),
+        new ContextSpanProcessorImplInstrumentation());
   }
 }

instrumentation/opentelemetry-instrumentation-api/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/ContextSpanProcessorTest.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.instrumentationapi;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.instrumenter.ContextSpanProcessor;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import io.opentelemetry.javaagent.instrumentation.testing.AgentSpanTesting;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class ContextSpanProcessorTest {
+
+  @RegisterExtension
+  static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+
+  @Test
+  void testUpdateSpanName() {
+    Context context =
+        Context.current()
+            .with(ContextSpanProcessor.wrap((context1, span) -> span.updateName("new span name")));
+
+    try (Scope scope = context.makeCurrent()) {
+      testing.runWithSpan("old span name", () -> {});
+    }
+
+    testing.waitAndAssertTraces(
+        trace -> trace.hasSpansSatisfyingExactly(span -> span.hasName("new span name")));
+  }
+
+  @Test
+  void testUpdateAgentSpanName() {
+    Context context =
+        Context.current()
+            .with(ContextSpanProcessor.wrap((context1, span) -> span.updateName("new span name")));
+
+    try (Scope scope = context.makeCurrent()) {
+      AgentSpanTesting.runWithAllSpanKeys("old span name", () -> {});
+    }
+
+    testing.waitAndAssertTraces(
+        trace -> trace.hasSpansSatisfyingExactly(span -> span.hasName("new span name")));
+  }
+}