Remove requirement that job of the pre-created JobCancelNode have to …

kotlinx-coroutines-core/common/src/JobSupport.kt

@@ -508,14 +508,17 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
         }
     }
 
-    private fun makeNode(handler: CompletionHandler, onCancelling: Boolean): JobNode<*> {
-        return if (onCancelling)
-            (handler as? JobCancellingNode<*>)?.also { assert { it.job === this } }
+    private fun makeNode(handler: CompletionHandler, onCancelling: Boolean): JobNode<*> =
+        if (onCancelling) {
+            (handler as? JobCancellingNode<*>)
+                ?.takeIf { it.job === this }
                 ?: InvokeOnCancelling(this, handler)
-        else
-            (handler as? JobNode<*>)?.also { assert { it.job === this && it !is JobCancellingNode } }
+        } else {
+            (handler as? JobNode<*>)
+                ?.also { assert { it !is JobCancellingNode } }
+                ?.takeIf { it.job === this }
                 ?: InvokeOnCompletion(this, handler)
-    }
+        }
 
     private fun addLastAtomic(expect: Any, list: NodeList, node: JobNode<*>) =
         list.addLastIf(node) { this.state === expect }

kotlinx-coroutines-core/common/test/AwaitTest.kt

@@ -351,4 +351,32 @@ class AwaitTest : TestBase() {
             async(NonCancellable) { throw TestException() }
         joinAll(job, job, job)
     }
+
+    @Test
+    fun testAwaitAllDelegates() = runTest {
+        expect(1)
+        val deferred = CompletableDeferred<String>()
+        val delegate = object : Deferred<String> by deferred {}
+        launch {
+            expect(3)
+            deferred.complete("OK")
+        }
+        expect(2)
+        awaitAll(delegate)
+        finish(4)
+    }
+
+    @Test
+    fun testCancelAwaitAllDelegate() = runTest {
+        expect(1)
+        val deferred = CompletableDeferred<String>()
+        val delegate = object : Deferred<String> by deferred {}
+        launch {
+            expect(3)
+            deferred.cancel()
+        }
+        expect(2)
+        assertFailsWith<CancellationException> { awaitAll(delegate) }
+        finish(4)
+    }
 }