From 2e8dc9d08f17c78ca992eeee40f67484bfc7f880 Mon Sep 17 00:00:00 2001
From: Nikolay Martynov <mar.kolya@gmail.com>
Date: Wed, 28 Nov 2018 22:14:38 -0800
Subject: [PATCH] Fix Ratpack tests that got broken by ExecutorInstrumentation
 refactoring

It turns out closing continuation also closes parent span. This is not
very good in cases when we end up not using continuation if
continuation in a state has already been setup.

This patch provides way to close continuation in a way that doesn't
affect parent scope.
---
 .../java/concurrent/ExecutorInstrumentation.java   |  3 ++-
 .../java/datadog/trace/context/TraceScope.java     | 14 +++++++++++++-
 .../opentracing/scopemanager/ContinuableScope.java | 11 ++++++++++-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java
index aeb6a4b2b2..d5bda26563 100644
--- a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java
+++ b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java
@@ -66,6 +66,7 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
       "scala.concurrent.impl.ExecutionContextImpl$$anon$3",
       "akka.dispatch.MessageDispatcher",
       "akka.dispatch.Dispatcher",
+      "akka.dispatch.Dispatcher$LazyExecutorServiceDelegate",
       "akka.actor.ActorSystemImpl$$anon$1",
       "akka.dispatch.ForkJoinExecutorConfigurator$AkkaForkJoinPool",
       "akka.dispatch.forkjoin.ForkJoinPool",
@@ -347,7 +348,7 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
       if (state.setContinuation(continuation)) {
         log.debug("created continuation {} from scope {}, state: {}", continuation, scope, state);
       } else {
-        continuation.close();
+        continuation.close(false);
       }
       return state;
     }
diff --git a/dd-trace-api/src/main/java/datadog/trace/context/TraceScope.java b/dd-trace-api/src/main/java/datadog/trace/context/TraceScope.java
index ff989a2844..866c81eeb6 100644
--- a/dd-trace-api/src/main/java/datadog/trace/context/TraceScope.java
+++ b/dd-trace-api/src/main/java/datadog/trace/context/TraceScope.java
@@ -32,7 +32,19 @@ public interface TraceScope {
      */
     TraceScope activate();
 
-    /** Cancel the continuation. */
+    /**
+     * Cancel the continuation. This also closes parent scope.
+     *
+     * <p>FIXME: the fact that this is closing parent scope is confusing, we should review this in
+     * new API.
+     */
     void close();
+
+    /**
+     * Close the continuation.
+     *
+     * @param closeContinuationScope true iff parent scope should also be closed
+     */
+    void close(boolean closeContinuationScope);
   }
 }
diff --git a/dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContinuableScope.java b/dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContinuableScope.java
index d2f1b29403..a17c97a0a9 100644
--- a/dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContinuableScope.java
+++ b/dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContinuableScope.java
@@ -128,9 +128,18 @@ public class ContinuableScope implements Scope, TraceScope {
 
     @Override
     public void close() {
+      close(true);
+    }
+
+    @Override
+    public void close(final boolean closeContinuationScope) {
       if (used.compareAndSet(false, true)) {
         trace.cancelContinuation(this);
-        ContinuableScope.this.close();
+        if (closeContinuationScope) {
+          ContinuableScope.this.close();
+        } else {
+          openCount.decrementAndGet();
+        }
       } else {
         log.debug("Failed to close continuation {}. Already used.", this);
       }