fixes onClose issue (#608)

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

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/internal/BaseDuplexConnection.java

@@ -0,0 +1,30 @@
+package io.rsocket.internal;
+
+import io.rsocket.DuplexConnection;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.MonoProcessor;
+
+public abstract class BaseDuplexConnection implements DuplexConnection {
+  private MonoProcessor<Void> onClose = MonoProcessor.create();
+
+  public BaseDuplexConnection() {
+    onClose.doFinally(s -> doOnClose()).subscribe();
+  }
+
+  protected abstract void doOnClose();
+
+  @Override
+  public final Mono<Void> onClose() {
+    return onClose;
+  }
+
+  @Override
+  public final void dispose() {
+    onClose.onComplete();
+  }
+
+  @Override
+  public final boolean isDisposed() {
+    return onClose.isDisposed();
+  }
+}

rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/TcpDuplexConnection.java

@@ -20,20 +20,18 @@
 import io.netty.buffer.ByteBufAllocator;
 import io.rsocket.DuplexConnection;
 import io.rsocket.frame.FrameLengthFlyweight;
+import io.rsocket.internal.BaseDuplexConnection;
 import java.util.Objects;
 import org.reactivestreams.Publisher;
-import reactor.core.Disposable;
 import reactor.core.Fuseable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.netty.Connection;
-import reactor.netty.FutureMono;
 
 /** An implementation of {@link DuplexConnection} that connects via TCP. */
-public final class TcpDuplexConnection implements DuplexConnection {
+public final class TcpDuplexConnection extends BaseDuplexConnection {
 
   private final Connection connection;
-  private final Disposable channelClosed;
   private final ByteBufAllocator allocator = ByteBufAllocator.DEFAULT;
   private final boolean encodeLength;
 
@@ -55,37 +53,21 @@ public TcpDuplexConnection(Connection connection) {
   public TcpDuplexConnection(Connection connection, boolean encodeLength) {
     this.encodeLength = encodeLength;
     this.connection = Objects.requireNonNull(connection, "connection must not be null");
-    this.channelClosed =
-        FutureMono.from(connection.channel().closeFuture())
-            .doFinally(
-                s -> {
-                  if (!isDisposed()) {
-                    dispose();
-                  }
-                })
-            .subscribe();
-  }
 
-  @Override
-  public void dispose() {
-    connection.dispose();
-  }
-
-  @Override
-  public boolean isDisposed() {
-    return connection.isDisposed();
+    connection
+        .channel()
+        .closeFuture()
+        .addListener(
+            future -> {
+              if (!isDisposed()) dispose();
+            });
   }
 
   @Override
-  public Mono<Void> onClose() {
-    return connection
-        .onDispose()
-        .doFinally(
-            s -> {
-              if (!channelClosed.isDisposed()) {
-                channelClosed.dispose();
-              }
-            });
+  protected void doOnClose() {
+    if (!connection.isDisposed()) {
+      connection.dispose();
+    }
   }
 
   @Override

rsocket-transport-netty/src/main/java/io/rsocket/transport/netty/WebsocketDuplexConnection.java

@@ -18,14 +18,13 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
 import io.rsocket.DuplexConnection;
+import io.rsocket.internal.BaseDuplexConnection;
 import java.util.Objects;
 import org.reactivestreams.Publisher;
-import reactor.core.Disposable;
 import reactor.core.Fuseable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.netty.Connection;
-import reactor.netty.FutureMono;
 import reactor.util.concurrent.Queues;
 
 /**
@@ -35,10 +34,9 @@
  * for message oriented transports so this must be specifically dropped from Frames sent and
  * stitched back on for frames received.
  */
-public final class WebsocketDuplexConnection implements DuplexConnection {
+public final class WebsocketDuplexConnection extends BaseDuplexConnection {
 
   private final Connection connection;
-  private final Disposable channelClosed;
 
   /**
    * Creates a new instance
@@ -47,37 +45,21 @@ public final class WebsocketDuplexConnection implements DuplexConnection {
    */
   public WebsocketDuplexConnection(Connection connection) {
     this.connection = Objects.requireNonNull(connection, "connection must not be null");
-    this.channelClosed =
-        FutureMono.from(connection.channel().closeFuture())
-            .doFinally(
-                s -> {
-                  if (!isDisposed()) {
-                    dispose();
-                  }
-                })
-            .subscribe();
-  }
-
-  @Override
-  public void dispose() {
-    connection.dispose();
-  }
 
-  @Override
-  public boolean isDisposed() {
-    return connection.isDisposed();
+    connection
+        .channel()
+        .closeFuture()
+        .addListener(
+            future -> {
+              if (!isDisposed()) dispose();
+            });
   }
 
   @Override
-  public Mono<Void> onClose() {
-    return connection
-        .onDispose()
-        .doFinally(
-            s -> {
-              if (!channelClosed.isDisposed()) {
-                channelClosed.dispose();
-              }
-            });
+  protected void doOnClose() {
+    if (!connection.isDisposed()) {
+      connection.dispose();
+    }
   }
 
   @Override