Roman Elizarov | 769d7dc | 2018-04-26 12:13:40 +0300 | [diff] [blame] | 1 | /* |
| 2 | * Copyright 2016-2017 JetBrains s.r.o. |
| 3 | * |
| 4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | * you may not use this file except in compliance with the License. |
| 6 | * You may obtain a copy of the License at |
| 7 | * |
| 8 | * http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | * |
| 10 | * Unless required by applicable law or agreed to in writing, software |
| 11 | * distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | * See the License for the specific language governing permissions and |
| 14 | * limitations under the License. |
| 15 | */ |
| 16 | |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 17 | package kotlinx.coroutines.experimental |
| 18 | |
| 19 | import kotlinx.atomicfu.atomic |
Roman Elizarov | e0b6db0 | 2018-04-28 19:22:29 +0300 | [diff] [blame] | 20 | import kotlinx.coroutines.experimental.internalAnnotations.Volatile |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 21 | |
| 22 | /** |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 23 | * Awaits for completion of given deferred values without blocking a thread and resumes normally with the list of values |
| 24 | * when all deferred computations are complete or resumes with the first thrown exception if any of computations |
| 25 | * complete exceptionally including cancellation. |
| 26 | * |
| 27 | * This function is **not** equivalent to `deferreds.map { it.await() }` which fails only when when it sequentially |
| 28 | * gets to wait the failing deferred, while this `awaitAll` fails immediately as soon as any of the deferreds fail. |
| 29 | * |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 30 | * This suspending function is cancellable. |
| 31 | * If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, |
| 32 | * this function immediately resumes with [CancellationException]. |
| 33 | */ |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 34 | public suspend fun <T> awaitAll(vararg deferreds: Deferred<T>): List<T> = |
| 35 | if (deferreds.isEmpty()) emptyList() else AwaitAll(deferreds).await() |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 36 | |
| 37 | /** |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 38 | * Awaits for completion of given deferred values without blocking a thread and resumes normally with the list of values |
| 39 | * when all deferred computations are complete or resumes with the first thrown exception if any of computations |
| 40 | * complete exceptionally including cancellation. |
| 41 | * |
| 42 | * This function is **not** equivalent to `this.map { it.await() }` which fails only when when it sequentially |
| 43 | * gets to wait the failing deferred, while this `awaitAll` fails immediately as soon as any of the deferreds fail. |
| 44 | * |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 45 | * This suspending function is cancellable. |
Roman Elizarov | 203abb0 | 2018-04-12 12:34:14 +0300 | [diff] [blame] | 46 | * If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 47 | * this function immediately resumes with [CancellationException]. |
| 48 | */ |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 49 | public suspend fun <T> Collection<Deferred<T>>.awaitAll(): List<T> = |
| 50 | if (isEmpty()) emptyList() else AwaitAll(toTypedArray()).await() |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 51 | |
| 52 | /** |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 53 | * Suspends current coroutine until all given jobs are complete. |
| 54 | * This method is semantically equivalent to joining all given jobs one by one with `jobs.forEach { it.join() }`. |
| 55 | * |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 56 | * This suspending function is cancellable. If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, |
| 57 | * this function immediately resumes with [CancellationException]. |
| 58 | */ |
Vsevolod Tolstopyatov | 9c69279 | 2018-04-11 18:22:35 +0300 | [diff] [blame] | 59 | public suspend fun joinAll(vararg jobs: Job): Unit = jobs.forEach { it.join() } |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 60 | |
| 61 | /** |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 62 | * Suspends current coroutine until all given jobs are complete. |
| 63 | * This method is semantically equivalent to joining all given jobs one by one with `forEach { it.join() }`. |
| 64 | * |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 65 | * This suspending function is cancellable. If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, |
| 66 | * this function immediately resumes with [CancellationException]. |
| 67 | */ |
Roman Elizarov | 203abb0 | 2018-04-12 12:34:14 +0300 | [diff] [blame] | 68 | public suspend fun Collection<Job>.joinAll(): Unit = forEach { it.join() } |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 69 | |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 70 | private class AwaitAll<T>(private val deferreds: Array<out Deferred<T>>) { |
| 71 | private val notCompletedCount = atomic(deferreds.size) |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 72 | |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 73 | suspend fun await(): List<T> = suspendCancellableCoroutine { cont -> |
Roman Elizarov | e0b6db0 | 2018-04-28 19:22:29 +0300 | [diff] [blame] | 74 | // Intricate dance here |
| 75 | // Step 1: Create nodes and install them as completion handlers, they may fire! |
| 76 | val nodes = Array<AwaitAllNode>(deferreds.size) { i -> |
| 77 | val deferred = deferreds[i] |
| 78 | deferred.start() // To properly await lazily started deferreds |
| 79 | AwaitAllNode(cont, deferred).apply { |
| 80 | handle = deferred.invokeOnCompletion(asHandler) |
| 81 | } |
Roman Elizarov | 6d9f40f | 2018-04-28 14:44:02 +0300 | [diff] [blame] | 82 | } |
Roman Elizarov | e0b6db0 | 2018-04-28 19:22:29 +0300 | [diff] [blame] | 83 | val disposer = DisposeHandlersOnCancel(nodes) |
| 84 | // Step 2: Set disposer to each node |
| 85 | nodes.forEach { it.disposer = disposer } |
| 86 | // Here we know that if any code the nodes complete, it will dipsose the rest |
| 87 | // Step 3: Now we can check if continuation is complete |
| 88 | if (cont.isCompleted) { |
| 89 | // it is already complete while handlers were being installed -- dispose them all |
| 90 | disposer.disposeAll() |
| 91 | } else { |
| 92 | cont.invokeOnCancellation(handler = disposer.asHandler) |
| 93 | } |
Vsevolod Tolstopyatov | 05d3823 | 2018-04-12 20:16:37 +0300 | [diff] [blame] | 94 | } |
| 95 | |
Roman Elizarov | e0b6db0 | 2018-04-28 19:22:29 +0300 | [diff] [blame] | 96 | private inner class DisposeHandlersOnCancel(private val nodes: Array<AwaitAllNode>) : CancelHandler() { |
| 97 | fun disposeAll() { |
| 98 | nodes.forEach { it.handle.dispose() } |
Roman Elizarov | 3e9f244 | 2018-04-28 17:38:22 +0300 | [diff] [blame] | 99 | } |
Roman Elizarov | e0b6db0 | 2018-04-28 19:22:29 +0300 | [diff] [blame] | 100 | |
| 101 | override fun invoke(cause: Throwable?) { disposeAll() } |
| 102 | override fun toString(): String = "DisposeHandlersOnCancel[$nodes]" |
Roman Elizarov | 3e9f244 | 2018-04-28 17:38:22 +0300 | [diff] [blame] | 103 | } |
| 104 | |
| 105 | private inner class AwaitAllNode(private val continuation: CancellableContinuation<List<T>>, job: Job) : JobNode<Job>(job) { |
Roman Elizarov | e0b6db0 | 2018-04-28 19:22:29 +0300 | [diff] [blame] | 106 | lateinit var handle: DisposableHandle |
| 107 | |
| 108 | @Volatile |
| 109 | var disposer: DisposeHandlersOnCancel? = null |
| 110 | |
Vsevolod Tolstopyatov | 05d3823 | 2018-04-12 20:16:37 +0300 | [diff] [blame] | 111 | override fun invoke(cause: Throwable?) { |
| 112 | if (cause != null) { |
| 113 | val token = continuation.tryResumeWithException(cause) |
| 114 | if (token != null) { |
| 115 | continuation.completeResume(token) |
Roman Elizarov | e0b6db0 | 2018-04-28 19:22:29 +0300 | [diff] [blame] | 116 | // volatile read of disposer AFTER continuation is complete |
| 117 | val disposer = this.disposer |
| 118 | // and if disposer was already set (all handlers where already installed, then dispose them all) |
| 119 | if (disposer != null) disposer.disposeAll() |
Vsevolod Tolstopyatov | 05d3823 | 2018-04-12 20:16:37 +0300 | [diff] [blame] | 120 | } |
| 121 | } else if (notCompletedCount.decrementAndGet() == 0) { |
Roman Elizarov | 189e995 | 2018-04-26 11:06:37 +0300 | [diff] [blame] | 122 | continuation.resume(deferreds.map { it.getCompleted() }) |
Roman Elizarov | e0b6db0 | 2018-04-28 19:22:29 +0300 | [diff] [blame] | 123 | // Note, that all deferreds are complete here, so we don't need to dispose their nodes |
Vsevolod Tolstopyatov | 4f0d48b | 2018-04-11 13:51:33 +0300 | [diff] [blame] | 124 | } |
| 125 | } |
| 126 | } |
| 127 | } |