Pass around context more instead of span (HttpServerTracer)

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/tracer/HttpServerTracer.java

@@ -12,7 +12,6 @@
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.api.trace.attributes.SemanticAttributes;
 import io.opentelemetry.context.Context;
-import io.opentelemetry.context.Scope;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import java.lang.reflect.Method;
 import java.util.concurrent.TimeUnit;
@@ -35,17 +34,29 @@ public HttpServerTracer(Tracer tracer) {
     super(tracer);
   }
 
-  public Context startSpan(REQUEST request, CONNECTION connection, Method origin) {
+  public Context startSpan(REQUEST request, CONNECTION connection, STORAGE storage, Method origin) {
     String spanName = spanNameForMethod(origin);
-    return startSpan(request, connection, spanName);
+    return startSpan(request, connection, storage, spanName);
   }
 
-  public Context startSpan(REQUEST request, CONNECTION connection, String spanName) {
-    return startSpan(request, connection, spanName, -1);
+  public Context startSpan(
+      REQUEST request, CONNECTION connection, STORAGE storage, String spanName) {
+    return startSpan(request, connection, storage, spanName, -1);
   }
 
   public Context startSpan(
-      REQUEST request, CONNECTION connection, String spanName, long startTimestamp) {
+      REQUEST request,
+      CONNECTION connection,
+      @Nullable STORAGE storage,
+      String spanName,
+      long startTimestamp) {
+
+    // not checking if inside of nested SERVER span because of concerns about context leaking
+    // and so always starting with a clean context here
+
+    // also we can't conditionally start a span in this method, because the caller won't know
+    // whether to call end() or not on the Span in the returned Context
+
     Context parentContext = extract(request, getGetter());
     SpanBuilder builder = tracer.spanBuilder(spanName).setSpanKind(SERVER).setParent(parentContext);
 
@@ -58,67 +69,51 @@ public Context startSpan(
     onRequest(span, request);
     onConnectionAndRequest(span, connection, request);
 
-    return parentContext.with(span);
-  }
+    Context context = parentContext.with(CONTEXT_SERVER_SPAN_KEY, span).with(span);
+    attachServerContext(context, storage);
 
-  /**
-   * Creates new scoped context, based on the current context, with the given span.
-   *
-   * <p>Attaches new context to the request to avoid creating duplicate server spans.
-   */
-  public Scope startScope(Span span, STORAGE storage) {
-    return startScope(span, storage, Context.current());
-  }
-
-  /**
-   * Creates new scoped context, based on the given context, with the given span.
-   *
-   * <p>Attaches new context to the request to avoid creating duplicate server spans.
-   */
-  public Scope startScope(Span span, STORAGE storage, Context context) {
-    // TODO we could do this in one go, but TracingContextUtils.CONTEXT_SPAN_KEY is private
-    Context newContext = context.with(CONTEXT_SERVER_SPAN_KEY, span).with(span);
-    attachServerContext(newContext, storage);
-    return newContext.makeCurrent();
+    return context;
   }
 
   /**
-   * Convenience method. Delegates to {@link #end(Span, Object, long)}, passing {@code timestamp}
+   * Convenience method. Delegates to {@link #end(Context, Object, long)}, passing {@code timestamp}
    * value of {@code -1}.
    */
   // TODO should end methods remove SPAN attribute from request as well?
-  public void end(Span span, RESPONSE response) {
-    end(span, response, -1);
+  public void end(Context context, RESPONSE response) {
+    end(context, response, -1);
   }
 
   // TODO should end methods remove SPAN attribute from request as well?
-  public void end(Span span, RESPONSE response, long timestamp) {
+  public void end(Context context, RESPONSE response, long timestamp) {
+    Span span = Span.fromContext(context);
     setStatus(span, responseStatus(response));
     endSpan(span, timestamp);
   }
 
   /**
-   * Convenience method. Delegates to {@link #endExceptionally(Span, Throwable, Object)}, passing
+   * Convenience method. Delegates to {@link #endExceptionally(Context, Throwable, Object)}, passing
    * {@code response} value of {@code null}.
    */
-  @Override
-  public void endExceptionally(Span span, Throwable throwable) {
-    endExceptionally(span, throwable, null);
+  public void endExceptionally(Context context, Throwable throwable) {
+    endExceptionally(context, throwable, null);
   }
 
   /**
-   * Convenience method. Delegates to {@link #endExceptionally(Span, Throwable, Object, long)},
+   * Convenience method. Delegates to {@link #endExceptionally(Context, Throwable, Object, long)},
    * passing {@code timestamp} value of {@code -1}.
    */
-  public void endExceptionally(Span span, Throwable throwable, RESPONSE response) {
-    endExceptionally(span, throwable, response, -1);
+  public void endExceptionally(Context context, Throwable throwable, RESPONSE response) {
+    endExceptionally(context, throwable, response, -1);
   }
 
   /**
    * If {@code response} is {@code null}, the {@code http.status_code} will be set to {@code 500}
    * and the {@link Span} status will be set to {@link io.opentelemetry.api.trace.StatusCode#ERROR}.
    */
-  public void endExceptionally(Span span, Throwable throwable, RESPONSE response, long timestamp) {
+  public void endExceptionally(
+      Context context, Throwable throwable, RESPONSE response, long timestamp) {
+    Span span = Span.fromContext(context);
     onError(span, unwrapThrowable(throwable));
     if (response == null) {
       setStatus(span, 500);

instrumentation-core/servlet-2.2/src/main/java/io/opentelemetry/instrumentation/servlet/ServletHttpServerTracer.java

@@ -8,7 +8,6 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.attributes.SemanticAttributes;
 import io.opentelemetry.context.Context;
-import io.opentelemetry.context.Scope;
 import io.opentelemetry.context.propagation.TextMapPropagator.Getter;
 import io.opentelemetry.instrumentation.api.servlet.ServletContextPath;
 import io.opentelemetry.instrumentation.api.tracer.HttpServerTracer;
@@ -26,19 +25,12 @@ public abstract class ServletHttpServerTracer<RESPONSE>
   private static final Logger log = LoggerFactory.getLogger(ServletHttpServerTracer.class);
 
   public Context startSpan(HttpServletRequest request) {
-    return startSpan(request, request, getSpanName(request));
-  }
-
-  @Override
-  public Scope startScope(Span span, HttpServletRequest request) {
-    Context context = Context.current();
-    // if we start returning Context from startSpan() then we can add ServletContextPath there
-    // (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1481)
+    Context context = startSpan(request, request, request, getSpanName(request));
     String contextPath = request.getContextPath();
     if (contextPath != null && !contextPath.isEmpty() && !contextPath.equals("/")) {
       context = context.with(ServletContextPath.CONTEXT_KEY, contextPath);
     }
-    return super.startScope(span, request, context);
+    return context;
   }
 
   @Override
@@ -109,10 +101,10 @@ protected Throwable unwrapThrowable(Throwable throwable) {
     return super.unwrapThrowable(result);
   }
 
-  public void setPrincipal(Span span, HttpServletRequest request) {
+  public void setPrincipal(Context context, HttpServletRequest request) {
     Principal principal = request.getUserPrincipal();
     if (principal != null) {
-      span.setAttribute(SemanticAttributes.ENDUSER_ID, principal.getName());
+      Span.fromContext(context).setAttribute(SemanticAttributes.ENDUSER_ID, principal.getName());
     }
   }
 

instrumentation/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/AkkaHttpServerInstrumentationModule.java

@@ -14,10 +14,8 @@
 import akka.http.scaladsl.model.HttpResponse;
 import akka.stream.Materializer;
 import com.google.auto.service.AutoService;
-import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
-import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.HashMap;
@@ -96,14 +94,13 @@ public SyncWrapper(Function1<HttpRequest, HttpResponse> userHandler) {
 
     @Override
     public HttpResponse apply(HttpRequest request) {
-      Context ctx = tracer().startSpan(request, request, "akka.request");
-      Span span = Java8BytecodeBridge.spanFromContext(ctx);
-      try (Scope ignored = tracer().startScope(span, null)) {
+      Context ctx = tracer().startSpan(request, request, null, "akka.request");
+      try (Scope ignored = ctx.makeCurrent()) {
         HttpResponse response = userHandler.apply(request);
-        tracer().end(span, response);
+        tracer().end(ctx, response);
         return response;
       } catch (Throwable t) {
-        tracer().endExceptionally(span, t);
+        tracer().endExceptionally(ctx, t);
         throw t;
       }
     }
@@ -122,29 +119,28 @@ public AsyncWrapper(
 
     @Override
     public Future<HttpResponse> apply(HttpRequest request) {
-      Context ctx = tracer().startSpan(request, request, "akka.request");
-      Span span = Java8BytecodeBridge.spanFromContext(ctx);
-      try (Scope ignored = tracer().startScope(span, null)) {
+      Context ctx = tracer().startSpan(request, request, null, "akka.request");
+      try (Scope ignored = ctx.makeCurrent()) {
         return userHandler
             .apply(request)
             .transform(
                 new AbstractFunction1<HttpResponse, HttpResponse>() {
                   @Override
                   public HttpResponse apply(HttpResponse response) {
-                    tracer().end(span, response);
+                    tracer().end(ctx, response);
                     return response;
                   }
                 },
                 new AbstractFunction1<Throwable, Throwable>() {
                   @Override
                   public Throwable apply(Throwable t) {
-                    tracer().endExceptionally(span, t);
+                    tracer().endExceptionally(ctx, t);
                     return t;
                   }
                 },
                 executionContext);
       } catch (Throwable t) {
-        tracer().endExceptionally(span, t);
+        tracer().endExceptionally(ctx, t);
         throw t;
       }
     }

instrumentation/armeria-1.0/library/src/main/java/io/opentelemetry/instrumentation/armeria/v1_0/server/OpenTelemetryService.java

@@ -59,25 +59,24 @@ public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exc
     long requestStartTimeMicros =
         ctx.log().ensureAvailable(RequestLogProperty.REQUEST_START_TIME).requestStartTimeMicros();
     long requestStartTimeNanos = TimeUnit.MICROSECONDS.toNanos(requestStartTimeMicros);
-    Context context = serverTracer.startSpan(req, ctx, spanName, requestStartTimeNanos);
-    Span span = Span.fromContext(context);
+    Context context = serverTracer.startSpan(req, ctx, null, spanName, requestStartTimeNanos);
 
-    if (span.isRecording()) {
+    if (Span.fromContext(context).isRecording()) {
       ctx.log()
           .whenComplete()
           .thenAccept(
               log -> {
                 long requestEndTimeNanos = requestStartTimeNanos + log.responseDurationNanos();
                 if (log.responseCause() != null) {
                   serverTracer.endExceptionally(
-                      span, log.responseCause(), log, requestEndTimeNanos);
+                      context, log.responseCause(), log, requestEndTimeNanos);
                 } else {
-                  serverTracer.end(span, log, requestEndTimeNanos);
+                  serverTracer.end(context, log, requestEndTimeNanos);
                 }
               });
     }
 
-    try (Scope ignored = span.makeCurrent()) {
+    try (Scope ignored = context.makeCurrent()) {
       return unwrap().serve(ctx, req);
     }
   }

instrumentation/grizzly-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/grizzly/HttpCodecFilterAdvice.java

@@ -7,9 +7,7 @@
 
 import static io.opentelemetry.javaagent.instrumentation.grizzly.GrizzlyHttpServerTracer.tracer;
 
-import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
-import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import java.lang.reflect.Method;
 import net.bytebuddy.asm.Advice;
 import org.glassfish.grizzly.filterchain.FilterChainContext;
@@ -31,11 +29,9 @@ public static void onExit(
       return;
     }
     HttpRequestPacket httpRequest = (HttpRequestPacket) httpHeader;
-    Context extractedContext = tracer().startSpan(httpRequest, httpRequest, method);
-    Span span = Java8BytecodeBridge.spanFromContext(extractedContext);
 
-    // We don't actually want to attach new context to this thread, as actual request will continue
-    // on some other thread. But we do want to attach that new context to FilterChainContext
-    tracer().startScope(span, ctx).close();
+    // We don't want to attach the new context to this thread, as actual request will continue on
+    // some other thread where we will read and attach it via tracer().getServerContext(ctx).
+    tracer().startSpan(httpRequest, httpRequest, ctx, method);
   }
 }

instrumentation/grizzly-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/grizzly/HttpCodecFilterOldAdvice.java

@@ -7,9 +7,7 @@
 
 import static io.opentelemetry.javaagent.instrumentation.grizzly.GrizzlyHttpServerTracer.tracer;
 
-import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
-import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import java.lang.reflect.Method;
 import net.bytebuddy.asm.Advice;
 import org.glassfish.grizzly.filterchain.FilterChainContext;
@@ -31,11 +29,9 @@ public static void onExit(
       return;
     }
     HttpRequestPacket httpRequest = (HttpRequestPacket) httpHeader;
-    Context extractedContext = tracer().startSpan(httpRequest, httpRequest, method);
-    Span span = Java8BytecodeBridge.spanFromContext(extractedContext);
 
-    // We don't actually want to attach new context to this thread, as actual request will continue
-    // on some other thread. But we do want to attach that new context to FilterChainContext
-    tracer().startScope(span, ctx).close();
+    // We don't want to attach the new context to this thread, as actual request will continue on
+    // some other thread where we will read and attach it via tracer().getServerContext(ctx).
+    tracer().startSpan(httpRequest, httpRequest, ctx, method);
   }
 }

instrumentation/grizzly-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/grizzly/HttpServerFilterAdvice.java

@@ -7,7 +7,7 @@
 
 import static io.opentelemetry.javaagent.instrumentation.grizzly.GrizzlyHttpServerTracer.tracer;
 
-import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
 import net.bytebuddy.asm.Advice;
 import org.glassfish.grizzly.filterchain.FilterChainContext;
 import org.glassfish.grizzly.http.HttpResponsePacket;
@@ -16,9 +16,9 @@ public class HttpServerFilterAdvice {
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void onExit(
       @Advice.Argument(0) FilterChainContext ctx, @Advice.Argument(2) HttpResponsePacket response) {
-    Span span = tracer().getServerSpan(ctx);
-    if (span != null) {
-      tracer().end(span, response);
+    Context context = tracer().getServerContext(ctx);
+    if (context != null) {
+      tracer().end(context, response);
     }
   }
 }