CoroutinesTimeout for JUnit5 (#2402)

gradle.properties

@@ -9,6 +9,7 @@ kotlin_version=1.5.0-RC
 
 # Dependencies
 junit_version=4.12
+junit5_version=5.7.0
 atomicfu_version=0.15.2
 knit_version=0.2.3
 html_version=0.7.2

kotlinx-coroutines-debug/api/kotlinx-coroutines-debug.api

@@ -61,3 +61,8 @@ public final class kotlinx/coroutines/debug/junit4/CoroutinesTimeout$Companion {
 	public static synthetic fun seconds$default (Lkotlinx/coroutines/debug/junit4/CoroutinesTimeout$Companion;JZZILjava/lang/Object;)Lkotlinx/coroutines/debug/junit4/CoroutinesTimeout;
 }
 
+public abstract interface annotation class kotlinx/coroutines/debug/junit5/CoroutinesTimeout : java/lang/annotation/Annotation {
+	public abstract fun cancelOnTimeout ()Z
+	public abstract fun testTimeoutMs ()J
+}
+

kotlinx-coroutines-debug/build.gradle

@@ -20,6 +20,9 @@ configurations {
 
 dependencies {
     compileOnly "junit:junit:$junit_version"
+    compileOnly "org.junit.jupiter:junit-jupiter-api:$junit5_version"
+    testCompile "org.junit.jupiter:junit-jupiter-engine:$junit5_version"
+    testCompile "org.junit.platform:junit-platform-testkit:1.7.0"
     shadowDeps "net.bytebuddy:byte-buddy:$byte_buddy_version"
     shadowDeps "net.bytebuddy:byte-buddy-agent:$byte_buddy_version"
     compileOnly "io.projectreactor.tools:blockhound:$blockhound_version"
@@ -38,6 +41,12 @@ if (rootProject.ext.jvm_ir_enabled) {
     }
 }
 
+java {
+    /* This is needed to be able to run JUnit5 tests. Otherwise, Gradle complains that it can't find the
+    JVM1.6-compatible version of the `junit-jupiter-api` artifact. */
+    disableAutoTargetJvm()
+}
+
 jar {
     manifest {
         attributes "Premain-Class": "kotlinx.coroutines.debug.AgentPremain"

kotlinx-coroutines-debug/src/junit/CoroutinesTimeoutImpl.kt

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.debug
+
+import java.util.concurrent.*
+
+/**
+ * Run [invocation] in a separate thread with the given timeout in ms, after which the coroutines info is dumped and, if
+ * [cancelOnTimeout] is set, the execution is interrupted.
+ *
+ * Assumes that [DebugProbes] are installed. Does not deinstall them.
+ */
+internal inline fun <T : Any?> runWithTimeoutDumpingCoroutines(
+    methodName: String,
+    testTimeoutMs: Long,
+    cancelOnTimeout: Boolean,
+    initCancellationException: () -> Throwable,
+    crossinline invocation: () -> T
+): T {
+    val testStartedLatch = CountDownLatch(1)
+    val testResult = FutureTask {
+        testStartedLatch.countDown()
+        invocation()
+    }
+    /*
+     * We are using hand-rolled thread instead of single thread executor
+     * in order to be able to safely interrupt thread in the end of a test
+     */
+    val testThread = Thread(testResult, "Timeout test thread").apply { isDaemon = true }
+    try {
+        testThread.start()
+        // Await until test is started to take only test execution time into account
+        testStartedLatch.await()
+        return testResult.get(testTimeoutMs, TimeUnit.MILLISECONDS)
+    } catch (e: TimeoutException) {
+        handleTimeout(testThread, methodName, testTimeoutMs, cancelOnTimeout, initCancellationException())
+    } catch (e: ExecutionException) {
+        throw e.cause ?: e
+    }
+}
+
+private fun handleTimeout(testThread: Thread, methodName: String, testTimeoutMs: Long, cancelOnTimeout: Boolean,
+                          cancellationException: Throwable): Nothing {
+    val units =
+        if (testTimeoutMs % 1000 == 0L)
+            "${testTimeoutMs / 1000} seconds"
+        else "$testTimeoutMs milliseconds"
+
+    System.err.println("\nTest $methodName timed out after $units\n")
+    System.err.flush()
+
+    DebugProbes.dumpCoroutines()
+    System.out.flush() // Synchronize serr/sout
+
+    /*
+     * Order is important:
+     * 1) Create exception with a stacktrace of hang test
+     * 2) Cancel all coroutines via debug agent API (changing system state!)
+     * 3) Throw created exception
+     */
+    cancellationException.attachStacktraceFrom(testThread)
+    testThread.interrupt()
+    cancelIfNecessary(cancelOnTimeout)
+    // If timed out test throws an exception, we can't do much except ignoring it
+    throw cancellationException
+}
+
+private fun cancelIfNecessary(cancelOnTimeout: Boolean) {
+    if (cancelOnTimeout) {
+        DebugProbes.dumpCoroutinesInfo().forEach {
+            it.job?.cancel()
+        }
+    }
+}
+
+private fun Throwable.attachStacktraceFrom(thread: Thread) {
+    val stackTrace = thread.stackTrace
+    this.stackTrace = stackTrace
+}

kotlinx-coroutines-debug/src/junit/junit4/CoroutinesTimeoutStatement.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.debug.junit4
+
+import kotlinx.coroutines.debug.*
+import org.junit.runner.*
+import org.junit.runners.model.*
+import java.util.concurrent.*
+
+internal class CoroutinesTimeoutStatement(
+    private val testStatement: Statement,
+    private val testDescription: Description,
+    private val testTimeoutMs: Long,
+    private val cancelOnTimeout: Boolean = false
+) : Statement() {
+
+    override fun evaluate() {
+        try {
+            runWithTimeoutDumpingCoroutines(testDescription.methodName, testTimeoutMs, cancelOnTimeout,
+                { TestTimedOutException(testTimeoutMs, TimeUnit.MILLISECONDS) })
+            {
+                testStatement.evaluate()
+            }
+        } finally {
+            DebugProbes.uninstall()
+        }
+    }
+}

kotlinx-coroutines-debug/src/junit/junit5/CoroutinesTimeout.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.debug.junit5
+import kotlinx.coroutines.debug.*
+import org.junit.jupiter.api.*
+import org.junit.jupiter.api.extension.*
+import org.junit.jupiter.api.parallel.*
+import java.lang.annotation.*
+
+/**
+ * Coroutines timeout annotation that is similar to JUnit5's [Timeout] annotation. It allows running test methods in a
+ * separate thread, failing them after the provided time limit and interrupting the thread.
+ *
+ * Additionally, it installs [DebugProbes] and dumps all coroutines at the moment of the timeout. It also cancels
+ * coroutines on timeout if [cancelOnTimeout] set to `true`. The dump contains the coroutine creation stack traces.
+ *
+ * This annotation has an effect on test, test factory, test template, and lifecycle methods and test classes that are
+ * annotated with it.
+ *
+ * Annotating a class is the same as annotating every test, test factory, and test template method (but not lifecycle
+ * methods) of that class and its inner test classes, unless any of them is annotated with [CoroutinesTimeout], in which
+ * case their annotation overrides the one on the containing class.
+ *
+ * Declaring [CoroutinesTimeout] on a test factory checks that it finishes in the specified time, but does not check
+ * whether the methods that it produces obey the timeout as well.
+ *
+ * Example usage:
+ * ```
+ * @CoroutinesTimeout(100)
+ * class CoroutinesTimeoutSimpleTest {
+ *     // does not time out, as the annotation on the method overrides the class-level one
+ *     @CoroutinesTimeout(1000)
+ *     @Test
+ *     fun classTimeoutIsOverridden() {
+ *         runBlocking {
+ *             delay(150)
+ *         }
+ *     }
+ *
+ *     // times out in 100 ms, timeout value is taken from the class-level annotation
+ *     @Test
+ *     fun classTimeoutIsUsed() {
+ *         runBlocking {
+ *             delay(150)
+ *         }
+ *     }
+ * }
+ * ```
+ *
+ * @see Timeout
+ */
+@ExtendWith(CoroutinesTimeoutExtension::class)
+@Inherited
+@MustBeDocumented
+@ResourceLock("coroutines timeout", mode = ResourceAccessMode.READ)
+@Retention(value = AnnotationRetention.RUNTIME)
+@Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
+public annotation class CoroutinesTimeout(
+    val testTimeoutMs: Long,
+    val cancelOnTimeout: Boolean = false
+)