introduces .connect method

Signed-off-by: Oleh Dokuka <[email protected]>
Signed-off-by: Oleh Dokuka <[email protected]>
Co-authored-by: Rossen Stoyanchev <[email protected]>
Signed-off-by: Oleh Dokuka <[email protected]>

Author: 2 people

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

@@ -543,6 +543,11 @@ void subscribe() {
       reconnectMono.subscribe(stressSubscriber);
     }
 
+    @Actor
+    void connect() {
+      reconnectMono.resolvingInner.connect();
+    }
+
     @Arbiter
     public void arbiter(IIIIII_Result r) {
       r.r1 = stressSubscription.subscribes;

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

@@ -77,6 +77,17 @@
  */
 public interface RSocketClient extends Closeable {
 
+  /**
+   * Connect to the remote rsocket endpoint, if not yet connected. This method is a shortcut for
+   * {@code RSocketClient#source().subscribe()}.
+   *
+   * @return {@code true} if an attempt to connect was triggered or if already connected, or {@code
+   *     false} if the client is terminated.
+   */
+  default boolean connect() {
+    throw new NotImplementedException();
+  }
+
   default Mono<Void> onClose() {
     return Mono.error(new NotImplementedException());
   }

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

@@ -41,6 +41,11 @@ public RSocket rsocket() {
     return rsocket;
   }
 
+  @Override
+  public boolean connect() {
+    throw new UnsupportedOperationException("Connect does not apply to a server side RSocket");
+  }
+
   @Override
   public Mono<RSocket> source() {
     return Mono.just(rsocket);

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

@@ -331,6 +331,30 @@ protected void doOnDispose() {
     // no ops
   }
 
+  public final boolean connect() {
+    for (; ; ) {
+      final BiConsumer<T, Throwable>[] a = this.subscribers;
+
+      if (a == TERMINATED) {
+        return false;
+      }
+
+      if (a == READY) {
+        return true;
+      }
+
+      if (a != EMPTY_UNSUBSCRIBED) {
+        // do nothing if already started
+        return true;
+      }
+
+      if (SUBSCRIBERS.compareAndSet(this, a, EMPTY_SUBSCRIBED)) {
+        this.doSubscribe();
+        return true;
+      }
+    }
+  }
+
   final int add(BiConsumer<T, Throwable> ps) {
     for (; ; ) {
       BiConsumer<T, Throwable>[] a = this.subscribers;

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

@@ -45,6 +45,11 @@ public Mono<Void> onClose() {
     return rSocketPool.onClose();
   }
 
+  @Override
+  public boolean connect() {
+    return rSocketPool.connect();
+  }
+
   /** Return {@code Mono} that selects an RSocket from the underlying pool. */
   @Override
   public Mono<RSocket> source() {

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

@@ -327,6 +327,30 @@ protected void doOnDispose() {
     // no ops
   }
 
+  public final boolean connect() {
+    for (; ; ) {
+      final BiConsumer<T, Throwable>[] a = this.subscribers;
+
+      if (a == TERMINATED) {
+        return false;
+      }
+
+      if (a == READY) {
+        return true;
+      }
+
+      if (a != EMPTY_UNSUBSCRIBED) {
+        // do nothing if already started
+        return true;
+      }
+
+      if (SUBSCRIBERS.compareAndSet(this, a, EMPTY_SUBSCRIBED)) {
+        this.doSubscribe();
+        return true;
+      }
+    }
+  }
+
   final int add(BiConsumer<T, Throwable> ps) {
     for (; ; ) {
       BiConsumer<T, Throwable>[] a = this.subscribers;

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

@@ -28,6 +28,7 @@
 import io.rsocket.frame.PayloadFrameCodec;
 import io.rsocket.frame.decoder.PayloadDecoder;
 import io.rsocket.internal.subscriber.AssertSubscriber;
+import io.rsocket.test.util.TestDuplexConnection;
 import io.rsocket.util.ByteBufPayload;
 import io.rsocket.util.RSocketProxy;
 import java.time.Duration;
@@ -495,6 +496,88 @@ public Mono<Void> onClose() {
     Assertions.assertThat(rule.socket.isDisposed()).isTrue();
   }
 
+  @Test
+  public void shouldResolveOnStartSource() {
+    AssertSubscriber<RSocket> assertSubscriber = AssertSubscriber.create();
+    Assertions.assertThat(rule.client.connect()).isTrue();
+    rule.client.source().subscribe(assertSubscriber);
+    rule.delayer.run();
+    assertSubscriber.assertTerminated().assertValueCount(1);
+
+    rule.client.dispose();
+
+    Assertions.assertThat(rule.client.isDisposed()).isTrue();
+
+    AssertSubscriber<Void> assertSubscriber1 = AssertSubscriber.create();
+
+    rule.client.onClose().subscribe(assertSubscriber1);
+
+    assertSubscriber1.assertTerminated().assertComplete();
+
+    Assertions.assertThat(rule.socket.isDisposed()).isTrue();
+  }
+
+  @Test
+  public void shouldNotStartIfAlreadyDisposed() {
+    Assertions.assertThat(rule.client.connect()).isTrue();
+    Assertions.assertThat(rule.client.connect()).isTrue();
+    rule.delayer.run();
+
+    rule.client.dispose();
+
+    Assertions.assertThat(rule.client.connect()).isFalse();
+
+    Assertions.assertThat(rule.client.isDisposed()).isTrue();
+
+    AssertSubscriber<Void> assertSubscriber1 = AssertSubscriber.create();
+
+    rule.client.onClose().subscribe(assertSubscriber1);
+
+    assertSubscriber1.assertTerminated().assertComplete();
+
+    Assertions.assertThat(rule.socket.isDisposed()).isTrue();
+  }
+
+  @Test
+  public void shouldBeRestartedIfSourceWasClosed() {
+    AssertSubscriber<RSocket> assertSubscriber = AssertSubscriber.create();
+    AssertSubscriber<Void> terminateSubscriber = AssertSubscriber.create();
+
+    Assertions.assertThat(rule.client.connect()).isTrue();
+    rule.client.source().subscribe(assertSubscriber);
+    rule.client.onClose().subscribe(terminateSubscriber);
+
+    rule.delayer.run();
+
+    assertSubscriber.assertTerminated().assertValueCount(1);
+
+    rule.socket.dispose();
+
+    terminateSubscriber.assertNotTerminated();
+    Assertions.assertThat(rule.client.isDisposed()).isFalse();
+
+    rule.connection = new TestDuplexConnection(rule.allocator);
+    rule.socket = rule.newRSocket();
+    rule.producer = Sinks.one();
+
+    AssertSubscriber<RSocket> assertSubscriber2 = AssertSubscriber.create();
+
+    Assertions.assertThat(rule.client.connect()).isTrue();
+    rule.client.source().subscribe(assertSubscriber2);
+
+    rule.delayer.run();
+
+    assertSubscriber2.assertTerminated().assertValueCount(1);
+
+    rule.client.dispose();
+
+    terminateSubscriber.assertTerminated().assertComplete();
+
+    Assertions.assertThat(rule.client.connect()).isFalse();
+
+    Assertions.assertThat(rule.socket.isDisposed()).isTrue();
+  }
+
   @Test
   public void shouldDisposeOriginalSourceIfRacing() {
     for (int i = 0; i < RaceTestConstants.REPEATS; i++) {
@@ -523,6 +606,37 @@ public void shouldDisposeOriginalSourceIfRacing() {
     }
   }
 
+  @Test
+  public void shouldStartOriginalSourceOnceIfRacing() {
+    for (int i = 0; i < RaceTestConstants.REPEATS; i++) {
+      ClientSocketRule rule = new ClientSocketRule();
+
+      rule.init();
+
+      AssertSubscriber<RSocket> assertSubscriber = AssertSubscriber.create();
+
+      RaceTestUtils.race(
+          () -> rule.client.source().subscribe(assertSubscriber), () -> rule.client.connect());
+
+      Assertions.assertThat(rule.producer.currentSubscriberCount()).isOne();
+
+      rule.delayer.run();
+
+      assertSubscriber.assertTerminated();
+
+      rule.client.dispose();
+
+      Assertions.assertThat(rule.client.isDisposed()).isTrue();
+      Assertions.assertThat(rule.socket.isDisposed()).isTrue();
+
+      AssertSubscriber<Void> assertSubscriber1 = AssertSubscriber.create();
+
+      rule.client.onClose().subscribe(assertSubscriber1);
+
+      assertSubscriber1.assertTerminated().assertComplete();
+    }
+  }
+
   public static class ClientSocketRule extends AbstractSocketRule<RSocket> {
 
     protected RSocketClient client;