add CoroutineScopeContext

utils/src/main/kotlin/cn/tursom/utils/coroutine/CoroutineLocalContext.kt

@@ -1,8 +1,6 @@
 package cn.tursom.utils.coroutine
 
-import kotlinx.coroutines.*
 import kotlin.coroutines.CoroutineContext
-import kotlin.coroutines.EmptyCoroutineContext
 
 open class CoroutineLocalContext(
   private val map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4)
@@ -10,34 +8,8 @@ open class CoroutineLocalContext(
   override val key: CoroutineContext.Key<*> get() = Companion
 
   override fun toString(): String {
-    return map.toString()
+    return "CoroutineLocalContext$map"
   }
 
   companion object : CoroutineContext.Key<CoroutineLocalContext>
-}
-
-fun CoroutineScope.launchWithCoroutineLocalContext(
-  context: CoroutineContext = EmptyCoroutineContext,
-  start: CoroutineStart = CoroutineStart.DEFAULT,
-  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
-  block: suspend CoroutineScope.() -> Unit
-): Job {
-  return launch(context + CoroutineLocalContext(map), start, block)
-}
-
-@Suppress("DeferredIsResult")
-fun <T> CoroutineScope.asyncWithCoroutineLocalContext(
-  context: CoroutineContext = EmptyCoroutineContext,
-  start: CoroutineStart = CoroutineStart.DEFAULT,
-  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
-  block: suspend CoroutineScope.() -> T
-): Deferred<T> {
-  return async(context + CoroutineLocalContext(map), start, block)
-}
-
-suspend fun <T> withCoroutineLocalContext(
-  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
-  block: suspend CoroutineScope.() -> T
-): T {
-  return withContext(CoroutineLocalContext(map), block)
 }

utils/src/main/kotlin/cn/tursom/utils/coroutine/CoroutineScopeContext.kt

@@ -0,0 +1,26 @@
+package cn.tursom.utils.coroutine
+
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.GlobalScope
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.coroutineContext
+
+class CoroutineScopeContext(
+  var coroutineScope: CoroutineScope
+) : CoroutineContext.Element {
+  override val key: CoroutineContext.Key<*> get() = Companion
+  override fun toString(): String = "CoroutineScopeContext(coroutineScope=$coroutineScope)"
+
+  companion object : CoroutineContext.Key<CoroutineScopeContext>, CoroutineLocal<CoroutineScope>() {
+    override suspend fun get(): CoroutineScope = coroutineContext[this]?.coroutineScope ?: super.get() ?: GlobalScope
+    override suspend fun set(value: CoroutineScope): Boolean {
+      val coroutineScopeContext = coroutineContext[this]
+      return if (coroutineScopeContext != null) {
+        coroutineScopeContext.coroutineScope = value
+        true
+      } else {
+        super.set(value)
+      }
+    }
+  }
+}

utils/src/main/kotlin/cn/tursom/utils/coroutine/utils.kt

@@ -0,0 +1,77 @@
+package cn.tursom.utils.coroutine
+
+import kotlinx.coroutines.*
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
+
+
+fun CoroutineScope.launchWithCoroutineLocalContext(
+  context: CoroutineContext = EmptyCoroutineContext,
+  start: CoroutineStart = CoroutineStart.DEFAULT,
+  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
+  block: suspend CoroutineScope.() -> Unit
+): Job {
+  return launch(context + CoroutineLocalContext(map), start, block)
+}
+
+@Suppress("DeferredIsResult")
+fun <T> CoroutineScope.asyncWithCoroutineLocalContext(
+  context: CoroutineContext = EmptyCoroutineContext,
+  start: CoroutineStart = CoroutineStart.DEFAULT,
+  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
+  block: suspend CoroutineScope.() -> T
+): Deferred<T> {
+  return async(context + CoroutineLocalContext(map), start, block)
+}
+
+suspend fun <T> withCoroutineLocalContext(
+  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
+  block: suspend CoroutineScope.() -> T
+): T {
+  return withContext(CoroutineLocalContext(map), block)
+}
+
+@Throws(InterruptedException::class)
+fun <T> runBlockingWithCoroutineLocalContext(
+  context: CoroutineContext = EmptyCoroutineContext,
+  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
+  block: suspend CoroutineScope.() -> T
+): T {
+  return runBlocking(context + CoroutineLocalContext(map), block)
+}
+
+fun CoroutineScope.launchWithCoroutineScopeContext(
+  context: CoroutineContext = EmptyCoroutineContext,
+  start: CoroutineStart = CoroutineStart.DEFAULT,
+  block: suspend CoroutineScope.() -> Unit
+): Job {
+  return launch(context + CoroutineScopeContext(this), start, block)
+}
+
+@Suppress("DeferredIsResult")
+fun <T> CoroutineScope.asyncWithCoroutineScopeContext(
+  context: CoroutineContext = EmptyCoroutineContext,
+  start: CoroutineStart = CoroutineStart.DEFAULT,
+  block: suspend CoroutineScope.() -> T
+): Deferred<T> {
+  return async(context + CoroutineScopeContext(this), start, block)
+}
+
+fun CoroutineScope.launchWithCoroutineLocalAndCoroutineScopeContext(
+  context: CoroutineContext = EmptyCoroutineContext,
+  start: CoroutineStart = CoroutineStart.DEFAULT,
+  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
+  block: suspend CoroutineScope.() -> Unit
+): Job {
+  return launch(context + CoroutineLocalContext(map) + CoroutineScopeContext(this), start, block)
+}
+
+@Suppress("DeferredIsResult")
+fun <T> CoroutineScope.asyncWithCoroutineLocalAndCoroutineScopeContext(
+  context: CoroutineContext = EmptyCoroutineContext,
+  start: CoroutineStart = CoroutineStart.DEFAULT,
+  map: MutableMap<CoroutineLocal<*>, Any?> = HashMap(4),
+  block: suspend CoroutineScope.() -> T
+): Deferred<T> {
+  return async(context + CoroutineLocalContext(map) + CoroutineScopeContext(this), start, block)
+}

utils/src/test/kotlin/cn/tursom/utils/coroutine/CoroutineLocalTest.kt

@@ -1,28 +1,25 @@
 package cn.tursom.utils.coroutine
 
 import cn.tursom.core.usingTime
-import kotlinx.coroutines.GlobalScope
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.delay
-import kotlinx.coroutines.launch
 
 val testCoroutineLocal = CoroutineLocal<Int>()
 val testCoroutineLocalList = Array(100000) {
   CoroutineLocal<Int>()
 }.asList()
 
-suspend fun main() {
-  println(testCoroutineLocal.set(1))
-  GlobalScope.launch {
-    coroutineScope { }
-    coroutineContext
+fun main() = runBlockingWithCoroutineLocalContext {
+  println(coroutineContext)
+  launchWithCoroutineLocalAndCoroutineScopeContext {
+    println(coroutineContext)
+    println(CoroutineScopeContext.get())
     testCoroutineLocal set 0
     testCoroutineLocal set 0
     testCoroutineLocalList.forEachIndexed { index, coroutineLocal ->
       coroutineLocal set index
     }
     println(usingTime {
-      repeat(1000000000) {
+      repeat(10000000) {
         testCoroutineLocal.get()
       }
     })