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