Fixed CancellationTest

Author: Mr3zee

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

@@ -9,7 +9,6 @@ import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.InternalCoroutinesApi
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.cancelAndJoin
@@ -77,7 +76,6 @@ public abstract class InitializedKrpcClient(
  * serializing data, tracking streams, processing exceptions, and other protocol responsibilities.
  * Leaves out the delivery of encoded messages to the specific implementations with [KrpcTransport].
  */
-@OptIn(InternalCoroutinesApi::class)
 public abstract class KrpcClient : RpcClient, KrpcEndpoint {
     /**
      * Called once to provide [KrpcTransport] for this client.
@@ -152,25 +150,22 @@ public abstract class KrpcClient : RpcClient, KrpcEndpoint {
 
         val context = SupervisorJob(transport.coroutineContext.job)
 
-        context.job.invokeOnCompletion(onCancelling = true) {
-            if (clientCancelled) {
-                return@invokeOnCompletion
-            }
-
-            clientCancelled = true
-
+        context.job.invokeOnCompletion {
             try {
-                if (!clientCancelledByServer) {
+                if (!clientCancelled && !clientCancelledByServer) {
                     sendCancellation(CancellationType.ENDPOINT, null, null, closeTransportAfterSending = true)
                 }
 
-                requestChannels.values.forEach {
-                    it.close(CancellationException("Client cancelled"))
-                    it.cancel()
-                }
+                clientCancelled = true
             } catch (_ : Exception) {
                 // ignore, we are already cancelled
             } finally {
+                requestChannels.values.forEach {
+                    val cause = CancellationException("Client cancelled")
+                    it.close(cause)
+                    it.cancel(cause)
+                }
+
                 requestChannels.clear()
             }
         }

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

@@ -31,7 +31,6 @@ import kotlin.reflect.KClass
  * @param transport [KrpcTransport] instance that will be used to send and receive RPC messages.
  * IMPORTANT: Must be exclusive to this server, otherwise unexpected behavior may occur.
  */
-@OptIn(InternalCoroutinesApi::class)
 public abstract class KrpcServer(
     private val config: KrpcConfig.Server,
     transport: KrpcTransport,
@@ -87,7 +86,7 @@ public abstract class KrpcServer(
     private var cancelledByClient = false
 
     init {
-        internalScope.coroutineContext.job.invokeOnCompletion(onCancelling = true) {
+        internalScope.coroutineContext.job.invokeOnCompletion {
             if (!cancelledByClient) {
                 sendCancellation(CancellationType.ENDPOINT, null, null, closeTransportAfterSending = true)
             }

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

@@ -112,6 +112,7 @@ class CancellationServiceImpl : CancellationService {
 
     override suspend fun outgoingStreamWithDelayedResponse(stream: Flow<Int>) {
         try {
+            waitCounter.increment()
             consume(stream)
 
             fence.await()

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

@@ -171,7 +171,7 @@ class CancellationTest {
         }
 
         val clientFlowJob = launch {
-            service.outgoingStream(flow {
+            service.outgoingStreamWithDelayedResponse(flow {
                 emit(0)
                 println("[testCancelClient] emit 0")
                 serverInstance().fence.await()
@@ -241,7 +241,7 @@ class CancellationTest {
         }
 
         val clientFlowJob = launch {
-            service.outgoingStream(flow {
+            service.outgoingStreamWithDelayedResponse(flow {
                 emit(0)
                 println("[testCancelServer] emit 0")
                 serverInstance().fence.await()