Merge pull request #1177 from DataDog/tyler/scope-depth-limit

Add limit to trace scope depth

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

@@ -76,6 +76,7 @@ public class Config {
   public static final String HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN = "trace.http.client.split-by-domain";
   public static final String DB_CLIENT_HOST_SPLIT_BY_INSTANCE = "trace.db.client.split-by-instance";
   public static final String SPLIT_BY_TAGS = "trace.split-by-tags";
+  public static final String SCOPE_DEPTH_LIMIT = "trace.scope.depth.limit";
   public static final String PARTIAL_FLUSH_MIN_SPANS = "trace.partial.flush.min.spans";
   public static final String RUNTIME_CONTEXT_FIELD_INJECTION =
       "trace.runtime.context.field.injection";
@@ -128,6 +129,7 @@ public class Config {
   private static final boolean DEFAULT_HTTP_CLIENT_SPLIT_BY_DOMAIN = false;
   private static final boolean DEFAULT_DB_CLIENT_HOST_SPLIT_BY_INSTANCE = false;
   private static final String DEFAULT_SPLIT_BY_TAGS = "";
+  private static final int DEFAULT_SCOPE_DEPTH_LIMIT = 100;
   private static final int DEFAULT_PARTIAL_FLUSH_MIN_SPANS = 1000;
   private static final String DEFAULT_PROPAGATION_STYLE_EXTRACT = PropagationStyle.DATADOG.name();
   private static final String DEFAULT_PROPAGATION_STYLE_INJECT = PropagationStyle.DATADOG.name();
@@ -188,6 +190,7 @@ public class Config {
   @Getter private final boolean httpClientSplitByDomain;
   @Getter private final boolean dbClientSplitByInstance;
   @Getter private final Set<String> splitByTags;
+  @Getter private final Integer scopeDepthLimit;
   @Getter private final Integer partialFlushMinSpans;
   @Getter private final boolean runtimeContextFieldInjection;
   @Getter private final Set<PropagationStyle> propagationStylesToExtract;
@@ -290,6 +293,9 @@ public class Config {
             new LinkedHashSet<>(
                 getListSettingFromEnvironment(SPLIT_BY_TAGS, DEFAULT_SPLIT_BY_TAGS)));
 
+    scopeDepthLimit =
+        getIntegerSettingFromEnvironment(SCOPE_DEPTH_LIMIT, DEFAULT_SCOPE_DEPTH_LIMIT);
+
     partialFlushMinSpans =
         getIntegerSettingFromEnvironment(PARTIAL_FLUSH_MIN_SPANS, DEFAULT_PARTIAL_FLUSH_MIN_SPANS);
 
@@ -416,6 +422,9 @@ public class Config {
                 getPropertyListValue(
                     properties, SPLIT_BY_TAGS, new ArrayList<>(parent.splitByTags))));
 
+    scopeDepthLimit =
+        getPropertyIntegerValue(properties, SCOPE_DEPTH_LIMIT, parent.scopeDepthLimit);
+
     partialFlushMinSpans =
         getPropertyIntegerValue(properties, PARTIAL_FLUSH_MIN_SPANS, parent.partialFlushMinSpans);
 

dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java

@@ -62,7 +62,7 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
   /** Sampler defines the sampling policy in order to reduce the number of traces for instance */
   final Sampler sampler;
   /** Scope manager is in charge of managing the scopes from which spans are created */
-  final ContextualScopeManager scopeManager = new ContextualScopeManager();
+  final ScopeManager scopeManager;
 
   /** A set of tags that are added only to the application's root span */
   private final Map<String, String> localRootSpanTags;
@@ -114,6 +114,7 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
       sampler(Sampler.Builder.forConfig(config));
       injector(HttpCodec.createInjector(config));
       extractor(HttpCodec.createExtractor(config, config.getHeaderTags()));
+      scopeManager(new ContextualScopeManager(config.getScopeDepthLimit()));
       localRootSpanTags(config.getLocalRootSpanTags());
       defaultSpanTags(config.getMergedSpanTags());
       serviceNameMappings(config.getServiceMapping());
@@ -258,6 +259,7 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
         sampler,
         HttpCodec.createInjector(Config.get()),
         HttpCodec.createExtractor(Config.get(), taggedHeaders),
+        new ContextualScopeManager(Config.get().getScopeDepthLimit()),
         localRootSpanTags,
         defaultSpanTags,
         serviceNameMappings,
@@ -274,6 +276,7 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
       final Sampler sampler,
       final HttpCodec.Injector injector,
       final HttpCodec.Extractor extractor,
+      final ScopeManager scopeManager,
       final Map<String, String> localRootSpanTags,
       final Map<String, String> defaultSpanTags,
       final Map<String, String> serviceNameMappings,
@@ -294,6 +297,7 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
     this.sampler = sampler;
     this.injector = injector;
     this.extractor = extractor;
+    this.scopeManager = scopeManager;
     this.localRootSpanTags = localRootSpanTags;
     this.defaultSpanTags = defaultSpanTags;
     this.serviceNameMappings = serviceNameMappings;
@@ -365,7 +369,9 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
 
   @Deprecated
   public void addScopeContext(final ScopeContext context) {
-    scopeManager.addScopeContext(context);
+    if (scopeManager instanceof ContextualScopeManager) {
+      ((ContextualScopeManager) scopeManager).addScopeContext(context);
+    }
   }
 
   /**
@@ -386,7 +392,7 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
   }
 
   @Override
-  public ContextualScopeManager scopeManager() {
+  public ScopeManager scopeManager() {
     return scopeManager;
   }
 
@@ -510,7 +516,9 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
 
   @Override
   public void addScopeListener(final ScopeListener listener) {
-    scopeManager.addScopeListener(listener);
+    if (scopeManager instanceof ContextualScopeManager) {
+      ((ContextualScopeManager) scopeManager).addScopeListener(listener);
+    }
   }
 
   @Override

dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContextualScopeManager.java

@@ -5,18 +5,35 @@ import datadog.trace.context.ScopeListener;
 import io.opentracing.Scope;
 import io.opentracing.ScopeManager;
 import io.opentracing.Span;
+import io.opentracing.noop.NoopScopeManager;
 import java.util.Deque;
 import java.util.List;
 import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.CopyOnWriteArrayList;
+import lombok.extern.slf4j.Slf4j;
 
+@Slf4j
 public class ContextualScopeManager implements ScopeManager {
   static final ThreadLocal<DDScope> tlsScope = new ThreadLocal<>();
   final Deque<ScopeContext> scopeContexts = new ConcurrentLinkedDeque<>();
   final List<ScopeListener> scopeListeners = new CopyOnWriteArrayList<>();
 
+  private final int depthLimit;
+
+  public ContextualScopeManager(final int depthLimit) {
+    this.depthLimit = depthLimit;
+  }
+
   @Override
   public Scope activate(final Span span, final boolean finishOnClose) {
+    final Scope active = active();
+    if (active instanceof DDScope) {
+      final int currentDepth = ((DDScope) active).depth();
+      if (depthLimit <= currentDepth) {
+        log.debug("Scope depth limit exceeded ({}).  Returning NoopScope.", currentDepth);
+        return NoopScopeManager.NoopScope.INSTANCE;
+      }
+    }
     for (final ScopeContext context : scopeContexts) {
       if (context.inContext()) {
         return context.activate(span, finishOnClose);

dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContinuableScope.java

@@ -29,6 +29,8 @@ public class ContinuableScope implements DDScope, TraceScope {
   private final Continuation continuation;
   /** Flag to propagate this scope across async boundaries. */
   private final AtomicBoolean isAsyncPropagating = new AtomicBoolean(false);
+  /** depth of scope on thread */
+  private final int depth;
 
   ContinuableScope(
       final ContextualScopeManager scopeManager,
@@ -51,6 +53,7 @@ public class ContinuableScope implements DDScope, TraceScope {
     this.finishOnClose = finishOnClose;
     toRestore = scopeManager.tlsScope.get();
     scopeManager.tlsScope.set(this);
+    depth = toRestore == null ? 0 : toRestore.depth() + 1;
     for (final ScopeListener listener : scopeManager.scopeListeners) {
       listener.afterScopeActivated();
     }
@@ -90,6 +93,11 @@ public class ContinuableScope implements DDScope, TraceScope {
     return spanUnderScope;
   }
 
+  @Override
+  public int depth() {
+    return depth;
+  }
+
   @Override
   public boolean isAsyncPropagating() {
     return isAsyncPropagating.get();

dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/DDScope.java

@@ -7,4 +7,6 @@ import io.opentracing.Span;
 interface DDScope extends Scope {
   @Override
   Span span();
+
+  int depth();
 }

dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/SimpleScope.java

@@ -9,6 +9,7 @@ public class SimpleScope implements DDScope {
   private final Span spanUnderScope;
   private final boolean finishOnClose;
   private final DDScope toRestore;
+  private final int depth;
 
   public SimpleScope(
       final ContextualScopeManager scopeManager,
@@ -20,6 +21,7 @@ public class SimpleScope implements DDScope {
     this.finishOnClose = finishOnClose;
     toRestore = scopeManager.tlsScope.get();
     scopeManager.tlsScope.set(this);
+    depth = toRestore == null ? 0 : toRestore.depth() + 1;
     for (final ScopeListener listener : scopeManager.scopeListeners) {
       listener.afterScopeActivated();
     }
@@ -48,4 +50,9 @@ public class SimpleScope implements DDScope {
   public Span span() {
     return spanUnderScope;
   }
+
+  @Override
+  public int depth() {
+    return depth;
+  }
 }

dd-trace-ot/src/test/groovy/datadog/opentracing/scopemanager/ScopeManagerTest.groovy

@@ -8,8 +8,11 @@ import datadog.trace.context.ScopeListener
 import datadog.trace.util.gc.GCUtils
 import datadog.trace.util.test.DDSpecification
 import io.opentracing.Scope
+import io.opentracing.ScopeManager
 import io.opentracing.Span
+import io.opentracing.noop.NoopScopeManager
 import io.opentracing.noop.NoopSpan
+import spock.lang.Shared
 import spock.lang.Subject
 import spock.lang.Timeout
 
@@ -22,14 +25,19 @@ import java.util.concurrent.atomic.AtomicReference
 import static java.util.concurrent.TimeUnit.SECONDS
 
 class ScopeManagerTest extends DDSpecification {
-  def latch
-  def writer
-  def tracer
 
+  @Shared
+  CountDownLatch latch
+  @Shared
+  ListWriter writer
+  @Shared
+  DDTracer tracer
+
+  @Shared
   @Subject
-  def scopeManager
+  ScopeManager scopeManager
 
-  def setup() {
+  def setupSpec() {
     latch = new CountDownLatch(1)
     final currentLatch = latch
     writer = new ListWriter() {
@@ -37,12 +45,15 @@ class ScopeManagerTest extends DDSpecification {
         currentLatch.countDown()
       }
     }
-    tracer = new DDTracer(writer)
+    tracer = DDTracer.builder().writer(writer).build()
     scopeManager = tracer.scopeManager()
   }
 
   def cleanup() {
     scopeManager.tlsScope.remove()
+    scopeManager.scopeContexts.clear()
+    scopeManager.scopeListeners.clear()
+    writer.clear()
   }
 
   def "non-ddspan activation results in a simple scope"() {
@@ -129,6 +140,36 @@ class ScopeManagerTest extends DDSpecification {
     false      | true
   }
 
+  def "scopemanager returns noop scope if depth exceeded"() {
+    when: "fill up the scope stack"
+    Scope scope = null
+    for (int i = 0; i <= depth; i++) {
+      scope = scopeManager.activate(NoopSpan.INSTANCE, false)
+      assert scope instanceof SimpleScope
+    }
+
+    then: "last scope is still valid"
+    (scope as SimpleScope).depth() == depth
+
+    when: "activate a scope over the limit"
+    scope = scopeManager.activate(NoopSpan.INSTANCE, false)
+
+    then: "a noop instance is returned"
+    scope instanceof NoopScopeManager.NoopScope
+
+    when: "try again for good measure"
+    scope = scopeManager.activate(NoopSpan.INSTANCE, false)
+
+    then: "still have a noop instance"
+    scope instanceof NoopScopeManager.NoopScope
+
+    and: "scope stack not effected."
+    (scopeManager.active() as SimpleScope).depth() == depth
+
+    where:
+    depth = scopeManager.depthLimit
+  }
+
   def "ContinuableScope only creates continuations when propagation is set"() {
     setup:
     def builder = tracer.buildSpan("test")