grpc: Start getRequestMetadata with coroutineContext

Jozott00 · Jozott00 · commit e9d54abaa3e6 · 2025-11-17T16:55:55.000+01:00
diff --git a/grpc/grpc-client/src/commonMain/kotlin/kotlinx/rpc/grpc/client/internal/GrpcChannel.kt b/grpc/grpc-client/src/commonMain/kotlin/kotlinx/rpc/grpc/client/internal/GrpcChannel.kt
@@ -4,10 +4,10 @@
 
 package kotlinx.rpc.grpc.client.internal
 
-import kotlinx.coroutines.CoroutineScope
 import kotlinx.rpc.grpc.client.GrpcCallOptions
 import kotlinx.rpc.grpc.descriptor.MethodDescriptor
 import kotlinx.rpc.internal.utils.InternalRpcApi
+import kotlin.coroutines.CoroutineContext
 
 @InternalRpcApi
 public expect abstract class GrpcChannel
@@ -16,4 +16,5 @@ public expect abstract class GrpcChannel
 public expect fun <RequestT, ResponseT> GrpcChannel.createCall(
     methodDescriptor: MethodDescriptor<RequestT, ResponseT>,
     callOptions: GrpcCallOptions,
+    coroutineContext: CoroutineContext,
 ): ClientCall<RequestT, ResponseT>
diff --git a/grpc/grpc-client/src/commonMain/kotlin/kotlinx/rpc/grpc/client/internal/suspendClientCalls.kt b/grpc/grpc-client/src/commonMain/kotlin/kotlinx/rpc/grpc/client/internal/suspendClientCalls.kt
@@ -6,6 +6,7 @@ package kotlinx.rpc.grpc.client.internal
 
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.channels.Channel
@@ -22,8 +23,6 @@ import kotlinx.rpc.grpc.Status
 import kotlinx.rpc.grpc.StatusCode
 import kotlinx.rpc.grpc.StatusException
 import kotlinx.rpc.grpc.client.ClientCallScope
-import kotlinx.rpc.grpc.client.EmptyCallCredentials
-import kotlinx.rpc.grpc.client.GrpcCallCredentials
 import kotlinx.rpc.grpc.client.GrpcCallOptions
 import kotlinx.rpc.grpc.client.GrpcClient
 import kotlinx.rpc.grpc.client.plus
@@ -202,7 +201,10 @@ private class ClientCallScopeImpl<Request, Response>(
         coroutineScope {
             // attach all call credentials set in the client to the call option ones.
             callOptions.callCredentials += client.callCredentials
-            val call = client.channel.platformApi.createCall(method, callOptions)
+            // pass this scope's context, so the suspend functions of call credentials launch in the same context.
+            // it will ensure that getRequestMetadata of the callCredential won't orphan if the call is cancelled
+            // by the user or by grpc-java
+            val call = client.channel.platformApi.createCall(method, callOptions, this.coroutineContext)
 
             /*
              * We maintain a buffer of size 1, so onMessage never has to block: it only gets called after
diff --git a/grpc/grpc-client/src/jvmMain/kotlin/kotlinx/rpc/grpc/client/GrpcCallOptions.jvm.kt b/grpc/grpc-client/src/jvmMain/kotlin/kotlinx/rpc/grpc/client/GrpcCallOptions.jvm.kt
@@ -6,11 +6,10 @@ package kotlinx.rpc.grpc.client
 
 import io.grpc.CallOptions
 import kotlinx.rpc.grpc.GrpcCompression
-import kotlinx.rpc.internal.utils.InternalRpcApi
 import java.util.concurrent.TimeUnit
+import kotlin.coroutines.CoroutineContext
 
-@InternalRpcApi
-public fun GrpcCallOptions.toJvm(): CallOptions {
+internal fun GrpcCallOptions.toJvm(coroutineContext: CoroutineContext): CallOptions {
     var default = CallOptions.DEFAULT
     if (timeout != null) {
         default = default.withDeadlineAfter(timeout!!.inWholeMilliseconds, TimeUnit.MILLISECONDS)
@@ -19,7 +18,7 @@ public fun GrpcCallOptions.toJvm(): CallOptions {
         default = default.withCompression(compression.name)
     }
     if (callCredentials !is EmptyCallCredentials) {
-        default = default.withCallCredentials(callCredentials.toJvm())
+        default = default.withCallCredentials(callCredentials.toJvm(coroutineContext))
     }
     return default
 }
diff --git a/grpc/grpc-client/src/jvmMain/kotlin/kotlinx/rpc/grpc/client/credentials.jvm.kt b/grpc/grpc-client/src/jvmMain/kotlin/kotlinx/rpc/grpc/client/credentials.jvm.kt
@@ -10,10 +10,10 @@ import io.grpc.InsecureChannelCredentials
 import io.grpc.SecurityLevel
 import io.grpc.TlsChannelCredentials
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.asCoroutineDispatcher
 import kotlinx.coroutines.launch
 import kotlinx.rpc.grpc.Status
 import java.util.concurrent.Executor
+import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.cancellation.CancellationException
 
 internal fun GrpcClientCredentials.toJvm(): ChannelCredentials {
@@ -46,32 +46,32 @@ private class JvmTlsCLientCredentialBuilder : GrpcTlsClientCredentialsBuilder {
     }
 }
 
-internal fun GrpcCallCredentials.toJvm(): CallCredentials {
+internal fun GrpcCallCredentials.toJvm(coroutineContext: CoroutineContext): CallCredentials {
     return object : CallCredentials() {
         override fun applyRequestMetadata(
             requestInfo: RequestInfo,
             appExecutor: Executor,
             applier: MetadataApplier
         ) {
-            val dispatcher = appExecutor.asCoroutineDispatcher()
-            CoroutineScope(dispatcher).launch {
-                try {
-                    check(!requiresTransportSecurity || requestInfo.securityLevel != SecurityLevel.NONE) {
+            CoroutineScope(coroutineContext).launch {
+                if (requiresTransportSecurity && requestInfo.securityLevel == SecurityLevel.NONE) {
+                    applier.fail(Status.UNAUTHENTICATED.withDescription(
                         "Established channel does not have a sufficient security level to transfer call credential."
-                    }
+                    ))
+                    return@launch
+                }
 
+                try {
                     val context = GrpcCallCredentials.Context(requestInfo.authority, requestInfo.methodDescriptor.fullMethodName)
                     val metadata = context.getRequestMetadata()
                     applier.apply(metadata)
-                } catch (e: Exception) {
+                } catch (err: Throwable) {
                     // we are not treating StatusExceptions separately, as currently there is no
                     // clean way to support the same feature on native. So for the sake of similar behavior,
                     // we always fail with Status.UNAVAILABLE. (KRPC-233)
-                    val description = "Getting metadata from call credentials failed with error: ${e.message}"
-                    applier.fail(Status.UNAVAILABLE.withDescription(description).withCause(e))
-                    if (e is CancellationException) {
-                        throw e
-                    }
+                    val description = "Getting metadata from call credentials failed with error: ${err.message}"
+                    applier.fail(Status.UNAVAILABLE.withDescription(description).withCause(err))
+                    if (err is CancellationException) throw err
                 }
             }
         }
diff --git a/grpc/grpc-client/src/jvmMain/kotlin/kotlinx/rpc/grpc/client/internal/GrpcChannel.jvm.kt b/grpc/grpc-client/src/jvmMain/kotlin/kotlinx/rpc/grpc/client/internal/GrpcChannel.jvm.kt
@@ -9,6 +9,7 @@ import kotlinx.rpc.grpc.client.GrpcCallOptions
 import kotlinx.rpc.grpc.client.toJvm
 import kotlinx.rpc.grpc.descriptor.MethodDescriptor
 import kotlinx.rpc.internal.utils.InternalRpcApi
+import kotlin.coroutines.CoroutineContext
 
 @InternalRpcApi
 public actual typealias GrpcChannel = Channel
@@ -17,6 +18,7 @@ public actual typealias GrpcChannel = Channel
 public actual fun <RequestT, ResponseT> GrpcChannel.createCall(
     methodDescriptor: MethodDescriptor<RequestT, ResponseT>,
     callOptions: GrpcCallOptions,
+    coroutineContext: CoroutineContext,
 ): ClientCall<RequestT, ResponseT> {
-    return this.newCall(methodDescriptor, callOptions.toJvm())
+    return this.newCall(methodDescriptor, callOptions.toJvm(coroutineContext))
 }
diff --git a/grpc/grpc-client/src/nativeMain/kotlin/kotlinx/rpc/grpc/client/GrpcCallCredentials.native.kt b/grpc/grpc-client/src/nativeMain/kotlin/kotlinx/rpc/grpc/client/GrpcCallCredentials.native.kt
@@ -20,12 +20,13 @@ import kotlinx.rpc.grpc.internal.toRaw
 import kotlinx.rpc.grpc.statusCode
 import libkgrpc.*
 import platform.posix.size_tVar
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.coroutineContext
 
 // Stable reference holder for Kotlin objects
 private class CredentialsPluginState(
     val kotlinCreds: GrpcCallCredentials,
-    val parentJob: Job,
-    val coroutineDispatcher: CoroutineDispatcher
+    val coroutineContext: CoroutineContext,
 )
 
 private fun getMetadataCallback(
@@ -62,7 +63,7 @@ private fun getMetadataCallback(
         }
     }
 
-    val scope = CoroutineScope(Job(pluginState.parentJob) + pluginState.coroutineDispatcher)
+    val scope = CoroutineScope(pluginState.coroutineContext)
 
     // Launch coroutine to call a suspend function asynchronously
     scope.launch {
@@ -120,11 +121,10 @@ private fun destroyCallback(state: COpaquePointer?) {
 }
 
 internal fun GrpcCallCredentials.createRaw(
-    parentJob: Job,
-    coroutineDispatcher: CoroutineDispatcher
+    coroutineContext: CoroutineContext,
 ): CPointer<grpc_call_credentials>? = memScoped {
     // Create a stable reference to keep the Kotlin object alive
-    val pluginState = CredentialsPluginState(this@createRaw, parentJob, coroutineDispatcher)
+    val pluginState = CredentialsPluginState(this@createRaw, coroutineContext)
     val stableRef = StableRef.create(pluginState)
 
     // Create plugin structure
diff --git a/grpc/grpc-client/src/nativeMain/kotlin/kotlinx/rpc/grpc/client/internal/GrpcChannel.native.kt b/grpc/grpc-client/src/nativeMain/kotlin/kotlinx/rpc/grpc/client/internal/GrpcChannel.native.kt
@@ -9,19 +9,22 @@ import kotlinx.rpc.grpc.GrpcMetadata
 import kotlinx.rpc.grpc.client.GrpcCallOptions
 import kotlinx.rpc.grpc.descriptor.MethodDescriptor
 import kotlinx.rpc.internal.utils.InternalRpcApi
+import kotlin.coroutines.CoroutineContext
 
 @InternalRpcApi
 public actual abstract class GrpcChannel {
     public abstract fun <RequestT, ResponseT> newCall(
         methodDescriptor: MethodDescriptor<RequestT, ResponseT>,
         callOptions: GrpcCallOptions,
+        coroutineContext: CoroutineContext
     ): ClientCall<RequestT, ResponseT>
 }
 
 @InternalRpcApi
 public actual fun <RequestT, ResponseT> GrpcChannel.createCall(
     methodDescriptor: MethodDescriptor<RequestT, ResponseT>,
     callOptions: GrpcCallOptions,
+    coroutineContext: CoroutineContext,
 ): ClientCall<RequestT, ResponseT> {
-    return this.newCall(methodDescriptor, callOptions)
+    return this.newCall(methodDescriptor, callOptions, coroutineContext)
 }
diff --git a/grpc/grpc-client/src/nativeMain/kotlin/kotlinx/rpc/grpc/client/internal/NativeClientCall.kt b/grpc/grpc-client/src/nativeMain/kotlin/kotlinx/rpc/grpc/client/internal/NativeClientCall.kt
@@ -24,10 +24,6 @@ import kotlinx.cinterop.value
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableJob
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.coroutineScope
-import kotlinx.coroutines.launch
 import kotlinx.rpc.grpc.GrpcMetadata
 import kotlinx.rpc.grpc.Status
 import kotlinx.rpc.grpc.StatusCode
@@ -43,14 +39,9 @@ import kotlinx.rpc.grpc.internal.toKotlin
 import kotlinx.rpc.protobuf.input.stream.asInputStream
 import kotlinx.rpc.protobuf.input.stream.asSource
 import kotlinx.rpc.grpc.GrpcCompression
-import kotlinx.rpc.grpc.StatusException
 import kotlinx.rpc.grpc.client.EmptyCallCredentials
-import kotlinx.rpc.grpc.client.GrpcCallCredentials
 import kotlinx.rpc.grpc.client.GrpcCallOptions
 import kotlinx.rpc.grpc.client.createRaw
-import kotlinx.rpc.grpc.internal.toRaw
-import kotlinx.rpc.grpc.merge
-import kotlinx.rpc.grpc.statusCode
 import libkgrpc.GRPC_OP_RECV_INITIAL_METADATA
 import libkgrpc.GRPC_OP_RECV_MESSAGE
 import libkgrpc.GRPC_OP_RECV_STATUS_ON_CLIENT
@@ -65,14 +56,14 @@ import libkgrpc.grpc_call_credentials_release
 import libkgrpc.grpc_call_error
 import libkgrpc.grpc_call_set_credentials
 import libkgrpc.grpc_call_unref
-import libkgrpc.grpc_channel_credentials_release
 import libkgrpc.grpc_metadata_array
 import libkgrpc.grpc_metadata_array_destroy
 import libkgrpc.grpc_metadata_array_init
 import libkgrpc.grpc_op
 import libkgrpc.grpc_slice
 import libkgrpc.grpc_slice_unref
 import libkgrpc.grpc_status_code
+import kotlin.coroutines.CoroutineContext
 import kotlin.experimental.ExperimentalNativeApi
 import kotlin.native.ref.createCleaner
 
@@ -83,6 +74,7 @@ internal class NativeClientCall<Request, Response>(
     private val methodDescriptor: MethodDescriptor<Request, Response>,
     private val callOptions: GrpcCallOptions,
     private val callJob: CompletableJob,
+    private val coroutineContext: CoroutineContext,
 ) : ClientCall<Request, Response>() {
 
     @Suppress("unused")
@@ -91,7 +83,7 @@ internal class NativeClientCall<Request, Response>(
     }
 
     private val rawCallCredentials = callOptions.callCredentials.let {
-        if (it is EmptyCallCredentials) null else it.createRaw(callJob, Dispatchers.Default)
+        if (it is EmptyCallCredentials) null else it.createRaw(coroutineContext)
     }
 
     @Suppress("unused")
diff --git a/grpc/grpc-client/src/nativeMain/kotlin/kotlinx/rpc/grpc/client/internal/NativeManagedChannel.kt b/grpc/grpc-client/src/nativeMain/kotlin/kotlinx/rpc/grpc/client/internal/NativeManagedChannel.kt
@@ -42,6 +42,7 @@ import libkgrpc.grpc_channel_create_call
 import libkgrpc.grpc_channel_credentials_release
 import libkgrpc.grpc_channel_destroy
 import libkgrpc.grpc_slice_unref
+import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.cancellation.CancellationException
 import kotlin.experimental.ExperimentalNativeApi
 import kotlin.native.ref.createCleaner
@@ -171,6 +172,7 @@ internal class NativeManagedChannel(
     override fun <RequestT, ResponseT> newCall(
         methodDescriptor: MethodDescriptor<RequestT, ResponseT>,
         callOptions: GrpcCallOptions,
+        coroutineContext: CoroutineContext
     ): ClientCall<RequestT, ResponseT> {
         check(!isShutdown) { internalError("Channel is shutdown") }
 
@@ -198,6 +200,7 @@ internal class NativeManagedChannel(
             methodDescriptor =methodDescriptor,
             callOptions = callOptions,
             callJob = Job(callJobSupervisor),
+            coroutineContext = coroutineContext,
         )
     }
 
diff --git a/grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/test/CoreClientTest.kt b/grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/test/CoreClientTest.kt
@@ -5,7 +5,7 @@ package kotlinx.rpc.grpc.test
 
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Job
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
@@ -54,7 +54,7 @@ class GrpcCoreClientTest {
         )
 
     private fun ManagedChannel.newHelloCall(fullName: String = "kotlinx.rpc.grpc.test.GreeterService/SayHello"): ClientCall<HelloRequest, HelloReply> =
-        platformApi.createCall(descriptorFor(fullName), GrpcCallOptions())
+        platformApi.createCall(descriptorFor(fullName), GrpcCallOptions(), Dispatchers.Default)
 
     private fun createChannel(): ManagedChannel = ManagedChannelBuilder(
         target = "localhost:$PORT",
diff --git a/grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/test/PlaygroundTest.kt b/grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/test/PlaygroundTest.kt
@@ -0,0 +1,46 @@
+/*
+ * 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.grpc.test
+
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.job
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import kotlin.test.Test
+
+class PlaygroundTest {
+
+    @Test
+    fun testCoroutineNew() {
+        runBlocking {
+
+            val scope = CoroutineScope(Dispatchers.Default)
+
+//            val job = Job(scope.coroutineContext.job)
+            val jobScope = CoroutineScope(scope.coroutineContext)
+
+            jobScope.launch {
+                println("Hello")
+            }.join()
+
+            scope.cancel("With this message")
+
+            jobScope.launch {
+                println("World")
+            }.apply {
+                invokeOnCompletion {
+                    println("Job is completed $it")
+                }
+            }.join()
+
+            println("Scope is active ${scope.coroutineContext.job.isActive}")
+        }
+    }
+
+}
