CoroutinesTimeout test rule in debug module

Fixes #938

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-debug.txt

@@ -38,3 +38,16 @@ public final class kotlinx/coroutines/debug/State : java/lang/Enum {
 	public static fun values ()[Lkotlinx/coroutines/debug/State;
 }
 
+public final class kotlinx/coroutines/debug/junit4/CoroutinesTimeout : org/junit/rules/TestRule {
+	public static final field Companion Lkotlinx/coroutines/debug/junit4/CoroutinesTimeout$Companion;
+	public fun <init> (JZ)V
+	public synthetic fun <init> (JZILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun apply (Lorg/junit/runners/model/Statement;Lorg/junit/runner/Description;)Lorg/junit/runners/model/Statement;
+	public static final fun seconds (IZ)Lkotlinx/coroutines/debug/junit4/CoroutinesTimeout;
+}
+
+public final class kotlinx/coroutines/debug/junit4/CoroutinesTimeout$Companion {
+	public final fun seconds (IZ)Lkotlinx/coroutines/debug/junit4/CoroutinesTimeout;
+	public static synthetic fun seconds$default (Lkotlinx/coroutines/debug/junit4/CoroutinesTimeout$Companion;IZILjava/lang/Object;)Lkotlinx/coroutines/debug/junit4/CoroutinesTimeout;
+}
+

kotlinx-coroutines-debug/build.gradle

@@ -5,6 +5,7 @@
 apply plugin: "com.github.johnrengelman.shadow"
 
 dependencies {
+    compileOnly "junit:junit:$junit_version"
     compile "net.bytebuddy:byte-buddy:$byte_buddy_version"
     compile "net.bytebuddy:byte-buddy-agent:$byte_buddy_version"
 }

kotlinx-coroutines-debug/src/junit4/CoroutinesTimeout.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2016-2019 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.rules.*
+import org.junit.runner.*
+import org.junit.runners.model.*
+import java.util.concurrent.*
+
+/**
+ * Coroutines timeout rule for JUnit4 that is applied to all methods in the class.
+ * This rule is very similar to [Timeout] rule: it runs tests in a separate thread,
+ * fails tests after the given timeout and interrupts test thread.
+ *
+ * Additionally, this rule installs [DebugProbes] and dumps all coroutines at the moment of the timeout.
+ * It may cancel coroutines on timeout if [cancelOnTimeout] set to `true`.
+ *
+ * Example of usage:
+ * ```
+ * class HangingTest {
+ *
+ *      @Rule
+ *   @JvmField
+ *   val timeout = CoroutinesTimeout.seconds(5)
+ *
+ *   @Test
+ *   fun testThatHangs() = runBlocking {
+ *     ...
+ *     delay(Long.MAX_VALUE) // somewhere deep in the stack
+ *     ...
+ *   }
+ * }
+ * ```
+ */
+public class CoroutinesTimeout(
+    private val testTimeoutMs: Long,
+    private val cancelOnTimeout: Boolean = false
+) : TestRule {
+
+    init {
+        require(testTimeoutMs > 0) { "Expected positive test timeout, but had $testTimeoutMs" }
+    }
+
+    companion object {
+        /**
+         * Creates [CoroutinesTimeout] rule with the given timeout in seconds.
+         */
+        @JvmStatic
+        public fun seconds(seconds: Int, cancelOnTimeout: Boolean = false): CoroutinesTimeout {
+            return CoroutinesTimeout(TimeUnit.SECONDS.toMillis(seconds.toLong()), cancelOnTimeout)
+        }
+    }
+
+    /**
+     * @suppress suppress from Dokka
+     */
+    override fun apply(base: Statement, description: Description): Statement {
+        return CoroutinesTimeoutStatement(base, description, testTimeoutMs, cancelOnTimeout)
+    }
+}

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

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2016-2019 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.*
+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() {
+
+    private val testExecutor = Executors.newSingleThreadExecutor {
+        Thread(it).apply {
+            name = "Timeout test executor"
+            isDaemon = true
+        }
+    }
+
+    // Thread to dump stack from, captured by testExecutor
+    private lateinit var testThread: Thread
+
+    override fun evaluate() {
+        DebugProbes.install() // Fail-fast if probes are unavailable
+        val latch = CountDownLatch(1)
+        val testFuture = CompletableFuture.runAsync(Runnable {
+            testThread = Thread.currentThread()
+            latch.countDown()
+            testStatement.evaluate()
+        }, testExecutor)
+
+        latch.await() // Await until test is started
+        try {
+            testFuture.get(testTimeoutMs, TimeUnit.MILLISECONDS)
+            return
+        } catch (e: TimeoutException) {
+            handleTimeout(testDescription)
+        } catch (e: ExecutionException) {
+            throw e.cause ?: e
+        } finally {
+            DebugProbes.uninstall()
+            testExecutor.shutdown()
+        }
+    }
+
+    private fun handleTimeout(description: Description) {
+        val units =
+            if (testTimeoutMs % 1000L == 0L)
+                "${testTimeoutMs / 1000} seconds"
+            else "$testTimeoutMs milliseconds"
+
+        val message = "Test ${description.methodName} timed out after $units"
+        System.err.println("\n$message\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
+         */
+        val exception = createTimeoutException(message, testThread)
+        cancelIfNecessary()
+        // If timed out test throws an exception, we can't do much except ignoring it
+        throw exception
+    }
+
+    private fun cancelIfNecessary() {
+        if (cancelOnTimeout) {
+            DebugProbes.dumpCoroutinesState().forEach {
+                it.jobOrNull?.cancel()
+            }
+        }
+    }
+
+    private fun createTimeoutException(message: String, thread: Thread): Exception {
+        val stackTrace = thread.stackTrace
+        val exception = TestTimedOutException(testTimeoutMs, TimeUnit.MILLISECONDS)
+        exception.stackTrace = stackTrace
+        thread.interrupt()
+        return exception
+    }
+}

kotlinx-coroutines-debug/test/junit4/CoroutinesTimeoutTest.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.debug.junit4
+
+import junit4.*
+import kotlinx.coroutines.*
+import org.junit.*
+import org.junit.runners.model.*
+
+class CoroutinesTimeoutTest : TestBase() {
+
+    @Rule
+    @JvmField
+    public val validation = TestFailureValidation(
+        1000, false,
+        TestResultSpec("throwingTest", error = RuntimeException::class.java),
+        TestResultSpec("successfulTest"),
+        TestResultSpec(
+            "hangingTest", expectedOutParts = listOf(
+                "Coroutines dump",
+                "Test hangingTest timed out after 1 seconds",
+                "BlockingCoroutine{Active}",
+                "runBlocking",
+                "at kotlinx.coroutines.debug.junit4.CoroutinesTimeoutTest.suspendForever",
+                "at kotlinx.coroutines.debug.junit4.CoroutinesTimeoutTest\$hangingTest\$1.invokeSuspend"),
+            notExpectedOutParts = listOf("delay", "throwingTest"),
+            error = TestTimedOutException::class.java)
+    )
+
+    @Test
+    fun hangingTest() = runBlocking<Unit> {
+        suspendForever()
+        expectUnreached()
+    }
+
+    private suspend fun suspendForever() {
+        delay(Long.MAX_VALUE)
+        expectUnreached()
+    }
+
+    @Test
+    fun throwingTest() = runBlocking<Unit> {
+        throw RuntimeException()
+    }
+
+    @Test
+    fun successfulTest() = runBlocking {
+        val job = launch {
+            yield()
+        }
+
+        job.join()
+    }
+}

kotlinx-coroutines-debug/test/junit4/TestFailureValidation.kt

@@ -0,0 +1,104 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package junit4
+
+import kotlinx.coroutines.debug.*
+import kotlinx.coroutines.debug.junit4.*
+import org.junit.rules.*
+import org.junit.runner.*
+import org.junit.runners.model.*
+import java.io.*
+import kotlin.test.*
+
+internal fun TestFailureValidation(timeoutMs: Long, cancelOnTimeout: Boolean, vararg specs: TestResultSpec): RuleChain =
+    RuleChain
+        .outerRule(TestFailureValidation(specs.associateBy { it.testName }))
+        .around(
+            CoroutinesTimeout(
+                timeoutMs,
+                cancelOnTimeout
+            )
+        )
+
+/**
+ * Rule that captures test result, serr and sout and validates it against provided [testsSpec]
+ */
+internal class TestFailureValidation(private val testsSpec: Map<String, TestResultSpec>) : TestRule {
+
+    companion object {
+        init {
+            DebugProbes.sanitizeStackTraces = false
+        }
+    }
+    override fun apply(base: Statement, description: Description): Statement {
+        return TestFailureStatement(base, description)
+    }
+
+    inner class TestFailureStatement(private val test: Statement, private val description: Description) : Statement() {
+        private lateinit var sout: PrintStream
+        private lateinit var serr: PrintStream
+        private val capturedOut = ByteArrayOutputStream()
+
+        override fun evaluate() {
+            try {
+                replaceOut()
+                test.evaluate()
+            } catch (e: Throwable) {
+                validateFailure(e)
+                return
+            } finally {
+                resetOut()
+            }
+
+            validateSuccess() // To avoid falling into catch
+        }
+
+        private fun validateSuccess() {
+            val spec = testsSpec[description.methodName] ?: error("Test spec not found: ${description.methodName}")
+            require(spec.error == null) { "Expected exception of type ${spec.error}, but test successfully passed" }
+
+            val captured = capturedOut.toString()
+            assertFalse(captured.contains("Coroutines dump"))
+            assertTrue(captured.isEmpty(), captured)
+        }
+
+        private fun validateFailure(e: Throwable) {
+            val spec = testsSpec[description.methodName] ?: error("Test spec not found: ${description.methodName}")
+            if (spec.error == null || !spec.error.isInstance(e)) {
+                throw IllegalStateException("Unexpected failure, expected ${spec.error}, had ${e::class}", e)
+            }
+
+            if (e !is TestTimedOutException) return
+
+            val captured = capturedOut.toString()
+            assertTrue(captured.contains("Coroutines dump"))
+            for (part in spec.expectedOutParts) {
+                assertTrue(captured.contains(part), "Expected $part to be part of the\n$captured")
+            }
+
+            for (part in spec.notExpectedOutParts) {
+                assertFalse(captured.contains(part), "Expected $part not to be part of the\n$captured")
+            }
+        }
+
+        private fun replaceOut() {
+            sout = System.out
+            serr = System.err
+
+            System.setOut(PrintStream(capturedOut))
+            System.setErr(PrintStream(capturedOut))
+        }
+
+        private fun resetOut() {
+            System.setOut(sout)
+            System.setErr(serr)
+        }
+    }
+}
+
+data class TestResultSpec(
+    val testName: String, val expectedOutParts: List<String> = listOf(), val notExpectedOutParts:
+    List<String> = listOf(), val error: Class<out Throwable>? = null
+)