propagate transport coroutine context tests

Author: yakivy

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

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+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
+import kotlinx.rpc.withService
+import kotlin.coroutines.CoroutineContext
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class CoroutineContextPropagationTest {
+    private val rpcServerConfig = rpcServerConfig {
+        serialization {
+            json()
+        }
+    }
+    private val rpcClientConfig = rpcClientConfig {
+        serialization {
+            json {
+                ignoreUnknownKeys = true
+            }
+        }
+    }
+
+    data class CoroutineElement(val value: String) : CoroutineContext.Element {
+        object Key : CoroutineContext.Key<CoroutineElement>
+
+        override val key: CoroutineContext.Key<*> = Key
+    }
+
+    @Test
+    fun test() = runTest {
+        var actualContext: CoroutineElement? = null
+        val transport = LocalTransport(CoroutineElement("transport"))
+        val server = KrpcTestServer(rpcServerConfig, transport.server)
+        val client = KrpcTestClient(rpcClientConfig, transport.client)
+        withContext(CoroutineElement("server")) {
+            server.registerService(Echo::class) {
+                object : Echo {
+                    override suspend fun echo(message: String): String = run {
+                        actualContext = currentCoroutineContext().get(CoroutineElement.Key)
+                        "response"
+                    }
+                }
+            }
+        }
+        withContext(CoroutineElement("client")) {
+            client.withService(Echo::class).echo("request")
+        }
+        assertEquals(CoroutineElement("transport"), actualContext)
+    }
+}

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

@@ -16,9 +16,10 @@ import kotlin.coroutines.CoroutineContext
 import kotlin.time.Clock
 import kotlin.time.ExperimentalTime
 
-class LocalTransport(parentScope: CoroutineScope? = null) : CoroutineScope {
-    override val coroutineContext = parentScope?.run { SupervisorJob(coroutineContext.job) }
-        ?: SupervisorJob()
+class LocalTransport(
+    parentContext: CoroutineContext? = null,
+) : CoroutineScope {
+    override val coroutineContext = SupervisorJob(parentContext?.job)
 
     private val clientIncoming = Channel<KrpcTransportMessage>()
     private val serverIncoming = Channel<KrpcTransportMessage>()
@@ -27,7 +28,10 @@ class LocalTransport(parentScope: CoroutineScope? = null) : CoroutineScope {
     val lastMessageSentOnServer = atomic(0L)
 
     val client: KrpcTransport = object : KrpcTransport {
-        override val coroutineContext: CoroutineContext = Job(this@LocalTransport.coroutineContext.job)
+        override val coroutineContext: CoroutineContext = Job(this@LocalTransport.coroutineContext.job).let {
+            if(parentContext != null) parentContext + it
+            else it
+        }
 
         @OptIn(ExperimentalTime::class)
         override suspend fun send(message: KrpcTransportMessage) {
@@ -41,7 +45,10 @@ class LocalTransport(parentScope: CoroutineScope? = null) : CoroutineScope {
     }
 
     val server: KrpcTransport = object : KrpcTransport {
-        override val coroutineContext: CoroutineContext = Job(this@LocalTransport.coroutineContext)
+        override val coroutineContext: CoroutineContext = Job(this@LocalTransport.coroutineContext.job).let {
+            if(parentContext != null) parentContext + it
+            else it
+        }
 
         @OptIn(ExperimentalTime::class)
         override suspend fun send(message: KrpcTransportMessage) {

krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/cancellation/CancellationToolkit.kt

@@ -50,7 +50,7 @@ class CancellationToolkit(scope: CoroutineScope) : CoroutineScope by scope {
         }
     }
 
-    val transport = LocalTransport(this)
+    val transport = LocalTransport(this.coroutineContext.job)
 
     val client by lazy {
         KrpcTestClient(rpcClientConfig {

krpc/krpc-test/src/jvmTest/kotlin/kotlinx/rpc/krpc/test/api/WireSamplingTestScope.kt

@@ -229,7 +229,7 @@ class WireSamplingTestScope(private val sampleName: String, scope: TestScope) :
 }
 
 private class WireToolkit(scope: CoroutineScope, format: SamplingFormat, val logger: RpcInternalCommonLogger? = null) {
-    val transport = LocalTransport(scope)
+    val transport = LocalTransport(scope.coroutineContext.job)
 
     val client by lazy {
         KrpcTestClient(rpcClientConfig {