Fixed cancellation tests

Author: Mr3zee

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

@@ -5,12 +5,27 @@
 package kotlinx.rpc.krpc.client
 
 import kotlinx.atomicfu.atomic
-import kotlinx.coroutines.*
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.DelicateCoroutinesApi
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.InternalCoroutinesApi
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.FlowCollector
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.job
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.supervisorScope
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import kotlinx.rpc.RpcCall
@@ -21,7 +36,8 @@ import kotlinx.rpc.descriptor.RpcInvokator
 import kotlinx.rpc.internal.utils.InternalRpcApi
 import kotlinx.rpc.internal.utils.getOrNull
 import kotlinx.rpc.internal.utils.map.RpcInternalConcurrentHashMap
-import kotlinx.rpc.krpc.*
+import kotlinx.rpc.krpc.KrpcConfig
+import kotlinx.rpc.krpc.KrpcTransport
 import kotlinx.rpc.krpc.client.internal.ClientStreamContext
 import kotlinx.rpc.krpc.client.internal.ClientStreamSerializer
 import kotlinx.rpc.krpc.client.internal.KrpcClientConnector
@@ -35,7 +51,6 @@ import kotlinx.serialization.StringFormat
 import kotlinx.serialization.modules.SerializersModule
 import kotlin.collections.first
 import kotlin.concurrent.Volatile
-import kotlin.coroutines.cancellation.CancellationException
 import kotlin.properties.Delegates
 
 /**
@@ -343,17 +358,8 @@ public abstract class KrpcClient : RpcClient, KrpcEndpoint {
             }
 
             is KrpcCallMessage.CallException -> {
-                val cause = runCatching {
-                    message.cause.deserialize()
-                }
-
-                val result = if (cause.isFailure) {
-                    cause.exceptionOrNull()!!
-                } else {
-                    cause.getOrNull()!!
-                }
-
-                channel.close(result)
+                val cause = message.cause.deserialize()
+                channel.close(cause)
             }
 
             is KrpcCallMessage.CallSuccess, is KrpcCallMessage.StreamMessage -> {
@@ -458,6 +464,21 @@ public abstract class KrpcClient : RpcClient, KrpcEndpoint {
                 serviceTypeString = serviceTypeString,
             )
         } catch (e: CancellationException) {
+            currentCoroutineContext().ensureActive()
+
+            val wrapped = ManualCancellationException(e)
+            val serializedReason = serializeException(wrapped)
+            val message = KrpcCallMessage.StreamCancel(
+                callId = outgoingStream.callId,
+                serviceType = serviceTypeString,
+                streamId = outgoingStream.streamId,
+                cause = serializedReason,
+                connectionId = outgoingStream.connectionId,
+                serviceId = outgoingStream.serviceId,
+            )
+            sender.sendMessage(message)
+
+            // stop the flow and its coroutine, other flows are not affected
             throw e
         } catch (@Suppress("detekt.TooGenericExceptionCaught") cause: Throwable) {
             val serializedReason = serializeException(cause)

krpc/krpc-core/src/commonMain/kotlin/kotlinx/rpc/krpc/internal/ExceptionUtils.kt

@@ -4,6 +4,7 @@
 
 package kotlinx.rpc.krpc.internal
 
+import kotlinx.coroutines.CancellationException
 import kotlinx.rpc.internal.rpcInternalTypeName
 import kotlinx.rpc.internal.utils.InternalRpcApi
 
@@ -19,8 +20,46 @@ public fun serializeException(cause: Throwable): SerializedException {
 
 internal expect fun Throwable.stackElements(): List<StackElement>
 
+internal expect fun SerializedException.deserializeUnsafe(): Throwable
+
+internal fun SerializedException.nonJvmManualCancellationExceptionDeserialize(): ManualCancellationException? {
+    if (className == ManualCancellationException::class.rpcInternalTypeName) {
+        val cancellation = cause?.deserializeUnsafe()
+            ?: error("ManualCancellationException must have a cause")
+
+        return ManualCancellationException(
+            CancellationException(
+                message = cancellation.message,
+                cause = cancellation.cause,
+            )
+        )
+    }
+
+    return null
+}
+
+@InternalRpcApi
+public fun SerializedException.deserialize(): Throwable {
+    val cause = runCatching {
+        deserializeUnsafe()
+    }
+
+    val result = if (cause.isFailure) {
+        cause.exceptionOrNull()!!
+    } else {
+        val ex = cause.getOrNull()!!
+        if (ex is ManualCancellationException) {
+            ex.cause
+        } else {
+            ex
+        }
+    }
+
+    return result
+}
+
 @InternalRpcApi
-public expect fun SerializedException.deserialize(): Throwable
+public class ManualCancellationException(override val cause: CancellationException): RuntimeException()
 
 internal expect class DeserializedException(
     toStringMessage: String,

krpc/krpc-core/src/commonMain/kotlin/kotlinx/rpc/krpc/internal/KrpcReceiveHandler.kt

@@ -182,7 +182,11 @@ internal class KrpcActingReceiveHandler(
 
     override fun close(key: HandlerKey<*>, e: Throwable?) {
         if (e != null) {
-            job.cancel(CancellationException(null, e))
+            if (e is CancellationException) {
+                job.cancel(e)
+            } else {
+                job.cancel(CancellationException(null, e))
+            }
         } else {
             job.cancel()
         }
@@ -191,24 +195,14 @@ internal class KrpcActingReceiveHandler(
     }
 
     private suspend fun tryHandle(message: KrpcMessage): HandlerResult {
-        val result = runCatching {
+        try {
             callHandler(message)
-        }
 
-        return when {
-            result.isFailure -> {
-                val exception = result.exceptionOrNull()
-
-                if (exception is CancellationException) {
-                    throw exception
-                }
-
-                HandlerResult.Failure(exception)
-            }
-
-            else -> {
-                HandlerResult.Success
-            }
+            return HandlerResult.Success
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: Throwable) {
+            return HandlerResult.Failure(e)
         }
     }
 

krpc/krpc-core/src/jsMain/kotlin/kotlinx/rpc/krpc/internal/ExceptionUtils.js.kt

@@ -1,13 +1,11 @@
 /*
- * 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.
  */
 
 @file:Suppress("detekt.MatchingDeclarationName")
 
 package kotlinx.rpc.krpc.internal
 
-import kotlinx.rpc.internal.utils.InternalRpcApi
-
 internal actual class DeserializedException actual constructor(
     private val toStringMessage: String,
     actual override val message: String,
@@ -24,7 +22,7 @@ internal actual class DeserializedException actual constructor(
 
 internal actual fun Throwable.stackElements(): List<StackElement> = emptyList()
 
-@InternalRpcApi
-public actual fun SerializedException.deserialize(): Throwable {
-    return DeserializedException(toStringMessage, message, stacktrace, cause, className)
+internal actual fun SerializedException.deserializeUnsafe(): Throwable {
+    return nonJvmManualCancellationExceptionDeserialize()
+        ?: DeserializedException(toStringMessage, message, stacktrace, cause, className)
 }

krpc/krpc-core/src/jvmMain/kotlin/kotlinx/rpc/krpc/internal/ExceptionUtils.jvm.kt

@@ -1,10 +1,9 @@
 /*
- * 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.internal
 
-import kotlinx.rpc.internal.utils.InternalRpcApi
 import java.lang.reflect.Constructor
 import java.lang.reflect.Modifier
 
@@ -37,8 +36,7 @@ internal actual fun Throwable.stackElements(): List<StackElement> = stackTrace.m
     )
 }
 
-@InternalRpcApi
-public actual fun SerializedException.deserialize(): Throwable {
+internal actual fun SerializedException.deserializeUnsafe(): Throwable {
     try {
         val clazz = Class.forName(className)
         val fieldsCount = clazz.fieldsCountOrDefault(throwableFields)

krpc/krpc-core/src/nativeMain/kotlin/kotlinx/rpc/krpc/internal/ExceptionUtils.native.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.
  */
 
 @file:Suppress("detekt.MatchingDeclarationName")
@@ -21,6 +21,7 @@ internal actual class DeserializedException actual constructor(
 
 internal actual fun Throwable.stackElements(): List<StackElement> = emptyList()
 
-public actual fun SerializedException.deserialize(): Throwable {
-    return DeserializedException(toStringMessage, message, stacktrace, cause, className)
+internal actual fun SerializedException.deserializeUnsafe(): Throwable {
+    return nonJvmManualCancellationExceptionDeserialize()
+        ?: DeserializedException(toStringMessage, message, stacktrace, cause, className)
 }

krpc/krpc-core/src/wasmJsMain/kotlin/kotlinx/rpc/krpc/internal/ExceptionUtils.wasm.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.
  */
 
 @file:Suppress("detekt.MatchingDeclarationName")
@@ -24,7 +24,7 @@ internal actual class DeserializedException actual constructor(
 
 internal actual fun Throwable.stackElements(): List<StackElement> = emptyList()
 
-@InternalRpcApi
-public actual fun SerializedException.deserialize(): Throwable {
-    return DeserializedException(toStringMessage, message, stacktrace, cause, className)
+internal actual fun SerializedException.deserializeUnsafe(): Throwable {
+    return nonJvmManualCancellationExceptionDeserialize()
+        ?: DeserializedException(toStringMessage, message, stacktrace, cause, className)
 }

krpc/krpc-server/src/commonMain/kotlin/kotlinx/rpc/krpc/server/KrpcServer.kt

@@ -133,7 +133,7 @@ public abstract class KrpcServer(
         }
     }
 
-    override fun <@Rpc Service : Any> deregisterService(serviceKClass: KClass<Service>) {
+    final override fun <@Rpc Service : Any> deregisterService(serviceKClass: KClass<Service>) {
         connector.unsubscribeFromServiceMessages(serviceDescriptorOf(serviceKClass).fqName)
         rpcServices.remove(serviceDescriptorOf(serviceKClass).fqName)
     }

krpc/krpc-server/src/commonMain/kotlin/kotlinx/rpc/krpc/server/internal/KrpcServerConnector.kt

@@ -11,7 +11,7 @@ import kotlinx.serialization.SerialFormat
 
 internal class KrpcServerConnector private constructor(
     private val connector: KrpcConnector,
-): KrpcMessageSender by connector {
+) : KrpcMessageSender by connector {
     constructor(
         serialFormat: SerialFormat,
         transport: KrpcTransport,