-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Introduce CoroutineContextThreadLocal API to integrate with thread-lo…
…cal sensitive code Fixes #119
- Loading branch information
Showing
8 changed files
with
225 additions
and
25 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
93 changes: 93 additions & 0 deletions
93
core/kotlinx-coroutines-core/src/CoroutineContextThreadLocal.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,93 @@ | ||
/* | ||
* Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license. | ||
*/ | ||
|
||
package kotlinx.coroutines.experimental | ||
|
||
import kotlin.coroutines.experimental.* | ||
|
||
/** | ||
* An extension point to define elements in [CoroutineContext] that are installed into thread local | ||
* variables every time the coroutine from the specified context in resumed on a thread. | ||
* | ||
* Implementations on this interface are looked up via [java.util.ServiceLoader]. | ||
* | ||
* Example usage looks like this: | ||
* | ||
* ``` | ||
* // declare custom coroutine context element | ||
* class MyElement : AbstractCoroutineContextElement(Key) { | ||
* companion object Key : CoroutineContext.Key<MyElement> | ||
* // some state is kept here | ||
* } | ||
* | ||
* // declare thread local variable | ||
* private val myThreadLocal = ThreadLocal<MyElement?>() | ||
* | ||
* // declare extension point implementation | ||
* class MyCoroutineContextThreadLocal : CoroutineContextThreadLocal<MyElement?> { | ||
* // this is invoked before coroutine is resumed on current thread | ||
* override fun updateThreadContext(context: CoroutineContext): MyElement? { | ||
* val oldValue = myThreadLocal.get() | ||
* myThreadLocal.set(context[MyElement]) | ||
* return oldValue | ||
* } | ||
* | ||
* // this is invoked after coroutine has suspended on current thread | ||
* override fun restoreThreadContext(context: CoroutineContext, oldValue: MyElement?) { | ||
* myThreadLocal.set(oldValue) | ||
* } | ||
* } | ||
* ``` | ||
* | ||
* Now, `MyCoroutineContextThreadLocal` fully qualified class named shall be registered via | ||
* `META-INF/services/kotlinx.coroutines.experimental.CoroutineContextThreadLocal` file. | ||
*/ | ||
public interface CoroutineContextThreadLocal<T> { | ||
/** | ||
* Updates context of the current thread. | ||
* This function is invoked before the coroutine in the specified [context] is resumed in the current thread. | ||
* The result of this function is the old value that will be passed to [restoreThreadContext]. | ||
*/ | ||
public fun updateThreadContext(context: CoroutineContext): T | ||
|
||
/** | ||
* Restores context of the current thread. | ||
* This function is invoked after the coroutine in the specified [context] is suspended in the current thread. | ||
* The value of [oldValue] is the result of the previous invocation of [updateThreadContext]. | ||
*/ | ||
public fun restoreThreadContext(context: CoroutineContext, oldValue: T) | ||
} | ||
|
||
/** | ||
* This class is used when multiple [CoroutineContextThreadLocal] are installed. | ||
*/ | ||
internal class CoroutineContextThreadLocalList( | ||
private val impls: Array<CoroutineContextThreadLocal<Any?>> | ||
) : CoroutineContextThreadLocal<Any?> { | ||
init { | ||
require(impls.size > 1) | ||
} | ||
|
||
private val threadLocalStack = ThreadLocal<ArrayList<Any?>?>() | ||
|
||
override fun updateThreadContext(context: CoroutineContext): Any? { | ||
val stack = threadLocalStack.get() ?: ArrayList<Any?>().also { | ||
threadLocalStack.set(it) | ||
} | ||
val lastIndex = impls.lastIndex | ||
for (i in 0 until lastIndex) { | ||
stack.add(impls[i].updateThreadContext(context)) | ||
} | ||
return impls[lastIndex].updateThreadContext(context) | ||
} | ||
|
||
override fun restoreThreadContext(context: CoroutineContext, oldValue: Any?) { | ||
val stack = threadLocalStack.get()!! // must be there | ||
val lastIndex = impls.lastIndex | ||
impls[lastIndex].restoreThreadContext(context, oldValue) | ||
for (i in lastIndex - 1 downTo 0) { | ||
impls[i].restoreThreadContext(context, stack.removeAt(stack.lastIndex)) | ||
} | ||
} | ||
} |
2 changes: 2 additions & 0 deletions
2
...t-resources/META-INF/services/kotlinx.coroutines.experimental.CoroutineContextThreadLocal
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
kotlinx.coroutines.experimental.MyCoroutineContextThreadLocal | ||
kotlinx.coroutines.experimental.ValidatingCoroutineContextThreadLocal |
59 changes: 59 additions & 0 deletions
59
core/kotlinx-coroutines-core/test/CoroutineContextThreadLocalTest.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,59 @@ | ||
/* | ||
* Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license. | ||
*/ | ||
|
||
package kotlinx.coroutines.experimental | ||
|
||
import org.junit.Test | ||
import kotlin.coroutines.experimental.* | ||
import kotlin.test.* | ||
|
||
class CoroutineContextThreadLocalTest : TestBase() { | ||
@Test | ||
fun testExample() = runTest { | ||
val mainDispatcher = coroutineContext[ContinuationInterceptor]!! | ||
val mainThread = Thread.currentThread() | ||
val element = MyElement() | ||
assertEquals(null, myThreadLocal.get()) | ||
val job = launch(element) { | ||
assertTrue(mainThread != Thread.currentThread()) | ||
assertSame(element, coroutineContext[MyElement]) | ||
assertSame(element, myThreadLocal.get()) | ||
withContext(mainDispatcher) { | ||
assertSame(mainThread, Thread.currentThread()) | ||
assertSame(element, coroutineContext[MyElement]) | ||
assertSame(element, myThreadLocal.get()) | ||
} | ||
assertTrue(mainThread != Thread.currentThread()) | ||
assertSame(element, coroutineContext[MyElement]) | ||
assertSame(element, myThreadLocal.get()) | ||
} | ||
assertEquals(null, myThreadLocal.get()) | ||
job.join() | ||
assertEquals(null, myThreadLocal.get()) | ||
} | ||
} | ||
|
||
// declare custom coroutine context element | ||
class MyElement : AbstractCoroutineContextElement(Key) { | ||
companion object Key : CoroutineContext.Key<MyElement> | ||
// some state is kept here | ||
} | ||
|
||
// declare thread local variable | ||
private val myThreadLocal = ThreadLocal<MyElement?>() | ||
|
||
// declare extension point implementation | ||
class MyCoroutineContextThreadLocal : CoroutineContextThreadLocal<MyElement?> { | ||
// this is invoked before coroutine is resumed on current thread | ||
override fun updateThreadContext(context: CoroutineContext): MyElement? { | ||
val oldValue = myThreadLocal.get() | ||
myThreadLocal.set(context[MyElement]) | ||
return oldValue | ||
} | ||
|
||
// this is invoked after coroutine has suspended on current thread | ||
override fun restoreThreadContext(context: CoroutineContext, oldValue: MyElement?) { | ||
myThreadLocal.set(oldValue) | ||
} | ||
} |
30 changes: 30 additions & 0 deletions
30
core/kotlinx-coroutines-core/test/ValidatingCoroutineContextThreadLocal.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,30 @@ | ||
/* | ||
* Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license. | ||
*/ | ||
|
||
package kotlinx.coroutines.experimental | ||
|
||
import kotlin.coroutines.experimental.* | ||
|
||
private val currentCoroutineId = ThreadLocal<CoroutineId?>() | ||
|
||
internal class ValidatingCoroutineContextThreadLocal : CoroutineContextThreadLocal<CoroutineId?> { | ||
override fun updateThreadContext(context: CoroutineContext): CoroutineId? { | ||
val id = context[CoroutineId] ?: error("Tests should be run in debug mode (enable assertions?)") | ||
val top = currentCoroutineId.get() | ||
require( top != id) { | ||
"Thread ${Thread.currentThread().name} already has coroutine context for coroutine $context" | ||
} | ||
currentCoroutineId.set(id) | ||
return top | ||
} | ||
|
||
override fun restoreThreadContext(context: CoroutineContext, oldValue: CoroutineId?) { | ||
val id = context[CoroutineId] | ||
val top = currentCoroutineId.get() | ||
require(top == id) { | ||
"Thread ${Thread.currentThread().name} does not have coroutine context for coroutine $context, but has for coroutine id $top" | ||
} | ||
currentCoroutineId.set(oldValue) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters