fragmentation and reassembly

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/core/RSocketConnector.kt

@@ -27,6 +27,7 @@ import kotlinx.coroutines.*
 @OptIn(TransportApi::class, RSocketLoggingApi::class)
 public class RSocketConnector internal constructor(
     private val loggerFactory: LoggerFactory,
+    private val maxFragmentSize: Int,
     private val interceptors: Interceptors,
     private val connectionConfigProvider: () -> ConnectionConfig,
     private val acceptor: ConnectionAcceptor,
@@ -61,6 +62,7 @@ public class RSocketConnector internal constructor(
         try {
             val requester = connection.connect(
                 isServer = false,
+                maxFragmentSize = maxFragmentSize,
                 interceptors = interceptors,
                 connectionConfig = connectionConfig,
                 acceptor = acceptor

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/core/RSocketConnectorBuilder.kt

@@ -26,6 +26,13 @@ import kotlinx.coroutines.*
 public class RSocketConnectorBuilder internal constructor() {
     @RSocketLoggingApi
     public var loggerFactory: LoggerFactory = DefaultLoggerFactory
+    public var maxFragmentSize: Int = 0
+        set(value) {
+            require(value == 0 || value >= 64) {
+                "maxFragmentSize should be zero (no fragmentation) or greater than or equal to 64, but was $value"
+            }
+            field = value
+        }
 
     private val connectionConfig: ConnectionConfigBuilder = ConnectionConfigBuilder()
     private val interceptors: InterceptorsBuilder = InterceptorsBuilder()
@@ -96,6 +103,7 @@ public class RSocketConnectorBuilder internal constructor() {
     @OptIn(RSocketLoggingApi::class)
     internal fun build(): RSocketConnector = RSocketConnector(
         loggerFactory,
+        maxFragmentSize,
         interceptors.build(),
         connectionConfig.producer(),
         acceptor ?: defaultAcceptor,

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/core/RSocketServer.kt

@@ -27,6 +27,7 @@ import kotlinx.coroutines.*
 @OptIn(TransportApi::class, RSocketLoggingApi::class)
 public class RSocketServer internal constructor(
     private val loggerFactory: LoggerFactory,
+    private val maxFragmentSize: Int,
     private val interceptors: Interceptors,
 ) {
 
@@ -44,6 +45,7 @@ public class RSocketServer internal constructor(
             else                                  -> try {
                 connect(
                     isServer = true,
+                    maxFragmentSize = maxFragmentSize,
                     interceptors = interceptors,
                     connectionConfig = ConnectionConfig(
                         keepAlive = setupFrame.keepAlive,

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/core/RSocketServerBuilder.kt

@@ -22,6 +22,13 @@ import io.rsocket.kotlin.logging.*
 public class RSocketServerBuilder internal constructor() {
     @RSocketLoggingApi
     public var loggerFactory: LoggerFactory = DefaultLoggerFactory
+    public var maxFragmentSize: Int = 0
+        set(value) {
+            require(value == 0 || value >= 64) {
+                "maxFragmentSize should be zero (no fragmentation) or greater than or equal to 64, but was $value"
+            }
+            field = value
+        }
 
     private val interceptors: InterceptorsBuilder = InterceptorsBuilder()
 
@@ -30,7 +37,7 @@ public class RSocketServerBuilder internal constructor() {
     }
 
     @OptIn(RSocketLoggingApi::class)
-    internal fun build(): RSocketServer = RSocketServer(loggerFactory, interceptors.build())
+    internal fun build(): RSocketServer = RSocketServer(loggerFactory, maxFragmentSize, interceptors.build())
 }
 
 public fun RSocketServer(configure: RSocketServerBuilder.() -> Unit = {}): RSocketServer {

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/Connect.kt

@@ -25,13 +25,15 @@ import kotlinx.coroutines.*
 @OptIn(TransportApi::class)
 internal suspend inline fun Connection.connect(
     isServer: Boolean,
+    maxFragmentSize: Int,
     interceptors: Interceptors,
     connectionConfig: ConnectionConfig,
     acceptor: ConnectionAcceptor
 ): RSocket {
     val keepAliveHandler = KeepAliveHandler(connectionConfig.keepAlive)
     val prioritizer = Prioritizer()
-    val streamsStorage = StreamsStorage(isServer)
+    val frameSender = FrameSender(prioritizer, pool, maxFragmentSize)
+    val streamsStorage = StreamsStorage(isServer, pool)
     val requestJob = SupervisorJob(job)
 
     requestJob.invokeOnCompletion {
@@ -43,7 +45,7 @@ internal suspend inline fun Connection.connect(
     val requestScope = CoroutineScope(requestJob + Dispatchers.Unconfined + CoroutineExceptionHandler { _, _ -> })
     val connectionScope = CoroutineScope(job + Dispatchers.Unconfined + CoroutineExceptionHandler { _, _ -> })
 
-    val requester = interceptors.wrapRequester(RSocketRequester(job, prioritizer, streamsStorage, requestScope))
+    val requester = interceptors.wrapRequester(RSocketRequester(job, frameSender, streamsStorage, requestScope, pool))
     val requestHandler = interceptors.wrapResponder(
         with(interceptors.wrapAcceptor(acceptor)) {
             ConnectionAcceptorContext(connectionConfig, requester).accept()
@@ -71,7 +73,7 @@ internal suspend inline fun Connection.connect(
 
     // start frame handling
     connectionScope.launch {
-        val rSocketResponder = RSocketResponder(prioritizer, requestHandler, requestScope)
+        val rSocketResponder = RSocketResponder(frameSender, requestHandler, requestScope)
         while (isActive) {
             receiveFrame().closeOnError { frame ->
                 when (frame.streamId) {

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/FrameSender.kt

@@ -0,0 +1,120 @@
+/*
+ * Copyright 2015-2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.rsocket.kotlin.internal
+
+import io.ktor.utils.io.core.*
+import io.ktor.utils.io.core.internal.*
+import io.ktor.utils.io.pool.*
+import io.rsocket.kotlin.frame.*
+import io.rsocket.kotlin.frame.io.*
+import io.rsocket.kotlin.payload.*
+import kotlinx.coroutines.*
+import kotlin.math.*
+
+private const val lengthSize = 3
+private const val headerSize = 6
+private const val fragmentOffset = lengthSize + headerSize
+private const val fragmentOffsetWithMetadata = fragmentOffset + lengthSize
+
+internal class FrameSender(
+    private val prioritizer: Prioritizer,
+    private val pool: ObjectPool<ChunkBuffer>,
+    private val maxFragmentSize: Int
+) {
+
+    suspend fun sendKeepAlive(respond: Boolean, lastPosition: Long, data: ByteReadPacket): Unit =
+        prioritizer.send(KeepAliveFrame(respond, lastPosition, data))
+
+    suspend fun sendMetadataPush(metadata: ByteReadPacket): Unit = prioritizer.send(MetadataPushFrame(metadata))
+
+    suspend fun sendCancel(id: Int): Unit = withContext(NonCancellable) { prioritizer.send(CancelFrame(id)) }
+    suspend fun sendError(id: Int, throwable: Throwable): Unit = withContext(NonCancellable) { prioritizer.send(ErrorFrame(id, throwable)) }
+    suspend fun sendRequestN(id: Int, n: Int): Unit = prioritizer.send(RequestNFrame(id, n))
+
+    suspend fun sendRequestPayload(type: FrameType, streamId: Int, payload: Payload, initialRequest: Int = 0) {
+        sendFragmented(type, streamId, payload, false, false, initialRequest)
+    }
+
+    suspend fun sendNextPayload(streamId: Int, payload: Payload) {
+        sendFragmented(FrameType.Payload, streamId, payload, false, true, 0)
+    }
+
+    suspend fun sendNextCompletePayload(streamId: Int, payload: Payload) {
+        sendFragmented(FrameType.Payload, streamId, payload, true, true, 0)
+    }
+
+    suspend fun sendCompletePayload(streamId: Int) {
+        sendFragmented(FrameType.Payload, streamId, Payload.Empty, true, false, 0)
+    }
+
+    private suspend fun sendFragmented(
+        type: FrameType,
+        streamId: Int,
+        payload: Payload,
+        complete: Boolean,
+        next: Boolean,
+        initialRequest: Int
+    ) {
+        //TODO release on fail ?
+        if (!payload.isFragmentable(type.hasInitialRequest)) {
+            prioritizer.send(RequestFrame(type, streamId, false, complete, next, initialRequest, payload))
+            return
+        }
+
+        val data = payload.data
+        val metadata = payload.metadata
+
+        val fragmentSize = maxFragmentSize - fragmentOffset - (if (type.hasInitialRequest) Int.SIZE_BYTES else 0)
+
+        var first = true
+        var remaining = fragmentSize
+        if (metadata != null) remaining -= lengthSize
+
+        do {
+            val metadataFragment = if (metadata != null && metadata.isNotEmpty) {
+                if (!first) remaining -= lengthSize
+                val length = min(metadata.remaining.toInt(), remaining)
+                remaining -= length
+                metadata.readPacket(pool, length)
+            } else null
+
+            val dataFragment = if (remaining > 0 && data.isNotEmpty) {
+                val length = min(data.remaining.toInt(), remaining)
+                remaining -= length
+                data.readPacket(pool, length)
+            } else {
+                ByteReadPacket.Empty
+            }
+
+            val fType = if (first && type.isRequestType) type else FrameType.Payload
+            val fragment = Payload(dataFragment, metadataFragment)
+            val follows = metadata != null && metadata.isNotEmpty || data.isNotEmpty
+            prioritizer.send(RequestFrame(fType, streamId, follows, (!follows && complete), !fType.isRequestType, initialRequest, fragment))
+            first = false
+            remaining = fragmentSize
+        } while (follows)
+    }
+
+    private fun Payload.isFragmentable(hasInitialRequest: Boolean) = when (maxFragmentSize) {
+        0    -> false
+        else -> when (val meta = metadata) {
+            null -> data.remaining > maxFragmentSize - fragmentOffset - (if (hasInitialRequest) Int.SIZE_BYTES else 0)
+            else -> data.remaining + meta.remaining > maxFragmentSize - fragmentOffsetWithMetadata - (if (hasInitialRequest) Int.SIZE_BYTES else 0)
+        }
+    }
+
+}

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/Prioritizer.kt

@@ -30,7 +30,7 @@ internal class Prioritizer {
     private val commonChannel = SafeChannel<Frame>(Channel.UNLIMITED)
 
     suspend fun send(frame: Frame) {
-        if (frame.type != FrameType.Cancel && frame.type != FrameType.Error) currentCoroutineContext().ensureActive()
+        currentCoroutineContext().ensureActive()
         val channel = if (frame.streamId == 0) priorityChannel else commonChannel
         channel.send(frame)
     }

rsocket-core/src/commonMain/kotlin/io/rsocket/kotlin/internal/RSocketRequester.kt

@@ -17,6 +17,8 @@
 package io.rsocket.kotlin.internal
 
 import io.ktor.utils.io.core.*
+import io.ktor.utils.io.core.internal.*
+import io.ktor.utils.io.pool.*
 import io.rsocket.kotlin.*
 import io.rsocket.kotlin.frame.*
 import io.rsocket.kotlin.internal.handler.*
@@ -28,83 +30,84 @@ import kotlinx.coroutines.flow.*
 @OptIn(ExperimentalStreamsApi::class)
 internal class RSocketRequester(
     connectionJob: Job,
-    private val prioritizer: Prioritizer,
+    private val sender: FrameSender,
     private val streamsStorage: StreamsStorage,
-    private val requestScope: CoroutineScope
+    private val requestScope: CoroutineScope,
+    private val pool: ObjectPool<ChunkBuffer>
 ) : RSocket {
     override val job: Job = connectionJob
 
     override suspend fun metadataPush(metadata: ByteReadPacket) {
         ensureActiveOrRelease(metadata)
         metadata.closeOnError {
-            prioritizer.send(MetadataPushFrame(metadata))
+            sender.sendMetadataPush(metadata)
         }
     }
 
     override suspend fun fireAndForget(payload: Payload) {
         ensureActiveOrRelease(payload)
 
-        val streamId = streamsStorage.nextId()
+        val id = streamsStorage.nextId()
         try {
-            prioritizer.send(RequestFireAndForgetFrame(streamId, payload))
+            sender.sendRequestPayload(FrameType.RequestFnF, id, payload)
         } catch (cause: Throwable) {
             payload.release()
-            if (job.isActive) prioritizer.send(CancelFrame(streamId)) //if cancelled during fragmentation
+            if (job.isActive) sender.sendCancel(id) //if cancelled during fragmentation
             throw cause
         }
     }
 
     override suspend fun requestResponse(payload: Payload): Payload {
         ensureActiveOrRelease(payload)
 
-        val streamId = streamsStorage.nextId()
+        val id = streamsStorage.nextId()
 
         val deferred = CompletableDeferred<Payload>()
-        val handler = RequesterRequestResponseFrameHandler(streamId, streamsStorage, deferred)
-        streamsStorage.save(streamId, handler)
+        val handler = RequesterRequestResponseFrameHandler(id, streamsStorage, deferred, pool)
+        streamsStorage.save(id, handler)
 
-        return handler.receiveOrCancel(streamId, payload) {
-            prioritizer.send(RequestResponseFrame(streamId, payload))
+        return handler.receiveOrCancel(id, payload) {
+            sender.sendRequestPayload(FrameType.RequestResponse, id, payload)
             deferred.await()
         }
     }
 
     override fun requestStream(payload: Payload): Flow<Payload> = requestFlow { strategy, initialRequest ->
         ensureActiveOrRelease(payload)
 
-        val streamId = streamsStorage.nextId()
+        val id = streamsStorage.nextId()
 
         val channel = SafeChannel<Payload>(Channel.UNLIMITED)
-        val handler = RequesterRequestStreamFrameHandler(streamId, streamsStorage, channel)
-        streamsStorage.save(streamId, handler)
+        val handler = RequesterRequestStreamFrameHandler(id, streamsStorage, channel, pool)
+        streamsStorage.save(id, handler)
 
-        handler.receiveOrCancel(streamId, payload) {
-            prioritizer.send(RequestStreamFrame(streamId, initialRequest, payload))
-            emitAllWithRequestN(channel, strategy) { prioritizer.send(RequestNFrame(streamId, it)) }
+        handler.receiveOrCancel(id, payload) {
+            sender.sendRequestPayload(FrameType.RequestStream, id, payload, initialRequest)
+            emitAllWithRequestN(channel, strategy) { sender.sendRequestN(id, it) }
         }
     }
 
     override fun requestChannel(initPayload: Payload, payloads: Flow<Payload>): Flow<Payload> = requestFlow { strategy, initialRequest ->
         ensureActiveOrRelease(initPayload)
 
-        val streamId = streamsStorage.nextId()
+        val id = streamsStorage.nextId()
 
         val channel = SafeChannel<Payload>(Channel.UNLIMITED)
         val limiter = Limiter(0)
-        val sender = Job(requestScope.coroutineContext.job)
-        val handler = RequesterRequestChannelFrameHandler(streamId, streamsStorage, limiter, sender, channel)
-        streamsStorage.save(streamId, handler)
+        val payloadsJob = Job(requestScope.coroutineContext.job)
+        val handler = RequesterRequestChannelFrameHandler(id, streamsStorage, limiter, payloadsJob, channel, pool)
+        streamsStorage.save(id, handler)
 
-        handler.receiveOrCancel(streamId, initPayload) {
-            prioritizer.send(RequestChannelFrame(streamId, initialRequest, initPayload))
+        handler.receiveOrCancel(id, initPayload) {
+            sender.sendRequestPayload(FrameType.RequestChannel, id, initPayload, initialRequest)
             //TODO lazy?
-            requestScope.launch(sender) {
-                handler.sendOrFail(streamId) {
-                    payloads.collectLimiting(limiter) { prioritizer.send(NextPayloadFrame(streamId, it)) }
-                    prioritizer.send(CompletePayloadFrame(streamId))
+            requestScope.launch(payloadsJob) {
+                handler.sendOrFail(id) {
+                    payloads.collectLimiting(limiter) { sender.sendNextPayload(id, it) }
+                    sender.sendCompletePayload(id)
                 }
             }
-            emitAllWithRequestN(channel, strategy) { prioritizer.send(RequestNFrame(streamId, it)) }
+            emitAllWithRequestN(channel, strategy) { sender.sendRequestN(id, it) }
         }
     }
 
@@ -114,7 +117,7 @@ internal class RSocketRequester(
             onSendComplete()
         } catch (cause: Throwable) {
             val isFailed = onSendFailed(cause)
-            if (job.isActive && isFailed) prioritizer.send(ErrorFrame(id, cause))
+            if (job.isActive && isFailed) sender.sendError(id, cause)
             throw cause
         }
     }
@@ -127,7 +130,7 @@ internal class RSocketRequester(
         } catch (cause: Throwable) {
             payload.release()
             val isCancelled = onReceiveCancelled(cause)
-            if (job.isActive && isCancelled) prioritizer.send(CancelFrame(id))
+            if (job.isActive && isCancelled) sender.sendCancel(id)
             throw cause
         }
     }