Fix span ordering inconsistencies in the tests

Also add comments for various scope closures.

dd-java-agent/instrumentation/elasticsearch/rest-5/src/main/java/datadog/trace/instrumentation/elasticsearch5/Elasticsearch5RestClientInstrumentation.java

@@ -79,6 +79,7 @@ public static void stopSpan(
         span.finish();
       }
       scope.close();
+      // span finished by RestResponseListener
     }
   }
 }

dd-java-agent/instrumentation/elasticsearch/rest-6.4/src/main/java/datadog/trace/instrumentation/elasticsearch6_4/Elasticsearch6RestClientInstrumentation.java

@@ -78,6 +78,7 @@ public static void stopSpan(
         span.finish();
       }
       scope.close();
+      // span finished by RestResponseListener
     }
   }
 }

dd-java-agent/instrumentation/elasticsearch/transport-2/src/main/java/datadog/trace/instrumentation/elasticsearch2/Elasticsearch2TransportClientInstrumentation.java

@@ -89,6 +89,7 @@ public static void stopSpan(
         span.finish();
       }
       scope.close();
+      // span finished by TransportActionListener
     }
   }
 }

dd-java-agent/instrumentation/elasticsearch/transport-5.3/src/main/java/datadog/trace/instrumentation/elasticsearch5_3/Elasticsearch53TransportClientInstrumentation.java

@@ -90,6 +90,7 @@ public static void stopSpan(
         span.finish();
       }
       scope.close();
+      // span finished by TransportActionListener
     }
   }
 }

dd-java-agent/instrumentation/elasticsearch/transport-5/src/main/java/datadog/trace/instrumentation/elasticsearch5/Elasticsearch5TransportClientInstrumentation.java

@@ -89,6 +89,7 @@ public static void stopSpan(
         span.finish();
       }
       scope.close();
+      // span finished by TransportActionListener
     }
   }
 }

dd-java-agent/instrumentation/elasticsearch/transport-6/src/main/java/datadog/trace/instrumentation/elasticsearch6/Elasticsearch6TransportClientInstrumentation.java

@@ -93,6 +93,7 @@ public static void stopSpan(
         span.finish();
       }
       scope.close();
+      // span finished by TransportActionListener
     }
   }
 }

dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java

@@ -101,6 +101,7 @@ public static void methodExit(
         span.finish();
       }
       scope.close();
+      // span finished by SpanClosingListener
     }
   }
 

dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/TracingServerInterceptor.java

@@ -55,6 +55,7 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
     } finally {
       scope.setAsyncPropagation(false);
       scope.close();
+      // span finished by TracingServerCall
     }
 
     // This ensures the server implementation can see the span in scope

dd-java-agent/instrumentation/hystrix-1.4/src/test/groovy/HystrixTest.groovy

@@ -34,7 +34,11 @@ class HystrixTest extends AgentTestRunner {
       }
     }
     def result = runUnderTrace("parent") {
-      operation(command)
+      try {
+        operation(command)
+      } finally {
+        blockUntilChildSpansFinished(2)
+      }
     }
     expect:
     TRANSFORMED_CLASSES_NAMES.contains("com.netflix.hystrix.strategy.concurrency.HystrixContextScheduler\$ThreadPoolWorker")
@@ -112,7 +116,11 @@ class HystrixTest extends AgentTestRunner {
       }
     }
     def result = runUnderTrace("parent") {
-      operation(command)
+      try {
+        return operation(command)
+      } finally {
+        blockUntilChildSpansFinished(2)
+      }
     }
     expect:
     TRANSFORMED_CLASSES_NAMES.contains("com.netflix.hystrix.strategy.concurrency.HystrixContextScheduler\$ThreadPoolWorker")

dd-java-agent/instrumentation/java-concurrent/akka-2.5-testing/src/test/groovy/AkkaExecutorInstrumentationTest.groovy

@@ -1,8 +1,8 @@
 import akka.dispatch.forkjoin.ForkJoinPool
 import akka.dispatch.forkjoin.ForkJoinTask
-import datadog.trace.core.DDSpan
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.Trace
+import datadog.trace.core.DDSpan
 import spock.lang.Shared
 
 import java.lang.reflect.InvocationTargetException
@@ -48,6 +48,7 @@ class AkkaExecutorInstrumentationTest extends AgentTestRunner {
         m(pool, new AkkaAsyncChild())
         // this child won't
         m(pool, new AkkaAsyncChild(false, false))
+        blockUntilChildSpansFinished(1)
       }
     }.run()
 

dd-java-agent/instrumentation/java-concurrent/akka-testing/src/test/scala/AkkaActors.scala

@@ -1,6 +1,7 @@
 import akka.actor.{Actor, ActorLogging, ActorRef, ActorSystem, Props}
 import akka.pattern.ask
 import akka.util.Timeout
+import datadog.trace.agent.test.AgentTestRunner.blockUntilChildSpansFinished
 import datadog.trace.api.Trace
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer.{activeScope, activeSpan}
 
@@ -33,23 +34,35 @@ class AkkaActors {
 
   @Trace
   def basicTell(): Unit = {
-    activeScope().setAsyncPropagation(true)
-    howdyGreeter ! WhoToGreet("Akka")
-    howdyGreeter ! Greet
+    try {
+      activeScope().setAsyncPropagation(true)
+      howdyGreeter ! WhoToGreet("Akka")
+      howdyGreeter ! Greet
+    } finally {
+      blockUntilChildSpansFinished(1)
+    }
   }
 
   @Trace
   def basicAsk(): Unit = {
-    activeScope().setAsyncPropagation(true)
-    howdyGreeter ! WhoToGreet("Akka")
-    howdyGreeter ? Greet
+    try {
+      activeScope().setAsyncPropagation(true)
+      howdyGreeter ! WhoToGreet("Akka")
+      howdyGreeter ? Greet
+    } finally {
+      blockUntilChildSpansFinished(1)
+    }
   }
 
   @Trace
   def basicForward(): Unit = {
-    activeScope().setAsyncPropagation(true)
-    helloGreeter ! WhoToGreet("Akka")
-    helloGreeter ? Greet
+    try {
+      activeScope().setAsyncPropagation(true)
+      helloGreeter ! WhoToGreet("Akka")
+      helloGreeter ? Greet
+    } finally {
+      blockUntilChildSpansFinished(1)
+    }
   }
 }
 

dd-java-agent/instrumentation/java-concurrent/scala-testing/src/test/groovy/ScalaExecutorInstrumentationTest.groovy

@@ -1,6 +1,6 @@
-import datadog.trace.core.DDSpan
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.Trace
+import datadog.trace.core.DDSpan
 import scala.concurrent.forkjoin.ForkJoinPool
 import scala.concurrent.forkjoin.ForkJoinTask
 import spock.lang.Shared
@@ -48,6 +48,7 @@ class ScalaExecutorInstrumentationTest extends AgentTestRunner {
         m(pool, new ScalaAsyncChild())
         // this child won't
         m(pool, new ScalaAsyncChild(false, false))
+        blockUntilChildSpansFinished(1)
       }
     }.run()
 

dd-java-agent/instrumentation/java-concurrent/scala-testing/src/test/scala/ScalaConcurrentTests.scala

@@ -1,3 +1,4 @@
+import datadog.trace.agent.test.AgentTestRunner.blockUntilChildSpansFinished
 import datadog.trace.api.Trace
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer.{activeScope, activeSpan}
 
@@ -8,8 +9,8 @@ import scala.concurrent.{Await, Future, Promise}
 class ScalaConcurrentTests {
 
   /**
-    * @return Number of expected spans in the trace
-    */
+   * @return Number of expected spans in the trace
+   */
   @Trace
   def traceWithFutureAndCallbacks(): Integer = {
     activeScope().setAsyncPropagation(true)
@@ -28,6 +29,7 @@ class ScalaConcurrentTests {
       case t: Throwable => tracedChild("failureCallback")
     }
 
+    blockUntilChildSpansFinished(4)
     return 5
   }
 
@@ -45,12 +47,13 @@ class ScalaConcurrentTests {
       }
     }
 
+    blockUntilChildSpansFinished(1)
     return 2
   }
 
   /**
-    * @return Number of expected spans in the trace
-    */
+   * @return Number of expected spans in the trace
+   */
   @Trace
   def traceWithPromises(): Integer = {
     activeScope().setAsyncPropagation(true)
@@ -78,12 +81,13 @@ class ScalaConcurrentTests {
       case t => tracedChild("brokenPromise")
     }
 
+    blockUntilChildSpansFinished(4)
     return 5
   }
 
   /**
-    * @return Number of expected spans in the trace
-    */
+   * @return Number of expected spans in the trace
+   */
   @Trace
   def tracedWithFutureFirstCompletions(): Integer = {
     activeScope().setAsyncPropagation(true)
@@ -102,12 +106,14 @@ class ScalaConcurrentTests {
           true
         }))
     Await.result(completedVal, 30 seconds)
+
+    blockUntilChildSpansFinished(3)
     return 4
   }
 
   /**
-    * @return Number of expected spans in the trace
-    */
+   * @return Number of expected spans in the trace
+   */
   @Trace
   def tracedTimeout(): Integer = {
     activeScope().setAsyncPropagation(true)
@@ -124,6 +130,8 @@ class ScalaConcurrentTests {
     } catch {
       case e: Exception => {}
     }
+
+    blockUntilChildSpansFinished(1)
     return 2
   }
 

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/AsyncPropagatingDisableInstrumentation.java

@@ -99,6 +99,7 @@ public static AgentScope enter() {
     public static void exit(@Advice.Enter final AgentScope scope) {
       if (scope != null) {
         scope.close();
+        // Don't need to finish noop span.
       }
     }
   }

dd-java-agent/instrumentation/java-concurrent/src/slickTest/groovy/CompletableFutureTest.groovy

@@ -1,6 +1,6 @@
-import datadog.trace.core.DDSpan
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.Trace
+import datadog.trace.core.DDSpan
 
 import java.util.concurrent.ArrayBlockingQueue
 import java.util.concurrent.CompletableFuture
@@ -42,10 +42,14 @@ class CompletableFutureTest extends AgentTestRunner {
       @Override
       @Trace(operationName = "parent")
       CompletableFuture<String> get() {
-        activeScope().setAsyncPropagation(true)
-        return CompletableFuture.supplyAsync(supplier, pool)
-          .thenCompose({ s -> CompletableFuture.supplyAsync(new AppendingSupplier(s), differentPool) })
-          .thenApply(function)
+        try {
+          activeScope().setAsyncPropagation(true)
+          return CompletableFuture.supplyAsync(supplier, pool)
+            .thenCompose({ s -> CompletableFuture.supplyAsync(new AppendingSupplier(s), differentPool) })
+            .thenApply(function)
+        } finally {
+          blockUntilChildSpansFinished(3)
+        }
       }
     }.get()
 

dd-java-agent/instrumentation/java-concurrent/src/slickTest/scala/SlickUtils.scala

@@ -1,3 +1,4 @@
+import datadog.trace.agent.test.AgentTestRunner.blockUntilChildSpansFinished
 import datadog.trace.api.Trace
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
 import slick.jdbc.H2Profile.api._
@@ -21,8 +22,12 @@ class SlickUtils {
 
   @Trace
   def startQuery(query: String): Future[Vector[Int]] = {
-    activeScope().setAsyncPropagation(true)
-    database.run(sql"#$query".as[Int])
+    try {
+      activeScope().setAsyncPropagation(true)
+      database.run(sql"#$query".as[Int])
+    } finally {
+      blockUntilChildSpansFinished(1)
+    }
   }
 
   def getResults(future: Future[Vector[Int]]): Int = {

dd-java-agent/instrumentation/java-concurrent/src/test/groovy/ExecutorInstrumentationTest.groovy

@@ -1,9 +1,9 @@
-import datadog.trace.core.DDSpan
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.agent.test.utils.ConfigUtils
 import datadog.trace.api.Trace
 import datadog.trace.bootstrap.instrumentation.java.concurrent.CallableWrapper
 import datadog.trace.bootstrap.instrumentation.java.concurrent.RunnableWrapper
+import datadog.trace.core.DDSpan
 import spock.lang.Shared
 
 import java.lang.reflect.InvocationTargetException
@@ -74,6 +74,7 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
         m(pool, new JavaAsyncChild())
         // this child won't
         m(pool, new JavaAsyncChild(false, false))
+        blockUntilChildSpansFinished(1)
       }
     }.run()
 

dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/AbstractRequestContextInstrumentation.java

@@ -101,8 +101,8 @@ public static void closeSpanAndScope(final AgentScope scope, final Throwable thr
       }
 
       DECORATE.beforeFinish(span);
-      span.finish();
       scope.close();
+      span.finish();
     }
   }
 }

dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java

@@ -148,6 +148,7 @@ public static void stopSpan(
         DECORATE.beforeFinish(span);
         span.finish();
       }
+      // else span finished by AsyncResponseAdvice
       scope.close();
     }
   }

dd-java-agent/instrumentation/jax-rs-client-2.0/src/test/groovy/JaxRsClientAsyncTest.groovy

@@ -15,6 +15,7 @@ import javax.ws.rs.client.InvocationCallback
 import javax.ws.rs.client.WebTarget
 import javax.ws.rs.core.MediaType
 import javax.ws.rs.core.Response
+import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 
 abstract class JaxRsClientAsyncTest extends HttpClientTest {
@@ -27,18 +28,21 @@ abstract class JaxRsClientAsyncTest extends HttpClientTest {
     headers.each { builder.header(it.key, it.value) }
     AsyncInvoker request = builder.async()
 
+    def latch = new CountDownLatch(1)
     def body = BODY_METHODS.contains(method) ? Entity.text("") : null
     Response response = request.method(method, (Entity) body, new InvocationCallback<Response>() {
       @Override
       void completed(Response s) {
         callback?.call()
+        latch.countDown()
       }
 
       @Override
       void failed(Throwable throwable) {
       }
     }).get()
 
+    latch.await()
     return response.status
   }
 

dd-java-agent/instrumentation/jetty-8/src/main/java/datadog/trace/instrumentation/jetty8/JettyHandlerAdvice.java

@@ -84,8 +84,9 @@ public static void stopSpan(
       if (!req.isAsyncStarted() && activated.compareAndSet(false, true)) {
         DECORATE.onResponse(span, resp);
         DECORATE.beforeFinish(span);
-        span.finish(); // Finish the span manually since finishSpanOnClose was false
+        span.finish();
       }
+      // else span finished by TagSettingAsyncListener
     }
     scope.close();
   }

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java

@@ -103,6 +103,7 @@ public static void stopSpan(
       PRODUCER_DECORATE.onError(scope, throwable);
       PRODUCER_DECORATE.beforeFinish(scope);
       scope.close();
+      // span finished by ProducerCallback
     }
   }