Skip to content

Commit c21114e

Browse files
zuevmaximzuevmaxim
committed
Move CopyableThreadContextElement to common
1 parent 877c70f commit c21114e

File tree

6 files changed

+190
-215
lines changed

6 files changed

+190
-215
lines changed

Diff for: ‎kotlinx-coroutines-core/common/src/CoroutineContext.common.kt

+79-14
Original file line numberDiff line numberDiff line change
@@ -3,26 +3,13 @@ package kotlinx.coroutines
33
import kotlinx.coroutines.internal.*
44
import kotlin.coroutines.*
55

6-
/**
7-
* Creates a context for a new coroutine. It installs [Dispatchers.Default] when no other dispatcher or
8-
* [ContinuationInterceptor] is specified and adds optional support for debugging facilities (when turned on)
9-
* and copyable-thread-local facilities on JVM.
10-
*/
11-
public expect fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext
12-
13-
/**
14-
* Creates a context for coroutine builder functions that do not launch a new coroutine, e.g. [withContext].
15-
* @suppress
16-
*/
17-
@InternalCoroutinesApi
18-
public expect fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext
19-
206
@PublishedApi // to have unmangled name when using from other modules via suppress
217
@Suppress("PropertyName")
228
internal expect val DefaultDelay: Delay
239

2410
internal expect fun Continuation<*>.toDebugString(): String
2511
internal expect val CoroutineContext.coroutineName: String?
12+
internal expect fun wrapContextWithDebug(context: CoroutineContext): CoroutineContext
2613

2714
/**
2815
* Executes a block using a given coroutine context.
@@ -98,3 +85,81 @@ internal object UndispatchedMarker: CoroutineContext.Element, CoroutineContext.K
9885
override val key: CoroutineContext.Key<*>
9986
get() = this
10087
}
88+
89+
/**
90+
* Creates a context for a new coroutine. It installs [Dispatchers.Default] when no other dispatcher or
91+
* [ContinuationInterceptor] is specified and
92+
*/
93+
@ExperimentalCoroutinesApi
94+
public fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
95+
val combined = foldCopies(coroutineContext, context, true)
96+
val debug = wrapContextWithDebug(combined)
97+
return if (combined !== Dispatchers.Default && combined[ContinuationInterceptor] == null)
98+
debug + Dispatchers.Default else debug
99+
}
100+
101+
/**
102+
* Creates a context for coroutine builder functions that do not launch a new coroutine, e.g. [withContext].
103+
* @suppress
104+
*/
105+
@InternalCoroutinesApi
106+
public fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext {
107+
/*
108+
* Fast-path: we only have to copy/merge if 'addedContext' (which typically has one or two elements)
109+
* contains copyable elements.
110+
*/
111+
if (!addedContext.hasCopyableElements()) return this + addedContext
112+
return foldCopies(this, addedContext, false)
113+
}
114+
115+
private fun CoroutineContext.hasCopyableElements(): Boolean =
116+
fold(false) { result, it -> result || it is CopyableThreadContextElement<*> }
117+
118+
/**
119+
* Folds two contexts properly applying [CopyableThreadContextElement] rules when necessary.
120+
* The rules are the following:
121+
* - If neither context has CTCE, the sum of two contexts is returned
122+
* - Every CTCE from the left-hand side context that does not have a matching (by key) element from right-hand side context
123+
* is [copied][CopyableThreadContextElement.copyForChild] if [isNewCoroutine] is `true`.
124+
* - Every CTCE from the left-hand side context that has a matching element in the right-hand side context is [merged][CopyableThreadContextElement.mergeForChild]
125+
* - Every CTCE from the right-hand side context that hasn't been merged is copied
126+
* - Everything else is added to the resulting context as is.
127+
*/
128+
private fun foldCopies(originalContext: CoroutineContext, appendContext: CoroutineContext, isNewCoroutine: Boolean): CoroutineContext {
129+
// Do we have something to copy left-hand side?
130+
val hasElementsLeft = originalContext.hasCopyableElements()
131+
val hasElementsRight = appendContext.hasCopyableElements()
132+
133+
// Nothing to fold, so just return the sum of contexts
134+
if (!hasElementsLeft && !hasElementsRight) {
135+
return originalContext + appendContext
136+
}
137+
138+
var leftoverContext = appendContext
139+
val folded = originalContext.fold<CoroutineContext>(EmptyCoroutineContext) { result, element ->
140+
if (element !is CopyableThreadContextElement<*>) return@fold result + element
141+
// Will this element be overwritten?
142+
val newElement = leftoverContext[element.key]
143+
// No, just copy it
144+
if (newElement == null) {
145+
// For 'withContext'-like builders we do not copy as the element is not shared
146+
return@fold result + if (isNewCoroutine) element.copyForChild() else element
147+
}
148+
// Yes, then first remove the element from append context
149+
leftoverContext = leftoverContext.minusKey(element.key)
150+
// Return the sum
151+
@Suppress("UNCHECKED_CAST")
152+
return@fold result + (element as CopyableThreadContextElement<Any?>).mergeForChild(newElement)
153+
}
154+
155+
if (hasElementsRight) {
156+
leftoverContext = leftoverContext.fold<CoroutineContext>(EmptyCoroutineContext) { result, element ->
157+
// We're appending new context element -- we have to copy it, otherwise it may be shared with others
158+
if (element is CopyableThreadContextElement<*>) {
159+
return@fold result + element.copyForChild()
160+
}
161+
return@fold result + element
162+
}
163+
}
164+
return folded + leftoverContext
165+
}

Diff for: ‎kotlinx-coroutines-core/common/src/ThreadContextElement.common.kt

+104
Original file line numberDiff line numberDiff line change
@@ -80,3 +80,107 @@ public interface ThreadContextElement<S> : CoroutineContext.Element {
8080
*/
8181
public fun restoreThreadContext(context: CoroutineContext, oldState: S)
8282
}
83+
84+
/**
85+
* A [ThreadContextElement] copied whenever a child coroutine inherits a context containing it.
86+
*
87+
* When an API uses a _mutable_ [ThreadLocal] for consistency, a [CopyableThreadContextElement]
88+
* can give coroutines "coroutine-safe" write access to that `ThreadLocal`.
89+
*
90+
* A write made to a `ThreadLocal` with a matching [CopyableThreadContextElement] by a coroutine
91+
* will be visible to _itself_ and any child coroutine launched _after_ that write.
92+
*
93+
* Writes will not be visible to the parent coroutine, peer coroutines, or coroutines that happen
94+
* to use the same thread. Writes made to the `ThreadLocal` by the parent coroutine _after_
95+
* launching a child coroutine will not be visible to that child coroutine.
96+
*
97+
* This can be used to allow a coroutine to use a mutable ThreadLocal API transparently and
98+
* correctly, regardless of the coroutine's structured concurrency.
99+
*
100+
* This example adapts a `ThreadLocal` method trace to be "coroutine local" while the method trace
101+
* is in a coroutine:
102+
*
103+
* ```
104+
* class TraceContextElement(private val traceData: TraceData?) : CopyableThreadContextElement<TraceData?> {
105+
* companion object Key : CoroutineContext.Key<TraceContextElement>
106+
*
107+
* override val key: CoroutineContext.Key<TraceContextElement> = Key
108+
*
109+
* override fun updateThreadContext(context: CoroutineContext): TraceData? {
110+
* val oldState = traceThreadLocal.get()
111+
* traceThreadLocal.set(traceData)
112+
* return oldState
113+
* }
114+
*
115+
* override fun restoreThreadContext(context: CoroutineContext, oldState: TraceData?) {
116+
* traceThreadLocal.set(oldState)
117+
* }
118+
*
119+
* override fun copyForChild(): TraceContextElement {
120+
* // Copy from the ThreadLocal source of truth at child coroutine launch time. This makes
121+
* // ThreadLocal writes between resumption of the parent coroutine and the launch of the
122+
* // child coroutine visible to the child.
123+
* return TraceContextElement(traceThreadLocal.get()?.copy())
124+
* }
125+
*
126+
* override fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext {
127+
* // Merge operation defines how to handle situations when both
128+
* // the parent coroutine has an element in the context and
129+
* // an element with the same key was also
130+
* // explicitly passed to the child coroutine.
131+
* // If merging does not require special behavior,
132+
* // the copy of the element can be returned.
133+
* return TraceContextElement(traceThreadLocal.get()?.copy())
134+
* }
135+
* }
136+
* ```
137+
*
138+
* A coroutine using this mechanism can safely call Java code that assumes the corresponding thread local element's
139+
* value is installed into the target thread local.
140+
*
141+
* ### Reentrancy and thread-safety
142+
*
143+
* Correct implementations of this interface must expect that calls to [restoreThreadContext]
144+
* may happen in parallel to the subsequent [updateThreadContext] and [restoreThreadContext] operations.
145+
*
146+
* Even though an element is copied for each child coroutine, an implementation should be able to handle the following
147+
* interleaving when a coroutine with the corresponding element is launched on a multithreaded dispatcher:
148+
*
149+
* ```
150+
* coroutine.updateThreadContext() // Thread #1
151+
* ... coroutine body ...
152+
* // suspension + immediate dispatch happen here
153+
* coroutine.updateThreadContext() // Thread #2, coroutine is already resumed
154+
* // ... coroutine body after suspension point on Thread #2 ...
155+
* coroutine.restoreThreadContext() // Thread #1, is invoked late because Thread #1 is slow
156+
* coroutine.restoreThreadContext() // Thread #2, may happen in parallel with the previous restore
157+
* ```
158+
*
159+
* All implementations of [CopyableThreadContextElement] should be thread-safe and guard their internal mutable state
160+
* within an element accordingly.
161+
*/
162+
@DelicateCoroutinesApi
163+
@ExperimentalCoroutinesApi
164+
public interface CopyableThreadContextElement<S> : ThreadContextElement<S> {
165+
166+
/**
167+
* Returns a [CopyableThreadContextElement] to replace `this` `CopyableThreadContextElement` in the child
168+
* coroutine's context that is under construction if the added context does not contain an element with the same [key].
169+
*
170+
* This function is called on the element each time a new coroutine inherits a context containing it,
171+
* and the returned value is folded into the context given to the child.
172+
*
173+
* Since this method is called whenever a new coroutine is launched in a context containing this
174+
* [CopyableThreadContextElement], implementations are performance-sensitive.
175+
*/
176+
public fun copyForChild(): CopyableThreadContextElement<S>
177+
178+
/**
179+
* Returns a [CopyableThreadContextElement] to replace `this` `CopyableThreadContextElement` in the child
180+
* coroutine's context that is under construction if the added context does contain an element with the same [key].
181+
*
182+
* This method is invoked on the original element, accepting as the parameter
183+
* the element that is supposed to overwrite it.
184+
*/
185+
public fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext
186+
}

Diff for: ‎kotlinx-coroutines-core/jsAndWasmShared/src/CoroutineContext.kt

+1-10
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,7 @@ import kotlin.coroutines.*
66
internal actual val DefaultDelay: Delay
77
get() = Dispatchers.Default as Delay
88

9-
public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
10-
val combined = coroutineContext + context
11-
return if (combined !== Dispatchers.Default && combined[ContinuationInterceptor] == null)
12-
combined + Dispatchers.Default else combined
13-
}
14-
15-
public actual fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext {
16-
return this + addedContext
17-
}
18-
199
// No debugging facilities on Wasm and JS
2010
internal actual fun Continuation<*>.toDebugString(): String = toString()
2111
internal actual val CoroutineContext.coroutineName: String? get() = null // not supported on Wasm and JS
12+
internal actual fun wrapContextWithDebug(context: CoroutineContext): CoroutineContext = context

Diff for: ‎kotlinx-coroutines-core/jvm/src/CoroutineContext.kt

+5-77
Original file line numberDiff line numberDiff line change
@@ -5,84 +5,12 @@ import kotlin.coroutines.*
55
import kotlin.coroutines.jvm.internal.CoroutineStackFrame
66

77
/**
8-
* Creates a context for a new coroutine. It installs [Dispatchers.Default] when no other dispatcher or
9-
* [ContinuationInterceptor] is specified and adds optional support for debugging facilities (when turned on)
10-
* and copyable-thread-local facilities on JVM.
11-
* See [DEBUG_PROPERTY_NAME] for description of debugging facilities on JVM.
8+
* Adds optional support for debugging facilities (when turned on)
9+
* and copyable-thread-local facilities on JVM.
10+
* See [DEBUG_PROPERTY_NAME] for description of debugging facilities on JVM.
1211
*/
13-
@ExperimentalCoroutinesApi
14-
public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
15-
val combined = foldCopies(coroutineContext, context, true)
16-
val debug = if (DEBUG) combined + CoroutineId(COROUTINE_ID.incrementAndGet()) else combined
17-
return if (combined !== Dispatchers.Default && combined[ContinuationInterceptor] == null)
18-
debug + Dispatchers.Default else debug
19-
}
20-
21-
/**
22-
* Creates a context for coroutine builder functions that do not launch a new coroutine, e.g. [withContext].
23-
* @suppress
24-
*/
25-
@InternalCoroutinesApi
26-
public actual fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext {
27-
/*
28-
* Fast-path: we only have to copy/merge if 'addedContext' (which typically has one or two elements)
29-
* contains copyable elements.
30-
*/
31-
if (!addedContext.hasCopyableElements()) return this + addedContext
32-
return foldCopies(this, addedContext, false)
33-
}
34-
35-
private fun CoroutineContext.hasCopyableElements(): Boolean =
36-
fold(false) { result, it -> result || it is CopyableThreadContextElement<*> }
37-
38-
/**
39-
* Folds two contexts properly applying [CopyableThreadContextElement] rules when necessary.
40-
* The rules are the following:
41-
* - If neither context has CTCE, the sum of two contexts is returned
42-
* - Every CTCE from the left-hand side context that does not have a matching (by key) element from right-hand side context
43-
* is [copied][CopyableThreadContextElement.copyForChild] if [isNewCoroutine] is `true`.
44-
* - Every CTCE from the left-hand side context that has a matching element in the right-hand side context is [merged][CopyableThreadContextElement.mergeForChild]
45-
* - Every CTCE from the right-hand side context that hasn't been merged is copied
46-
* - Everything else is added to the resulting context as is.
47-
*/
48-
private fun foldCopies(originalContext: CoroutineContext, appendContext: CoroutineContext, isNewCoroutine: Boolean): CoroutineContext {
49-
// Do we have something to copy left-hand side?
50-
val hasElementsLeft = originalContext.hasCopyableElements()
51-
val hasElementsRight = appendContext.hasCopyableElements()
52-
53-
// Nothing to fold, so just return the sum of contexts
54-
if (!hasElementsLeft && !hasElementsRight) {
55-
return originalContext + appendContext
56-
}
57-
58-
var leftoverContext = appendContext
59-
val folded = originalContext.fold<CoroutineContext>(EmptyCoroutineContext) { result, element ->
60-
if (element !is CopyableThreadContextElement<*>) return@fold result + element
61-
// Will this element be overwritten?
62-
val newElement = leftoverContext[element.key]
63-
// No, just copy it
64-
if (newElement == null) {
65-
// For 'withContext'-like builders we do not copy as the element is not shared
66-
return@fold result + if (isNewCoroutine) element.copyForChild() else element
67-
}
68-
// Yes, then first remove the element from append context
69-
leftoverContext = leftoverContext.minusKey(element.key)
70-
// Return the sum
71-
@Suppress("UNCHECKED_CAST")
72-
return@fold result + (element as CopyableThreadContextElement<Any?>).mergeForChild(newElement)
73-
}
74-
75-
if (hasElementsRight) {
76-
leftoverContext = leftoverContext.fold<CoroutineContext>(EmptyCoroutineContext) { result, element ->
77-
// We're appending new context element -- we have to copy it, otherwise it may be shared with others
78-
if (element is CopyableThreadContextElement<*>) {
79-
return@fold result + element.copyForChild()
80-
}
81-
return@fold result + element
82-
}
83-
}
84-
return folded + leftoverContext
85-
}
12+
internal actual fun wrapContextWithDebug(context: CoroutineContext): CoroutineContext =
13+
if (DEBUG) context + CoroutineId(COROUTINE_ID.incrementAndGet()) else context
8614

8715
internal actual val CoroutineContext.coroutineName: String? get() {
8816
if (!DEBUG) return null

0 commit comments

Comments
 (0)