Fix bidirectional flows in non-suspend streams (#316)

Author: Mr3zee

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

@@ -311,7 +311,7 @@ public abstract class KrpcClient(
         connector.sendMessage(firstMessage)
     }
 
-    private val nonSuspendingSerialFormat = config.serialFormatInitializer.build()
+    private val noFlowSerialFormat = config.serialFormatInitializer.build()
 
     @Suppress("detekt.CyclomaticComplexMethod")
     override fun <T> callServerStreaming(call: RpcCall): Flow<T> {
@@ -326,21 +326,33 @@ public abstract class KrpcClient(
 
             val channel = Channel<T>()
 
-            val request = serializeRequest(
-                callId = callId,
-                call = call,
-                callable = callable,
-                serialFormat = nonSuspendingSerialFormat,
-                pluginParams = mapOf(KrpcPluginKey.NON_SUSPENDING_SERVER_FLOW_MARKER to ""),
-            )
-
-            connector.sendMessage(request)
+            val streamScope = StreamScope(currentCoroutineContext())
 
             try {
+                val streamContext = LazyKrpcStreamContext(streamScope, null) {
+                    KrpcStreamContext(callId, config, connectionId, call.serviceId, it)
+                }
+
+                val serialFormat = prepareSerialFormat(streamContext)
+
+                val request = serializeRequest(
+                    callId = callId,
+                    call = call,
+                    callable = callable,
+                    serialFormat = serialFormat,
+                    pluginParams = mapOf(KrpcPluginKey.NON_SUSPENDING_SERVER_FLOW_MARKER to ""),
+                )
+
+                connector.sendMessage(request)
+
                 connector.subscribeToCallResponse(call.descriptor.fqName, callId) { message ->
                     handleServerStreamingMessage(message, channel, callable, call, callId)
                 }
 
+                streamContext.valueOrNull?.launchIf({ outgoingStreamsAvailable }) {
+                    handleOutgoingStreams(it, serialFormat, call.descriptor.fqName)
+                }
+
                 while (true) {
                     val element = channel.receiveCatching()
                     if (element.isClosed) {
@@ -358,11 +370,14 @@ public abstract class KrpcClient(
                 connector.unsubscribeFromMessages(call.descriptor.fqName, callId)
 
                 throw e
+            } finally {
+                streamScope.close()
+                channel.close()
             }
         }
     }
 
-    private suspend fun <T, @Rpc R: Any> handleServerStreamingMessage(
+    private suspend fun <T, @Rpc R : Any> handleServerStreamingMessage(
         message: KrpcCallMessage,
         channel: Channel<T>,
         callable: RpcCallable<R>,
@@ -390,10 +405,10 @@ public abstract class KrpcClient(
 
             is KrpcCallMessage.CallSuccess, is KrpcCallMessage.StreamMessage -> {
                 val value = runCatching {
-                    val serializerResult = nonSuspendingSerialFormat.serializersModule
+                    val serializerResult = noFlowSerialFormat.serializersModule
                         .rpcSerializerForType(callable.returnType)
 
-                    decodeMessageData(nonSuspendingSerialFormat, serializerResult, message)
+                    decodeMessageData(noFlowSerialFormat, serializerResult, message)
                 }
 
                 @Suppress("UNCHECKED_CAST")

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

@@ -49,6 +49,9 @@ interface KrpcTestService : RemoteService {
     fun nonSuspendFlow(): Flow<Int>
     fun nonSuspendFlowErrorOnEmit(): Flow<Int>
     fun nonSuspendFlowErrorOnReturn(): Flow<Int>
+    fun nonSuspendBidirectional(flow: Flow<Int>): Flow<Int>
+    fun nonSuspendBidirectionalPayload(payloadWithStream: PayloadWithStream): Flow<Int>
+
     suspend fun empty()
     suspend fun returnType(): String
     suspend fun simpleWithParams(name: String): String

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

@@ -37,6 +37,14 @@ class KrpcTestServiceBackend(override val coroutineContext: CoroutineContext) :
         error("nonSuspendFlowErrorOnReturn")
     }
 
+    override fun nonSuspendBidirectional(flow: Flow<Int>): Flow<Int> {
+        return flow.map { it * 2 }
+    }
+
+    override fun nonSuspendBidirectionalPayload(payloadWithStream: PayloadWithStream): Flow<Int> {
+        return payloadWithStream.stream.map { it.length }
+    }
+
     @Suppress("detekt.EmptyFunctionBlock")
     override suspend fun empty() {}
 

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

@@ -136,6 +136,22 @@ abstract class KrpcTransportTestBase {
         }
     }
 
+    @Test
+    fun nonSuspendBidirectional() = runTest {
+        assertEquals(
+            expected = List(10) { it * 2 },
+            actual = client.nonSuspendBidirectional(List(10) { it }.asFlow()).toList(),
+        )
+        print(1)
+    }
+
+    @Test
+    fun nonSuspendBidirectionalPayload() = runTest {
+        assertEquals(
+            expected = List(3) { 2 },
+            actual = client.nonSuspendBidirectionalPayload(payload(0)).toList(),
+        )
+    }
 
     @Test
     fun empty() {