optimizes FragmentationDuplexConnection by extending ReassemblyDuplexConnection

Signed-off-by: Oleh Dokuka <[email protected]>

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/fragmentation/FragmentationDuplexConnection.java

@@ -34,13 +34,14 @@
 import reactor.core.publisher.Mono;
 
 /**
- * A {@link DuplexConnection} implementation that fragments {@link ByteBuf}s.
+ * A {@link DuplexConnection} implementation that fragments and reassembly {@link ByteBuf}s.
  *
  * @see <a
  *     href="https://github.com/rsocket/rsocket/blob/master/Protocol.md#fragmentation-and-reassembly">Fragmentation
  *     and Reassembly</a>
  */
-public final class FragmentationDuplexConnection implements DuplexConnection {
+public final class FragmentationDuplexConnection extends ReassemblyDuplexConnection
+    implements DuplexConnection {
   private static final int MIN_MTU_SIZE = 64;
   private static final Logger logger = LoggerFactory.getLogger(FragmentationDuplexConnection.class);
   private final DuplexConnection delegate;
@@ -54,11 +55,13 @@ public FragmentationDuplexConnection(
       DuplexConnection delegate,
       ByteBufAllocator allocator,
       int mtu,
-      boolean encodeLength,
+      boolean encodeAndEncodeLength,
       String type) {
+    super(delegate, allocator, encodeAndEncodeLength);
+
     Objects.requireNonNull(delegate, "delegate must not be null");
     Objects.requireNonNull(allocator, "byteBufAllocator must not be null");
-    this.encodeLength = encodeLength;
+    this.encodeLength = encodeAndEncodeLength;
     this.allocator = allocator;
     this.delegate = delegate;
     this.mtu = assertMtu(mtu);
@@ -137,19 +140,4 @@ private ByteBuf encode(ByteBuf frame) {
       return frame;
     }
   }
-
-  @Override
-  public Flux<ByteBuf> receive() {
-    return delegate.receive();
-  }
-
-  @Override
-  public Mono<Void> onClose() {
-    return delegate.onClose();
-  }
-
-  @Override
-  public void dispose() {
-    delegate.dispose();
-  }
 }

rsocket-core/src/main/java/io/rsocket/fragmentation/ReassemblyDuplexConnection.java

@@ -32,7 +32,7 @@
  *     href="https://github.com/rsocket/rsocket/blob/master/Protocol.md#fragmentation-and-reassembly">Fragmentation
  *     and Reassembly</a>
  */
-public final class ReassemblyDuplexConnection implements DuplexConnection {
+public class ReassemblyDuplexConnection implements DuplexConnection {
   private final DuplexConnection delegate;
   private final FrameReassembler frameReassembler;
   private final boolean decodeLength;

rsocket-transport-local/src/main/java/io/rsocket/transport/local/LocalClientTransport.java

@@ -66,32 +66,26 @@ private Mono<DuplexConnection> connect() {
           UnboundedProcessor<ByteBuf> out = new UnboundedProcessor<>();
           MonoProcessor<Void> closeNotifier = MonoProcessor.create();
 
-          server.accept(
-              new ReassemblyDuplexConnection(
-                  new LocalDuplexConnection(out, in, closeNotifier),
-                  ByteBufAllocator.DEFAULT,
-                  false));
+          server.accept(new LocalDuplexConnection(out, in, closeNotifier));
 
-          return Mono.just(
-              (DuplexConnection)
-                  new ReassemblyDuplexConnection(
-                      new LocalDuplexConnection(in, out, closeNotifier),
-                      ByteBufAllocator.DEFAULT,
-                      false));
+          return Mono.just((DuplexConnection) new LocalDuplexConnection(in, out, closeNotifier));
         });
   }
 
   @Override
   public Mono<DuplexConnection> connect(int mtu) {
     Mono<DuplexConnection> isError = FragmentationDuplexConnection.checkMtu(mtu);
     Mono<DuplexConnection> connect = isError != null ? isError : connect();
-    if (mtu > 0) {
-      return connect.map(
-          duplexConnection ->
-              new FragmentationDuplexConnection(
-                  duplexConnection, ByteBufAllocator.DEFAULT, mtu, false, "client"));
-    } else {
-      return connect;
-    }
+
+    return connect.map(
+        duplexConnection -> {
+          if (mtu > 0) {
+            return new FragmentationDuplexConnection(
+                duplexConnection, ByteBufAllocator.DEFAULT, mtu, false, "client");
+          } else {
+            return new ReassemblyDuplexConnection(
+                duplexConnection, ByteBufAllocator.DEFAULT, false);
+          }
+        });
   }
 }

rsocket-transport-local/src/main/java/io/rsocket/transport/local/LocalServerTransport.java

@@ -20,6 +20,7 @@
 import io.rsocket.Closeable;
 import io.rsocket.DuplexConnection;
 import io.rsocket.fragmentation.FragmentationDuplexConnection;
+import io.rsocket.fragmentation.ReassemblyDuplexConnection;
 import io.rsocket.transport.ClientTransport;
 import io.rsocket.transport.ServerTransport;
 import java.util.Objects;
@@ -168,6 +169,9 @@ public void accept(DuplexConnection duplexConnection) {
         duplexConnection =
             new FragmentationDuplexConnection(
                 duplexConnection, ByteBufAllocator.DEFAULT, mtu, false, "server");
+      } else {
+        duplexConnection =
+            new ReassemblyDuplexConnection(duplexConnection, ByteBufAllocator.DEFAULT, false);
       }
 
       acceptor.apply(duplexConnection).subscribe();

rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/client/TcpClientTransport.java

@@ -105,8 +105,7 @@ public Mono<DuplexConnection> connect(int mtu) {
                 c -> {
                   if (mtu > 0) {
                     return new FragmentationDuplexConnection(
-                        new ReassemblyDuplexConnection(
-                            new TcpDuplexConnection(c, false), ByteBufAllocator.DEFAULT, true),
+                        new TcpDuplexConnection(c, false),
                         ByteBufAllocator.DEFAULT,
                         mtu,
                         true,

rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/client/WebsocketClientTransport.java

@@ -161,13 +161,14 @@ public Mono<DuplexConnection> connect(int mtu) {
             .connect()
             .map(
                 c -> {
-                  DuplexConnection connection =
-                      new ReassemblyDuplexConnection(
-                          new WebsocketDuplexConnection(c), ByteBufAllocator.DEFAULT, false);
+                  DuplexConnection connection = new WebsocketDuplexConnection(c);
                   if (mtu > 0) {
                     connection =
                         new FragmentationDuplexConnection(
                             connection, ByteBufAllocator.DEFAULT, mtu, false, "client");
+                  } else {
+                    connection =
+                        new ReassemblyDuplexConnection(connection, ByteBufAllocator.DEFAULT, false);
                   }
                   return connection;
                 });

rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/server/TcpServerTransport.java

@@ -106,8 +106,7 @@ public Mono<CloseableChannel> start(ConnectionAcceptor acceptor, int mtu) {
                   if (mtu > 0) {
                     connection =
                         new FragmentationDuplexConnection(
-                            new ReassemblyDuplexConnection(
-                                new TcpDuplexConnection(c, false), ByteBufAllocator.DEFAULT, true),
+                            new TcpDuplexConnection(c, false),
                             ByteBufAllocator.DEFAULT,
                             mtu,
                             true,

rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/server/WebsocketRouteTransport.java

@@ -103,13 +103,13 @@ public static BiFunction<WebsocketInbound, WebsocketOutbound, Publisher<Void>> n
   public static BiFunction<WebsocketInbound, WebsocketOutbound, Publisher<Void>> newHandler(
       ConnectionAcceptor acceptor, int mtu) {
     return (in, out) -> {
-      DuplexConnection connection =
-          new ReassemblyDuplexConnection(
-              new WebsocketDuplexConnection((Connection) in), ByteBufAllocator.DEFAULT, false);
+      DuplexConnection connection = new WebsocketDuplexConnection((Connection) in);
       if (mtu > 0) {
         connection =
             new FragmentationDuplexConnection(
                 connection, ByteBufAllocator.DEFAULT, mtu, false, "server");
+      } else {
+        connection = new ReassemblyDuplexConnection(connection, ByteBufAllocator.DEFAULT, false);
       }
       return acceptor.apply(connection).then(out.neverComplete());
     };

rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/server/WebsocketServerTransport.java

@@ -126,14 +126,15 @@ public Mono<CloseableChannel> start(ConnectionAcceptor acceptor, int mtu) {
                   return response.sendWebsocket(
                       (in, out) -> {
                         DuplexConnection connection =
-                            new ReassemblyDuplexConnection(
-                                new WebsocketDuplexConnection((Connection) in),
-                                ByteBufAllocator.DEFAULT,
-                                false);
+                            new WebsocketDuplexConnection((Connection) in);
                         if (mtu > 0) {
                           connection =
                               new FragmentationDuplexConnection(
                                   connection, ByteBufAllocator.DEFAULT, mtu, false, "server");
+                        } else {
+                          connection =
+                              new ReassemblyDuplexConnection(
+                                  connection, ByteBufAllocator.DEFAULT, false);
                         }
                         return acceptor.apply(connection).then(out.neverComplete());
                       },