More refactoring for ScopeManager

.palantir/revapi.yml

@@ -256,6 +256,12 @@ acceptedBreaks:
     - code: "java.class.removed"
       old: "class datadog.trace.common.writer.ddagent.BatchWritingDisruptor"
       justification: "refactor trace serialization pipeline"
+    - code: "java.class.removed"
+      old: "class datadog.trace.core.scopemanager.ContextualScopeManager"
+      justification: "Refactor and split up ContextualScopeManager"
+    - code: "java.class.removed"
+      old: "class datadog.trace.core.scopemanager.ContinuableScope"
+      justification: "Refactor and split up ContextualScopeManager"
     - code: "java.class.removed"
       old: "class datadog.trace.core.scopemanager.SimpleScope"
       justification: "SimpleScope shouldn't be used and is now removed"

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/log/LogContextScopeListener.java

@@ -1,5 +1,7 @@
 package datadog.trace.agent.tooling.log;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+
 import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.context.ScopeListener;
 import java.lang.reflect.Method;
@@ -37,6 +39,10 @@ public void afterScopeActivated() {
 
   @Override
   public void afterScopeClosed() {
+    if (activeSpan() != null) {
+      afterScopeActivated();
+      return;
+    }
     try {
       removeMethod.invoke(null, CorrelationIdentifier.getTraceIdKey());
       removeMethod.invoke(null, CorrelationIdentifier.getSpanIdKey());

dd-java-agent/instrumentation/slf4j-mdc/src/main/java/datadog/trace/instrumentation/slf4j/mdc/MDCInjectionInstrumentation.java

@@ -15,6 +15,7 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.BooleanMatcher;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.utility.JavaModule;
 
@@ -26,6 +27,8 @@ public class MDCInjectionInstrumentation extends Instrumenter.Default {
   // mdcClassName = org.slf4j.MDC
   private static final String mdcClassName = "org.TMP.MDC".replaceFirst("TMP", "slf4j");
 
+  private boolean initialized = false;
+
   public MDCInjectionInstrumentation() {
     super(MDC_INSTRUMENTATION_NAME);
   }
@@ -47,15 +50,22 @@ public void postMatch(
       final JavaModule module,
       final Class<?> classBeingRedefined,
       final ProtectionDomain protectionDomain) {
-    if (classBeingRedefined != null) {
+    if (classBeingRedefined != null && !initialized) {
       MDCAdvice.mdcClassInitialized(classBeingRedefined);
     }
+    initialized = true;
   }
 
   @Override
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
     return singletonMap(
-        isTypeInitializer(), MDCInjectionInstrumentation.class.getName() + "$MDCAdvice");
+        new BooleanMatcher<MethodDescription>(false) {
+          @Override
+          public boolean matches(final MethodDescription target) {
+            return initialized;
+          }
+        }.and(isTypeInitializer()),
+        MDCInjectionInstrumentation.class.getName() + "$MDCAdvice");
   }
 
   @Override

dd-trace-core/dd-trace-core.gradle

@@ -13,7 +13,8 @@ excludedClassesCoverage += [
   // This code is copied from okHttp samples and we have integration tests to verify that it works.
   'datadog.trace.common.writer.unixdomainsockets.TunnelingUnixSocket',
   'datadog.trace.common.writer.unixdomainsockets.UnixDomainSocketFactory',
-  'datadog.trace.core.StringCachingBigInteger'
+  'datadog.trace.core.StringCachingBigInteger',
+  'datadog.trace.core.scopemanager.ScopeInterceptor.DelegatingScope',
 ]
 
 apply plugin: 'org.unbroken-dome.test-sets'

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -22,7 +22,7 @@
 import datadog.trace.core.propagation.ExtractedContext;
 import datadog.trace.core.propagation.HttpCodec;
 import datadog.trace.core.propagation.TagContext;
-import datadog.trace.core.scopemanager.ContextualScopeManager;
+import datadog.trace.core.scopemanager.ContinuableScopeManager;
 import datadog.trace.core.scopemanager.DDScopeManager;
 import java.io.Closeable;
 import java.lang.ref.WeakReference;
@@ -123,7 +123,7 @@ public CoreTracerBuilder config(final Config config) {
       injector(HttpCodec.createInjector(config));
       extractor(HttpCodec.createExtractor(config, config.getHeaderTags()));
       scopeManager(
-          new ContextualScopeManager(config.getScopeDepthLimit(), createScopeEventFactory()));
+          new ContinuableScopeManager(config.getScopeDepthLimit(), createScopeEventFactory()));
       localRootSpanTags(config.getLocalRootSpanTags());
       defaultSpanTags(config.getMergedSpanTags());
       serviceNameMappings(config.getServiceMapping());
@@ -413,8 +413,8 @@ public boolean addTraceInterceptor(final TraceInterceptor interceptor) {
 
   @Override
   public void addScopeListener(final ScopeListener listener) {
-    if (scopeManager instanceof ContextualScopeManager) {
-      ((ContextualScopeManager) scopeManager).addScopeListener(listener);
+    if (scopeManager instanceof ContinuableScopeManager) {
+      ((ContinuableScopeManager) scopeManager).addScopeListener(listener);
     }
   }