propagate transport coroutine context

Author: yakivy

krpc/krpc-server/src/commonMain/kotlin/kotlinx/rpc/krpc/server/KrpcServer.kt

@@ -60,7 +60,9 @@ public abstract class KrpcServer(
      */
 
     @InternalRpcApi
-    public val internalScope: CoroutineScope = CoroutineScope(SupervisorJob(transport.coroutineContext.job))
+    public val internalScope: CoroutineScope = CoroutineScope(
+        transport.coroutineContext + SupervisorJob(transport.coroutineContext.job)
+    )
 
     private val logger = RpcInternalCommonLogger.logger(rpcInternalObjectId())
 

krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/CoroutineContextPropagationTest.kt

@@ -7,6 +7,7 @@ package kotlinx.rpc.krpc.test
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.withContext
 import kotlinx.rpc.krpc.rpcClientConfig
 import kotlinx.rpc.krpc.rpcServerConfig
 import kotlinx.rpc.krpc.serialization.json.json
@@ -29,7 +30,7 @@ class CoroutineContextPropagationTest {
         }
     }
 
-    object CoroutineElement : CoroutineContext.Element {
+    data class CoroutineElement(val value: String) : CoroutineContext.Element {
         object Key : CoroutineContext.Key<CoroutineElement>
 
         override val key: CoroutineContext.Key<*> = Key
@@ -38,20 +39,22 @@ class CoroutineContextPropagationTest {
     @Test
     fun test() = runTest {
         var actualContext: CoroutineElement? = null
-        val transport = LocalTransport(transportContext = CoroutineElement)
+        val transport = LocalTransport(transportContext = CoroutineElement("transport"))
         val server = KrpcTestServer(rpcServerConfig, transport.server)
         val client = KrpcTestClient(rpcClientConfig, transport.client)
-        server.registerService(Echo::class) {
-            object : Echo {
-                override suspend fun echo(message: String): String = run {
-                    actualContext = currentCoroutineContext().get(CoroutineElement.Key)
-                    "response"
+        withContext(CoroutineElement("server")) {
+            server.registerService(Echo::class) {
+                object : Echo {
+                    override suspend fun echo(message: String): String = run {
+                        actualContext = currentCoroutineContext().get(CoroutineElement.Key)
+                        "response"
+                    }
                 }
-
-                override val coroutineContext: CoroutineContext = it
             }
         }
-        client.withService(Echo::class).echo("request")
-        assertEquals(actualContext, CoroutineElement)
+        withContext(CoroutineElement("client")) {
+            client.withService(Echo::class).echo("request")
+        }
+        assertEquals(CoroutineElement("transport"), actualContext)
     }
 }