polishes logging and refactor Client/ServerRSocketSession. adds tests…

…. Improves KeepAliveSupport Signed-off-by: Oleh Dokuka <[email protected]>
rsocket · Jun 8, 2021 · f521a6a · f521a6a
1 parent de60762
commit f521a6a
Show file tree

Hide file tree

Showing 8 changed files with 900 additions and 200 deletions.
diff --git a/rsocket-core/src/main/java/io/rsocket/core/ServerSetup.java b/rsocket-core/src/main/java/io/rsocket/core/ServerSetup.java
@@ -114,10 +114,10 @@ public Mono<Void> acceptRSocketSetup(
         final ServerRSocketSession serverRSocketSession =
             new ServerRSocketSession(
                 resumeToken,
-                duplexConnection,
                 resumableDuplexConnection,
-                resumeSessionDuration,
+                duplexConnection,
                 resumableFramesStore,
+                resumeSessionDuration,
                 cleanupStoreOnKeepAlive);
 
         sessionManager.save(serverRSocketSession, resumeToken);

diff --git a/rsocket-core/src/main/java/io/rsocket/keepalive/KeepAliveSupport.java b/rsocket-core/src/main/java/io/rsocket/keepalive/KeepAliveSupport.java
@@ -23,7 +23,7 @@
 import io.rsocket.resume.ResumeStateHolder;
 import java.time.Duration;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 import java.util.function.Consumer;
 import reactor.core.Disposable;
 import reactor.core.publisher.Flux;
@@ -38,11 +38,19 @@ public abstract class KeepAliveSupport implements KeepAliveFramesAcceptor {
   final Duration keepAliveTimeout;
   final long keepAliveTimeoutMillis;
 
-  final AtomicBoolean started = new AtomicBoolean();
+  volatile int state;
+  static final AtomicIntegerFieldUpdater<KeepAliveSupport> STATE =
+      AtomicIntegerFieldUpdater.newUpdater(KeepAliveSupport.class, "state");
+
+  static final int STOPPED_STATE = 0;
+  static final int STARTING_STATE = 1;
+  static final int STARTED_STATE = 2;
+  static final int DISPOSED_STATE = -1;
 
   volatile Consumer<KeepAlive> onTimeout;
   volatile Consumer<ByteBuf> onFrameSent;
-  volatile Disposable ticksDisposable;
+
+  Disposable ticksDisposable;
 
   volatile ResumeStateHolder resumeStateHolder;
   volatile long lastReceivedMillis;
@@ -57,25 +65,30 @@ private KeepAliveSupport(
   }
 
   public KeepAliveSupport start() {
-    this.lastReceivedMillis = scheduler.now(TimeUnit.MILLISECONDS);
-    if (started.compareAndSet(false, true)) {
-      ticksDisposable =
+    if (this.state == STOPPED_STATE && STATE.compareAndSet(this, STOPPED_STATE, STARTING_STATE)) {
+      this.lastReceivedMillis = scheduler.now(TimeUnit.MILLISECONDS);
+
+      final Disposable disposable =
           Flux.interval(keepAliveInterval, scheduler).subscribe(v -> onIntervalTick());
+      this.ticksDisposable = disposable;
+
+      if (this.state != STARTING_STATE
+          || !STATE.compareAndSet(this, STARTING_STATE, STARTED_STATE)) {
+        disposable.dispose();
+      }
     }
     return this;
   }
 
   public void stop() {
-    if (started.compareAndSet(true, false)) {
-      ticksDisposable.dispose();
-    }
+    terminate(STOPPED_STATE);
   }
 
   @Override
   public void receive(ByteBuf keepAliveFrame) {
     this.lastReceivedMillis = scheduler.now(TimeUnit.MILLISECONDS);
     if (resumeStateHolder != null) {
-      long remoteLastReceivedPos = remoteLastReceivedPosition(keepAliveFrame);
+      final long remoteLastReceivedPos = KeepAliveFrameCodec.lastPosition(keepAliveFrame);
       resumeStateHolder.onImpliedPosition(remoteLastReceivedPos);
     }
     if (KeepAliveFrameCodec.respondFlag(keepAliveFrame)) {
@@ -104,6 +117,16 @@ public KeepAliveSupport onTimeout(Consumer<KeepAlive> onTimeout) {
     return this;
   }
 
+  @Override
+  public void dispose() {
+    terminate(DISPOSED_STATE);
+  }
+
+  @Override
+  public boolean isDisposed() {
+    return ticksDisposable.isDisposed();
+  }
+
   abstract void onIntervalTick();
 
   void send(ByteBuf frame) {
@@ -122,22 +145,24 @@ void tryTimeout() {
     }
   }
 
-  long localLastReceivedPosition() {
-    return resumeStateHolder != null ? resumeStateHolder.impliedPosition() : 0;
-  }
+  void terminate(int terminationState) {
+    for (; ; ) {
+      final int state = this.state;
 
-  long remoteLastReceivedPosition(ByteBuf keepAliveFrame) {
-    return KeepAliveFrameCodec.lastPosition(keepAliveFrame);
-  }
+      if (state == STOPPED_STATE || state == DISPOSED_STATE) {
+        return;
+      }
 
-  @Override
-  public void dispose() {
-    stop();
+      final Disposable disposable = this.ticksDisposable;
+      if (STATE.compareAndSet(this, state, terminationState)) {
+        disposable.dispose();
+        return;
+      }
+    }
   }
 
-  @Override
-  public boolean isDisposed() {
-    return ticksDisposable.isDisposed();
+  long localLastReceivedPosition() {
+    return resumeStateHolder != null ? resumeStateHolder.impliedPosition() : 0;
   }
 
   public static final class ClientKeepAliveSupport extends KeepAliveSupport {