Add option to provide Mono<Payload> for the SETUP payload

Signed-off-by: Rossen Stoyanchev <[email protected]>

Author: rstoyanchev

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

@@ -47,6 +47,7 @@
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 import reactor.util.annotation.Nullable;
+import reactor.util.function.Tuples;
 import reactor.util.retry.Retry;
 
 /**
@@ -77,7 +78,7 @@ public class RSocketConnector {
   private static final BiConsumer<RSocket, Invalidatable> INVALIDATE_FUNCTION =
       (r, i) -> r.onClose().subscribe(null, __ -> i.invalidate(), i::invalidate);
 
-  private Payload setupPayload = EmptyPayload.INSTANCE;
+  private Mono<Payload> setupPayloadMono = Mono.empty();
   private String metadataMimeType = "application/binary";
   private String dataMimeType = "application/binary";
   private Duration keepAliveInterval = Duration.ofSeconds(20);
@@ -118,22 +119,38 @@ public static Mono<RSocket> connectWith(ClientTransport transport) {
   }
 
   /**
-   * Provide a {@code Payload} with data and/or metadata for the initial {@code SETUP} frame. Data
-   * and metadata should be formatted according to the MIME types specified via {@link
+   * Provide a {@code Mono} from which to obtain the {@code Payload} for the initial SETUP frame.
+   * Data and metadata should be formatted according to the MIME types specified via {@link
    * #dataMimeType(String)} and {@link #metadataMimeType(String)}.
    *
-   * @param payload the payload containing data and/or metadata for the {@code SETUP} frame. Note,
-   *     if the instance of the given payload is not a {@link DefaultPayload}, its content will be
-   *     copied
+   * @param setupPayloadMono the payload with data and/or metadata for the {@code SETUP} frame.
+   * @return the same instance for method chaining
+   * @since 1.0.2
+   * @see <a href="https://github.com/rsocket/rsocket/blob/master/Protocol.md#frame-setup">SETUP
+   *     Frame</a>
+   */
+  public RSocketConnector setupPayload(Mono<Payload> setupPayloadMono) {
+    this.setupPayloadMono = setupPayloadMono;
+    return this;
+  }
+
+  /**
+   * Variant of {@link #setupPayload(Mono)} that accepts a {@code Payload} instance.
+   *
+   * <p>Note: if the given payload is {@link io.rsocket.util.ByteBufPayload}, it is copied to a
+   * {@link DefaultPayload} and released immediately. This ensures it can re-used to obtain a
+   * connection more than once.
+   *
+   * @param payload the payload with data and/or metadata for the {@code SETUP} frame.
    * @return the same instance for method chaining
    * @see <a href="https://github.com/rsocket/rsocket/blob/master/Protocol.md#frame-setup">SETUP
    *     Frame</a>
    */
   public RSocketConnector setupPayload(Payload payload) {
     if (payload instanceof DefaultPayload) {
-      this.setupPayload = payload;
+      this.setupPayloadMono = Mono.just(payload);
     } else {
-      this.setupPayload = DefaultPayload.create(Objects.requireNonNull(payload));
+      this.setupPayloadMono = Mono.just(DefaultPayload.create(Objects.requireNonNull(payload)));
       payload.release();
     }
     return this;
@@ -479,9 +496,20 @@ public Mono<RSocket> connect(Supplier<ClientTransport> transportSupplier) {
                     mtu > 0
                         ? new FragmentationDuplexConnection(connection, mtu, "client")
                         : new ReassemblyDuplexConnection(connection));
+
     return connectionMono
         .flatMap(
-            connection -> {
+            connection ->
+                setupPayloadMono
+                    .defaultIfEmpty(EmptyPayload.INSTANCE)
+                    .map(setupPayload -> Tuples.of(connection, setupPayload))
+                    .doOnError(ex -> connection.dispose())
+                    .doOnCancel(connection::dispose))
+        .flatMap(
+            tuple -> {
+              DuplexConnection connection = tuple.getT1();
+              Payload setupPayload = tuple.getT2();
+
               ByteBuf resumeToken;
               KeepAliveHandler keepAliveHandler;
               DuplexConnection wrappedConnection;

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

@@ -9,6 +9,8 @@
 import io.rsocket.test.util.TestClientTransport;
 import io.rsocket.util.ByteBufPayload;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
 import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
@@ -99,4 +101,50 @@ public void ensuresThatMonoFromRSocketConnectorCanBeUsedForMultipleSubscriptions
         .allMatch(ReferenceCounted::release);
     Assertions.assertThat(setupPayload.refCnt()).isZero();
   }
+
+  @Test
+  public void ensuresThatSetupPayloadProvidedAsMonoIsReleased() {
+    List<Payload> saved = new ArrayList<>();
+    Mono<Payload> setupPayloadMono =
+        Mono.create(
+            sink -> {
+              Payload payload = ByteBufPayload.create("TestData", "TestMetadata");
+              saved.add(payload);
+              sink.success(payload);
+            });
+
+    TestClientTransport testClientTransport = new TestClientTransport();
+    Mono<RSocket> connectionMono =
+        RSocketConnector.create().setupPayload(setupPayloadMono).connect(testClientTransport);
+
+    connectionMono
+        .as(StepVerifier::create)
+        .expectNextCount(1)
+        .expectComplete()
+        .verify(Duration.ofMillis(100));
+
+    connectionMono
+        .as(StepVerifier::create)
+        .expectNextCount(1)
+        .expectComplete()
+        .verify(Duration.ofMillis(100));
+
+    Assertions.assertThat(testClientTransport.testConnection().getSent())
+        .hasSize(2)
+        .allMatch(
+            bb -> {
+              DefaultConnectionSetupPayload payload = new DefaultConnectionSetupPayload(bb);
+              return payload.getDataUtf8().equals("TestData")
+                  && payload.getMetadataUtf8().equals("TestMetadata");
+            })
+        .allMatch(ReferenceCounted::release);
+
+    Assertions.assertThat(saved)
+        .as("Metadata and data were consumed and released as slices")
+        .allMatch(
+            payload ->
+                payload.refCnt() == 1
+                    && payload.data().refCnt() == 0
+                    && payload.metadata().refCnt() == 0);
+  }
 }