com.firefly.kotlin.ext.common.CoroutineLocal.kt Maven / Gradle / Ivy
package com.firefly.kotlin.ext.common
import kotlinx.coroutines.*
import java.util.concurrent.ConcurrentHashMap
import kotlin.coroutines.ContinuationInterceptor
/**
* Retain a value in the coroutine scope.
*
* @author Pengtao Qiu
*/
class CoroutineLocal {
private val threadLocal = ThreadLocal()
/**
* Convert a value to the coroutine context element.
*
* @param value A value runs through in the coroutine scope.
* @return The coroutine context element.
*/
fun asElement(value: T) = threadLocal.asContextElement(value)
/**
* Get the value in the coroutine scope.
*
* @return The value in the coroutine scope.
*/
fun get(): T? = threadLocal.get()
/**
* Set the value in the coroutine scope.
*
* @param value The value in the coroutine scope.
*/
fun set(value: T?) = threadLocal.set(value)
}
/**
* Retain the attributes in the coroutine scope.
*
* @author Pengtao Qiu
*/
object CoroutineLocalContext {
private val ctx: CoroutineLocal> by lazy {
CoroutineLocal>()
}
/**
* Convert the attributes to the coroutine context element.
*
* @param attributes The attributes runs through in the coroutine scope.
* @return The coroutine context element.
*/
fun asElement(attributes: MutableMap): ThreadContextElement> {
return if (attributes is ConcurrentHashMap) {
ctx.asElement(attributes)
} else {
ctx.asElement(ConcurrentHashMap(attributes))
}
}
/**
* Merge the attributes into the parent coroutine context element.
*
* @param attributes The attributes are merged into the parent coroutine context element.
* @return The coroutine context element.
*/
fun inheritParentElement(attributes: MutableMap? = null): ThreadContextElement> {
val parentAttributes = getAttributes()
val attrs = if (parentAttributes.isNullOrEmpty()) {
attributes ?: ConcurrentHashMap()
} else {
if (!attributes.isNullOrEmpty()) {
parentAttributes.putAll(attributes)
}
parentAttributes
}
return asElement(attrs)
}
/**
* Get the current attributes.
*
* @return The current attributes.
*/
fun getAttributes(): MutableMap? = ctx.get()
/**
* Get a attribute in the current coroutine scope.
*
* @param name The attribute name.
* @return A attribute in the current coroutine scope.
*/
inline fun getAttr(name: String): T? {
val value = getAttributes()?.get(name)
return if (value == null) null else value as T
}
/**
* Get a attribute in the current coroutine scope, if the value is null return the default value.
*
* @param name The attribute name.
* @param func Get the default value lazily.
* @return A attribute in the current coroutine scope or the default value.
*/
inline fun getAttrOrDefault(name: String, func: (String) -> T): T {
val value = getAttributes()?.get(name)
return if (value == null) func.invoke(name) else value as T
}
/**
* Set a attribute in the current coroutine scope.
*
* @param name The attribute name.
* @param value The attribute value.
* @return The old value in the current coroutine scope.
*/
inline fun setAttr(name: String, value: T): T? {
val oldValue = getAttributes()?.put(name, value)
return if (oldValue == null) null else oldValue as T
}
/**
* If the specified attribute name is not already associated with a value (or is mapped
* to null), attempts to compute its value using the given mapping
* function and enters it into this map unless null.
*
* @param name The attribute name.
* @param func The mapping function.
* @return The value in the current coroutine scope.
*/
inline fun computeIfAbsent(name: String, crossinline func: (String) -> T): T? {
val value = getAttributes()?.computeIfAbsent(name) {
func.invoke(it)
}
return if (value == null) null else value as T
}
}
/**
* Starts a asynchronous task and inherits parent coroutine local attributes.
*
* @param context Additional to [CoroutineScope.coroutineContext] context of the coroutine.
* @param attributes The attributes are merged into the parent coroutine context element.
* @param block The coroutine code.
* @return The deferred task result.
*/
fun asyncTraceable(
context: ContinuationInterceptor = CoroutineDispatchers.computation,
attributes: MutableMap? = null,
block: suspend CoroutineScope.() -> T
): Deferred {
return GlobalScope.async(context + CoroutineLocalContext.inheritParentElement(attributes)) { block.invoke(this) }
}
/**
* Starts a asynchronous task without the return value and inherits parent coroutine local attributes.
*
* @param context Additional to [CoroutineScope.coroutineContext] context of the coroutine.
* @param attributes The attributes are merged into the parent coroutine context element.
* @param block The coroutine code.
* @return The current job.
*/
fun launchTraceable(
context: ContinuationInterceptor = CoroutineDispatchers.computation,
attributes: MutableMap? = null,
block: suspend CoroutineScope.() -> Unit
): Job {
return GlobalScope.launch(context + CoroutineLocalContext.inheritParentElement(attributes)) { block.invoke(this) }
}
/**
* Starts a asynchronous task waiting the result and inherits parent coroutine local attributes.
*
* @param context Additional to [CoroutineScope.coroutineContext] context of the coroutine.
* @param attributes The attributes are merged into the parent coroutine context element.
* @param block The coroutine code.
* @return The task result.
*/
suspend fun withContextTraceable(
context: ContinuationInterceptor = CoroutineDispatchers.computation,
attributes: MutableMap? = null,
block: suspend CoroutineScope.() -> T
): T {
return withContext(context + CoroutineLocalContext.inheritParentElement(attributes)) { block.invoke(this) }
}