Document withContext and some more pitfalls

Author: dkhalanskyjb

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

@@ -320,28 +320,131 @@ private class LazyDeferredCoroutine<T>(
 // --------------- withContext ---------------
 
 /**
- * Calls the specified suspending block with a given coroutine context, suspends until it completes, and returns
+ * Calls the specified suspending [block] with an updated coroutine context, suspends until it completes, and returns
  * the result.
  *
- * The resulting context for the [block] is derived by merging the current [coroutineContext] with the
- * specified [context] using `coroutineContext + context` (see [CoroutineContext.plus]).
- * This suspending function is cancellable. It immediately checks for cancellation of
- * the resulting context and throws [CancellationException] if it is not [active][CoroutineContext.isActive].
+ * [context] specifies the additional context elements for the coroutine to combine with
+ * the elements already present in the [CoroutineScope.coroutineContext].
+ * It is incorrect to pass a [Job] element there, as this breaks structured concurrency,
+ * unless it is [NonCancellable].
+ *
+ * If the resulting [CoroutineScope.coroutineContext] is cancelled before the [block] starts running,
+ * [block] will immediately finish with a [CancellationException],
+ * possibly without even being scheduled for execution.
+ *
+ * ## Structured Concurrency
+ *
+ * The behavior of [withContext] is similar to [coroutineScope], as it, too, creates a new *scoped child coroutine*.
+ * Refer to the documentation of that function for details.
+ *
+ * The difference is that [withContext] does not simply call the [block] in a new coroutine
+ * but updates the [currentCoroutineContext] used for running it.
+ *
+ * The context of the new scope is created like this:
+ * - First, [currentCoroutineContext] is combined with the [context] argument.
+ *   In most cases, this means that elements from [context] simply override
+ *   the elements in the [currentCoroutineContext],
+ *   but if they are `CopyableThreadContextElement`s, they are copied and combined as needed.
+ * - Then, the [Job] in the [currentCoroutineContext], if any, is used as the *parent* of the new scope,
+ *   unless overridden.
+ *   Overriding the [Job] is forbidden with the notable exception of [NonCancellable];
+ *   see a separate subsection below for details.
+ *   The new scope's [Job] is added to the resulting context.
+ *
+ * The context of the new scope is obtained by combining the [currentCoroutineContext] with a new [Job]
+ * whose parent is the [Job] of the caller [currentCoroutineContext] (if any).
+ * The [Job] of the new scope is not a normal child of the caller coroutine but a lexically scoped one,
+ * meaning that the failure of the [Job] will not affect the parent [Job].
+ * Instead, the exception leading to the failure will be rethrown to the caller of this function.
+ *
+ * ### Overriding the parent job
+ *
+ * #### [NonCancellable]
+ *
+ * Passing [NonCancellable] in the [context] argument is a special case that allows
+ * the [block] to run even if the parent coroutine is cancelled.
+ *
+ * This is useful in particular for performing cleanup operations
+ * if the cleanup procedure is itself a `suspend` function.
+ *
+ * Example:
+ *
+ * ```
+ * class Connection {
+ *     suspend fun terminate()
+ * }
+ *
+ * val connection = Connection()
+ * try {
+ *     // some cancellable operations...
+ * } finally {
+ *     withContext(NonCancellable) {
+ *         // this block will run even if the parent coroutine is cancelled
+ *         connection.terminate()
+ *     }
+ * }
+ * ```
+ *
+ * Beware that combining [NonCancellable] with context elements that change the dispatcher
+ * will make this cleanup code incorrect. See the [NonCancellable] documentation for details.
+ *
+ * #### Other [Job] elements
+ *
+ * Passing a [Job] in the [context] argument breaks structured concurrency and is not a supported pattern.
+ * It does not throw an exception only for backward compatibility reasons, as a lot of code was written this way.
+ * Always structure your coroutines such that the lifecycle of the child coroutine is
+ * contained in the lifecycle of the [CoroutineScope] it is launched in.
+ *
+ * To help with migrating to structured concurrency, the specific behaviour of passing a [Job] in the [context] argument
+ * is described here.
+ * **Do not rely on this behaviour in new code.**
+ *
+ * If [context] contains a [Job] element, it will be the *parent* of the new coroutine,
+ * and the lifecycle of the new coroutine will not be tied to the [CoroutineScope] at all.
  *
- * Calls to [withContext] whose [context] argument provides a [CoroutineDispatcher] that is
- * different from the current one, by necessity, perform additional dispatches: the [block]
- * can not be executed immediately and needs to be dispatched for execution on
- * the passed [CoroutineDispatcher], and then when the [block] completes, the execution
- * has to shift back to the original dispatcher.
+ * In specific terms:
  *
- * Note that the result of `withContext` invocation is dispatched into the original context in a cancellable way
+ * - If the [Job] passed to the [context] is cancelled, the new coroutine will be cancelled.
+ * - However, because [withContext] creates a scoped child coroutine, the failure of [block]
+ *   will not cancel the parent [Job].
+ * - If the [currentCoroutineContext] is cancelled, the new coroutine will not be affected.
+ * - In particular, if [withContext] avoided a dispatch (see "Dispatching behavior" below)
+ *   and its block finished without an exception, [withContext] itself will not throw a [CancellationException].
+ *   This means that the block execution will continue even if the parent coroutine was cancelled.
+ *
+ * ## Dispatching behavior
+ *
+ * If [context] provides a [ContinuationInterceptor] other than the one used by the caller,
+ * the [block] can not simply be called inline with the updated context and has to go through a dispatch,
+ * that is, get scheduled on the new [ContinuationInterceptor].
+ * It is up to the [ContinuationInterceptor] to actually run the [block],
+ * and it may take arbitrarily long for that to happen.
+ * After the [block] has completed, the computation has to be dispatched back to the original
+ * [ContinuationInterceptor].
+ *
+ * If the resulting context in which [block] should run has the same [ContinuationInterceptor]
+ * as the caller, no dispatch is performed.
+ * In that case, no dispatch happens on exiting the [block] either, unless child coroutines have to be awaited.
+ *
+ * Note that the result of [withContext] invocation is dispatched into the original context in a cancellable way
  * with a **prompt cancellation guarantee**, which means that if the original [coroutineContext]
- * in which `withContext` was invoked is cancelled by the time its dispatcher starts to execute the code,
- * it discards the result of `withContext` and throws [CancellationException].
+ * in which [withContext] was invoked is cancelled by the time its dispatcher starts to execute the code,
+ * it discards the result of [withContext] and throws a [CancellationException].
+ *
+ * Note that if the dispatch from [withContext] back to the original context does not need to happen
+ * (because of having the same dispatcher and not having to wait for the children)
+ * *and* the [context] passed to [withContext] contains [NonCancellable],
+ * then cancellation of the caller will not prevent a value from being successfully returned.
+ *
+ * ## Pitfalls
+ *
+ * ### Returning closeable resources
+ *
+ * Values returned from [withContext] will typically be lost if the caller is cancelled.
+ * The exception is the `withContext(NonCancellable)` pattern.
  *
- * The cancellation behaviour described above is enabled if and only if the dispatcher is being changed.
- * For example, when using `withContext(NonCancellable) { ... }` there is no change in dispatcher and
- * this call will not be cancelled neither on entry to the block inside `withContext` nor on exit from it.
+ * See the corresponding section in the [coroutineScope] documentation for details,
+ * as well as the [NonCancellable] documentation.
  */
 public suspend fun <T> withContext(
     context: CoroutineContext,

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

@@ -797,6 +797,61 @@ public object GlobalScope : CoroutineScope {
  *
  * There is a **prompt cancellation guarantee**: even if this function is ready to return the result, but was cancelled
  * while suspended, [CancellationException] will be thrown. See [suspendCancellableCoroutine] for low-level details.
+ *
+ * ## Pitfall: returning closeable resources from a scoped coroutine
+ *
+ * [R] must be a value that can safely be dropped. For example, this code is incorrect:
+ *
+ * ```
+ * // DO NOT DO THIS
+ * val closeableResource = coroutineScope {
+ *     // calculate the closeable resource somehow
+ *     obtainResource()
+ * }
+ * closeableResource.use { resource ->
+ *     // use the resource
+ * }
+ * ```
+ *
+ * The problem is that, if the caller gets cancelled before [coroutineScope] completes,
+ * then even if the calculation of the closeable resource does not suspend at all,
+ * [coroutineScope] will throw [CancellationException] instead of returning any value.
+ *
+ * This pitfall applies to all [coroutineScope]-like functions, like [withContext], [withTimeout], or [supervisorScope].
+ * For this discussion, we call them collectively `myLexicalScope`.
+ *
+ * If it is necessary to process a value returned from a lexical coroutine scope, the following pattern should be used:
+ *
+ * ```
+ * var resource: MyResource? = null
+ * try {
+ *     // note: do not simply return the resource here!
+ *     myLexicalScope {
+ *         resource = obtainResource()
+ *     }
+ *     // the resource is available here
+ * } finally {
+ *     resource?.close()
+ * }
+ * ```
+ *
+ * If cancellation during the acquisition of the resource is also undesired, the following pattern can be used
+ * ([coroutineScope] is not needed here, as [withContext] defines its own [CoroutineScope], but :
+ *
+ * ```
+ * withContext(NonCancellable) {
+ *     myLexicalScope {
+ *         obtainResource()
+ *     }
+ * }.use { resource ->
+ * }
+ * ```
+ *
+ * Be aware, however, that like any [NonCancellable] usage, this creates the risk of accessing values past the point
+ * where they are valid.
+ * For example, if the caller coroutine scope is tied to the lifecycle of a UI element, with cancellation meaning
+ * that the UI element was already disposed of, accessing it during the acquisition of a resource or
+ * before the first suspension point in [use] is not allowed and may lead to crashes.
  */
 public suspend fun <R> coroutineScope(block: suspend CoroutineScope.() -> R): R {
     contract {

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

@@ -20,6 +20,150 @@ import kotlin.coroutines.*
  * if you write `launch(NonCancellable) { ... }` then not only the newly launched job will not be cancelled
  * when the parent is cancelled, the whole parent-child relation between parent and child is severed.
  * The parent will not wait for the child's completion, nor will be cancelled when the child crashed.
+ *
+ * ## Pitfalls
+ *
+ * ### Overriding the exception with a [CancellationException] in a finalizer
+ *
+ * #### Combining [NonCancellable] with a [ContinuationInterceptor]
+ *
+ * The typical usage of [NonCancellable] is to ensure that cleanup code is executed even if the parent job is cancelled.
+ * Example:
+ *
+ * ```
+ * try {
+ *     // some code using a resource
+ * } finally {
+ *     withContext(NonCancellable) {
+ *         // cleanup code that should not be cancelled
+ *     }
+ * }
+ * ```
+ *
+ * However, it is easy to get this pattern wrong if the cleanup code needs to run on some specific dispatcher:
+ *
+ * ```
+ * // DO NOT DO THIS
+ * withContext(Dispatchers.Main) {
+ *     try {
+ *         // some code using a resource
+ *     } finally {
+ *         // THIS IS INCORRECT
+ *         withContext(NonCancellable + Dispatchers.Default) {
+ *             // cleanup code that should not be cancelled
+ *         } // this line may throw a `CancellationException`!
+ *     }
+ * }
+ * ```
+ *
+ * In this case, if the parent job is cancelled, [withContext] will throw a [CancellationException] as soon
+ * as it tries to switch back from the [Dispatchers.Default] dispatcher back to the original one.
+ * The reason for this is that [withContext] obeys the **prompt cancellation** principle,
+ * which means that dispatching back from it to the original context will fail with a [CancellationException]
+ * even if the block passed to [withContext] finished successfully,
+ * overriding the original exception thrown by the `try` block, if any.
+ *
+ * To avoid this, you should use [NonCancellable] as the only element in the context of the `withContext` call,
+ * and then inside the block, you can switch to any dispatcher you need:
+ *
+ * ```
+ * withContext(Dispatchers.Main) {
+ *     try {
+ *         // some code using a resource
+ *     } finally {
+ *         withContext(NonCancellable) {
+ *             withContext(Dispatchers.Default) {
+ *                 // cleanup code that should not be cancelled
+ *             }
+ *         }
+ *     }
+ * }
+ * ```
+ *
+ * #### Launching child coroutines
+ *
+ * Child coroutines should not be started in `withContext(NonCancellable)` blocks in resource cleanup handlers directly.
+ *
+ * ```
+ * // DO NOT DO THIS
+ * withContext(Dispatchers.Main) {
+ *     try {
+ *         // some code using a resource
+ *     } finally {
+ *         // THIS IS INCORRECT
+ *         withContext(NonCancellable) {
+ *             // cleanup code that should not be cancelled
+ *             launch { delay(100.milliseconds) }
+ *         } // this line may throw a `CancellationException`!
+ *     }
+ * }
+ * ```
+ *
+ * Similarly to the case of specifying a dispatcher alongside [NonCancellable] in a [withContext] argument,
+ * having to wait for child coroutines can lead to a dispatch at the end of the [withContext] call,
+ * which will lead to it throwing a [CancellationException] due to the prompt cancellation guarantee.
+ *
+ * The solution to this is also similar:
+ *
+ * ```
+ * withContext(Dispatchers.Main) {
+ *     try {
+ *         // some code using a resource
+ *     } finally {
+ *         withContext(NonCancellable) {
+ *             // note: `coroutineScope` here is required
+ *             // to prevent a sporadic CancellationException
+ *             coroutineScope {
+ *                 // cleanup code that should not be cancelled
+ *                 launch { delay(100.milliseconds) }
+ *             }
+ *         }
+ *     }
+ * }
+ * ```
+ *
+ * Because now [coroutineScope] and not [withContext] has to wait for the children, there is once again no dispatch
+ * between the last line of the [withContext] block and getting back to the caller.
+ *
+ * ### Not reacting to cancellations right outside the [withContext]
+ *
+ * Just like combining [NonCancellable] with other elements is incorrect because cancellation may override
+ * the original exception, the opposite can also be incorrect, depending on the context:
+ *
+ * ```
+ * // DO NOT DO THIS
+ * withContext(Dispatchers.Main) {
+ *     withContext(NonCancellable) {
+ *         withContext(Dispatchers.Default) {
+ *             // do something
+ *         }
+ *     } // will not react to the caller's cancellation!
+ *     // BUG HERE
+ *     updateUi() // may be invoked when the caller is already cancelled
+ * }
+ * ```
+ *
+ * Here, the following may happen:
+ * 1. The `do something` block gets entered, and the main thread gets released and is free to perform other tasks.
+ * 2. Some other task updates the UI and cancels this coroutine, which is no longer needed.
+ * 3. `do something` finishes, and the computation is dispatched back to the main thread.
+ * 4. `updateUi()` is called, even though the coroutine was already cancelled and the UI is no longer in a valid state
+ *    for this update operation, potentially leading to a crash.
+ *
+ * [ensureActive] can be used to manually ensure that cancelled code no longer runs:
+ *
+ * ```
+ * withContext(Dispatchers.Main) {
+ *     withContext(NonCancellable) {
+ *         withContext(Dispatchers.Default) {
+ *             // do something
+ *         }
+ *     }
+ *     ensureActive() // check if we are still allowed to run the code
+ *     updateUi()
+ * }
+ * ```
+ *
  */
 @OptIn(InternalForInheritanceCoroutinesApi::class)
 public object NonCancellable : AbstractCoroutineContextElement(Job), Job {

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

@@ -77,7 +77,11 @@ public fun SupervisorJob0(parent: Job? = null) : Job = SupervisorJob(parent)
  * There is a **prompt cancellation guarantee**: even if this function is ready to return the result, but was cancelled
  * while suspended, [CancellationException] will be thrown. See [suspendCancellableCoroutine] for low-level details.
  *
- * **Pitfall**: [supervisorScope] does not install a [CoroutineExceptionHandler] in the new scope.
+ * ## Pitfalls
+ *
+ * ### Uncaught exceptions in child coroutines
+ *
+ * [supervisorScope] does not install a [CoroutineExceptionHandler] in the new scope.
  * This means that if a child coroutine started with [launch] fails, its exception will be unhandled,
  * possibly crashing the program. Use the following pattern to avoid this:
  *
@@ -92,6 +96,11 @@ public fun SupervisorJob0(parent: Job? = null) : Job = SupervisorJob(parent)
  * ```
  *
  * Alternatively, the [CoroutineExceptionHandler] can be supplied to the newly launched coroutines themselves.
+ *
+ * ### Returning closeable resources
+ *
+ * Values returned from [supervisorScope] will be lost if the caller is cancelled.
+ * See the corresponding section in the [coroutineScope] documentation for details.
  */
 public suspend fun <R> supervisorScope(block: suspend CoroutineScope.() -> R): R {
     contract {