Added cancellation tests and GC tests, checked exceptions

Author: Mr3zee

krpc/krpc-client/src/commonMain/kotlin/kotlinx/rpc/krpc/client/KrpcClient.kt

@@ -340,65 +340,72 @@ public abstract class KrpcClient(
 
             connector.sendMessage(request)
 
-            connector.subscribeToCallResponse(call.descriptor.fqName, callId) { message ->
-                when (message) {
-                    is KrpcCallMessage.CallData -> {
-                        error("Unexpected message")
-                    }
-
-                    is KrpcCallMessage.CallException -> {
-                        val cause = runCatching {
-                            message.cause.deserialize()
+            try {
+                connector.subscribeToCallResponse(call.descriptor.fqName, callId) { message ->
+                    when (message) {
+                        is KrpcCallMessage.CallData -> {
+                            error("Unexpected message")
                         }
 
-                        val result = if (cause.isFailure) {
-                            cause.exceptionOrNull()!!
-                        } else {
-                            cause.getOrNull()!!
+                        is KrpcCallMessage.CallException -> {
+                            val cause = runCatching {
+                                message.cause.deserialize()
+                            }
+
+                            val result = if (cause.isFailure) {
+                                cause.exceptionOrNull()!!
+                            } else {
+                                cause.getOrNull()!!
+                            }
+
+                            channel.close(result)
                         }
 
-                        channel.close(result)
-                    }
+                        is KrpcCallMessage.CallSuccess, is KrpcCallMessage.StreamMessage -> {
+                            val value = runCatching {
+                                val serializerResult =
+                                    nonSuspendingSerialFormat.serializersModule.rpcSerializerForType(callable.returnType)
 
-                    is KrpcCallMessage.CallSuccess, is KrpcCallMessage.StreamMessage -> {
-                        val value = runCatching {
-                            val serializerResult =
-                                nonSuspendingSerialFormat.serializersModule.rpcSerializerForType(callable.returnType)
+                                decodeMessageData(nonSuspendingSerialFormat, serializerResult, message)
+                            }
 
-                            decodeMessageData(nonSuspendingSerialFormat, serializerResult, message)
+                            @Suppress("UNCHECKED_CAST")
+                            channel.send(value.getOrNull() as T)
                         }
 
-                        @Suppress("UNCHECKED_CAST")
-                        channel.send(value.getOrNull() as T)
-                    }
-
-                    is KrpcCallMessage.StreamFinished -> {
-                        connector.unsubscribeFromMessages(call.descriptor.fqName, callId)
-                        channel.close()
-                    }
+                        is KrpcCallMessage.StreamFinished -> {
+                            connector.unsubscribeFromMessages(call.descriptor.fqName, callId)
+                            channel.close()
+                        }
 
-                    is KrpcCallMessage.StreamCancel -> {
-                        connector.unsubscribeFromMessages(call.descriptor.fqName, callId)
-                        val cause = message.cause.deserialize()
-                        channel.close(cause)
+                        is KrpcCallMessage.StreamCancel -> {
+                            connector.unsubscribeFromMessages(call.descriptor.fqName, callId)
+                            val cause = message.cause.deserialize()
+                            channel.close(cause)
+                        }
                     }
                 }
-            }
 
-            try {
-                while (true) {
-                    val element = channel.receiveCatching()
-                    if (element.isClosed) {
-                        val ex = element.exceptionOrNull() ?: break
-                        error(ex)
-                    }
+                try {
+                    while (true) {
+                        val element = channel.receiveCatching()
+                        if (element.isClosed) {
+                            val ex = element.exceptionOrNull() ?: break
+                            error(ex)
+                        }
 
-                    if (!element.isFailure) {
-                        emit(element.getOrThrow())
+                        if (!element.isFailure) {
+                            emit(element.getOrThrow())
+                        }
                     }
+                } catch (_: ClosedReceiveChannelException) {
+                    // ignore
                 }
-            } catch (_: ClosedReceiveChannelException) {
-                // ignore
+            } catch (e: CancellationException) {
+                // sendCancellation is not suspending, so no need for NonCancellable
+                sendCancellation(CancellationType.REQUEST, call.serviceId.toString(), callId)
+
+                throw e
             }
         }
     }

krpc/krpc-test/src/jvmMain/kotlin/kotlinx/rpc/krpc/test/KrpcTestService.kt

@@ -18,6 +18,8 @@ import java.time.LocalDateTime
 @Rpc
 interface KrpcTestService : RemoteService {
     fun nonSuspendFlow(): Flow<Int>
+    fun nonSuspendFlowErrorOnEmit(): Flow<Int>
+    fun nonSuspendFlowErrorOnReturn(): Flow<Int>
     suspend fun empty()
     suspend fun returnType(): String
     suspend fun simpleWithParams(name: String): String

krpc/krpc-test/src/jvmMain/kotlin/kotlinx/rpc/krpc/test/KrpcTestServiceBackend.kt

@@ -30,6 +30,16 @@ class KrpcTestServiceBackend(override val coroutineContext: CoroutineContext) :
         }
     }
 
+    override fun nonSuspendFlowErrorOnEmit(): Flow<Int> {
+        return flow {
+            error("nonSuspendFlowErrorOnEmit")
+        }
+    }
+
+    override fun nonSuspendFlowErrorOnReturn(): Flow<Int> {
+        error("nonSuspendFlowErrorOnReturn")
+    }
+
     @Suppress("detekt.EmptyFunctionBlock")
     override suspend fun empty() {}
 

krpc/krpc-test/src/jvmMain/kotlin/kotlinx/rpc/krpc/test/KrpcTransportTestBase.kt

@@ -128,6 +128,25 @@ abstract class KrpcTransportTestBase {
         }
     }
 
+    @Test
+    fun nonSuspendErrorOnEmit() {
+        runBlocking {
+            val flow = client.nonSuspendFlowErrorOnReturn()
+            assertFailsWith<IllegalStateException> {
+                flow.toList()
+            }
+        }
+    }
+
+    @Test
+    fun nonSuspendErrorOnReturn() {
+        runBlocking {
+            assertFailsWith<IllegalStateException> {
+                client.nonSuspendFlowErrorOnReturn().toList()
+            }
+        }
+    }
+
     @Test
     fun empty() {
         backend.cancel()

krpc/krpc-test/src/jvmTest/kotlin/kotlinx/rpc/krpc/test/cancellation/CancellationService.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ * 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.cancellation
@@ -41,6 +41,8 @@ interface CancellationService : RemoteService {
     suspend fun closedStreamScopeCallback()
 
     suspend fun closedStreamScopeCallbackWithStream(): Flow<Int>
+
+    fun nonSuspendable(): Flow<Int>
 }
 
 class CancellationServiceImpl(override val coroutineContext: CoroutineContext) : CancellationService {
@@ -169,6 +171,31 @@ class CancellationServiceImpl(override val coroutineContext: CoroutineContext) :
             }
         }
     }
+
+    var nonSuspendableSecond = false
+    val nonSuspendableFinished = CompletableDeferred<Unit>()
+
+    override fun nonSuspendable(): Flow<Int> {
+        return flow {
+            try {
+                repeat(2) {
+                    if (it == 1) {
+                        nonSuspendableSecond = true
+                    }
+
+                    emit(it)
+
+                    if (it == 0) {
+                        fence.await()
+                    }
+                }
+            } catch (e: CancellationException) {
+                nonSuspendableFinished.complete(Unit)
+
+                throw e
+            }
+        }
+    }
 }
 
 fun resumableFlow(fence: Deferred<Unit>, onEmit: (Int) -> Unit = {}): Flow<Int> = flow {

krpc/krpc-test/src/jvmTest/kotlin/kotlinx/rpc/krpc/test/cancellation/CancellationTest.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ * 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.cancellation
@@ -671,6 +671,61 @@ class CancellationTest {
         stopAllAndJoin()
     }
 
+    @Test
+    fun testCancellingNonSuspendable() = runCancellationTest {
+        val flow = service.nonSuspendable()
+        val firstDone = CompletableDeferred<Unit>()
+        val requestJob = launch {
+            flow.collect {
+                if (it == 0) {
+                    firstDone.complete(Unit)
+                }
+            }
+        }
+
+        firstDone.await()
+        requestJob.cancel("Cancelled by test")
+        requestJob.join()
+        serverInstance().nonSuspendableFinished.await()
+
+        assertEquals(false, serverInstance().nonSuspendableSecond)
+
+        checkAlive()
+        stopAllAndJoin()
+    }
+
+    @Test
+    fun testGCNonSuspendable() = runCancellationTest {
+        val firstDone = CompletableDeferred<Unit>()
+        val latch = CompletableDeferred<Unit>()
+        val requestJob = processFlowAndLeaveUnusedForGC(firstDone, latch)
+
+        firstDone.await()
+        System.gc() // hint GC to collect the flow
+        serverInstance().nonSuspendableFinished.await()
+
+        assertEquals(false, serverInstance().nonSuspendableSecond)
+        latch.complete(Unit)
+        requestJob.join()
+
+        checkAlive()
+        stopAllAndJoin()
+    }
+
+    private fun CancellationToolkit.processFlowAndLeaveUnusedForGC(
+        firstDone: CompletableDeferred<Unit>,
+        latch: CompletableDeferred<Unit>
+    ): Job {
+        val flow = service.nonSuspendable()
+        val requestJob = launch {
+            flow.first()
+            firstDone.complete(Unit)
+            latch.await()
+        }
+
+        return requestJob
+    }
+
     private fun CancellationToolkit.checkAlive(
         serviceAlive: Boolean = true,
         clientAlive: Boolean = true,