Migrates to using a context when span is started (#1525)

* Return Context when starting span to preserve baggage.

Co-authored-by: Marcin Grzejszczak <marcin.grzejszczak@gmail.com>

* Use bridge

Co-authored-by: Marcin Grzejszczak <marcin.grzejszczak@gmail.com>

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

@@ -44,19 +44,20 @@ public abstract class HttpServerTracer<REQUEST, RESPONSE, CONNECTION, STORAGE> e
     super(tracer);
   }
 
-  public Span startSpan(REQUEST request, CONNECTION connection, Method origin) {
+  public Context startSpan(REQUEST request, CONNECTION connection, Method origin) {
     String spanName = spanNameForMethod(origin);
     return startSpan(request, connection, spanName);
   }
 
-  public Span startSpan(REQUEST request, CONNECTION connection, String spanName) {
+  public Context startSpan(REQUEST request, CONNECTION connection, String spanName) {
     return startSpan(request, connection, spanName, -1);
   }
 
-  public Span startSpan(
+  public Context startSpan(
       REQUEST request, CONNECTION connection, String spanName, long startTimestamp) {
+    Context parentContext = extract(request, getGetter());
     Span.Builder builder =
-        tracer.spanBuilder(spanName).setSpanKind(SERVER).setParent(extract(request, getGetter()));
+        tracer.spanBuilder(spanName).setSpanKind(SERVER).setParent(parentContext);
 
     if (startTimestamp >= 0) {
       builder.setStartTimestamp(startTimestamp);
@@ -67,7 +68,7 @@ public abstract class HttpServerTracer<REQUEST, RESPONSE, CONNECTION, STORAGE> e
     onRequest(span, request);
     onConnectionAndRequest(span, connection, request);
 
-    return span;
+    return parentContext.with(span);
   }
 
   /**

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

@@ -25,7 +25,7 @@ public abstract class ServletHttpServerTracer<RESPONSE>
 
   private static final Logger log = LoggerFactory.getLogger(ServletHttpServerTracer.class);
 
-  public Span startSpan(HttpServletRequest request) {
+  public Context startSpan(HttpServletRequest request) {
     return startSpan(request, request, getSpanName(request));
   }
 

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

@@ -13,7 +13,9 @@ import akka.http.scaladsl.model.HttpRequest;
 import akka.http.scaladsl.model.HttpResponse;
 import akka.stream.Materializer;
 import com.google.auto.service.AutoService;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.tooling.Instrumenter;
 import io.opentelemetry.trace.Span;
 import java.util.HashMap;
@@ -96,7 +98,8 @@ public final class AkkaHttpServerInstrumentation extends Instrumenter.Default {
 
     @Override
     public HttpResponse apply(HttpRequest request) {
-      Span span = TRACER.startSpan(request, request, "akka.request");
+      Context ctx = TRACER.startSpan(request, request, "akka.request");
+      Span span = Java8BytecodeBridge.spanFromContext(ctx);
       try (Scope ignored = TRACER.startScope(span, null)) {
         HttpResponse response = userHandler.apply(request);
         TRACER.end(span, response);
@@ -121,7 +124,8 @@ public final class AkkaHttpServerInstrumentation extends Instrumenter.Default {
 
     @Override
     public Future<HttpResponse> apply(HttpRequest request) {
-      Span span = TRACER.startSpan(request, request, "akka.request");
+      Context ctx = TRACER.startSpan(request, request, "akka.request");
+      Span span = Java8BytecodeBridge.spanFromContext(ctx);
       try (Scope ignored = TRACER.startScope(span, null)) {
         return userHandler
             .apply(request)

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

@@ -12,6 +12,7 @@ import com.linecorp.armeria.server.HttpService;
 import com.linecorp.armeria.server.Route;
 import com.linecorp.armeria.server.ServiceRequestContext;
 import com.linecorp.armeria.server.SimpleDecoratingHttpService;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.trace.Span;
 import io.opentelemetry.trace.Tracer;
@@ -58,7 +59,8 @@ public class OpenTelemetryService extends SimpleDecoratingHttpService {
     long requestStartTimeMicros =
         ctx.log().ensureAvailable(RequestLogProperty.REQUEST_START_TIME).requestStartTimeMicros();
     long requestStartTimeNanos = TimeUnit.MICROSECONDS.toNanos(requestStartTimeMicros);
-    Span span = serverTracer.startSpan(req, ctx, spanName, requestStartTimeNanos);
+    Context context = serverTracer.startSpan(req, ctx, spanName, requestStartTimeNanos);
+    Span span = Span.fromContext(context);
 
     if (span.isRecording()) {
       ctx.log()

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

@@ -8,6 +8,7 @@ package io.opentelemetry.javaagent.instrumentation.grizzly;
 import static io.opentelemetry.javaagent.instrumentation.grizzly.GrizzlyHttpServerTracer.TRACER;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.trace.Span;
 import java.lang.reflect.Method;
 import net.bytebuddy.asm.Advice;
@@ -30,7 +31,8 @@ public class HttpCodecFilterAdvice {
       return;
     }
     HttpRequestPacket httpRequest = (HttpRequestPacket) httpHeader;
-    Span span = TRACER.startSpan(httpRequest, httpRequest, method);
+    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

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

@@ -8,6 +8,7 @@ package io.opentelemetry.javaagent.instrumentation.grizzly;
 import static io.opentelemetry.javaagent.instrumentation.grizzly.GrizzlyHttpServerTracer.TRACER;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.trace.Span;
 import java.lang.reflect.Method;
 import net.bytebuddy.asm.Advice;
@@ -30,7 +31,8 @@ public class HttpCodecFilterOldAdvice {
       return;
     }
     HttpRequestPacket httpRequest = (HttpRequestPacket) httpHeader;
-    Span span = TRACER.startSpan(httpRequest, httpRequest, method);
+    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

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

@@ -9,6 +9,7 @@ import static io.opentelemetry.javaagent.instrumentation.jetty.JettyHttpServerTr
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.servlet.v3_0.TagSettingAsyncListener;
 import io.opentelemetry.trace.Span;
 import java.lang.reflect.Method;
@@ -33,7 +34,8 @@ public class JettyHandlerAdvice {
       return;
     }
 
-    span = TRACER.startSpan(request, request, method);
+    Context ctx = TRACER.startSpan(request, request, method);
+    span = Java8BytecodeBridge.spanFromContext(ctx);
     scope = TRACER.startScope(span, request);
   }
 

instrumentation/netty/netty-3.8/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v3_8/server/HttpServerRequestTracingHandler.java

@@ -10,6 +10,7 @@ import static io.opentelemetry.javaagent.instrumentation.netty.v3_8.server.Netty
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.netty.v3_8.ChannelTraceContext;
 import io.opentelemetry.trace.Span;
 import org.jboss.netty.channel.Channel;
@@ -45,7 +46,8 @@ public class HttpServerRequestTracingHandler extends SimpleChannelUpstreamHandle
 
     HttpRequest request = (HttpRequest) msg.getMessage();
 
-    Span span = TRACER.startSpan(request, ctx.getChannel(), "netty.request");
+    Context context = TRACER.startSpan(request, ctx.getChannel(), "netty.request");
+    Span span = Java8BytecodeBridge.spanFromContext(context);
     try (Scope ignored = TRACER.startScope(span, channelTraceContext)) {
       ctx.sendUpstream(msg);
     } catch (Throwable throwable) {

instrumentation/netty/netty-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_0/server/HttpServerRequestTracingHandler.java

@@ -13,6 +13,7 @@ import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.handler.codec.http.HttpRequest;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.trace.Span;
 
 public class HttpServerRequestTracingHandler extends ChannelInboundHandlerAdapter {
@@ -33,7 +34,8 @@ public class HttpServerRequestTracingHandler extends ChannelInboundHandlerAdapte
       return;
     }
 
-    Span span = TRACER.startSpan((HttpRequest) msg, channel, "netty.request");
+    Context context = TRACER.startSpan((HttpRequest) msg, channel, "netty.request");
+    Span span = Java8BytecodeBridge.spanFromContext(context);
     try (Scope ignored = TRACER.startScope(span, channel)) {
       ctx.fireChannelRead(msg);
     } catch (Throwable throwable) {

instrumentation/netty/netty-4.1/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_1/server/HttpServerRequestTracingHandler.java

@@ -13,6 +13,7 @@ import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.handler.codec.http.HttpRequest;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.trace.Span;
 
 public class HttpServerRequestTracingHandler extends ChannelInboundHandlerAdapter {
@@ -33,7 +34,8 @@ public class HttpServerRequestTracingHandler extends ChannelInboundHandlerAdapte
       return;
     }
 
-    Span span = TRACER.startSpan((HttpRequest) msg, channel, "netty.request");
+    Context context = TRACER.startSpan((HttpRequest) msg, channel, "netty.request");
+    Span span = Java8BytecodeBridge.spanFromContext(context);
     try (Scope ignored = TRACER.startScope(span, channel)) {
       ctx.fireChannelRead(msg);
     } catch (Throwable throwable) {

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

@@ -7,8 +7,10 @@ package io.opentelemetry.javaagent.instrumentation.servlet.v2_2;
 
 import static io.opentelemetry.javaagent.instrumentation.servlet.v2_2.Servlet2HttpServerTracer.TRACER;
 
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.trace.Span;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
@@ -35,7 +37,8 @@ public class Servlet2Advice {
       return;
     }
 
-    span = TRACER.startSpan(httpServletRequest);
+    Context ctx = TRACER.startSpan(httpServletRequest);
+    span = Java8BytecodeBridge.spanFromContext(ctx);
     scope = TRACER.startScope(span, httpServletRequest);
   }
 

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

@@ -9,6 +9,7 @@ import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.Servlet3Ht
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.trace.Span;
 import java.util.concurrent.atomic.AtomicBoolean;
 import javax.servlet.ServletRequest;
@@ -41,7 +42,8 @@ public class Servlet3Advice {
       return;
     }
 
-    span = TRACER.startSpan(httpServletRequest);
+    Context ctx = TRACER.startSpan(httpServletRequest);
+    span = Java8BytecodeBridge.spanFromContext(ctx);
     scope = TRACER.startScope(span, httpServletRequest);
   }
 

instrumentation/spring/spring-webmvc-3.1/library/src/main/java/io/opentelemetry/instrumentation/spring/webmvc/WebMvcTracingFilter.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.instrumentation.spring.webmvc;
 
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.trace.Span;
 import io.opentelemetry.trace.Tracer;
@@ -30,8 +31,8 @@ public class WebMvcTracingFilter extends OncePerRequestFilter implements Ordered
   public void doFilterInternal(
       HttpServletRequest request, HttpServletResponse response, FilterChain filterChain)
       throws ServletException, IOException {
-    Span serverSpan = tracer.startSpan(request, request, FILTER_CLASS + "." + FILTER_METHOD);
+    Context ctx = tracer.startSpan(request, request, FILTER_CLASS + "." + FILTER_METHOD);
-
+    Span serverSpan = Span.fromContext(ctx);
     try (Scope ignored = tracer.startScope(serverSpan, request)) {
       filterChain.doFilter(request, response);
       tracer.end(serverSpan, response);