Leaking handles in awaitAll are fixed

Author: elizarov

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

@@ -17,6 +17,7 @@
 package kotlinx.coroutines.experimental
 
 import kotlinx.atomicfu.atomic
+import kotlinx.coroutines.experimental.internalAnnotations.Volatile
 
 /**
  * Awaits for completion of given deferred values without blocking a thread and resumes normally with the list of values
@@ -70,29 +71,56 @@ private class AwaitAll<T>(private val deferreds: Array<out Deferred<T>>) {
     private val notCompletedCount = atomic(deferreds.size)
 
     suspend fun await(): List<T> = suspendCancellableCoroutine { cont ->
-        val handlers = deferreds.map {
-            it.start() // To properly await lazily started deferreds
-            it.invokeOnCompletion(AwaitAllNode(cont, it).asHandler)
+        // Intricate dance here
+        // Step 1: Create nodes and install them as completion handlers, they may fire!
+        val nodes = Array<AwaitAllNode>(deferreds.size) { i ->
+            val deferred = deferreds[i]
+            deferred.start() // To properly await lazily started deferreds
+            AwaitAllNode(cont, deferred).apply {
+                handle = deferred.invokeOnCompletion(asHandler)
+            }
+        }
+        val disposer = DisposeHandlersOnCancel(nodes)
+        // Step 2: Set disposer to each node
+        nodes.forEach { it.disposer = disposer }
+        // Here we know that if any code the nodes complete, it will dipsose the rest
+        // Step 3: Now we can check if continuation is complete
+        if (cont.isCompleted) {
+            // it is already complete while handlers were being installed -- dispose them all
+            disposer.disposeAll()
+        } else {
+            cont.invokeOnCancellation(handler = disposer.asHandler)
         }
-        cont.invokeOnCancellation(handler = DisposeHandlersOnCancel(handlers).asHandler)
     }
 
-    private class DisposeHandlersOnCancel(private val handlers: List<DisposableHandle>) : CancelHandler() {
-        override fun invoke(cause: Throwable?) {
-            handlers.forEach { it.dispose() }
+    private inner class DisposeHandlersOnCancel(private val nodes: Array<AwaitAllNode>) : CancelHandler() {
+        fun disposeAll() {
+            nodes.forEach { it.handle.dispose() }
         }
-        override fun toString(): String = "DisposeHandlersOnCancel[$handlers]"
+
+        override fun invoke(cause: Throwable?) { disposeAll() }
+        override fun toString(): String = "DisposeHandlersOnCancel[$nodes]"
     }
 
     private inner class AwaitAllNode(private val continuation: CancellableContinuation<List<T>>, job: Job) : JobNode<Job>(job) {
+        lateinit var handle: DisposableHandle
+
+        @Volatile
+        var disposer: DisposeHandlersOnCancel? = null
+        
         override fun invoke(cause: Throwable?) {
             if (cause != null) {
                 val token = continuation.tryResumeWithException(cause)
                 if (token != null) {
                     continuation.completeResume(token)
+                    // volatile read of disposer AFTER continuation is complete
+                    val disposer = this.disposer
+                    // and if disposer was already set (all handlers where already installed, then dispose them all)
+                    if (disposer != null) disposer.disposeAll()
                 }
             } else if (notCompletedCount.decrementAndGet() == 0) {
                 continuation.resume(deferreds.map { it.getCompleted() })
+                // Note, that all deferreds are complete here, so we don't need to dispose their nodes
             }
         }
     }

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2016-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kotlinx.coroutines.experimental
+
+import org.junit.*
+
+class AwaitJvmTest : TestBase() {
+    @Test
+    public fun testSecondLeak() = runTest {
+        // This test is to make sure that handlers installed on the second deferred do not leak
+        val d1 = CompletableDeferred<Int>()
+        val d2 = CompletableDeferred<Int>()
+        d1.completeExceptionally(TestException()) // first is crashed
+        val iterations = 3_000_000 * stressTestMultiplier
+        for (iter in 1..iterations) {
+            try {
+                awaitAll(d1, d2)
+                expectUnreached()
+            } catch (e: TestException) {
+                expect(iter)
+            }
+        }
+        finish(iterations + 1)
+    }
+
+    private class TestException : Exception()
+}