Add context customizer hook to Instrumenter API

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.instrumenter;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.context.Context;
+
+/**
+ * Customizer of the {@link Context}. Instrumented libraries will call {@link #start(Context,
+ * Object, Attributes)} during {@link Instrumenter#start(Context, Object)}, allowing customization
+ * of the {@link Context} that is returned from that method.
+ */
+public interface ContextCustomizer<REQUEST> {
+
+  /**
+   * Context customizer method that is called during {@link Instrumenter#start(Context, Object)},
+   * allowing customization of the {@link Context} that is returned from that method.
+   */
+  Context start(Context context, REQUEST request, Attributes startAttributes);
+}

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

@@ -69,6 +69,7 @@ public static <REQUEST, RESPONSE> InstrumenterBuilder<REQUEST, RESPONSE> newBuil
   private final List<? extends SpanLinksExtractor<? super REQUEST>> spanLinksExtractors;
   private final List<? extends AttributesExtractor<? super REQUEST, ? super RESPONSE>>
       attributesExtractors;
+  private final List<? extends ContextCustomizer<? super REQUEST>> contextCustomizers;
   private final List<? extends RequestListener> requestListeners;
   private final ErrorCauseExtractor errorCauseExtractor;
   @Nullable private final StartTimeExtractor<REQUEST> startTimeExtractor;
@@ -85,6 +86,7 @@ public static <REQUEST, RESPONSE> InstrumenterBuilder<REQUEST, RESPONSE> newBuil
     this.spanStatusExtractor = builder.spanStatusExtractor;
     this.spanLinksExtractors = new ArrayList<>(builder.spanLinksExtractors);
     this.attributesExtractors = new ArrayList<>(builder.attributesExtractors);
+    this.contextCustomizers = new ArrayList<>(builder.contextCustomizers);
     this.requestListeners = new ArrayList<>(builder.requestListeners);
     this.errorCauseExtractor = builder.errorCauseExtractor;
     this.startTimeExtractor = builder.startTimeExtractor;
@@ -144,6 +146,10 @@ public Context start(Context parentContext, REQUEST request) {
 
     Context context = parentContext;
 
+    for (ContextCustomizer<? super REQUEST> contextCustomizer : contextCustomizers) {
+      context = contextCustomizer.start(context, request, attributes);
+    }
+
     for (RequestListener requestListener : requestListeners) {
       context = requestListener.start(context, attributes);
     }

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

@@ -12,6 +12,7 @@
 import io.opentelemetry.api.metrics.Meter;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.context.propagation.TextMapSetter;
 import io.opentelemetry.instrumentation.api.annotations.UnstableApi;
@@ -46,6 +47,7 @@ public final class InstrumenterBuilder<REQUEST, RESPONSE> {
   final List<SpanLinksExtractor<? super REQUEST>> spanLinksExtractors = new ArrayList<>();
   final List<AttributesExtractor<? super REQUEST, ? super RESPONSE>> attributesExtractors =
       new ArrayList<>();
+  final List<ContextCustomizer<? super REQUEST>> contextCustomizers = new ArrayList<>();
   final List<RequestListener> requestListeners = new ArrayList<>();
 
   SpanKindExtractor<? super REQUEST> spanKindExtractor = SpanKindExtractor.alwaysInternal();
@@ -106,6 +108,16 @@ public InstrumenterBuilder<REQUEST, RESPONSE> addSpanLinksExtractor(
     return this;
   }
 
+  /**
+   * Adds a {@link ContextCustomizer} to customize the context during {@link
+   * Instrumenter#start(Context, Object)}.
+   */
+  public InstrumenterBuilder<REQUEST, RESPONSE> addContextCustomizer(
+      ContextCustomizer<? super REQUEST> contextCustomizer) {
+    contextCustomizers.add(contextCustomizer);
+    return this;
+  }
+
   /** Adds a {@link RequestMetrics} whose metrics will be recorded for request start and stop. */
   @UnstableApi
   public InstrumenterBuilder<REQUEST, RESPONSE> addRequestMetrics(RequestMetrics factory) {

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/servlet/MappingResolver.java

@@ -11,6 +11,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Helper class for finding a mapping that matches current request from a collection of mappings.
@@ -55,7 +56,8 @@ public static MappingResolver build(Collection<String> mappings) {
   }
 
   /** Find mapping for requested path. */
-  public String resolve(String servletPath, String pathInfo) {
+  @Nullable
+  public String resolve(@Nullable String servletPath, @Nullable String pathInfo) {
     if (servletPath == null) {
       return null;
     }

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterTest.java

@@ -20,6 +20,7 @@
 import io.opentelemetry.api.trace.TraceState;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.instrumentation.api.instrumenter.db.DbAttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.http.HttpAttributesExtractor;
@@ -498,6 +499,24 @@ void shouldNotAddInvalidLink() {
                     span -> span.hasName("test span").hasTotalRecordedLinks(0)));
   }
 
+  @Test
+  void shouldUseContextCustomizer() {
+    // given
+    ContextKey<String> testKey = ContextKey.named("test");
+    Instrumenter<String, String> instrumenter =
+        Instrumenter.<String, String>newBuilder(
+                otelTesting.getOpenTelemetry(), "test", request -> "test span")
+            .addContextCustomizer(
+                (context, request, attributes) -> context.with(testKey, "testVal"))
+            .newInstrumenter();
+
+    // when
+    Context context = instrumenter.start(Context.root(), "request");
+
+    // then
+    assertThat(context.get(testKey)).isEqualTo("testVal");
+  }
+
   @Test
   void extractForwarded() {
     assertThat(ServerInstrumenter.extractForwarded("for=1.1.1.1")).isEqualTo("1.1.1.1");

instrumentation/jetty/jetty-11.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v11_0/Jetty11Singletons.java

@@ -5,7 +5,11 @@
 
 package io.opentelemetry.javaagent.instrumentation.jetty.v11_0;
 
+import static io.opentelemetry.instrumentation.api.servlet.ServerSpanNaming.Source.CONTAINER;
+
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.api.servlet.AppServerBridge;
+import io.opentelemetry.instrumentation.api.servlet.ServerSpanNaming;
 import io.opentelemetry.javaagent.instrumentation.jetty.common.JettyHelper;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletInstrumenterBuilder;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
@@ -20,8 +24,14 @@ public final class Jetty11Singletons {
   private static final Instrumenter<
           ServletRequestContext<HttpServletRequest>, ServletResponseContext<HttpServletResponse>>
       INSTRUMENTER =
-          ServletInstrumenterBuilder.newInstrumenter(
-              INSTRUMENTATION_NAME, Servlet5Accessor.INSTANCE);
+          ServletInstrumenterBuilder.<HttpServletRequest, HttpServletResponse>create()
+              .addContextCustomizer(
+                  (context, request, attributes) -> {
+                    context = ServerSpanNaming.init(context, CONTAINER);
+                    return AppServerBridge.init(context, /* shouldRecordException= */ false);
+                  })
+              .build(INSTRUMENTATION_NAME, Servlet5Accessor.INSTANCE);
+
   private static final JettyHelper<HttpServletRequest, HttpServletResponse> HELPER =
       new JettyHelper<>(INSTRUMENTER, Servlet5Accessor.INSTANCE);
 

instrumentation/jetty/jetty-8.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v8_0/Jetty8Singletons.java

@@ -5,7 +5,11 @@
 
 package io.opentelemetry.javaagent.instrumentation.jetty.v8_0;
 
+import static io.opentelemetry.instrumentation.api.servlet.ServerSpanNaming.Source.CONTAINER;
+
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.api.servlet.AppServerBridge;
+import io.opentelemetry.instrumentation.api.servlet.ServerSpanNaming;
 import io.opentelemetry.instrumentation.servlet.v3_0.Servlet3Accessor;
 import io.opentelemetry.javaagent.instrumentation.jetty.common.JettyHelper;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletInstrumenterBuilder;
@@ -20,8 +24,14 @@ public final class Jetty8Singletons {
   private static final Instrumenter<
           ServletRequestContext<HttpServletRequest>, ServletResponseContext<HttpServletResponse>>
       INSTRUMENTER =
-          ServletInstrumenterBuilder.newInstrumenter(
-              INSTRUMENTATION_NAME, Servlet3Accessor.INSTANCE);
+          ServletInstrumenterBuilder.<HttpServletRequest, HttpServletResponse>create()
+              .addContextCustomizer(
+                  (context, request, attributes) -> {
+                    context = ServerSpanNaming.init(context, CONTAINER);
+                    return AppServerBridge.init(context, /* shouldRecordException= */ false);
+                  })
+              .build(INSTRUMENTATION_NAME, Servlet3Accessor.INSTANCE);
+
   private static final JettyHelper<HttpServletRequest, HttpServletResponse> HELPER =
       new JettyHelper<>(INSTRUMENTER, Servlet3Accessor.INSTANCE);
 

instrumentation/jetty/jetty-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/common/JettyHelper.java

@@ -5,12 +5,9 @@
 
 package io.opentelemetry.javaagent.instrumentation.jetty.common;
 
-import static io.opentelemetry.instrumentation.api.servlet.ServerSpanNaming.Source.CONTAINER;
-
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
-import io.opentelemetry.instrumentation.api.servlet.AppServerBridge;
 import io.opentelemetry.instrumentation.servlet.ServletAccessor;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletHelper;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
@@ -24,15 +21,6 @@ public JettyHelper(
     super(instrumenter, accessor);
   }
 
-  public Context start(Context parentContext, ServletRequestContext<REQUEST> requestContext) {
-    return start(parentContext, requestContext, CONTAINER);
-  }
-
-  @Override
-  protected Context customizeContext(Context context, REQUEST httpServletRequest) {
-    return AppServerBridge.init(context, /* shouldRecordException= */ false);
-  }
-
   public void end(
       ServletRequestContext<REQUEST> requestContext,
       REQUEST request,

instrumentation/servlet/servlet-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v2_2/Servlet2Advice.java

@@ -58,7 +58,7 @@ public static void onEnter(
       return;
     }
 
-    context = helper().startSpan(parentContext, requestContext);
+    context = helper().start(parentContext, requestContext);
     scope = context.makeCurrent();
     // reset response status from previous request
     // (some servlet containers reuse response objects to reduce memory allocations)

instrumentation/servlet/servlet-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v2_2/Servlet2Helper.java

@@ -5,8 +5,6 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v2_2;
 
-import static io.opentelemetry.instrumentation.api.servlet.ServerSpanNaming.Source.SERVLET;
-
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.javaagent.instrumentation.servlet.BaseServletHelper;
@@ -25,11 +23,6 @@ public class Servlet2Helper extends BaseServletHelper<HttpServletRequest, HttpSe
     super(instrumenter, Servlet2Accessor.INSTANCE);
   }
 
-  public Context startSpan(
-      Context parentContext, ServletRequestContext<HttpServletRequest> requestContext) {
-    return start(parentContext, requestContext, SERVLET);
-  }
-
   public void stopSpan(
       Context context,
       ServletRequestContext<HttpServletRequest> requestContext,

instrumentation/servlet/servlet-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v2_2/Servlet2Singletons.java

@@ -5,9 +5,12 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v2_2;
 
+import static io.opentelemetry.instrumentation.api.servlet.ServerSpanNaming.Source.SERVLET;
+
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.http.HttpAttributesExtractor;
+import io.opentelemetry.instrumentation.api.servlet.ServerSpanNaming;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletInstrumenterBuilder;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletResponseContext;
@@ -29,11 +32,15 @@ public final class Servlet2Singletons {
     Instrumenter<
             ServletRequestContext<HttpServletRequest>, ServletResponseContext<HttpServletResponse>>
         instrumenter =
-            ServletInstrumenterBuilder.newInstrumenter(
-                INSTRUMENTATION_NAME,
-                Servlet2Accessor.INSTANCE,
-                spanNameExtractor,
-                httpAttributesExtractor);
+            ServletInstrumenterBuilder.<HttpServletRequest, HttpServletResponse>create()
+                .addContextCustomizer(
+                    (context, request, attributes) -> ServerSpanNaming.init(context, SERVLET))
+                .build(
+                    INSTRUMENTATION_NAME,
+                    Servlet2Accessor.INSTANCE,
+                    spanNameExtractor,
+                    httpAttributesExtractor);
+
     HELPER = new Servlet2Helper(instrumenter);
   }
 

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/Servlet3Advice.java

@@ -13,10 +13,8 @@
 import io.opentelemetry.instrumentation.api.servlet.MappingResolver;
 import io.opentelemetry.instrumentation.api.tracer.ServerSpan;
 import io.opentelemetry.javaagent.instrumentation.api.CallDepth;
-import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
-import javax.servlet.Filter;
 import javax.servlet.Servlet;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
@@ -47,21 +45,11 @@ public static void onEnter(
 
     HttpServletRequest httpServletRequest = (HttpServletRequest) request;
 
-    boolean servlet = servletOrFilter instanceof Servlet;
-    MappingResolver mappingResolver;
-    if (servlet) {
-      mappingResolver =
-          InstrumentationContext.get(Servlet.class, MappingResolver.class)
-              .get((Servlet) servletOrFilter);
-    } else {
-      mappingResolver =
-          InstrumentationContext.get(Filter.class, MappingResolver.class)
-              .get((Filter) servletOrFilter);
-    }
-
     Context currentContext = Java8BytecodeBridge.currentContext();
     Context attachedContext = helper().getServerContext(httpServletRequest);
     if (attachedContext != null && helper().needsRescoping(currentContext, attachedContext)) {
+      MappingResolver mappingResolver = Servlet3Singletons.getMappingResolver(servletOrFilter);
+      boolean servlet = servletOrFilter instanceof Servlet;
       attachedContext =
           helper().updateContext(attachedContext, httpServletRequest, mappingResolver, servlet);
       scope = attachedContext.makeCurrent();
@@ -75,6 +63,8 @@ public static void onEnter(
       // In case server span was created by app server instrumentations calling updateContext
       // returns a new context that contains servlet context path that is used in other
       // instrumentations for naming server span.
+      MappingResolver mappingResolver = Servlet3Singletons.getMappingResolver(servletOrFilter);
+      boolean servlet = servletOrFilter instanceof Servlet;
       Context updatedContext =
           helper().updateContext(currentContext, httpServletRequest, mappingResolver, servlet);
       if (currentContext != updatedContext) {
@@ -85,13 +75,13 @@ public static void onEnter(
       return;
     }
 
-    requestContext = new ServletRequestContext<>(httpServletRequest, mappingResolver);
+    requestContext = new ServletRequestContext<>(httpServletRequest, servletOrFilter);
 
     if (!helper().shouldStart(currentContext, requestContext)) {
       return;
     }
 
-    context = helper().start(currentContext, requestContext, servlet);
+    context = helper().start(currentContext, requestContext);
     scope = context.makeCurrent();
 
     helper().setAsyncListenerResponse(httpServletRequest, (HttpServletResponse) response);