grpc: Add client interceptor support

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

Author: Jozott00

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

@@ -307,7 +307,7 @@ private class ClientCallScopeImpl<Request, Response>(
             } catch (exception: Throwable) {
                 cause = exception
                 if (exception !is StatusException) {
-                    val status = Status(StatusCode.INTERNAL, "Interceptor threw an error", exception)
+                    val status = Status(StatusCode.CANCELLED, "Interceptor threw an error", exception)
                     cause = StatusException(status)
                 }
             }

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

@@ -67,6 +67,7 @@ class ClientInterceptorTest: GrpcProtoTest() {
             runGrpcTest(clientInterceptors = interceptor, test = ::unaryCall)
         }
 
+        assertEquals(StatusCode.CANCELLED, error.getStatus().statusCode)
         assertIs<IllegalStateException>(error.cause)
         assertEquals("Failing in onHeader", error.cause?.message)
     }
@@ -83,6 +84,7 @@ class ClientInterceptorTest: GrpcProtoTest() {
             runGrpcTest(clientInterceptors = interceptor, test = ::unaryCall)
         }
 
+        assertEquals(StatusCode.CANCELLED, error.getStatus().statusCode)
         assertIs<IllegalStateException>(error.cause)
         assertEquals("Failing in onClose", error.cause?.message)
     }
@@ -129,13 +131,13 @@ class ClientInterceptorTest: GrpcProtoTest() {
     }
 
     @Test
-    fun `append a response message once closed`() = repeat(1000) {
+    fun `append a response message once closed`() {
         val interceptor = interceptor { scope, req -> channelFlow {
             scope.proceed(req).collect {
                 trySend(it)
             }
-            scope.onClose { _, _ ->
-                trySend(EchoResponse { message = "Appended-after-close" })
+            scope.onClose { status, _ ->
+                trySend(EchoResponse { message = "Appended-after-close-with-${status.statusCode}" })
             }
         } }
 
@@ -148,10 +150,8 @@ class ClientInterceptorTest: GrpcProtoTest() {
                             emit(EchoRequest { message = "Eccchhooo" })
                         }
                     }).toList()
-
-            println("Respone messages: ${responses.map { it.message }}")
             assertEquals(6, responses.size)
-            assertTrue(responses.any { it.message == "Appended-after-close" })
+            assertTrue(responses.any { it.message == "Appended-after-close-with-OK" })
         }
     }
 

grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/Status.native.kt

@@ -12,6 +12,10 @@ public actual class Status internal constructor(
     public actual fun getDescription(): String? = description
 
     public actual fun getCause(): Throwable? = cause
+
+    override fun toString(): String {
+        return "Status(description=$description, statusCode=$statusCode, cause=$cause)"
+    }
 }
 
 public actual val Status.statusCode: StatusCode

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

@@ -81,7 +81,7 @@ internal class NativeClientCall<Request, Response>(
     private var listener: Listener<Response>? = null
     private var halfClosed = false
     private var cancelled = false
-    private var closed = atomic(false)
+    private val closed = atomic(false)
 
     // tracks how many operations are in flight (not yet completed by the listener).
     // if 0 and we got a closeInfo (containing the status), there are no more ongoing operations.
@@ -133,7 +133,9 @@ internal class NativeClientCall<Request, Response>(
             val lst = checkNotNull(listener) { internalError("Not yet started") }
             // allows the managed channel to join for the call to finish.
             callJob.complete()
-            lst.onClose(info.first, info.second)
+            safeUserCode("Failed to call onClose.") {
+                lst.onClose(info.first, info.second)
+            }
         }
     }
 
@@ -142,9 +144,8 @@ internal class NativeClientCall<Request, Response>(
      * This is called as soon as the RECV_STATUS_ON_CLIENT batch (started with [startRecvStatus]) finished.
      */
     private fun markClosePending(status: Status, trailers: GrpcTrailers) {
-        if (closeInfo.compareAndSet(null, Pair(status, trailers))) {
-            tryToCloseCall()
-        }
+        closeInfo.compareAndSet(null, Pair(status, trailers))
+        tryToCloseCall()
     }
 
     /**
@@ -153,7 +154,9 @@ internal class NativeClientCall<Request, Response>(
      */
     private fun turnReady() {
         if (ready.compareAndSet(expect = false, update = true)) {
-            listener?.onReady()
+            safeUserCode("Failed to call onClose.") {
+                listener?.onReady()
+            }
         }
     }
 
@@ -163,7 +166,6 @@ internal class NativeClientCall<Request, Response>(
         headers: GrpcTrailers,
     ) {
         check(listener == null) { internalError("Already started") }
-        check(!cancelled) { internalError("Already cancelled.") }
 
         listener = responseListener
 
@@ -254,7 +256,8 @@ internal class NativeClientCall<Request, Response>(
             is BatchResult.Submitted -> {
                 callResult.future.onComplete {
                     val details = statusDetails.toByteArray().toKString()
-                    val status = Status(statusCode.value.toKotlin(), details, null)
+                    val kStatusCode = statusCode.value.toKotlin()
+                    val status = Status(kStatusCode, details, null)
                     val trailers = GrpcTrailers()
 
                     // cleanup
@@ -306,7 +309,9 @@ internal class NativeClientCall<Request, Response>(
             grpc_metadata_array_destroy(meta.ptr)
             arena.clear()
         }) {
-            // TODO: Send headers to listener
+            safeUserCode("Failed to call onHeaders.") {
+                listener?.onHeaders(GrpcTrailers())
+            }
         }
     }
 
@@ -319,7 +324,10 @@ internal class NativeClientCall<Request, Response>(
         // limit numMessages to prevent potential stack overflows
         check(numMessages <= 16) { internalError("numMessages must be <= 16") }
         val listener = checkNotNull(listener) { internalError("Not yet started") }
-        check(!cancelled) { internalError("Already cancelled") }
+        if (cancelled) {
+            // no need to send message if the call got already cancelled.
+            return
+        }
 
         var remainingMessages = numMessages
 
@@ -342,7 +350,9 @@ internal class NativeClientCall<Request, Response>(
                 val buf = recvPtr.value ?: return@runBatch
                 val msg = methodDescriptor.getResponseMarshaller()
                     .parse(buf.toKotlin().asInputStream())
-                listener.onMessage(msg)
+                safeUserCode("Failed to call onClose.") {
+                    listener.onMessage(msg)
+                }
                 post()
             }
         }
@@ -353,8 +363,11 @@ internal class NativeClientCall<Request, Response>(
 
     override fun cancel(message: String?, cause: Throwable?) {
         cancelled = true
-        val message = if (cause != null) "$message: ${cause.message}" else message
-        cancelInternal(grpc_status_code.GRPC_STATUS_CANCELLED, message ?: "Call cancelled")
+        val status = Status(StatusCode.CANCELLED, message ?: "Call cancelled", cause)
+        // user side cancellation must always win over any other status (even if the call is already completed).
+        // this will also preserve the cancellation cause, which cannot be passed to the grpc-core.
+        closeInfo.value = Pair(status, GrpcTrailers())
+        cancelInternal(grpc_status_code.GRPC_STATUS_CANCELLED, message ?: "Call cancelled with cause: ${cause?.message}")
     }
 
     private fun cancelInternal(statusCode: grpc_status_code, message: String) {
@@ -366,7 +379,7 @@ internal class NativeClientCall<Request, Response>(
 
     override fun halfClose() {
         check(!halfClosed) { internalError("Already half closed.") }
-        check(!cancelled) { internalError("Already cancelled.") }
+        if (cancelled) return
         halfClosed = true
 
         val arena = Arena()
@@ -384,9 +397,10 @@ internal class NativeClientCall<Request, Response>(
     override fun sendMessage(message: Request) {
         checkNotNull(listener) { internalError("Not yet started") }
         check(!halfClosed) { internalError("Already half closed.") }
-        check(!cancelled) { internalError("Already cancelled.") }
         check(isReady()) { internalError("Not yet ready.") }
 
+        if (cancelled) return
+
         // set ready false, as only one message can be sent at a time.
         ready.value = false
 
@@ -408,6 +422,18 @@ internal class NativeClientCall<Request, Response>(
             turnReady()
         }
     }
+
+    /**
+     * Safely executes the provided block of user code, catching any thrown exceptions or errors.
+     * If an exception is caught, it cancels the operation with the specified message and cause.
+     */
+    private fun safeUserCode(cancelMsg: String, block: () -> Unit) {
+        try {
+            block()
+        } catch (e: Throwable) {
+            cancel(cancelMsg, e)
+        }
+    }
 }