fixes performance degradation when fragmentation is used (#995)

rsocket · Mar 4, 2021 · e4d62b6 · e4d62b6
1 parent 765abb7
commit e4d62b6
Showing 1 changed file with 18 additions and 2 deletions.
diff --git a/rsocket-core/src/main/java/io/rsocket/fragmentation/FragmentationDuplexConnection.java b/rsocket-core/src/main/java/io/rsocket/fragmentation/FragmentationDuplexConnection.java
@@ -84,7 +84,18 @@ public static int assertMtu(int mtu) {
 
   @Override
   public Mono<Void> send(Publisher<ByteBuf> frames) {
-    return Flux.from(frames).concatMap(this::sendOne).then();
+    return delegate.send(
+        Flux.from(frames)
+            .concatMap(
+                frame -> {
+                  FrameType frameType = FrameHeaderCodec.frameType(frame);
+                  int readableBytes = frame.readableBytes();
+                  if (!shouldFragment(frameType, readableBytes)) {
+                    return Flux.just(frame);
+                  }
+
+                  return logFragments(Flux.from(fragmentFrame(alloc(), mtu, frame, frameType)));
+                }));
   }
 
   @Override
@@ -95,6 +106,11 @@ public Mono<Void> sendOne(ByteBuf frame) {
       return delegate.sendOne(frame);
     }
     Flux<ByteBuf> fragments = Flux.from(fragmentFrame(alloc(), mtu, frame, frameType));
+    fragments = logFragments(fragments);
+    return delegate.send(fragments);
+  }
+
+  protected Flux<ByteBuf> logFragments(Flux<ByteBuf> fragments) {
     if (logger.isDebugEnabled()) {
       fragments =
           fragments.doOnNext(
@@ -107,6 +123,6 @@ public Mono<Void> sendOne(ByteBuf frame) {
                     ByteBufUtil.prettyHexDump(byteBuf));
               });
     }
-    return delegate.send(fragments);
+    return fragments;
   }
 }