diff --git a/krpc/krpc-server/src/commonMain/kotlin/kotlinx/rpc/krpc/server/KrpcServer.kt b/krpc/krpc-server/src/commonMain/kotlin/kotlinx/rpc/krpc/server/KrpcServer.kt index 898f3a266..1e45590a6 100644 --- a/krpc/krpc-server/src/commonMain/kotlin/kotlinx/rpc/krpc/server/KrpcServer.kt +++ b/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()) diff --git a/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/CoroutineContextPropagationTest.kt b/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/CoroutineContextPropagationTest.kt new file mode 100644 index 000000000..a0a3704a5 --- /dev/null +++ b/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/CoroutineContextPropagationTest.kt @@ -0,0 +1,59 @@ +/* + * 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.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 + + 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) + } +} diff --git a/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/LocalTransport.kt b/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/LocalTransport.kt index fb7c0b565..5f9073d46 100644 --- a/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/LocalTransport.kt +++ b/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?.get(Job)) private val clientIncoming = Channel() private val serverIncoming = Channel() @@ -27,7 +28,9 @@ 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 +44,9 @@ 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) { diff --git a/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/cancellation/CancellationToolkit.kt b/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/cancellation/CancellationToolkit.kt index 187df520c..3dd848f09 100644 --- a/krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/cancellation/CancellationToolkit.kt +++ b/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 { diff --git a/krpc/krpc-test/src/jvmTest/kotlin/kotlinx/rpc/krpc/test/api/WireSamplingTestScope.kt b/krpc/krpc-test/src/jvmTest/kotlin/kotlinx/rpc/krpc/test/api/WireSamplingTestScope.kt index df05355b0..8e0210f83 100644 --- a/krpc/krpc-test/src/jvmTest/kotlin/kotlinx/rpc/krpc/test/api/WireSamplingTestScope.kt +++ b/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 {