More orderly shutdown of executors in tests

Author: elizarov

kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/CommonPool.kt

@@ -16,10 +16,7 @@
 
 package kotlinx.coroutines.experimental
 
-import java.util.concurrent.ExecutorService
-import java.util.concurrent.Executors
-import java.util.concurrent.RejectedExecutionException
-import java.util.concurrent.TimeUnit
+import java.util.concurrent.*
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.coroutines.experimental.CoroutineContext
 
@@ -33,7 +30,7 @@ object CommonPool : CoroutineDispatcher() {
     private var usePrivatePool = false
 
     @Volatile
-    private var _pool: ExecutorService? = null
+    private var _pool: Executor? = null
 
     private inline fun <T> Try(block: () -> T) = try { block() } catch (e: Throwable) { null }
 
@@ -59,7 +56,7 @@ object CommonPool : CoroutineDispatcher() {
     private fun defaultParallelism() = (Runtime.getRuntime().availableProcessors() - 1).coerceAtLeast(1)
 
     @Synchronized
-    private fun getOrCreatePoolSync(): ExecutorService =
+    private fun getOrCreatePoolSync(): Executor =
         _pool ?: createPool().also { _pool = it }
 
     override fun dispatch(context: CoroutineContext, block: Runnable) =
@@ -69,20 +66,29 @@ object CommonPool : CoroutineDispatcher() {
     // used for tests
     @Synchronized
     internal fun usePrivatePool() {
-        shutdownAndRelease(0)
+        shutdown(0)
         usePrivatePool = true
+        _pool = null
     }
 
     // used for tests
     @Synchronized
-    internal fun shutdownAndRelease(timeout: Long) {
-        _pool?.apply {
+    internal fun shutdown(timeout: Long) {
+        (_pool as? ExecutorService)?.apply {
             shutdown()
             if (timeout > 0)
                 awaitTermination(timeout, TimeUnit.MILLISECONDS)
-            _pool = null
+            shutdownNow().forEach { defaultExecutor.execute(it) }
         }
+        _pool = Executor { throw RejectedExecutionException("CommonPool was shutdown") }
+    }
+
+    // used for tests
+    @Synchronized
+    internal fun restore() {
+        shutdown(0)
         usePrivatePool = false
+        _pool = null
     }
 
     override fun toString(): String = "CommonPool"

kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/DefaultExecutor.kt

@@ -50,14 +50,11 @@ private fun getOrCreateExecutorSync(): ScheduledExecutorService =
 
 // used for tests
 @Synchronized
-internal fun defaultExecutorShutdownNow() {
-    _executor?.shutdownNow()
-}
-
-@Synchronized
-internal fun defaultExecutorShutdownNowAndRelease() {
+internal fun shutdownDefaultExecutor(timeout: Long) {
     _executor?.apply {
-        shutdownNow()
+        shutdown()
+        awaitTermination(timeout, TimeUnit.MILLISECONDS)
+        shutdownNow() // ignore all remaining
         _executor = null
     }
 }

kotlinx-coroutines-core/src/test/kotlin/guide/test/TestUtil.kt

@@ -22,6 +22,7 @@ import org.junit.Assert.assertTrue
 import java.io.ByteArrayInputStream
 import java.io.ByteArrayOutputStream
 import java.io.PrintStream
+import java.util.concurrent.TimeUnit
 
 fun test(block: () -> Unit): List<String> {
     val oldOut = System.out
@@ -42,24 +43,29 @@ fun test(block: () -> Unit): List<String> {
         // capture output
         bytes = bytesOut.toByteArray()
         // the shutdown
-        defaultExecutorShutdownNow()
-        shutdownDispatcherPools()
-        CommonPool.shutdownAndRelease(10000L) // wait at most 10 sec
-        defaultExecutorShutdownNowAndRelease()
+        val timeout = 5000L // 5 sec at most to wait
+        CommonPool.shutdown(timeout)
+        shutdownDispatcherPools(timeout)
+        shutdownDefaultExecutor(timeout) // the last man standing -- kill it too
+        CommonPool.restore()
         System.setOut(oldOut)
         System.setErr(oldErr)
 
     }
     return ByteArrayInputStream(bytes).bufferedReader().readLines()
 }
 
-private fun shutdownDispatcherPools() {
+private fun shutdownDispatcherPools(timeout: Long) {
     val threads = arrayOfNulls<Thread>(Thread.activeCount())
     val n = Thread.enumerate(threads)
     for (i in 0 until n) {
         val thread = threads[i]
         if (thread is PoolThread)
-            thread.dispatcher.executor.shutdownNow()
+            thread.dispatcher.executor.apply {
+                shutdown()
+                awaitTermination(timeout, TimeUnit.MILLISECONDS)
+                shutdownNow().forEach { defaultExecutor.execute(it) }
+            }
     }
 }