awaitAll improvements: create defensive copy of jobs to consistently handle concurrent mutations, make handler JobNode to avoid extra object allocation

Author: qwwdfsad

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

@@ -20,7 +20,8 @@ public suspend fun awaitAll(vararg jobs: Job): Unit = jobs.asList().awaitAll()
  */
 public suspend fun Collection<Job>.awaitAll() {
     if (isEmpty()) return
-    AwaitAll(this).await()
+    val snapshot = ArrayList(this)
+    AwaitAll(snapshot).await()
 }
 
 /**
@@ -43,21 +44,23 @@ private class AwaitAll(private val jobs: Collection<Job>) {
 
     suspend fun await() {
         suspendCancellableCoroutine<Unit> { cont ->
-            // todo: create a separate named class instance of JobNode to avoid extra object
-            val handler: (Throwable?) -> Unit = {
-                if (it != null) {
-                    val token = cont.tryResumeWithException(it)
-                    if (token != null) {
-                        cont.completeResume(token)
-                    }
-                } else if (notCompletedCount.decrementAndGet() == 0) {
-                    cont.resume(Unit)
-                }
-            }
-
             jobs.forEach {
                 it.start() // To properly await lazily started jobs
-                cont.disposeOnCompletion(it.invokeOnCompletion(handler))
+                cont.disposeOnCompletion(it.invokeOnCompletion(AwaitAllNode(cont, it)))
+            }
+        }
+    }
+
+    inner class AwaitAllNode(private val continuation: CancellableContinuation<Unit>, job: Job) : JobNode<Job>(job), CompletionHandler {
+
+        override fun invoke(cause: Throwable?) {
+            if (cause != null) {
+                val token = continuation.tryResumeWithException(cause)
+                if (token != null) {
+                    continuation.completeResume(token)
+                }
+            } else if (notCompletedCount.decrementAndGet() == 0) {
+                continuation.resume(Unit)
             }
         }
     }

core/kotlinx-coroutines-core/src/test/kotlin/kotlinx/coroutines/experimental/AwaitStressTest.kt

@@ -1,8 +1,8 @@
 package kotlinx.coroutines.experimental
 
-import org.junit.After
-import java.util.concurrent.CyclicBarrier
-import kotlin.test.Test
+import org.junit.*
+import org.junit.Test
+import java.util.concurrent.*
 
 class AwaitStressTest : TestBase() {
 
@@ -100,4 +100,34 @@ class AwaitStressTest : TestBase() {
 
         require(cancelledOnce) { "Cancellation exception wasn't properly caught" }
     }
+
+    @Test
+    fun testMutatingCollection() = runTest {
+        val barrier = CyclicBarrier(4)
+
+        repeat(iterations) {
+            val jobs = mutableListOf<Job>()
+
+            jobs += async(pool) {
+                barrier.await()
+                1L
+            }
+
+            jobs += async(pool) {
+                barrier.await()
+                2L
+            }
+
+            jobs += async(pool) {
+                barrier.await()
+                jobs.removeAt(2)
+            }
+
+            val allJobs = ArrayList(jobs)
+            barrier.await()
+            jobs.awaitAll() // shouldn't hang
+            allJobs.awaitAll()
+            barrier.reset()
+        }
+    }
 }