Add span processing function from application code

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-api/src/main/java/io/opentelemetry/instrumentation/api/util/ContextSpanProcessorUtil.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.util;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import java.util.function.BiConsumer;
+import javax.annotation.Nullable;
+
+/**
+ * A helper class that allow setting a processing function that will be applied to the spans created
+ * while the context with processing function is active.
+ */
+public final class ContextSpanProcessorUtil {
+  private static final ContextKey<BiConsumer<Context, Span>> KEY =
+      ContextKey.named("opentelemetry-context-span-processor");
+
+  private ContextSpanProcessorUtil() {}
+
+  /** Returns currently set span processor function from context or null. */
+  // instrumented by ContextSpanProcessorUtilInstrumentation
+  @Nullable
+  public static BiConsumer<Context, Span> fromContextOrNull(Context context) {
+    return context.get(KEY);
+  }
+
+  /**
+   * Store a span processor function in context. This processor will be applied to all spans created
+   * while the returned context is active. Processing function is called from {@code
+   * io.opentelemetry.sdk.trace.SpanProcessor#onStart} and should adhere to the same rules as the
+   * onStart method.
+   */
+  // instrumented by ContextSpanProcessorUtilInstrumentation
+  public static Context storeInContext(Context context, BiConsumer<Context, Span> processor) {
+    return context.with(KEY, processor);
+  }
+}

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

@@ -0,0 +1,86 @@
+/*
+ * 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.util.ContextSpanProcessorUtil;
+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 ContextSpanProcessorUtilInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("application.io.opentelemetry.instrumentation.api.util.ContextSpanProcessorUtil");
+  }
+
+  @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 =
+          ContextSpanProcessorUtil.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 =
+          ContextSpanProcessorUtil.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 ContextSpanProcessorUtilInstrumentation());
   }
 }

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

@@ -0,0 +1,38 @@
+/*
+ * 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.util.ContextSpanProcessorUtil;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class ContextSpanProcessorTest {
+
+  @RegisterExtension
+  static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+
+  @Test
+  void tesUpdateSpanName() {
+    Context context = Context.current();
+    context =
+        ContextSpanProcessorUtil.storeInContext(
+            context,
+            (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")));
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentTracerProviderConfigurer.java

@@ -39,6 +39,7 @@ private static SdkTracerProviderBuilder configure(
       sdkTracerProviderBuilder.addSpanProcessor(new AddThreadDetailsSpanProcessor());
     }
 
+    sdkTracerProviderBuilder.addSpanProcessor(new ContextSpanProcessor());
     maybeEnableLoggingExporter(sdkTracerProviderBuilder);
 
     return sdkTracerProviderBuilder;

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/ContextSpanProcessor.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.util.ContextSpanProcessorUtil;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+import java.util.function.BiConsumer;
+
+/** A span processor that retrieves the actual processing function from context. */
+final class ContextSpanProcessor implements SpanProcessor {
+
+  @Override
+  public void onStart(Context context, ReadWriteSpan readWriteSpan) {
+    BiConsumer<Context, Span> spanProcessor = ContextSpanProcessorUtil.fromContextOrNull(context);
+    if (spanProcessor != null) {
+      spanProcessor.accept(context, readWriteSpan);
+    }
+  }
+
+  @Override
+  public boolean isStartRequired() {
+    return true;
+  }
+
+  @Override
+  public void onEnd(ReadableSpan readableSpan) {}
+
+  @Override
+  public boolean isEndRequired() {
+    return false;
+  }
+}