adds first frame handling timeout

Signed-off-by: Oleh Dokuka <[email protected]> Signed-off-by: Oleh Dokuka <[email protected]> Signed-off-by: Oleh Dokuka <[email protected]>
rsocket · Nov 9, 2021 · 93c30b6 · 93c30b6
1 parent 1f71914
commit 93c30b6
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 5 deletions.
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,22 @@ public RSocketServer maxInboundPayloadSize(int maxInboundPayloadSize) {
     return this;
   }
 
+  /**
+   * Specifies timeout for the first incoming frame from the accepted connection.
+   *
+   * <p>By default this is set to 1 minute.
+   *
+   * @param timeout duration
+   * @return the same instance for method chaining
+   */
+  public RSocketServer setupHandlingTimeout(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 +305,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 +344,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 +487,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().setupHandlingTimeout(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()