refactors to convenient Retry

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

Author: OlegDokuka

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

@@ -48,10 +48,9 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
-import org.reactivestreams.Publisher;
 import reactor.core.Disposable;
-import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.retry.Retry;
 
 /** Factory for creating RSocket clients and servers. */
 public class RSocketFactory {
@@ -133,7 +132,7 @@ public static class ClientRSocketFactory implements ClientTransportAcceptor {
     private boolean leaseEnabled;
     private Supplier<Leases<?>> leasesSupplier = Leases::new;
     private boolean reconnectEnabled;
-    private Function<Flux<Throwable>, ? extends Publisher<?>> whenFactory;
+    private Retry retrySpec;
 
     private ByteBufAllocator allocator = ByteBufAllocator.DEFAULT;
 
@@ -355,9 +354,8 @@ public ClientRSocketFactory reconnect() {
      * @param whenFactory a retry factory applied for {@link Mono.retryWhen(whenFactory)}
      * @return a shared instance of {@code Mono<RSocket>}.
      */
-    public ClientRSocketFactory reconnect(
-        Function<Flux<Throwable>, ? extends Publisher<?>> whenFactory) {
-      this.whenFactory = Objects.requireNonNull(whenFactory);
+    public ClientRSocketFactory reconnect(Retry retrySpec) {
+      this.retrySpec = Objects.requireNonNull(retrySpec);
       this.reconnectEnabled = true;
       return this;
     }
@@ -529,7 +527,7 @@ public Mono<RSocket> start() {
                 source -> {
                   if (reconnectEnabled) {
                     return new ReconnectMono<>(
-                        whenFactory == null ? source : source.retryWhen(whenFactory),
+                        retrySpec == null ? source : source.retryWhen(retrySpec),
                         Disposable::dispose,
                         INVALIDATE_FUNCTION);
                   } else {
@@ -841,9 +839,11 @@ public Mono<T> start() {
 
               @Override
               public Mono<T> get() {
-                return transportServer
-                    .get()
-                    .start(duplexConnection -> acceptor(serverSetup, duplexConnection), mtu)
+                return Mono.fromSupplier(transportServer)
+                    .flatMap(
+                        transport ->
+                            transport.start(
+                                duplexConnection -> acceptor(serverSetup, duplexConnection), mtu))
                     .doOnNext(c -> c.onClose().doFinally(v -> serverSetup.dispose()).subscribe());
               }
             });

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

@@ -470,3 +470,8 @@ public Object scanUnsafe(Attr key) {
     }
   }
 }
+
+interface Invalidatable {
+
+  void invalidate();
+}

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

@@ -53,7 +53,8 @@
 public class ReconnectMonoTests {
 
   private Queue<RetryContext<?>> retries = new ConcurrentLinkedQueue<>();
-  private Queue<Tuple2<Object, Invalidatable>> received = new ConcurrentLinkedQueue<>();
+  private Queue<Tuple2<Object, ReconnectMono.Invalidatable>> received =
+      new ConcurrentLinkedQueue<>();
   private Queue<Object> expired = new ConcurrentLinkedQueue<>();
 
   @Test
@@ -906,7 +907,7 @@ Consumer<? super RetryContext<?>> onRetry() {
     return context -> retries.add(context);
   }
 
-  <T> BiConsumer<T, Invalidatable> onValue() {
+  <T> BiConsumer<T, ReconnectMono.Invalidatable> onValue() {
     return (v, __) -> received.add(Tuples.of(v, __));
   }