R&D week: blackhole span

dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java

@@ -7,7 +7,6 @@
 import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient.HttpHeadersInjectAdapter.SETTER;
 
-import datadog.trace.api.Config;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -19,46 +18,30 @@
 import org.apache.http.client.methods.HttpUriRequest;
 
 public class HelperMethods {
-  private static final boolean AWS_LEGACY_TRACING =
-      Config.get().isLegacyTracingEnabled(false, "aws-sdk");
-
   public static AgentScope doMethodEnter(final HttpUriRequest request) {
-    boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
-    if (!AWS_LEGACY_TRACING && awsClientCall) {
-      // avoid creating an extra HTTP client span beneath the AWS client call
-      return null;
-    }
-
     final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
     if (callDepth > 0) {
       return null;
     }
-
-    return activateHttpSpan(request, awsClientCall);
+    return activateHttpSpan(request);
   }
 
   public static AgentScope doMethodEnter(HttpHost host, HttpRequest request) {
-    boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
-    if (!AWS_LEGACY_TRACING && awsClientCall) {
-      // avoid creating an extra HTTP client span beneath the AWS client call
-      return null;
-    }
-
     final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
     if (callDepth > 0) {
       return null;
     }
 
-    return activateHttpSpan(new HostAndRequestAsHttpUriRequest(host, request), awsClientCall);
+    return activateHttpSpan(new HostAndRequestAsHttpUriRequest(host, request));
   }
 
-  private static AgentScope activateHttpSpan(
-      final HttpUriRequest request, final boolean awsClientCall) {
+  private static AgentScope activateHttpSpan(final HttpUriRequest request) {
     final AgentSpan span = startSpan(HTTP_REQUEST);
     final AgentScope scope = activateSpan(span);
 
     DECORATE.afterStart(span);
     DECORATE.onRequest(span, request);
+    final boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
 
     // AWS calls are often signed, so we can't add headers without breaking the signature.
     if (!awsClientCall) {

dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java

@@ -7,7 +7,6 @@
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
 
-import datadog.trace.api.Config;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -18,47 +17,32 @@
 import org.apache.hc.core5.http.HttpResponse;
 
 public class HelperMethods {
-  private static final boolean AWS_LEGACY_TRACING =
-      Config.get().isLegacyTracingEnabled(false, "aws-sdk");
-
   public static AgentScope doMethodEnter(final HttpRequest request) {
-    boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
-    if (!AWS_LEGACY_TRACING && awsClientCall) {
-      // avoid creating an extra HTTP client span beneath the AWS client call
-      return null;
-    }
-
     final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
     if (callDepth > 0) {
       return null;
     }
 
-    return activateHttpSpan(request, awsClientCall);
+    return activateHttpSpan(request);
   }
 
   public static AgentScope doMethodEnter(HttpHost host, HttpRequest request) {
-    boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
-    if (!AWS_LEGACY_TRACING && awsClientCall) {
-      // avoid creating an extra HTTP client span beneath the AWS client call
-      return null;
-    }
-
     final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
     if (callDepth > 0) {
       return null;
     }
 
-    return activateHttpSpan(new HostAndRequestAsHttpUriRequest(host, request), awsClientCall);
+    return activateHttpSpan(new HostAndRequestAsHttpUriRequest(host, request));
   }
 
-  private static AgentScope activateHttpSpan(
-      final HttpRequest request, final boolean awsClientCall) {
+  private static AgentScope activateHttpSpan(final HttpRequest request) {
     final AgentSpan span = startSpan(HTTP_REQUEST);
     final AgentScope scope = activateSpan(span);
 
     DECORATE.afterStart(span);
     DECORATE.onRequest(span, request);
 
+    final boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
     // AWS calls are often signed, so we can't add headers without breaking the signature.
     if (!awsClientCall) {
       propagate().inject(span, request, SETTER);

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java

@@ -1,7 +1,7 @@
 package datadog.trace.instrumentation.aws.v0;
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_IN;
@@ -53,7 +53,7 @@ public void beforeRequest(final Request<?> request) {
     if (!AWS_LEGACY_TRACING && isPollingRequest(request.getOriginalRequest())) {
       // SQS messages spans are created by aws-java-sqs-1.0 - replace client scope with no-op,
       // so we can tell when receive call is complete without affecting the rest of the trace
-      span = noopSpan();
+      activateSpan(blackholeSpan());
     } else {
       span = requestSpanStore.remove(request.getOriginalRequest());
       if (span != null) {
@@ -74,10 +74,13 @@ public void beforeRequest(final Request<?> request) {
           log.warn("Unable to inject trace header", e);
         }
       }
+      // This scope will be closed by AwsHttpClientInstrumentation
+      if (AWS_LEGACY_TRACING) {
+        activateSpan(span);
+      } else {
+        activateSpan(blackholeSpan());
+      }
     }
-
-    // This scope will be closed by AwsHttpClientInstrumentation
-    activateSpan(span);
   }
 
   @Override

dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java

@@ -1,7 +1,7 @@
 package datadog.trace.instrumentation.aws.v2;
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.AWS_LEGACY_TRACING;
@@ -85,10 +85,8 @@ public SdkHttpRequest modifyHttpRequest(
   public void beforeTransmission(
       final Context.BeforeTransmission context, final ExecutionAttributes executionAttributes) {
     final AgentSpan span;
-    if (!AWS_LEGACY_TRACING && isPollingRequest(context.request())) {
-      // SQS messages spans are created by aws-java-sqs-2.0 - replace client scope with no-op,
-      // so we can tell when receive call is complete without affecting the rest of the trace
-      span = noopSpan();
+    if (!AWS_LEGACY_TRACING) {
+      span = blackholeSpan();
     } else {
       span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
     }

dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/DoNotTraceAdvice.java

@@ -0,0 +1,26 @@
+package datadog.trace.instrumentation.trace_annotation;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
+import static datadog.trace.instrumentation.trace_annotation.TraceDecorator.DECORATE;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+
+public class DoNotTraceAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope before() {
+    return activateSpan(blackholeSpan());
+  }
+
+  @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+  public static void after(
+      @Advice.Enter final AgentScope scope,
+      @Advice.Return(typing = Assigner.Typing.DYNAMIC, readOnly = false) Object result) {
+    if (scope != null) {
+      scope.close();
+      result = DECORATE.wrapAsyncResultOrFinishSpan(result, scope.span());
+    }
+  }
+}

dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/DoNotTraceAnnotationInstrumentation.java

@@ -0,0 +1,44 @@
+package datadog.trace.instrumentation.trace_annotation;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresMethod;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.isAnnotatedWith;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class DoNotTraceAnnotationInstrumentation extends Instrumenter.Tracing
+    implements Instrumenter.ForTypeHierarchy {
+
+  @SuppressForbidden
+  public DoNotTraceAnnotationInstrumentation() {
+    super("not-not-trace", "do-not-trace-annotation");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "datadog.trace.api.DoNotTrace";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return declaresMethod(isAnnotatedWith(named(hierarchyMarkerType())));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".TraceDecorator",
+    };
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isAnnotatedWith(named(hierarchyMarkerType())), packageName + ".DoNotTraceAdvice");
+  }
+}

dd-java-agent/instrumentation/trace-annotation/src/test/groovy/TraceAnnotationsTest.groovy

@@ -2,6 +2,7 @@ import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.agent.test.utils.TraceUtils
 import datadog.trace.api.Trace
 import datadog.trace.bootstrap.instrumentation.api.Tags
+import dd.test.trace.annotation.DontTraceClass
 import dd.test.trace.annotation.SayTracedHello
 import dd.test.trace.annotation.TracedSubClass
 
@@ -510,5 +511,25 @@ class TraceAnnotationsTest extends AgentTestRunner {
       }
     }
   }
+  def "@DoNotTrace should mute tracing"() {
+    setup:
+    TraceUtils.runUnderTrace("parent", () -> {
+      new DontTraceClass().muted()
+    })
+    expect:
+    assertTraces(1) {
+      trace(1) {
+        span {
+          hasServiceName()
+          resourceName "parent"
+          operationName "parent"
+          parent()
+          errored false
+          tags {
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
 }
-

dd-java-agent/instrumentation/trace-annotation/src/test/java/dd/test/trace/annotation/DontTraceClass.java

@@ -0,0 +1,14 @@
+package dd.test.trace.annotation;
+
+import datadog.trace.api.DoNotTrace;
+import datadog.trace.api.Trace;
+
+public class DontTraceClass {
+  @DoNotTrace
+  public void muted() {
+    normallyTraced();
+  }
+
+  @Trace
+  public void normallyTraced() {}
+}

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.groovy

@@ -327,14 +327,14 @@ abstract class AgentTestRunner extends DDSpecification implements AgentBuilder.L
 
     boolean enabledFinishTimingChecks = this.enabledFinishTimingChecks()
     TEST_TRACER.startSpan(*_) >> {
-      DDSpan agentSpan = callRealMethod()
+      AgentSpan agentSpan = callRealMethod()
       TEST_SPANS.add(agentSpan.spanId)
       if (!enabledFinishTimingChecks) {
         return agentSpan
       }
 
       // rest of closure if for checking duplicate finishes and tags set after finish
-      DDSpan spiedAgentSpan = Spy(agentSpan)
+      AgentSpan spiedAgentSpan = Spy(agentSpan)
       originalToSpySpan[agentSpan] = spiedAgentSpan
       def handleFinish = { MockInvocation mi ->
         def depth = CallDepthThreadLocalMap.incrementCallDepth(DDSpan)

dd-trace-api/build.gradle

@@ -35,6 +35,8 @@ excludedClassesCoverage += [
   'datadog.trace.api.experimental.DataStreamsContextCarrier',
   'datadog.trace.api.experimental.DataStreamsContextCarrier.NoOp',
   'datadog.appsec.api.blocking.*',
+  'datadog.trace.context.NoopTraceScope.NoopContinuation',
+  'datadog.trace.context.NoopTraceScope',
 ]
 
 description = 'dd-trace-api'

dd-trace-api/src/main/java/datadog/trace/api/DoNotTrace.java

@@ -0,0 +1,12 @@
+package datadog.trace.api;
+
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+/** Set this annotation to a method to mute tracing until its scope is closed */
+@Retention(RUNTIME)
+@Target(METHOD)
+public @interface DoNotTrace {}

dd-trace-api/src/main/java/datadog/trace/api/GlobalTracer.java

@@ -2,6 +2,8 @@
 
 import datadog.trace.api.interceptor.TraceInterceptor;
 import datadog.trace.api.internal.InternalTracer;
+import datadog.trace.context.NoopTraceScope;
+import datadog.trace.context.TraceScope;
 import java.util.ArrayList;
 import java.util.Collection;
 
@@ -28,6 +30,11 @@ public String getSpanId() {
         public boolean addTraceInterceptor(TraceInterceptor traceInterceptor) {
           return false;
         }
+
+        @Override
+        public TraceScope muteTracing() {
+          return NoopTraceScope.INSTANCE;
+        }
       };
 
   private static final Collection<Callback> installationCallbacks = new ArrayList<>();

dd-trace-api/src/main/java/datadog/trace/api/Tracer.java

@@ -1,6 +1,7 @@
 package datadog.trace.api;
 
 import datadog.trace.api.interceptor.TraceInterceptor;
+import datadog.trace.context.TraceScope;
 
 /** A class with Datadog tracer features. */
 public interface Tracer {
@@ -21,4 +22,6 @@ public interface Tracer {
    * @return false if an interceptor with same priority exists.
    */
   boolean addTraceInterceptor(TraceInterceptor traceInterceptor);
+
+  TraceScope muteTracing();
 }