Move client span creation to decorator and automatically suppress creation of neste…

agent-bootstrap/src/main/java/io/opentelemetry/auto/bootstrap/instrumentation/decorator/ClientDecorator.java

@@ -15,10 +15,59 @@
  */
 package io.opentelemetry.auto.bootstrap.instrumentation.decorator;
 
+import io.grpc.Context;
+import io.opentelemetry.context.ContextUtils;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.trace.DefaultSpan;
 import io.opentelemetry.trace.Span;
+import io.opentelemetry.trace.Span.Kind;
+import io.opentelemetry.trace.Tracer;
+import io.opentelemetry.trace.TracingContextUtils;
 
 public abstract class ClientDecorator extends BaseDecorator {
 
+  // Keeps track of the client span in a subtree corresponding to a client request.
+  // Visible for testing
+  static final Context.Key<Span> CONTEXT_CLIENT_SPAN_KEY =
+      Context.key("opentelemetry-trace-auto-client-span-key");
+
+  /**
+   * Creates a {@link Context} with the provided {@link Span} and sets it as the current {@link
+   * Context}
+   *
+   * @return a {@link Scope} that will restore the previous context. All callers of this method must
+   *     also call {@link Scope#close()} when this next context is no longer needed.
+   */
+  public static Scope currentContextWith(final Span clientSpan) {
+    return ContextUtils.withScopedContext(withSpan(clientSpan, Context.current()));
+  }
+
+  /** Returns a new {@link Context} forked from {@code context} with the {@link Span} set. */
+  public static Context withSpan(final Span clientSpan, final Context context) {
+    if (!clientSpan.getContext().isValid()) {
+      return TracingContextUtils.withSpan(clientSpan, context);
+    }
+    return TracingContextUtils.withSpan(
+        clientSpan, context.withValue(CONTEXT_CLIENT_SPAN_KEY, clientSpan));
+  }
+
+  /**
+   * Returns a new client {@link Span} if there is no client {@link Span} in the current {@link
+   * Context}, or an invalid {@link Span} otherwise.
+   */
+  public static Span getOrCreateSpan(String name, Tracer tracer) {
+    final Context context = Context.current();
+    final Span clientSpan = CONTEXT_CLIENT_SPAN_KEY.get(context);
+
+    if (clientSpan != null) {
+      // We don't want to create two client spans for a given client call, suppress inner spans.
+      return DefaultSpan.getInvalid();
+    }
+
+    final Span current = TracingContextUtils.getSpan(context);
+    return tracer.spanBuilder(name).setSpanKind(Kind.CLIENT).setParent(current).startSpan();
+  }
+
   @Override
   public Span afterStart(final Span span) {
     assert span != null;

agent-bootstrap/src/main/java/io/opentelemetry/auto/bootstrap/instrumentation/decorator/HttpClientDecorator.java

@@ -20,6 +20,7 @@
 import io.opentelemetry.auto.instrumentation.api.Tags;
 import io.opentelemetry.trace.Span;
 import io.opentelemetry.trace.Status;
+import io.opentelemetry.trace.Tracer;
 import java.net.URI;
 import java.net.URISyntaxException;
 import lombok.extern.slf4j.Slf4j;
@@ -35,6 +36,10 @@ public abstract class HttpClientDecorator<REQUEST, RESPONSE> extends ClientDecor
 
   protected abstract Integer status(RESPONSE response);
 
+  public Span getOrCreateSpan(REQUEST request, Tracer tracer) {
+    return getOrCreateSpan(spanNameForRequest(request), tracer);
+  }
+
   public String spanNameForRequest(final REQUEST request) {
     if (request == null) {
       return DEFAULT_SPAN_NAME;

agent-bootstrap/src/test/groovy/io/opentelemetry/auto/bootstrap/instrumentation/decorator/ClientDecoratorTest.groovy

@@ -15,10 +15,16 @@
  */
 package io.opentelemetry.auto.bootstrap.instrumentation.decorator
 
+import io.grpc.Context
+import io.opentelemetry.OpenTelemetry
 import io.opentelemetry.trace.Span
+import io.opentelemetry.trace.Tracer
+import io.opentelemetry.trace.TracingContextUtils
 
 class ClientDecoratorTest extends BaseDecoratorTest {
 
+  private static final Tracer TRACER = OpenTelemetry.getTracerProvider().get("io.opentelemetry.auto")
+
   def span = Mock(Span)
 
   def "test afterStart"() {
@@ -44,6 +50,75 @@ class ClientDecoratorTest extends BaseDecoratorTest {
     0 * _
   }
 
+  def "test getOrCreateSpan when no existing client span"() {
+    when:
+    def span = ClientDecorator.getOrCreateSpan("test", TRACER)
+
+    then:
+    assert span.getContext().isValid()
+  }
+
+  def "test getOrCreateSpan when existing client span"() {
+    setup:
+    def existing = ClientDecorator.getOrCreateSpan("existing", TRACER)
+    def scope = ClientDecorator.currentContextWith(existing)
+
+    when:
+    def span = ClientDecorator.getOrCreateSpan("test", TRACER)
+
+    then:
+    assert !span.getContext().isValid()
+
+    cleanup:
+    scope.close()
+  }
+
+  def "test getOrCreateSpan internal after client span"() {
+    setup:
+    def client = ClientDecorator.getOrCreateSpan("existing", TRACER)
+    def scope = ClientDecorator.currentContextWith(client)
+
+    when:
+    def internal = TRACER.spanBuilder("internal").setSpanKind(Span.Kind.INTERNAL).startSpan()
+    def scope2 = TracingContextUtils.currentContextWith(internal)
+
+    then:
+    assert internal.getContext().isValid()
+    assert ClientDecorator.CONTEXT_CLIENT_SPAN_KEY.get(Context.current()) == client
+    assert TracingContextUtils.getSpan(Context.current()) == internal
+
+    cleanup:
+    scope2.close()
+    scope.close()
+  }
+
+  def "test withSpan"() {
+    setup:
+    def span = ClientDecorator.getOrCreateSpan("test", TRACER)
+
+    when:
+    def context = ClientDecorator.withSpan(span, Context.ROOT)
+
+    then:
+    assert ClientDecorator.CONTEXT_CLIENT_SPAN_KEY.get(context) == span
+    assert TracingContextUtils.getSpan(context) == span
+  }
+
+  def "test currentContextWith"() {
+    setup:
+    def span = ClientDecorator.getOrCreateSpan("test", TRACER)
+
+    when:
+    def scope = ClientDecorator.currentContextWith(span)
+
+    then:
+    assert ClientDecorator.CONTEXT_CLIENT_SPAN_KEY.get(Context.current()) == span
+    assert TracingContextUtils.getSpan(Context.current()) == span
+
+    cleanup:
+    scope.close()
+  }
+
   @Override
   def newDecorator() {
     return newDecorator("test-service")

instrumentation-core/aws-sdk/aws-sdk-2.2-core/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/TracingExecutionInterceptor.java

@@ -17,6 +17,7 @@
 
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkClientDecorator.DECORATE;
 
+import io.opentelemetry.auto.bootstrap.instrumentation.decorator.ClientDecorator;
 import io.opentelemetry.trace.Span;
 import io.opentelemetry.trace.Span.Kind;
 import software.amazon.awssdk.core.interceptor.Context;
@@ -40,10 +41,7 @@ final class TracingExecutionInterceptor implements ExecutionInterceptor {
   public void beforeExecution(
       final Context.BeforeExecution context, final ExecutionAttributes executionAttributes) {
     final Span span =
-        AwsSdk.tracer()
-            .spanBuilder(DECORATE.spanName(executionAttributes))
-            .setSpanKind(kind)
-            .startSpan();
+        ClientDecorator.getOrCreateSpan(DECORATE.spanName(executionAttributes), AwsSdk.tracer());
     DECORATE.afterStart(span);
     executionAttributes.putAttribute(SPAN_ATTRIBUTE, span);
   }

instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/io/opentelemetry/auto/instrumentation/apachehttpclient/v4_0/ApacheHttpClientInstrumentation.java

@@ -21,8 +21,6 @@
 import static io.opentelemetry.auto.tooling.ClassLoaderMatcher.hasClassesNamed;
 import static io.opentelemetry.auto.tooling.bytebuddy.matcher.AgentElementMatchers.implementsInterface;
 import static io.opentelemetry.context.ContextUtils.withScopedContext;
-import static io.opentelemetry.trace.Span.Kind.CLIENT;
-import static io.opentelemetry.trace.TracingContextUtils.withSpan;
 import static net.bytebuddy.matcher.ElementMatchers.isAbstract;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -34,6 +32,7 @@
 import io.grpc.Context;
 import io.opentelemetry.OpenTelemetry;
 import io.opentelemetry.auto.bootstrap.CallDepthThreadLocalMap;
+import io.opentelemetry.auto.bootstrap.instrumentation.decorator.ClientDecorator;
 import io.opentelemetry.auto.instrumentation.api.SpanWithScope;
 import io.opentelemetry.auto.tooling.Instrumenter;
 import io.opentelemetry.context.Scope;
@@ -172,16 +171,14 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
 
   public static class HelperMethods {
     public static SpanWithScope doMethodEnter(final HttpUriRequest request) {
-      final Span span =
-          TRACER.spanBuilder(DECORATE.spanNameForRequest(request)).setSpanKind(CLIENT).startSpan();
+      final Span span = DECORATE.getOrCreateSpan(request, TRACER);
 
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, request);
 
-      final Context context = withSpan(span, Context.current());
-      final boolean awsClientCall = request.getHeaders("amz-sdk-invocation-id").length > 0;
-      // AWS calls are often signed, so we can't add headers without breaking the signature.
-      if (!awsClientCall) {
+      final Context context = ClientDecorator.withSpan(span, Context.current());
+      // TODO(anuraaga): Seems like a bug that invalid context still gets injected by the injector.
+      if (span.getContext().isValid()) {
         OpenTelemetry.getPropagators().getHttpTextFormat().inject(context, request, SETTER);
       }
       final Scope scope = withScopedContext(context);

instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/auto/instrumentation/awssdk/v1_11/TracingRequestHandler.java

@@ -16,14 +16,14 @@
 package io.opentelemetry.auto.instrumentation.awssdk.v1_11;
 
 import static io.opentelemetry.auto.instrumentation.awssdk.v1_11.RequestMeta.SPAN_SCOPE_PAIR_CONTEXT_KEY;
-import static io.opentelemetry.trace.TracingContextUtils.currentContextWith;
 
 import com.amazonaws.AmazonWebServiceRequest;
 import com.amazonaws.Request;
 import com.amazonaws.Response;
 import com.amazonaws.handlers.RequestHandler2;
 import io.opentelemetry.OpenTelemetry;
 import io.opentelemetry.auto.bootstrap.ContextStore;
+import io.opentelemetry.auto.bootstrap.instrumentation.decorator.ClientDecorator;
 import io.opentelemetry.auto.instrumentation.api.SpanWithScope;
 import io.opentelemetry.trace.Span;
 import io.opentelemetry.trace.Tracer;
@@ -47,11 +47,12 @@ public AmazonWebServiceRequest beforeMarshalling(final AmazonWebServiceRequest r
 
   @Override
   public void beforeRequest(final Request<?> request) {
-    final Span span = TRACER.spanBuilder(decorate.spanNameForRequest(request)).startSpan();
+    final Span span = decorate.getOrCreateSpan(request, TRACER);
     decorate.afterStart(span);
     decorate.onRequest(span, request);
     request.addHandlerContext(
-        SPAN_SCOPE_PAIR_CONTEXT_KEY, new SpanWithScope(span, currentContextWith(span)));
+        SPAN_SCOPE_PAIR_CONTEXT_KEY,
+        new SpanWithScope(span, ClientDecorator.currentContextWith(span)));
   }
 
   @Override