From fd43210e0787175c44db7cf3cec7962cc2a734a7 Mon Sep 17 00:00:00 2001
From: Tyler Benson <tyler.benson@datadoghq.com>
Date: Thu, 23 May 2019 16:27:53 -0700
Subject: [PATCH] Add task timeout for tests.

---
 dd-java-agent/dd-java-agent.gradle            |  3 ++-
 .../agent/test/IntegrationTestUtils.java      | 26 +++++++++++++++++--
 gradle/java.gradle                            |  7 ++++-
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/dd-java-agent/dd-java-agent.gradle b/dd-java-agent/dd-java-agent.gradle
index 43c3c0482a..4a8567f3fa 100644
--- a/dd-java-agent/dd-java-agent.gradle
+++ b/dd-java-agent/dd-java-agent.gradle
@@ -107,7 +107,8 @@ dependencies {
 }
 
 tasks.withType(Test).configureEach {
-  jvmArgs "-Ddd.writer.type=LogWriter", "-Ddd.service.name=java-app"
+  jvmArgs "-Ddd.service.name=java-agent-tests"
+  jvmArgs "-Ddd.writer.type=LoggingWriter"
   jvmArgs "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=debug"
   jvmArgs "-Dorg.slf4j.simpleLogger.defaultLogLevel=debug"
 
diff --git a/dd-java-agent/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java b/dd-java-agent/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java
index c106179ead..1a85694388 100644
--- a/dd-java-agent/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java
+++ b/dd-java-agent/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java
@@ -17,6 +17,8 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.jar.Attributes;
 import java.util.jar.JarEntry;
 import java.util.jar.JarOutputStream;
@@ -201,7 +203,8 @@ public class IntegrationTestUtils {
     final ProcessBuilder processBuilder = new ProcessBuilder(commands.toArray(new String[0]));
     processBuilder.environment().putAll(envVars);
     final Process process = processBuilder.start();
-    final int result = process.waitFor();
+
+    waitFor(process, 30, TimeUnit.SECONDS);
 
     if (printOutputStreams) {
       final BufferedReader stdInput =
@@ -221,6 +224,25 @@ public class IntegrationTestUtils {
       }
       System.out.println("--- stderr end ---");
     }
-    return result;
+    return process.exitValue();
+  }
+
+  private static void waitFor(final Process process, final long timeout, final TimeUnit unit)
+      throws InterruptedException, TimeoutException {
+    final long startTime = System.nanoTime();
+    long rem = unit.toNanos(timeout);
+
+    do {
+      try {
+        process.exitValue();
+        return;
+      } catch (final IllegalThreadStateException ex) {
+        if (rem > 0) {
+          Thread.sleep(Math.min(TimeUnit.NANOSECONDS.toMillis(rem) + 1, 100));
+        }
+      }
+      rem = unit.toNanos(timeout) - (System.nanoTime() - startTime);
+    } while (rem > 0);
+    throw new TimeoutException();
   }
 }
diff --git a/gradle/java.gradle b/gradle/java.gradle
index a0c1af6cf9..9f6b50da8f 100644
--- a/gradle/java.gradle
+++ b/gradle/java.gradle
@@ -1,3 +1,5 @@
+import java.time.Duration
+
 apply plugin: 'java'
 apply plugin: 'groovy'
 
@@ -301,8 +303,11 @@ for (def env : System.getenv().entrySet()) {
   }
 }
 
-// Disable all tests if skipTests property was specified
 tasks.withType(Test).configureEach {
+  // All tests must complete within 2 minutes.
+  timeout = Duration.ofMinutes(2)
+  
+  // Disable all tests if skipTests property was specified
   onlyIf { !project.rootProject.hasProperty("skipTests") }
 }