1 /*
2 * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
3 */
4
5 package kotlinx.coroutines
6
7 import kotlinx.coroutines.internal.*
8 import kotlin.coroutines.*
9
10 /**
11 * Defines elements in [CoroutineContext] that are installed into thread context
12 * every time the coroutine with this element in the context is resumed on a thread.
13 *
14 * Implementations of this interface define a type [S] of the thread-local state that they need to store on
15 * resume of a coroutine and restore later on suspend. The infrastructure provides the corresponding storage.
16 *
17 * Example usage looks like this:
18 *
19 * ```
20 * // Appends "name" of a coroutine to a current thread name when coroutine is executed
21 * class CoroutineName(val name: String) : ThreadContextElement<String> {
22 * // declare companion object for a key of this element in coroutine context
23 * companion object Key : CoroutineContext.Key<CoroutineName>
24 *
25 * // provide the key of the corresponding context element
26 * override val key: CoroutineContext.Key<CoroutineName>
27 * get() = Key
28 *
29 * // this is invoked before coroutine is resumed on current thread
30 * override fun updateThreadContext(context: CoroutineContext): String {
31 * val previousName = Thread.currentThread().name
32 * Thread.currentThread().name = "$previousName # $name"
33 * return previousName
34 * }
35 *
36 * // this is invoked after coroutine has suspended on current thread
37 * override fun restoreThreadContext(context: CoroutineContext, oldState: String) {
38 * Thread.currentThread().name = oldState
39 * }
40 * }
41 *
42 * // Usage
43 * launch(Dispatchers.Main + CoroutineName("Progress bar coroutine")) { ... }
44 * ```
45 *
46 * Every time this coroutine is resumed on a thread, UI thread name is updated to
47 * "UI thread original name # Progress bar coroutine" and the thread name is restored to the original one when
48 * this coroutine suspends.
49 *
50 * To use [ThreadLocal] variable within the coroutine use [ThreadLocal.asContextElement][asContextElement] function.
51 *
52 * ### Reentrancy and thread-safety
53 *
54 * Correct implementations of this interface must expect that calls to [restoreThreadContext]
55 * may happen in parallel to the subsequent [updateThreadContext] and [restoreThreadContext] operations.
56 * See [CopyableThreadContextElement] for advanced interleaving details.
57 *
58 * All implementations of [ThreadContextElement] should be thread-safe and guard their internal mutable state
59 * within an element accordingly.
60 */
61 public interface ThreadContextElement<S> : CoroutineContext.Element {
62 /**
63 * Updates context of the current thread.
64 * This function is invoked before the coroutine in the specified [context] is resumed in the current thread
65 * when the context of the coroutine this element.
66 * The result of this function is the old value of the thread-local state that will be passed to [restoreThreadContext].
67 * This method should handle its own exceptions and do not rethrow it. Thrown exceptions will leave coroutine which
68 * context is updated in an undefined state and may crash an application.
69 *
70 * @param context the coroutine context.
71 */
updateThreadContextnull72 public fun updateThreadContext(context: CoroutineContext): S
73
74 /**
75 * Restores context of the current thread.
76 * This function is invoked after the coroutine in the specified [context] is suspended in the current thread
77 * if [updateThreadContext] was previously invoked on resume of this coroutine.
78 * The value of [oldState] is the result of the previous invocation of [updateThreadContext] and it should
79 * be restored in the thread-local state by this function.
80 * This method should handle its own exceptions and do not rethrow it. Thrown exceptions will leave coroutine which
81 * context is updated in an undefined state and may crash an application.
82 *
83 * @param context the coroutine context.
84 * @param oldState the value returned by the previous invocation of [updateThreadContext].
85 */
86 public fun restoreThreadContext(context: CoroutineContext, oldState: S)
87 }
88
89 /**
90 * A [ThreadContextElement] copied whenever a child coroutine inherits a context containing it.
91 *
92 * When an API uses a _mutable_ [ThreadLocal] for consistency, a [CopyableThreadContextElement]
93 * can give coroutines "coroutine-safe" write access to that `ThreadLocal`.
94 *
95 * A write made to a `ThreadLocal` with a matching [CopyableThreadContextElement] by a coroutine
96 * will be visible to _itself_ and any child coroutine launched _after_ that write.
97 *
98 * Writes will not be visible to the parent coroutine, peer coroutines, or coroutines that happen
99 * to use the same thread. Writes made to the `ThreadLocal` by the parent coroutine _after_
100 * launching a child coroutine will not be visible to that child coroutine.
101 *
102 * This can be used to allow a coroutine to use a mutable ThreadLocal API transparently and
103 * correctly, regardless of the coroutine's structured concurrency.
104 *
105 * This example adapts a `ThreadLocal` method trace to be "coroutine local" while the method trace
106 * is in a coroutine:
107 *
108 * ```
109 * class TraceContextElement(private val traceData: TraceData?) : CopyableThreadContextElement<TraceData?> {
110 * companion object Key : CoroutineContext.Key<TraceContextElement>
111 *
112 * override val key: CoroutineContext.Key<TraceContextElement> = Key
113 *
114 * override fun updateThreadContext(context: CoroutineContext): TraceData? {
115 * val oldState = traceThreadLocal.get()
116 * traceThreadLocal.set(traceData)
117 * return oldState
118 * }
119 *
120 * override fun restoreThreadContext(context: CoroutineContext, oldState: TraceData?) {
121 * traceThreadLocal.set(oldState)
122 * }
123 *
124 * override fun copyForChild(): TraceContextElement {
125 * // Copy from the ThreadLocal source of truth at child coroutine launch time. This makes
126 * // ThreadLocal writes between resumption of the parent coroutine and the launch of the
127 * // child coroutine visible to the child.
128 * return TraceContextElement(traceThreadLocal.get()?.copy())
129 * }
130 *
131 * override fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext {
132 * // Merge operation defines how to handle situations when both
133 * // the parent coroutine has an element in the context and
134 * // an element with the same key was also
135 * // explicitly passed to the child coroutine.
136 * // If merging does not require special behavior,
137 * // the copy of the element can be returned.
138 * return TraceContextElement(traceThreadLocal.get()?.copy())
139 * }
140 * }
141 * ```
142 *
143 * A coroutine using this mechanism can safely call Java code that assumes the corresponding thread local element's
144 * value is installed into the target thread local.
145 *
146 * ### Reentrancy and thread-safety
147 *
148 * Correct implementations of this interface must expect that calls to [restoreThreadContext]
149 * may happen in parallel to the subsequent [updateThreadContext] and [restoreThreadContext] operations.
150 *
151 * Even though an element is copied for each child coroutine, an implementation should be able to handle the following
152 * interleaving when a coroutine with the corresponding element is launched on a multithreaded dispatcher:
153 *
154 * ```
155 * coroutine.updateThreadContext() // Thread #1
156 * ... coroutine body ...
157 * // suspension + immediate dispatch happen here
158 * coroutine.updateThreadContext() // Thread #2, coroutine is already resumed
159 * // ... coroutine body after suspension point on Thread #2 ...
160 * coroutine.restoreThreadContext() // Thread #1, is invoked late because Thread #1 is slow
161 * coroutine.restoreThreadContext() // Thread #2, may happen in parallel with the previous restore
162 * ```
163 *
164 * All implementations of [CopyableThreadContextElement] should be thread-safe and guard their internal mutable state
165 * within an element accordingly.
166 */
167 @DelicateCoroutinesApi
168 @ExperimentalCoroutinesApi
169 public interface CopyableThreadContextElement<S> : ThreadContextElement<S> {
170
171 /**
172 * Returns a [CopyableThreadContextElement] to replace `this` `CopyableThreadContextElement` in the child
173 * coroutine's context that is under construction if the added context does not contain an element with the same [key].
174 *
175 * This function is called on the element each time a new coroutine inherits a context containing it,
176 * and the returned value is folded into the context given to the child.
177 *
178 * Since this method is called whenever a new coroutine is launched in a context containing this
179 * [CopyableThreadContextElement], implementations are performance-sensitive.
180 */
181 public fun copyForChild(): CopyableThreadContextElement<S>
182
183 /**
184 * Returns a [CopyableThreadContextElement] to replace `this` `CopyableThreadContextElement` in the child
185 * coroutine's context that is under construction if the added context does contain an element with the same [key].
186 *
187 * This method is invoked on the original element, accepting as the parameter
188 * the element that is supposed to overwrite it.
189 */
190 public fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext
191 }
192
193 /**
194 * Wraps [ThreadLocal] into [ThreadContextElement]. The resulting [ThreadContextElement]
195 * maintains the given [value] of the given [ThreadLocal] for coroutine regardless of the actual thread its is resumed on.
196 * By default [ThreadLocal.get] is used as a value for the thread-local variable, but it can be overridden with [value] parameter.
197 * Beware that context element **does not track** modifications of the thread-local and accessing thread-local from coroutine
198 * without the corresponding context element returns **undefined** value. See the examples for a detailed description.
199 *
200 *
201 * Example usage:
202 * ```
203 * val myThreadLocal = ThreadLocal<String?>()
204 * ...
205 * println(myThreadLocal.get()) // Prints "null"
206 * launch(Dispatchers.Default + myThreadLocal.asContextElement(value = "foo")) {
207 * println(myThreadLocal.get()) // Prints "foo"
208 * withContext(Dispatchers.Main) {
209 * println(myThreadLocal.get()) // Prints "foo", but it's on UI thread
210 * }
211 * }
212 * println(myThreadLocal.get()) // Prints "null"
213 * ```
214 *
215 * The context element does not track modifications of the thread-local variable, for example:
216 *
217 * ```
218 * myThreadLocal.set("main")
219 * withContext(Dispatchers.Main) {
220 * println(myThreadLocal.get()) // Prints "main"
221 * myThreadLocal.set("UI")
222 * }
223 * println(myThreadLocal.get()) // Prints "main", not "UI"
224 * ```
225 *
226 * Use `withContext` to update the corresponding thread-local variable to a different value, for example:
227 * ```
228 * withContext(myThreadLocal.asContextElement("foo")) {
229 * println(myThreadLocal.get()) // Prints "foo"
230 * }
231 * ```
232 *
233 * Accessing the thread-local without corresponding context element leads to undefined value:
234 * ```
235 * val tl = ThreadLocal.withInitial { "initial" }
236 *
237 * runBlocking {
238 * println(tl.get()) // Will print "initial"
239 * // Change context
240 * withContext(tl.asContextElement("modified")) {
241 * println(tl.get()) // Will print "modified"
242 * }
243 * // Context is changed again
244 * println(tl.get()) // <- WARN: can print either "modified" or "initial"
245 * }
246 * ```
247 * to fix this behaviour use `runBlocking(tl.asContextElement())`
248 */
asContextElementnull249 public fun <T> ThreadLocal<T>.asContextElement(value: T = get()): ThreadContextElement<T> =
250 ThreadLocalElement(value, this)
251
252 /**
253 * Return `true` when current thread local is present in the coroutine context, `false` otherwise.
254 * Thread local can be present in the context only if it was added via [asContextElement] to the context.
255 *
256 * Example of usage:
257 * ```
258 * suspend fun processRequest() {
259 * if (traceCurrentRequestThreadLocal.isPresent()) { // Probabilistic tracing
260 * // Do some heavy-weight tracing
261 * }
262 * // Process request regularly
263 * }
264 * ```
265 */
266 public suspend inline fun ThreadLocal<*>.isPresent(): Boolean = coroutineContext[ThreadLocalKey(this)] !== null
267
268 /**
269 * Checks whether current thread local is present in the coroutine context and throws [IllegalStateException] if it is not.
270 * It is a good practice to validate that thread local is present in the context, especially in large code-bases,
271 * to avoid stale thread-local values and to have a strict invariants.
272 *
273 * E.g. one may use the following method to enforce proper use of the thread locals with coroutines:
274 * ```
275 * public suspend inline fun <T> ThreadLocal<T>.getSafely(): T {
276 * ensurePresent()
277 * return get()
278 * }
279 *
280 * // Usage
281 * withContext(...) {
282 * val value = threadLocal.getSafely() // Fail-fast in case of improper context
283 * }
284 * ```
285 */
286 public suspend inline fun ThreadLocal<*>.ensurePresent(): Unit =
287 check(isPresent()) { "ThreadLocal $this is missing from context $coroutineContext" }
288