improves BaseDuplexConnection and related subclasses

Signed-off-by: Oleh Dokuka <odokuka@vmware.com> Signed-off-by: Oleh Dokuka <oleh.dokuka@icloud.com>
rsocket · Sep 14, 2022 · 000f6da · 000f6da
1 parent 9804688
commit 000f6da
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 25 deletions.
diff --git a/rsocket-core/src/main/java/io/rsocket/internal/BaseDuplexConnection.java b/rsocket-core/src/main/java/io/rsocket/internal/BaseDuplexConnection.java
@@ -39,7 +39,7 @@ public void sendFrame(int streamId, ByteBuf frame) {
   protected abstract void doOnClose();
 
   @Override
-  public final Mono<Void> onClose() {
+  public Mono<Void> onClose() {
     return onClose.asMono();
   }
 

diff --git a/rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/TcpDuplexConnection.java b/rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/TcpDuplexConnection.java
@@ -26,6 +26,7 @@
 import java.net.SocketAddress;
 import java.util.Objects;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 import reactor.netty.Connection;
 
 /** An implementation of {@link DuplexConnection} that connects via TCP. */
@@ -67,24 +68,19 @@ protected void doOnClose() {
     connection.dispose();
   }
 
+  @Override
+  public Mono<Void> onClose() {
+    return super.onClose().and(connection.onDispose());
+  }
+
   @Override
   public void sendErrorAndClose(RSocketErrorException e) {
     final ByteBuf errorFrame = ErrorFrameCodec.encode(alloc(), 0, e);
     connection
         .outbound()
         .sendObject(FrameLengthCodec.encode(alloc(), errorFrame.readableBytes(), errorFrame))
-        .then()
-        .subscribe(
-            null,
-            t -> onClose.tryEmitError(t),
-            () -> {
-              final Throwable cause = e.getCause();
-              if (cause == null) {
-                onClose.tryEmitEmpty();
-              } else {
-                onClose.tryEmitError(cause);
-              }
-            });
+        .subscribe(connection.disposeSubscriber());
+    sender.onComplete();
   }
 
   @Override

diff --git a/...t-transport-netty/src/main/java/io/rsocket/transport/netty/WebsocketDuplexConnection.java b/...t-transport-netty/src/main/java/io/rsocket/transport/netty/WebsocketDuplexConnection.java
@@ -25,6 +25,7 @@
 import java.net.SocketAddress;
 import java.util.Objects;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 import reactor.netty.Connection;
 
 /**
@@ -72,6 +73,11 @@ protected void doOnClose() {
     connection.dispose();
   }
 
+  @Override
+  public Mono<Void> onClose() {
+    return super.onClose().and(connection.onDispose());
+  }
+
   @Override
   public Flux<ByteBuf> receive() {
     return connection.inbound().receive();
@@ -83,17 +89,7 @@ public void sendErrorAndClose(RSocketErrorException e) {
     connection
         .outbound()
         .sendObject(new BinaryWebSocketFrame(errorFrame))
-        .then()
-        .subscribe(
-            null,
-            t -> onClose.tryEmitError(t),
-            () -> {
-              final Throwable cause = e.getCause();
-              if (cause == null) {
-                onClose.tryEmitEmpty();
-              } else {
-                onClose.tryEmitError(cause);
-              }
-            });
+        .subscribe(connection.disposeSubscriber());
+    sender.onComplete();
   }
 }