adds first frame handling timeout (rsocket#1027)

OlegDokuka · rstoyanchev · web-flow · commit efd1269a149e · 2021-11-09T23:27:46.000+02:00
Co-authored-by: Rossen Stoyanchev &lt;rstoyanchev@users.noreply.github.com&gt;
diff --git a/.github/workflows/gradle-all.yml b/.github/workflows/gradle-all.yml
@@ -142,7 +142,10 @@ jobs:
         run: chmod +x gradlew
       - name: Publish Packages to Artifactory
         if: ${{ matrix.jdk == '1.8' }}
-        run: ./gradlew -PversionSuffix="-${githubRef#refs/heads/}-SNAPSHOT" -PbuildNumber="${buildNumber}" publishMavenPublicationToGitHubPackagesRepository --no-daemon --stacktrace
+        run: |
+          githubRef="${githubRef#refs/heads/}"
+          githubRef="${githubRef////-}"
+          ./gradlew -PversionSuffix="-${githubRef}-SNAPSHOT" -PbuildNumber="${buildNumber}" publishMavenPublicationToGitHubPackagesRepository --no-daemon --stacktrace
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           githubRef: ${{ github.ref }}
diff --git a/rsocket-core/src/main/java/io/rsocket/core/RSocketServer.java b/rsocket-core/src/main/java/io/rsocket/core/RSocketServer.java
@@ -41,6 +41,7 @@
 import io.rsocket.plugins.RequestInterceptor;
 import io.rsocket.resume.SessionManager;
 import io.rsocket.transport.ServerTransport;
+import java.time.Duration;
 import java.util.Objects;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
@@ -70,6 +71,7 @@ public final class RSocketServer {
   private int mtu = 0;
   private int maxInboundPayloadSize = Integer.MAX_VALUE;
   private PayloadDecoder payloadDecoder = PayloadDecoder.DEFAULT;
+  private Duration timeout = Duration.ofMinutes(1);
 
   private RSocketServer() {}
 
@@ -223,6 +225,23 @@ public RSocketServer maxInboundPayloadSize(int maxInboundPayloadSize) {
     return this;
   }
 
+  /**
+   * Specify the max time to wait for the first frame (e.g. {@code SETUP}) on an accepted
+   * connection.
+   *
+   * <p>By default this is set to 1 minute.
+   *
+   * @param timeout duration
+   * @return the same instance for method chaining
+   */
+  public RSocketServer maxTimeToFirstFrame(Duration timeout) {
+    if (timeout.isNegative() || timeout.isZero()) {
+      throw new IllegalArgumentException("Setup Handling Timeout should be greater than zero");
+    }
+    this.timeout = timeout;
+    return this;
+  }
+
   /**
    * When this is set, frames larger than the given maximum transmission unit (mtu) size value are
    * fragmented.
@@ -287,7 +306,7 @@ public RSocketServer payloadDecoder(PayloadDecoder decoder) {
   public <T extends Closeable> Mono<T> bind(ServerTransport<T> transport) {
     return Mono.defer(
         new Supplier<Mono<T>>() {
-          final ServerSetup serverSetup = serverSetup();
+          final ServerSetup serverSetup = serverSetup(timeout);
 
           @Override
           public Mono<T> get() {
@@ -326,7 +345,7 @@ public ServerTransport.ConnectionAcceptor asConnectionAcceptor() {
   public ServerTransport.ConnectionAcceptor asConnectionAcceptor(int maxFrameLength) {
     assertValidateSetup(maxFrameLength, maxInboundPayloadSize, mtu);
     return new ServerTransport.ConnectionAcceptor() {
-      private final ServerSetup serverSetup = serverSetup();
+      private final ServerSetup serverSetup = serverSetup(timeout);
 
       @Override
       public Mono<Void> apply(DuplexConnection connection) {
@@ -469,12 +488,13 @@ private Mono<Void> acceptSetup(
         });
   }
 
-  private ServerSetup serverSetup() {
-    return resume != null ? createSetup() : new ServerSetup.DefaultServerSetup();
+  private ServerSetup serverSetup(Duration timeout) {
+    return resume != null ? createSetup(timeout) : new ServerSetup.DefaultServerSetup(timeout);
   }
 
-  ServerSetup createSetup() {
+  ServerSetup createSetup(Duration timeout) {
     return new ServerSetup.ResumableServerSetup(
+        timeout,
         new SessionManager(),
         resume.getSessionDuration(),
         resume.getStreamTimeout(),
diff --git a/rsocket-core/src/main/java/io/rsocket/core/ServerSetup.java b/rsocket-core/src/main/java/io/rsocket/core/ServerSetup.java
@@ -36,9 +36,16 @@
 
 abstract class ServerSetup {
 
+  final Duration timeout;
+
+  protected ServerSetup(Duration timeout) {
+    this.timeout = timeout;
+  }
+
   Mono<Tuple2<ByteBuf, DuplexConnection>> init(DuplexConnection connection) {
     return Mono.<Tuple2<ByteBuf, DuplexConnection>>create(
             sink -> sink.onRequest(__ -> new SetupHandlingDuplexConnection(connection, sink)))
+        .timeout(this.timeout)
         .or(connection.onClose().then(Mono.error(ClosedChannelException::new)));
   }
 
@@ -57,6 +64,10 @@ void sendError(DuplexConnection duplexConnection, RSocketErrorException exceptio
 
   static class DefaultServerSetup extends ServerSetup {
 
+    DefaultServerSetup(Duration timeout) {
+      super(timeout);
+    }
+
     @Override
     public Mono<Void> acceptRSocketSetup(
         ByteBuf frame,
@@ -86,11 +97,13 @@ static class ResumableServerSetup extends ServerSetup {
     private final boolean cleanupStoreOnKeepAlive;
 
     ResumableServerSetup(
+        Duration timeout,
         SessionManager sessionManager,
         Duration resumeSessionDuration,
         Duration resumeStreamTimeout,
         Function<? super ByteBuf, ? extends ResumableFramesStore> resumeStoreFactory,
         boolean cleanupStoreOnKeepAlive) {
+      super(timeout);
       this.sessionManager = sessionManager;
       this.resumeSessionDuration = resumeSessionDuration;
       this.resumeStreamTimeout = resumeStreamTimeout;
diff --git a/rsocket-core/src/main/java/io/rsocket/core/SetupHandlingDuplexConnection.java b/rsocket-core/src/main/java/io/rsocket/core/SetupHandlingDuplexConnection.java
@@ -96,6 +96,7 @@ public void request(long n) {
 
   @Override
   public void cancel() {
+    source.dispose();
     s.cancel();
   }
 
diff --git a/rsocket-core/src/test/java/io/rsocket/core/RSocketServerTest.java b/rsocket-core/src/test/java/io/rsocket/core/RSocketServerTest.java
@@ -29,11 +29,13 @@
 import io.rsocket.test.util.TestServerTransport;
 import java.time.Duration;
 import java.util.Random;
+import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 import reactor.core.Scannable;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
 import reactor.test.StepVerifier;
+import reactor.test.scheduler.VirtualTimeScheduler;
 
 public class RSocketServerTest {
 
@@ -60,6 +62,32 @@ public void unexpectedFramesBeforeSetupFrame() {
         .hasNoLeaks();
   }
 
+  @Test
+  public void timeoutOnNoFirstFrame() {
+    final VirtualTimeScheduler scheduler = VirtualTimeScheduler.getOrSet();
+    try {
+      TestServerTransport transport = new TestServerTransport();
+      RSocketServer.create().maxTimeToFirstFrame(Duration.ofMinutes(2)).bind(transport).block();
+
+      final TestDuplexConnection duplexConnection = transport.connect();
+
+      scheduler.advanceTimeBy(Duration.ofMinutes(1));
+
+      Assertions.assertThat(duplexConnection.isDisposed()).isFalse();
+
+      scheduler.advanceTimeBy(Duration.ofMinutes(1));
+
+      StepVerifier.create(duplexConnection.onClose())
+          .expectSubscription()
+          .expectComplete()
+          .verify(Duration.ofSeconds(10));
+
+      FrameAssert.assertThat(duplexConnection.pollFrame()).isNull();
+    } finally {
+      VirtualTimeScheduler.reset();
+    }
+  }
+
   @Test
   public void ensuresMaxFrameLengthCanNotBeLessThenMtu() {
     RSocketServer.create()
diff --git a/rsocket-test/src/main/java/io/rsocket/test/TransportTest.java b/rsocket-test/src/main/java/io/rsocket/test/TransportTest.java
@@ -39,6 +39,7 @@
 import java.io.InputStreamReader;
 import java.net.SocketAddress;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadLocalRandom;
@@ -96,7 +97,18 @@ default void close() {
     getTransportPair().responder.awaitAllInteractionTermination(getTimeout());
     getTransportPair().dispose();
     getTransportPair().awaitClosed();
-    RuntimeException throwable = new RuntimeException();
+    RuntimeException throwable =
+        new RuntimeException() {
+          @Override
+          public synchronized Throwable fillInStackTrace() {
+            return this;
+          }
+
+          @Override
+          public String getMessage() {
+            return Arrays.toString(getSuppressed());
+          }
+        };
 
     try {
       getTransportPair().byteBufAllocator2.assertHasNoLeaks();
@@ -776,8 +788,11 @@ public void onSubscribe(Subscription s) {
 
       @Override
       public void onNext(ByteBuf buf) {
-        actual.onNext(buf);
-        buf.release();
+        try {
+          actual.onNext(buf);
+        } finally {
+          buf.release();
+        }
       }
 
       Mono<Void> onClose() {

Original file line number	Diff line number	Diff line change
`@@ -96,6 +96,7 @@ public void request(long n) {`
`96`	`96`
`97`	`97`	`@Override`
`98`	`98`	`public void cancel() {`
	`99`	`+ source.dispose();`
`99`	`100`	`s.cancel();`
`100`	`101`	`}`
`101`	`102`