improves KeepAlive to use Scheduler time, fixes tests to be stable

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

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/keepalive/KeepAliveSupport.java

@@ -22,36 +22,45 @@
 import io.rsocket.frame.KeepAliveFrameCodec;
 import io.rsocket.resume.ResumeStateHolder;
 import java.time.Duration;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import reactor.core.Disposable;
 import reactor.core.publisher.Flux;
+import reactor.core.scheduler.Scheduler;
+import reactor.core.scheduler.Schedulers;
 
 public abstract class KeepAliveSupport implements KeepAliveFramesAcceptor {
+
   final ByteBufAllocator allocator;
-  private final Duration keepAliveInterval;
-  private final Duration keepAliveTimeout;
-  private final long keepAliveTimeoutMillis;
-  private volatile Consumer<KeepAlive> onTimeout;
-  private volatile Consumer<ByteBuf> onFrameSent;
-  private volatile Disposable ticksDisposable;
-  private final AtomicBoolean started = new AtomicBoolean();
+  final Scheduler scheduler;
+  final Duration keepAliveInterval;
+  final Duration keepAliveTimeout;
+  final long keepAliveTimeoutMillis;
+
+  final AtomicBoolean started = new AtomicBoolean();
+
+  volatile Consumer<KeepAlive> onTimeout;
+  volatile Consumer<ByteBuf> onFrameSent;
+  volatile Disposable ticksDisposable;
 
-  private volatile ResumeStateHolder resumeStateHolder;
-  private volatile long lastReceivedMillis;
+  volatile ResumeStateHolder resumeStateHolder;
+  volatile long lastReceivedMillis;
 
   private KeepAliveSupport(
       ByteBufAllocator allocator, int keepAliveInterval, int keepAliveTimeout) {
     this.allocator = allocator;
+    this.scheduler = Schedulers.parallel();
     this.keepAliveInterval = Duration.ofMillis(keepAliveInterval);
     this.keepAliveTimeout = Duration.ofMillis(keepAliveTimeout);
     this.keepAliveTimeoutMillis = keepAliveTimeout;
   }
 
   public KeepAliveSupport start() {
-    this.lastReceivedMillis = System.currentTimeMillis();
+    this.lastReceivedMillis = scheduler.now(TimeUnit.MILLISECONDS);
     if (started.compareAndSet(false, true)) {
-      ticksDisposable = Flux.interval(keepAliveInterval).subscribe(v -> onIntervalTick());
+      ticksDisposable =
+          Flux.interval(keepAliveInterval, scheduler).subscribe(v -> onIntervalTick());
     }
     return this;
   }
@@ -64,7 +73,7 @@ public void stop() {
 
   @Override
   public void receive(ByteBuf keepAliveFrame) {
-    this.lastReceivedMillis = System.currentTimeMillis();
+    this.lastReceivedMillis = scheduler.now(TimeUnit.MILLISECONDS);
     if (resumeStateHolder != null) {
       long remoteLastReceivedPos = remoteLastReceivedPosition(keepAliveFrame);
       resumeStateHolder.onImpliedPosition(remoteLastReceivedPos);
@@ -104,7 +113,7 @@ void send(ByteBuf frame) {
   }
 
   void tryTimeout() {
-    long now = System.currentTimeMillis();
+    long now = scheduler.now(TimeUnit.MILLISECONDS);
     if (now - lastReceivedMillis >= keepAliveTimeoutMillis) {
       if (onTimeout != null) {
         onTimeout.accept(new KeepAlive(keepAliveInterval, keepAliveTimeout));

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

@@ -37,17 +37,32 @@
 import io.rsocket.util.DefaultPayload;
 import java.time.Duration;
 import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import reactor.core.Disposable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
+import reactor.test.scheduler.VirtualTimeScheduler;
 
 public class KeepAliveTest {
   private static final int KEEP_ALIVE_INTERVAL = 100;
   private static final int KEEP_ALIVE_TIMEOUT = 1000;
   private static final int RESUMABLE_KEEP_ALIVE_TIMEOUT = 200;
 
+  VirtualTimeScheduler virtualTimeScheduler;
+
+  @BeforeEach
+  public void setUp() {
+    virtualTimeScheduler = VirtualTimeScheduler.getOrSet();
+  }
+
+  @AfterEach
+  public void tearDown() {
+    VirtualTimeScheduler.reset();
+  }
+
   static RSocketState requester(int tickPeriod, int timeout) {
     LeaksTrackingByteBufAllocator allocator =
         LeaksTrackingByteBufAllocator.instrument(ByteBufAllocator.DEFAULT);
@@ -95,7 +110,7 @@ static ResumableRSocketState resumableRequester(int tickPeriod, int timeout) {
   }
 
   @Test
-  void rSocketNotDisposedOnPresentKeepAlives() throws InterruptedException {
+  void rSocketNotDisposedOnPresentKeepAlives() {
     RSocketState requesterState = requester(KEEP_ALIVE_INTERVAL, KEEP_ALIVE_TIMEOUT);
 
     TestDuplexConnection connection = requesterState.connection();
@@ -108,21 +123,19 @@ void rSocketNotDisposedOnPresentKeepAlives() throws InterruptedException {
                         KeepAliveFrameCodec.encode(
                             requesterState.allocator, true, 0, Unpooled.EMPTY_BUFFER)));
 
-    Thread.sleep(KEEP_ALIVE_TIMEOUT * 2);
+    virtualTimeScheduler.advanceTimeBy(Duration.ofMillis(KEEP_ALIVE_TIMEOUT * 2));
 
     RSocket rSocket = requesterState.rSocket();
 
     Assertions.assertThat(rSocket.isDisposed()).isFalse();
 
     disposable.dispose();
-    Thread.sleep(KEEP_ALIVE_INTERVAL);
 
     requesterState.connection.dispose();
     requesterState.rSocket.dispose();
 
-    Thread.sleep(KEEP_ALIVE_TIMEOUT);
-
     Assertions.assertThat(requesterState.connection.getSent()).allMatch(ByteBuf::release);
+
     requesterState.allocator.assertHasNoLeaks();
   }
 
@@ -132,11 +145,12 @@ void noKeepAlivesSentAfterRSocketDispose() {
 
     requesterState.rSocket().dispose();
 
-    StepVerifier.create(
-            Flux.from(requesterState.connection().getSentAsPublisher())
-                .take(Duration.ofMillis(500)))
+    Duration duration = Duration.ofMillis(500);
+    StepVerifier.create(Flux.from(requesterState.connection().getSentAsPublisher()).take(duration))
+        .then(() -> virtualTimeScheduler.advanceTimeBy(duration))
         .expectComplete()
         .verify(Duration.ofSeconds(1));
+
     requesterState.allocator.assertHasNoLeaks();
   }
 
@@ -146,7 +160,7 @@ void rSocketDisposedOnMissingKeepAlives() {
 
     RSocket rSocket = requesterState.rSocket();
 
-    Mono.delay(Duration.ofMillis(KEEP_ALIVE_TIMEOUT * 2)).block();
+    virtualTimeScheduler.advanceTimeBy(Duration.ofMillis(KEEP_ALIVE_TIMEOUT * 2));
 
     Assertions.assertThat(rSocket.isDisposed()).isTrue();
     rSocket
@@ -161,72 +175,73 @@ void rSocketDisposedOnMissingKeepAlives() {
   }
 
   @Test
-  void clientRequesterSendsKeepAlives() throws InterruptedException {
+  void clientRequesterSendsKeepAlives() {
     RSocketState RSocketState = requester(KEEP_ALIVE_INTERVAL, KEEP_ALIVE_TIMEOUT);
     TestDuplexConnection connection = RSocketState.connection();
 
     StepVerifier.create(Flux.from(connection.getSentAsPublisher()).take(3))
+        .then(() -> virtualTimeScheduler.advanceTimeBy(Duration.ofMillis(KEEP_ALIVE_INTERVAL)))
         .expectNextMatches(this::keepAliveFrameWithRespondFlag)
+        .then(() -> virtualTimeScheduler.advanceTimeBy(Duration.ofMillis(KEEP_ALIVE_INTERVAL)))
         .expectNextMatches(this::keepAliveFrameWithRespondFlag)
+        .then(() -> virtualTimeScheduler.advanceTimeBy(Duration.ofMillis(KEEP_ALIVE_INTERVAL)))
         .expectNextMatches(this::keepAliveFrameWithRespondFlag)
         .expectComplete()
         .verify(Duration.ofSeconds(5));
 
     RSocketState.rSocket.dispose();
     RSocketState.connection.dispose();
 
-    Thread.sleep(KEEP_ALIVE_INTERVAL);
-
     RSocketState.allocator.assertHasNoLeaks();
   }
 
   @Test
-  void requesterRespondsToKeepAlives() throws InterruptedException {
+  void requesterRespondsToKeepAlives() {
     RSocketState rSocketState = requester(100_000, 100_000);
     TestDuplexConnection connection = rSocketState.connection();
-    Mono.delay(Duration.ofMillis(100))
+    Duration duration = Duration.ofMillis(100);
+    Mono.delay(duration)
         .subscribe(
             l ->
                 connection.addToReceivedBuffer(
                     KeepAliveFrameCodec.encode(
                         rSocketState.allocator, true, 0, Unpooled.EMPTY_BUFFER)));
 
     StepVerifier.create(Flux.from(connection.getSentAsPublisher()).take(1))
+        .then(() -> virtualTimeScheduler.advanceTimeBy(duration))
         .expectNextMatches(this::keepAliveFrameWithoutRespondFlag)
         .expectComplete()
         .verify(Duration.ofSeconds(5));
 
     rSocketState.rSocket.dispose();
     rSocketState.connection.dispose();
 
-    Thread.sleep(KEEP_ALIVE_TIMEOUT);
-
     rSocketState.allocator.assertHasNoLeaks();
   }
 
   @Test
-  void resumableRequesterNoKeepAlivesAfterDisconnect() throws InterruptedException {
+  void resumableRequesterNoKeepAlivesAfterDisconnect() {
     ResumableRSocketState rSocketState =
         resumableRequester(KEEP_ALIVE_INTERVAL, KEEP_ALIVE_TIMEOUT);
     TestDuplexConnection testConnection = rSocketState.connection();
     ResumableDuplexConnection resumableDuplexConnection = rSocketState.resumableDuplexConnection();
 
     resumableDuplexConnection.disconnect();
 
-    StepVerifier.create(Flux.from(testConnection.getSentAsPublisher()).take(Duration.ofMillis(500)))
+    Duration duration = Duration.ofMillis(500);
+    StepVerifier.create(Flux.from(testConnection.getSentAsPublisher()).take(duration))
+        .then(() -> virtualTimeScheduler.advanceTimeBy(duration))
         .expectComplete()
         .verify(Duration.ofSeconds(5));
 
     rSocketState.rSocket.dispose();
     rSocketState.connection.dispose();
 
-    Thread.sleep(KEEP_ALIVE_INTERVAL);
-
     rSocketState.allocator.assertHasNoLeaks();
   }
 
   @Test
-  void resumableRequesterKeepAlivesAfterReconnect() throws InterruptedException {
+  void resumableRequesterKeepAlivesAfterReconnect() {
     ResumableRSocketState rSocketState =
         resumableRequester(KEEP_ALIVE_INTERVAL, KEEP_ALIVE_TIMEOUT);
     ResumableDuplexConnection resumableDuplexConnection = rSocketState.resumableDuplexConnection();
@@ -236,33 +251,31 @@ void resumableRequesterKeepAlivesAfterReconnect() throws InterruptedException {
     resumableDuplexConnection.resume(0, 0, ignored -> Mono.empty());
 
     StepVerifier.create(Flux.from(newTestConnection.getSentAsPublisher()).take(1))
+        .then(() -> virtualTimeScheduler.advanceTimeBy(Duration.ofMillis(KEEP_ALIVE_INTERVAL)))
         .expectNextMatches(frame -> keepAliveFrame(frame) && frame.release())
         .expectComplete()
         .verify(Duration.ofSeconds(5));
 
     rSocketState.rSocket.dispose();
     rSocketState.connection.dispose();
 
-    Thread.sleep(KEEP_ALIVE_INTERVAL);
-
     rSocketState.allocator.assertHasNoLeaks();
   }
 
   @Test
-  void resumableRequesterNoKeepAlivesAfterDispose() throws InterruptedException {
+  void resumableRequesterNoKeepAlivesAfterDispose() {
     ResumableRSocketState rSocketState =
         resumableRequester(KEEP_ALIVE_INTERVAL, KEEP_ALIVE_TIMEOUT);
     rSocketState.rSocket().dispose();
-    StepVerifier.create(
-            Flux.from(rSocketState.connection().getSentAsPublisher()).take(Duration.ofMillis(500)))
+    Duration duration = Duration.ofMillis(500);
+    StepVerifier.create(Flux.from(rSocketState.connection().getSentAsPublisher()).take(duration))
+        .then(() -> virtualTimeScheduler.advanceTimeBy(duration))
         .expectComplete()
         .verify(Duration.ofSeconds(5));
 
     rSocketState.rSocket.dispose();
     rSocketState.connection.dispose();
 
-    Thread.sleep(KEEP_ALIVE_INTERVAL);
-
     rSocketState.allocator.assertHasNoLeaks();
   }
 
@@ -273,7 +286,7 @@ void resumableRSocketsNotDisposedOnMissingKeepAlives() throws InterruptedExcepti
     RSocket rSocket = resumableRequesterState.rSocket();
     TestDuplexConnection connection = resumableRequesterState.connection();
 
-    Mono.delay(Duration.ofMillis(500)).block();
+    virtualTimeScheduler.advanceTimeBy(Duration.ofMillis(500));
 
     Assertions.assertThat(rSocket.isDisposed()).isFalse();
     Assertions.assertThat(connection.isDisposed()).isTrue();
@@ -283,8 +296,6 @@ void resumableRSocketsNotDisposedOnMissingKeepAlives() throws InterruptedExcepti
     resumableRequesterState.connection.dispose();
     resumableRequesterState.rSocket.dispose();
 
-    Thread.sleep(KEEP_ALIVE_INTERVAL);
-
     resumableRequesterState.allocator.assertHasNoLeaks();
   }