grpc: Refactor cancel API in ClientCallScope to return Nothing

Signed-off-by: Johannes Zottele <[email protected]>

Author: Jozott00

grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/ClientInterceptor.kt

@@ -14,7 +14,7 @@ public interface ClientCallScope<Request, Response> {
     public val callOptions: GrpcCallOptions
     public fun onHeaders(block: (GrpcMetadata) -> Unit)
     public fun onClose(block: (Status, GrpcMetadata) -> Unit)
-    public fun cancel(message: String, cause: Throwable? = null)
+    public fun cancel(message: String, cause: Throwable? = null): Nothing
     public fun proceed(request: Flow<Request>): Flow<Response>
 }
 

grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/internal/suspendClientCalls.kt

@@ -224,8 +224,8 @@ private class ClientCallScopeImpl<Request, Response>(
         onCloseFuture.onComplete { block(it.first, it.second) }
     }
 
-    override fun cancel(message: String, cause: Throwable?) {
-        call.cancel(message, cause)
+    override fun cancel(message: String, cause: Throwable?): Nothing {
+        throw StatusException(Status(StatusCode.CANCELLED, message, cause))
     }
 
     override fun proceed(request: Flow<Request>): Flow<Response> {
@@ -302,7 +302,15 @@ private class ClientCallScopeImpl<Request, Response>(
         responses: Channel<Response>,
         ready: Ready,
     ) = clientCallListener(
-        onHeaders = { onHeadersFuture.complete(it) },
+        onHeaders = {
+            try {
+                onHeadersFuture.complete(it)
+            } catch (e: StatusException) {
+                // if a client interceptor called cancel, we throw a StatusException.
+                // as the JVM implementation treats them differently, we need to catch them here.
+                call.cancel(e.message, e.cause)
+            }
+        },
         onMessage = { message: Response ->
             responses.trySend(message).onFailure { e ->
                 throw e ?: AssertionError("onMessage should never be called until responses is ready")

grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/test/proto/ClientInterceptorTest.kt

@@ -87,7 +87,6 @@ class ClientInterceptorTest : GrpcProtoTest() {
         val error = assertFailsWith<StatusException> {
             val interceptor = interceptor {
                 cancel("Canceling in interceptor", IllegalStateException("Cancellation cause"))
-                proceed(it)
             }
             runGrpcTest(clientInterceptors = interceptor, test = ::unaryCall)
         }
@@ -98,6 +97,85 @@ class ClientInterceptorTest : GrpcProtoTest() {
         assertEquals("Cancellation cause", error.cause?.message)
     }
 
+    @Test
+    fun `cancel in request flow - should fail with cancellation`() {
+        val error = assertFailsWith<StatusException> {
+            val interceptor = interceptor {
+                proceed(it.map {
+                    val msg = it as EchoRequest
+                    if (msg.message == "Echo-3") {
+                        cancel("Canceling in request flow", IllegalStateException("Cancellation cause"))
+                    }
+                    it
+                })
+            }
+            runGrpcTest(clientInterceptors = interceptor, test = ::bidiStream)
+        }
+
+        assertEquals(StatusCode.CANCELLED, error.getStatus().statusCode)
+        assertContains(error.message!!, "Canceling in request flow")
+        assertIs<IllegalStateException>(error.cause)
+        assertEquals("Cancellation cause", error.cause?.message)
+    }
+
+    @Test
+    fun `cancel in response flow - should fail with cancellation`() {
+        val error = assertFailsWith<StatusException> {
+            val interceptor = interceptor {
+                flow {
+                    proceed(it).collect { resp ->
+                        val msg = resp as EchoResponse
+                        if (msg.message == "Echo-3") {
+                            cancel("Canceling in response flow", IllegalStateException("Cancellation cause"))
+                        }
+                        emit(resp)
+                    }
+                }
+            }
+            runGrpcTest(clientInterceptors = interceptor, test = ::bidiStream)
+        }
+
+        assertEquals(StatusCode.CANCELLED, error.getStatus().statusCode)
+        assertContains(error.message!!, "Canceling in response flow")
+        assertIs<IllegalStateException>(error.cause)
+        assertEquals("Cancellation cause", error.cause?.message)
+    }
+
+    @Test
+    fun `cancel onHeaders - should fail with cancellation`() {
+        val error = assertFailsWith<StatusException> {
+            val interceptor = interceptor {
+                this.onHeaders {
+                    cancel("Canceling in headers", IllegalStateException("Cancellation cause"))
+                }
+                proceed(it)
+            }
+            runGrpcTest(clientInterceptors = interceptor, test = ::bidiStream)
+        }
+
+        assertEquals(StatusCode.CANCELLED, error.getStatus().statusCode)
+        assertContains(error.message!!, "Canceling in headers")
+        assertIs<IllegalStateException>(error.cause)
+        assertEquals("Cancellation cause", error.cause?.message)
+    }
+
+    @Test
+    fun `cancel onClose - should fail with cancellation`() {
+        val error = assertFailsWith<StatusException> {
+            val interceptor = interceptor {
+                this.onClose { _, _ ->
+                    cancel("Canceling in onClose", IllegalStateException("Cancellation cause"))
+                }
+                proceed(it)
+            }
+            runGrpcTest(clientInterceptors = interceptor, test = ::bidiStream)
+        }
+        assertEquals(StatusCode.CANCELLED, error.getStatus().statusCode)
+        assertContains(error.message!!, "Canceling in onClose")
+        assertIs<IllegalStateException>(error.cause)
+        assertEquals("Cancellation cause", error.cause?.message)
+    }
+
     @Test
     fun `modify request message - should return modified message`() {
         val interceptor = interceptor {
@@ -156,6 +234,19 @@ class ClientInterceptorTest : GrpcProtoTest() {
         assertEquals("Hello", response.message)
     }
 
+    private suspend fun bidiStream(grpcClient: GrpcClient) {
+        val service = grpcClient.withService<EchoService>()
+        val responses = service.BidirectionalStreamingEcho(flow {
+            repeat(5) {
+                emit(EchoRequest { message = "Echo-$it" })
+            }
+        }).toList()
+        assertEquals(5, responses.size)
+        repeat(5) {
+            assertEquals("Echo-$it", responses[it].message)
+        }
+    }
+
 }
 
 private fun interceptor(

grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/test/proto/ServerInterceptorTest.kt

@@ -148,6 +148,27 @@ class ServerInterceptorTest : GrpcProtoTest() {
         }
     }
 
+    @Test
+    fun `proceedFlow - should succeed on client`() {
+        val interceptor = interceptor {
+            kotlinx.coroutines.flow.flow {
+                proceedFlow(it)
+            }
+        }
+        runGrpcTest(serverInterceptors = interceptor, test = ::unaryCall)
+    }
+
+    @Test
+    fun `method descriptor - full method name is exposed`() {
+        var methodName: String? = null
+        val interceptor = interceptor {
+            methodName = method.getFullMethodName()
+            proceed(it)
+        }
+        runGrpcTest(serverInterceptors = interceptor, test = ::unaryCall)
+        assertContains(methodName!!, "EchoService/UnaryEcho")
+    }
+
     private suspend fun unaryCall(grpcClient: GrpcClient) {
         val service = grpcClient.withService<EchoService>()
         val response = service.UnaryEcho(EchoRequest { message = "Hello" })

grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/NativeServerCall.kt

@@ -101,6 +101,7 @@ internal class NativeServerCall<Request, Response>(
     }
 
     private fun initialize() {
+        println("Initializing native server call")
         // finishes if the whole connection is closed.
         // this triggers onClose()/onCanceled() callback.
         val arena = Arena()