improves BaseDuplexConnection and fixes PingClient impl (#1062)

rsocket · Aug 20, 2022 · d330a32 · d330a32
1 parent 1fe2d64
commit d330a32
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 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
@@ -22,13 +22,10 @@
 import reactor.core.publisher.Sinks;
 
 public abstract class BaseDuplexConnection implements DuplexConnection {
-  protected Sinks.Empty<Void> onClose = Sinks.empty();
+  protected final Sinks.Empty<Void> onClose = Sinks.empty();
+  protected final UnboundedProcessor sender = new UnboundedProcessor(onClose::tryEmitEmpty);
 
-  protected UnboundedProcessor sender = new UnboundedProcessor();
-
-  public BaseDuplexConnection() {
-    onClose().doFinally(s -> doOnClose()).subscribe();
-  }
+  public BaseDuplexConnection() {}
 
   @Override
   public void sendFrame(int streamId, ByteBuf frame) {
@@ -48,7 +45,7 @@ public final Mono<Void> onClose() {
 
   @Override
   public final void dispose() {
-    onClose.tryEmitEmpty();
+    doOnClose();
   }
 
   @Override

diff --git a/rsocket-test/src/main/java/io/rsocket/test/PingClient.java b/rsocket-test/src/main/java/io/rsocket/test/PingClient.java
@@ -63,8 +63,8 @@ Flux<Payload> pingPong(
       BiFunction<RSocket, ? super Payload, ? extends Publisher<Payload>> interaction,
       int count,
       final Recorder histogram) {
-    return client
-        .flatMapMany(
+    return Flux.usingWhen(
+            client,
             rsocket ->
                 Flux.range(1, count)
                     .flatMap(
@@ -78,7 +78,11 @@ Flux<Payload> pingPong(
                                     histogram.recordValue(diff);
                                   });
                         },
-                        64))
+                        64),
+            rsocket -> {
+              rsocket.dispose();
+              return rsocket.onClose();
+            })
         .doOnError(Throwable::printStackTrace);
   }
 }