1 files changed, 81 insertions, 0 deletions
diff --git a/kotlinx-coroutines-debug/src/junit/CoroutinesTimeoutImpl.kt b/kotlinx-coroutines-debug/src/junit/CoroutinesTimeoutImpl.kt
new file mode 100644
index 00000000..06a84a5b
--- /dev/null
+++ b/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
+}