introduces onClose listener for RSocketClient

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

Author: Oleh Dokuka

rsocket-core/src/jcstress/java/io/rsocket/core/ReconnectMonoStressTest.java

@@ -517,7 +517,7 @@ public void arbiter(IIIIII_Result r) {
       id = {"1, 0, 1, 0, 1, 2"},
       expect = ACCEPTABLE)
   @State
-  public static class SubscribeBlockRace extends BaseStressTest {
+  public static class SubscribeBlockConnectRace extends BaseStressTest {
 
     String receivedValue;
 

rsocket-core/src/main/java/io/rsocket/core/DefaultRSocketClient.java

@@ -35,6 +35,7 @@
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoOperator;
 import reactor.core.publisher.Operators;
+import reactor.core.publisher.Sinks;
 import reactor.util.annotation.Nullable;
 import reactor.util.context.Context;
 
@@ -65,19 +66,27 @@ class DefaultRSocketClient extends ResolvingOperator<RSocket>
 
   final Mono<RSocket> source;
 
+  final Sinks.Empty<Void> onDisposeSink;
+
   volatile Subscription s;
 
   static final AtomicReferenceFieldUpdater<DefaultRSocketClient, Subscription> S =
       AtomicReferenceFieldUpdater.newUpdater(DefaultRSocketClient.class, Subscription.class, "s");
 
   DefaultRSocketClient(Mono<RSocket> source) {
     this.source = unwrapReconnectMono(source);
+    this.onDisposeSink = Sinks.empty();
   }
 
   private Mono<RSocket> unwrapReconnectMono(Mono<RSocket> source) {
     return source instanceof ReconnectMono ? ((ReconnectMono<RSocket>) source).getSource() : source;
   }
 
+  @Override
+  public Mono<Void> onClose() {
+    return this.onDisposeSink.asMono();
+  }
+
   @Override
   public Mono<RSocket> source() {
     return Mono.fromDirect(this);
@@ -194,6 +203,12 @@ protected void doOnValueExpired(RSocket value) {
   @Override
   protected void doOnDispose() {
     Operators.terminate(S, this);
+    final RSocket value = this.value;
+    if (value != null) {
+      value.onClose().subscribe(null, onDisposeSink::tryEmitError, onDisposeSink::tryEmitEmpty);
+    } else {
+      onDisposeSink.tryEmitEmpty();
+    }
   }
 
   static final class FlatMapMain<R> implements CoreSubscriber<Payload>, Context, Scannable {

rsocket-core/src/main/java/io/rsocket/core/RSocketClient.java

@@ -15,12 +15,13 @@
  */
 package io.rsocket.core;
 
+import io.rsocket.Closeable;
 import io.rsocket.Payload;
 import io.rsocket.RSocket;
 import org.reactivestreams.Publisher;
-import reactor.core.Disposable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import sun.reflect.generics.reflectiveObjects.NotImplementedException;
 
 /**
  * Contract for performing RSocket requests.
@@ -74,7 +75,11 @@
  * @since 1.1
  * @see io.rsocket.loadbalance.LoadbalanceRSocketClient
  */
-public interface RSocketClient extends Disposable {
+public interface RSocketClient extends Closeable {
+
+  default Mono<Void> onClose() {
+    return Mono.error(new NotImplementedException());
+  }
 
   /** Return the underlying source used to obtain a shared {@link RSocket} connection. */
   Mono<RSocket> source();

rsocket-core/src/main/java/io/rsocket/core/RSocketClientAdapter.java

@@ -46,6 +46,11 @@ public Mono<RSocket> source() {
     return Mono.just(rsocket);
   }
 
+  @Override
+  public Mono<Void> onClose() {
+    return rsocket.onClose();
+  }
+
   @Override
   public Mono<Void> fireAndForget(Mono<Payload> payloadMono) {
     return payloadMono.flatMap(rsocket::fireAndForget);

rsocket-core/src/main/java/io/rsocket/loadbalance/LoadbalanceRSocketClient.java

@@ -40,6 +40,11 @@ private LoadbalanceRSocketClient(RSocketPool rSocketPool) {
     this.rSocketPool = rSocketPool;
   }
 
+  @Override
+  public Mono<Void> onClose() {
+    return rSocketPool.onClose();
+  }
+
   /** Return {@code Mono} that selects an RSocket from the underlying pool. */
   @Override
   public Mono<RSocket> source() {

rsocket-core/src/main/java/io/rsocket/loadbalance/PooledRSocket.java

@@ -26,6 +26,7 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Operators;
+import reactor.core.publisher.Sinks;
 import reactor.util.context.Context;
 
 /** Default implementation of {@link RSocket} stored in {@link RSocketPool} */
@@ -35,6 +36,7 @@ final class PooledRSocket extends ResolvingOperator<RSocket>
   final RSocketPool parent;
   final Mono<RSocket> rSocketSource;
   final LoadbalanceTarget loadbalanceTarget;
+  final Sinks.Empty<Void> onCloseSink;
 
   volatile Subscription s;
 
@@ -46,6 +48,7 @@ final class PooledRSocket extends ResolvingOperator<RSocket>
     this.parent = parent;
     this.rSocketSource = rSocketSource;
     this.loadbalanceTarget = loadbalanceTarget;
+    this.onCloseSink = Sinks.unsafe().empty();
   }
 
   @Override
@@ -155,6 +158,12 @@ void doCleanup(Throwable t) {
         break;
       }
     }
+
+    if (t == ON_DISPOSE) {
+      this.onCloseSink.tryEmitEmpty();
+    } else {
+      this.onCloseSink.tryEmitError(t);
+    }
   }
 
   @Override
@@ -165,6 +174,13 @@ protected void doOnValueExpired(RSocket value) {
   @Override
   protected void doOnDispose() {
     Operators.terminate(S, this);
+
+    final RSocket value = this.value;
+    if (value != null) {
+      value.onClose().subscribe(null, onCloseSink::tryEmitError, onCloseSink::tryEmitEmpty);
+    } else {
+      onCloseSink.tryEmitEmpty();
+    }
   }
 
   @Override
@@ -193,7 +209,12 @@ public Mono<Void> metadataPush(Payload payload) {
   }
 
   LoadbalanceTarget target() {
-    return loadbalanceTarget;
+    return this.loadbalanceTarget;
+  }
+
+  @Override
+  public Mono<Void> onClose() {
+    return this.onCloseSink.asMono();
   }
 
   @Override

rsocket-core/src/main/java/io/rsocket/loadbalance/RSocketPool.java

@@ -16,6 +16,7 @@
 package io.rsocket.loadbalance;
 
 import io.netty.util.ReferenceCountUtil;
+import io.rsocket.Closeable;
 import io.rsocket.Payload;
 import io.rsocket.RSocket;
 import io.rsocket.core.RSocketConnector;
@@ -28,16 +29,18 @@
 import java.util.ListIterator;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
+import java.util.stream.Collectors;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Operators;
+import reactor.core.publisher.Sinks;
 import reactor.util.annotation.Nullable;
 
 class RSocketPool extends ResolvingOperator<Object>
-    implements CoreSubscriber<List<LoadbalanceTarget>> {
+    implements CoreSubscriber<List<LoadbalanceTarget>>, Closeable {
 
   static final AtomicReferenceFieldUpdater<RSocketPool, PooledRSocket[]> ACTIVE_SOCKETS =
       AtomicReferenceFieldUpdater.newUpdater(
@@ -49,6 +52,7 @@ class RSocketPool extends ResolvingOperator<Object>
   final DeferredResolutionRSocket deferredResolutionRSocket = new DeferredResolutionRSocket(this);
   final RSocketConnector connector;
   final LoadbalanceStrategy loadbalanceStrategy;
+  final Sinks.Empty<Void> onAllClosedSink = Sinks.unsafe().empty();
   volatile PooledRSocket[] activeSockets;
   volatile Subscription s;
 
@@ -64,6 +68,11 @@ public RSocketPool(
     targetPublisher.subscribe(this);
   }
 
+  @Override
+  public Mono<Void> onClose() {
+    return onAllClosedSink.asMono();
+  }
+
   @Override
   protected void doOnDispose() {
     Operators.terminate(S, this);
@@ -72,6 +81,14 @@ protected void doOnDispose() {
     for (RSocket rSocket : activeSockets) {
       rSocket.dispose();
     }
+
+    if (activeSockets.length > 0) {
+      Mono.whenDelayError(
+              Arrays.stream(activeSockets).map(RSocket::onClose).collect(Collectors.toList()))
+          .subscribe(null, onAllClosedSink::tryEmitError, onAllClosedSink::tryEmitEmpty);
+    } else {
+      onAllClosedSink.tryEmitEmpty();
+    }
   }
 
   @Override

rsocket-core/src/test/java/io/rsocket/core/DefaultRSocketClientTests.java

@@ -29,6 +29,7 @@
 import io.rsocket.frame.decoder.PayloadDecoder;
 import io.rsocket.internal.subscriber.AssertSubscriber;
 import io.rsocket.util.ByteBufPayload;
+import io.rsocket.util.RSocketProxy;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -457,6 +458,43 @@ public void shouldDisposeOriginalSource() {
     Assertions.assertThat(rule.socket.isDisposed()).isTrue();
   }
 
+  @Test
+  public void shouldReceiveOnCloseNotificationOnDisposeOriginalSource() {
+    Sinks.Empty<Void> onCloseDelayer = Sinks.empty();
+    ClientSocketRule rule =
+        new ClientSocketRule() {
+          @Override
+          protected RSocket newRSocket() {
+            return new RSocketProxy(super.newRSocket()) {
+              @Override
+              public Mono<Void> onClose() {
+                return super.onClose().and(onCloseDelayer.asMono());
+              }
+            };
+          }
+        };
+    rule.init();
+    AssertSubscriber<RSocket> assertSubscriber = AssertSubscriber.create();
+    rule.client.source().subscribe(assertSubscriber);
+    rule.delayer.run();
+    assertSubscriber.assertTerminated().assertValueCount(1);
+
+    rule.client.dispose();
+
+    Assertions.assertThat(rule.client.isDisposed()).isTrue();
+
+    AssertSubscriber<Void> onCloseSubscriber = AssertSubscriber.create();
+
+    rule.client.onClose().subscribe(onCloseSubscriber);
+    onCloseSubscriber.assertNotTerminated();
+
+    onCloseDelayer.tryEmitEmpty();
+
+    onCloseSubscriber.assertTerminated().assertComplete();
+
+    Assertions.assertThat(rule.socket.isDisposed()).isTrue();
+  }
+
   @Test
   public void shouldDisposeOriginalSourceIfRacing() {
     for (int i = 0; i < RaceTestConstants.REPEATS; i++) {
@@ -485,7 +523,7 @@ public void shouldDisposeOriginalSourceIfRacing() {
     }
   }
 
-  public static class ClientSocketRule extends AbstractSocketRule<RSocketRequester> {
+  public static class ClientSocketRule extends AbstractSocketRule<RSocket> {
 
     protected RSocketClient client;
     protected Runnable delayer;
@@ -498,14 +536,16 @@ protected void doInit() {
       producer = Sinks.one();
       client =
           new DefaultRSocketClient(
-              producer
-                  .asMono()
-                  .doOnCancel(() -> socket.dispose())
-                  .doOnDiscard(Disposable.class, Disposable::dispose));
+              Mono.defer(
+                  () ->
+                      producer
+                          .asMono()
+                          .doOnCancel(() -> socket.dispose())
+                          .doOnDiscard(Disposable.class, Disposable::dispose)));
     }
 
     @Override
-    protected RSocketRequester newRSocket() {
+    protected RSocket newRSocket() {
       return new RSocketRequester(
           connection,
           PayloadDecoder.ZERO_COPY,