Fix sending redundant RequestN frame (#101)

whyoleg · web-flow · commit e8de0d935ddb · 2020-09-23T15:55:11.000+03:00
diff --git a/rsocket-core/build.gradle.kts b/rsocket-core/build.gradle.kts
@@ -36,6 +36,7 @@ kotlin {
         }
         val commonTest by getting {
             dependencies {
+                implementation("app.cash.turbine:turbine:0.2.1")
                 implementation("io.ktor:ktor-utils:1.4.0")
                 implementation(project(":rsocket-transport-local"))
             }
diff --git a/rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/flow/RequestChannelRequesterFlow.kt b/rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/flow/RequestChannelRequesterFlow.kt
@@ -35,8 +35,7 @@ internal class RequestChannelRequesterFlow(
     override fun create(context: CoroutineContext, capacity: Int): RequestChannelRequesterFlow =
         RequestChannelRequesterFlow(payloads, requester, state, context, capacity)
 
-    @OptIn(ExperimentalCoroutinesApi::class)
-    override suspend fun collectTo(scope: ProducerScope<Payload>): Unit = with(state) {
+    override suspend fun collectImpl(collectContext: CoroutineContext, collector: FlowCollector<Payload>): Unit = with(state) {
         val streamId = requester.createStream()
         val receiverDeferred = CompletableDeferred<ReceiveChannel<RequestFrame>?>()
         val request = launchCancelable(streamId) {
@@ -47,6 +46,7 @@ internal class RequestChannelRequesterFlow(
         }
         request.invokeOnCompletion {
             if (receiverDeferred.isCompleted) {
+                @OptIn(ExperimentalCoroutinesApi::class)
                 if (it != null && it !is CancellationException) receiverDeferred.getCompleted()?.cancelConsumed(it)
             } else {
                 if (it == null) receiverDeferred.complete(null)
@@ -55,7 +55,7 @@ internal class RequestChannelRequesterFlow(
         }
         try {
             val receiver = receiverDeferred.await() ?: return
-            collectStream(streamId, receiver, scope)
+            collectStream(streamId, receiver, collectContext, collector)
         } catch (e: Throwable) {
             if (e is CancellationException) request.cancel(e)
             else request.cancel("Receiver failed", e)
diff --git a/rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/flow/RequestChannelResponderFlow.kt b/rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/flow/RequestChannelResponderFlow.kt
@@ -19,8 +19,8 @@ package io.rsocket.kotlin.internal.flow
 import io.rsocket.kotlin.frame.*
 import io.rsocket.kotlin.internal.*
 import io.rsocket.kotlin.payload.*
-import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.flow.*
 import kotlin.coroutines.*
 
 //TODO prevent consuming more then one time - add atomic ?
@@ -35,9 +35,8 @@ internal class RequestChannelResponderFlow(
     override fun create(context: CoroutineContext, capacity: Int): RequestChannelResponderFlow =
         RequestChannelResponderFlow(streamId, receiver, state, context, capacity)
 
-    @OptIn(ExperimentalCoroutinesApi::class)
-    override suspend fun collectTo(scope: ProducerScope<Payload>): Unit = with(state) {
+    override suspend fun collectImpl(collectContext: CoroutineContext, collector: FlowCollector<Payload>): Unit = with(state) {
         send(RequestNFrame(streamId, requestSize - 1)) //-1 because first payload received
-        collectStream(streamId, receiver, scope)
+        collectStream(streamId, receiver, collectContext, collector)
     }
 }
diff --git a/rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/flow/RequestStreamRequesterFlow.kt b/rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/flow/RequestStreamRequesterFlow.kt
@@ -19,8 +19,8 @@ package io.rsocket.kotlin.internal.flow
 import io.rsocket.kotlin.frame.*
 import io.rsocket.kotlin.internal.*
 import io.rsocket.kotlin.payload.*
-import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.flow.*
 import kotlin.coroutines.*
 
 internal class RequestStreamRequesterFlow(
@@ -33,11 +33,10 @@ internal class RequestStreamRequesterFlow(
     override fun create(context: CoroutineContext, capacity: Int): RequestStreamRequesterFlow =
         RequestStreamRequesterFlow(payload, requester, state, context, capacity)
 
-    @OptIn(ExperimentalCoroutinesApi::class)
-    override suspend fun collectTo(scope: ProducerScope<Payload>): Unit = with(state) {
+    override suspend fun collectImpl(collectContext: CoroutineContext, collector: FlowCollector<Payload>): Unit = with(state) {
         val streamId = requester.createStream()
         val receiver = createReceiverFor(streamId)
         send(RequestStreamFrame(streamId, requestSize, payload))
-        collectStream(streamId, receiver, scope)
+        collectStream(streamId, receiver, collectContext, collector)
     }
 }
diff --git a/rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/flow/StreamFlow.kt b/rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/flow/StreamFlow.kt
@@ -21,6 +21,7 @@ import io.rsocket.kotlin.internal.*
 import io.rsocket.kotlin.payload.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.flow.internal.*
 import kotlin.coroutines.*
 
@@ -40,18 +41,33 @@ internal abstract class StreamFlow(
             else               -> capacity.also { check(it >= 1) }
         }
 
+    protected abstract suspend fun collectImpl(collectContext: CoroutineContext, collector: FlowCollector<Payload>)
+
+    final override suspend fun collect(collector: FlowCollector<Payload>) {
+        val collectContext = context + coroutineContext
+        withContext(coroutineContext + context) {
+            collectImpl(collectContext, collector)
+        }
+    }
+
+    final override suspend fun collectTo(scope: ProducerScope<Payload>): Unit =
+        collectImpl(scope.coroutineContext, SendingCollector(scope.channel))
+
     protected suspend fun collectStream(
         streamId: Int,
         receiver: ReceiveChannel<RequestFrame>,
-        scope: ProducerScope<Payload>,
+        collectContext: CoroutineContext,
+        collector: FlowCollector<Payload>,
     ): Unit = with(state) {
-        val collector = SendingCollector(scope.channel)
         consumeReceiverFor(streamId) {
             var consumed = 0
             //TODO fragmentation
             for (frame in receiver) {
                 if (frame.complete) return //TODO check next flag
-                collector.emit(frame.payload)
+                //emit in collectContext to prevent `Flow invariant is violated`
+                withContext(collectContext) {
+                    collector.emit(frame.payload)
+                }
                 if (++consumed == requestSize) {
                     consumed = 0
                     send(RequestNFrame(streamId, requestSize))
diff --git a/rsocket-core/src/commonTest/kotlin/io/rsocket/kotlin/Test.common.kt b/rsocket-core/src/commonTest/kotlin/io/rsocket/kotlin/Test.common.kt
@@ -20,3 +20,5 @@ import kotlinx.coroutines.*
 import kotlin.time.*
 
 expect fun test(timeout: Duration? = 10.seconds, block: suspend CoroutineScope.() -> Unit)
+
+expect val anotherDispatcher: CoroutineDispatcher
diff --git a/rsocket-core/src/commonTest/kotlin/io/rsocket/kotlin/internal/RSocketRequesterTest.kt b/rsocket-core/src/commonTest/kotlin/io/rsocket/kotlin/internal/RSocketRequesterTest.kt
@@ -16,6 +16,7 @@
 
 package io.rsocket.kotlin.internal
 
+import app.cash.turbine.*
 import io.rsocket.kotlin.*
 import io.rsocket.kotlin.error.*
 import io.rsocket.kotlin.frame.*
@@ -24,6 +25,7 @@ import io.rsocket.kotlin.payload.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.flow.*
+import kotlin.coroutines.*
 import kotlin.test.*
 import kotlin.time.*
 
@@ -57,6 +59,110 @@ class RSocketRequesterTest {
         assertEquals(5, frame.initialRequest)
     }
 
+    @Test
+    fun testStreamBuffer() = test {
+        val flow =
+            requester.requestStream(Payload.Empty)
+                .buffer(2)
+                .take(2)
+
+        assertEquals(0, connection.sentFrames.size)
+
+        flow.launchIn(CoroutineScope(connection.job))
+
+        connection.sentAsFlow().test {
+            expectItem().let { frame ->
+                assertTrue(frame is RequestFrame)
+                assertEquals(FrameType.RequestStream, frame.type)
+                assertEquals(2, frame.initialRequest)
+            }
+            delay(200)
+            expectNoEvents()
+            connection.sendToReceiver(NextPayloadFrame(1, Payload.Empty))
+            delay(200)
+            expectNoEvents()
+            connection.sendToReceiver(NextPayloadFrame(1, Payload.Empty))
+            delay(200)
+            expectItem().let { frame ->
+                assertTrue(frame is CancelFrame)
+            }
+            delay(200)
+            expectNoEvents()
+        }
+    }
+
+    class SomeContext(val context: Int) : AbstractCoroutineContextElement(SomeContext) {
+        companion object Key : CoroutineContext.Key<SomeContext>
+    }
+
+    @Test
+    fun testStreamBufferWithAdditionalContext() = test {
+        val flow =
+            requester.requestStream(Payload.Empty)
+                .buffer(2)
+                .flowOn(SomeContext(2))
+                .take(2)
+
+        assertEquals(0, connection.sentFrames.size)
+
+        flow.launchIn(CoroutineScope(connection.job))
+
+        connection.sentAsFlow().test {
+            expectItem().let { frame ->
+                assertTrue(frame is RequestFrame)
+                assertEquals(FrameType.RequestStream, frame.type)
+                assertEquals(2, frame.initialRequest)
+            }
+            delay(200)
+            expectNoEvents()
+            connection.sendToReceiver(NextPayloadFrame(1, Payload.Empty))
+            delay(200)
+            expectNoEvents()
+            connection.sendToReceiver(NextPayloadFrame(1, Payload.Empty))
+            delay(200)
+            expectItem().let { frame ->
+                assertTrue(frame is CancelFrame)
+            }
+            delay(200)
+            expectNoEvents()
+        }
+    }
+
+    @Test
+    fun testStreamBufferWithAnotherDispatcher() = test {
+        val flow =
+            requester.requestStream(Payload.Empty)
+                .buffer(2)
+                .flowOn(anotherDispatcher) //change dispatcher before take
+                .take(2)
+                .transform { emit(it) } //force using SafeCollector to check that `Flow invariant is violated` not happens
+
+        assertEquals(0, connection.sentFrames.size)
+
+        flow.launchIn(CoroutineScope(connection.job))
+
+        connection.sentAsFlow().test {
+            expectItem().let { frame ->
+                assertTrue(frame is RequestFrame)
+                assertEquals(FrameType.RequestStream, frame.type)
+                assertEquals(2, frame.initialRequest)
+            }
+            delay(200)
+            expectNoEvents()
+            connection.sendToReceiver(NextPayloadFrame(1, Payload.Empty))
+            delay(200)
+            expectNoEvents() //will fail here if `Flow invariant is violated`
+            connection.sendToReceiver(NextPayloadFrame(1, Payload.Empty))
+            delay(200)
+            expectItem().let { frame ->
+                println(frame)
+                assertTrue(frame is CancelFrame)
+            }
+            delay(200)
+            expectNoEvents()
+        }
+    }
+
     @Test
     fun testHandleSetupException() = test {
         val errorMessage = "error"
diff --git a/rsocket-core/src/jsTest/kotlin/io/rsocket/kotlin/Test.kt b/rsocket-core/src/jsTest/kotlin/io/rsocket/kotlin/Test.kt
@@ -25,3 +25,6 @@ actual fun test(timeout: Duration?, block: suspend CoroutineScope.() -> Unit): d
         else -> withTimeout(timeout) { block() }
     }
 }
+
+//JS is single threaded, so it have only one dispatcher backed by one threed
+actual val anotherDispatcher: CoroutineDispatcher get() = Dispatchers.Default
diff --git a/rsocket-core/src/jvmTest/kotlin/io/rsocket/kotlin/Test.kt b/rsocket-core/src/jvmTest/kotlin/io/rsocket/kotlin/Test.kt
@@ -31,3 +31,5 @@ actual fun test(timeout: Duration?, block: suspend CoroutineScope.() -> Unit): U
         else -> withTimeout(timeout) { block() }
     }
 }
+
+actual val anotherDispatcher: CoroutineDispatcher get() = Dispatchers.IO

Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@ kotlin {`
`36`	`36`	`}`
`37`	`37`	`val commonTest by getting {`
`38`	`38`	`dependencies {`
	`39`	`+ implementation("app.cash.turbine:turbine:0.2.1")`
`39`	`40`	`implementation("io.ktor:ktor-utils:1.4.0")`
`40`	`41`	`implementation(project(":rsocket-transport-local"))`
`41`	`42`	`}`
Original file line number	Diff line number	Diff line change
`@@ -25,3 +25,6 @@ actual fun test(timeout: Duration?, block: suspend CoroutineScope.() -> Unit): d`
`25`	`25`	`else -> withTimeout(timeout) { block() }`
`26`	`26`	`}`
`27`	`27`	`}`
	`28`	`+`
	`29`	`+//JS is single threaded, so it have only one dispatcher backed by one threed`
	`30`	`+actual val anotherDispatcher: CoroutineDispatcher get() = Dispatchers.Default`
Original file line number	Diff line number	Diff line change
`@@ -31,3 +31,5 @@ actual fun test(timeout: Duration?, block: suspend CoroutineScope.() -> Unit): U`
`31`	`31`	`else -> withTimeout(timeout) { block() }`
`32`	`32`	`}`
`33`	`33`	`}`
	`34`	`+`
	`35`	`+actual val anotherDispatcher: CoroutineDispatcher get() = Dispatchers.IO`