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