resume: fix race between send / release of frames on received implied position from keepalives

replace state enum with constants

remove unnecessary allocation

make cleanup of store on keep-alive optional

formatter

Signed-off-by: Maksym Ostroverkhov <[email protected]>

Author: mostroverkhov

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

@@ -104,6 +104,7 @@ public static class ClientRSocketFactory implements ClientTransportAcceptor {
     private String dataMimeType = "application/binary";
 
     private boolean resumeEnabled;
+    private boolean resumeCleanupStoreOnKeepAlive;
     private Supplier<ByteBuf> resumeTokenSupplier = ResumeFrameFlyweight::generateResumeToken;
     private Function<? super ByteBuf, ? extends ResumableFramesStore> resumeStoreFactory =
         token -> new InMemoryResumableFramesStore("client", 100_000);
@@ -216,6 +217,11 @@ public ClientRSocketFactory resumeStrategy(Supplier<ResumeStrategy> resumeStrate
       return this;
     }
 
+    public ClientRSocketFactory resumeCleanupOnKeepAlive() {
+      resumeCleanupStoreOnKeepAlive = true;
+      return this;
+    }
+
     @Override
     public Start<RSocket> transport(Supplier<ClientTransport> transportClient) {
       return new StartClient(transportClient);
@@ -326,7 +332,8 @@ private ClientSetup clientSetup() {
               resumeStoreFactory.apply(resumeToken),
               resumeSessionDuration,
               resumeStreamTimeout,
-              resumeStrategySupplier);
+              resumeStrategySupplier,
+              resumeCleanupStoreOnKeepAlive);
         } else {
           return new ClientSetup.DefaultClientSetup();
         }
@@ -360,6 +367,7 @@ public static class ServerRSocketFactory {
         token -> new InMemoryResumableFramesStore("server", 100_000);
 
     private ByteBufAllocator allocator = ByteBufAllocator.DEFAULT;
+    private boolean resumeCleanupStoreOnKeepAlive;
 
     private ServerRSocketFactory() {}
 
@@ -425,6 +433,11 @@ public ServerRSocketFactory resumeStreamTimeout(Duration resumeStreamTimeout) {
       return this;
     }
 
+    public ServerRSocketFactory resumeCleanupOnKeepAlive() {
+      resumeCleanupStoreOnKeepAlive = true;
+      return this;
+    }
+
     private class ServerStart<T extends Closeable> implements Start<T> {
       private final Supplier<ServerTransport<T>> transportServer;
 
@@ -543,7 +556,8 @@ private ServerSetup serverSetup() {
                 new SessionManager(),
                 resumeSessionDuration,
                 resumeStreamTimeout,
-                resumeStoreFactory)
+                resumeStoreFactory,
+                resumeCleanupStoreOnKeepAlive)
             : new ServerSetup.DefaultServerSetup(allocator);
       }
 

rsocket-core/src/main/java/io/rsocket/internal/ClientSetup.java

@@ -56,6 +56,7 @@ class ResumableClientSetup implements ClientSetup {
     private final Supplier<ResumeStrategy> resumeStrategySupplier;
     private final ResumableFramesStore resumableFramesStore;
     private final Duration resumeStreamTimeout;
+    private final boolean cleanupStoreOnKeepAlive;
 
     public ResumableClientSetup(
         ByteBufAllocator allocator,
@@ -64,14 +65,16 @@ public ResumableClientSetup(
         ResumableFramesStore resumableFramesStore,
         Duration resumeSessionDuration,
         Duration resumeStreamTimeout,
-        Supplier<ResumeStrategy> resumeStrategySupplier) {
+        Supplier<ResumeStrategy> resumeStrategySupplier,
+        boolean cleanupStoreOnKeepAlive) {
       this.allocator = allocator;
       this.newConnectionFactory = newConnectionFactory;
       this.resumeToken = resumeToken;
       this.resumeSessionDuration = resumeSessionDuration;
       this.resumeStrategySupplier = resumeStrategySupplier;
       this.resumableFramesStore = resumableFramesStore;
       this.resumeStreamTimeout = resumeStreamTimeout;
+      this.cleanupStoreOnKeepAlive = cleanupStoreOnKeepAlive;
     }
 
     @Override
@@ -83,7 +86,8 @@ public DuplexConnection wrappedConnection(KeepAliveConnection connection) {
                   resumeSessionDuration,
                   resumeStrategySupplier,
                   resumableFramesStore,
-                  resumeStreamTimeout)
+                  resumeStreamTimeout,
+                  cleanupStoreOnKeepAlive)
               .continueWith(newConnectionFactory)
               .resumeToken(resumeToken);
 

rsocket-core/src/main/java/io/rsocket/internal/ServerSetup.java

@@ -106,18 +106,21 @@ class ResumableServerSetup implements ServerSetup {
     private final Duration resumeSessionDuration;
     private final Duration resumeStreamTimeout;
     private final Function<? super ByteBuf, ? extends ResumableFramesStore> resumeStoreFactory;
+    private final boolean cleanupStoreOnKeepAlive;
 
     public ResumableServerSetup(
         ByteBufAllocator allocator,
         SessionManager sessionManager,
         Duration resumeSessionDuration,
         Duration resumeStreamTimeout,
-        Function<? super ByteBuf, ? extends ResumableFramesStore> resumeStoreFactory) {
+        Function<? super ByteBuf, ? extends ResumableFramesStore> resumeStoreFactory,
+        boolean cleanupStoreOnKeepAlive) {
       this.allocator = allocator;
       this.sessionManager = sessionManager;
       this.resumeSessionDuration = resumeSessionDuration;
       this.resumeStreamTimeout = resumeStreamTimeout;
       this.resumeStoreFactory = resumeStoreFactory;
+      this.cleanupStoreOnKeepAlive = cleanupStoreOnKeepAlive;
     }
 
     @Override
@@ -144,7 +147,8 @@ public Mono<Void> acceptRSocketSetup(
                         resumeStreamTimeout,
                         resumeStoreFactory,
                         resumeToken,
-                        keepAliveData))
+                        keepAliveData,
+                        cleanupStoreOnKeepAlive))
                 .resumableConnection();
         return then.apply(new ClientServerInputMultiplexer(resumableConnection));
       } else {

rsocket-core/src/main/java/io/rsocket/resume/ClientRSocketSession.java

@@ -46,11 +46,16 @@ public ClientRSocketSession(
       Duration resumeSessionDuration,
       Supplier<ResumeStrategy> resumeStrategy,
       ResumableFramesStore resumableFramesStore,
-      Duration resumeStreamTimeout) {
+      Duration resumeStreamTimeout,
+      boolean cleanupStoreOnKeepAlive) {
     this.allocator = allocator;
     this.resumableConnection =
         new ResumableDuplexConnection(
-            "client", duplexConnection, resumableFramesStore, resumeStreamTimeout);
+            "client",
+            duplexConnection,
+            resumableFramesStore,
+            resumeStreamTimeout,
+            cleanupStoreOnKeepAlive);
 
     /*session completed: release token initially retained in resumeToken(ByteBuf)*/
     onClose().doFinally(s -> resumeToken.release()).subscribe();

rsocket-core/src/main/java/io/rsocket/resume/ResumableDuplexConnection.java

@@ -40,6 +40,7 @@ class ResumableDuplexConnection implements DuplexConnection, ResumeStateHolder {
   private final String tag;
   private final ResumableFramesStore resumableFramesStore;
   private final Duration resumeStreamTimeout;
+  private final boolean cleanupOnKeepAlive;
 
   private final ReplayProcessor<DuplexConnection> connections = ReplayProcessor.create(1);
   private final EmitterProcessor<Throwable> connectionErrors = EmitterProcessor.create();
@@ -57,23 +58,25 @@ class ResumableDuplexConnection implements DuplexConnection, ResumeStateHolder {
   private final UnicastProcessor<Flux<ByteBuf>> upstreams = UnicastProcessor.create();
   private final UpstreamFramesSubscriber upstreamSubscriber =
       new UpstreamFramesSubscriber(
-          128,
+          Queues.SMALL_BUFFER_SIZE,
           downStreamRequestListener.requests(),
           resumeSaveStreamRequestListener.requests(),
           this::dispatch);
 
-  private volatile State state;
+  private volatile int state;
   private volatile Disposable resumedStreamDisposable = Disposables.disposed();
   private final AtomicBoolean disposed = new AtomicBoolean();
 
   ResumableDuplexConnection(
       String tag,
       ResumePositionsConnection duplexConnection,
       ResumableFramesStore resumableFramesStore,
-      Duration resumeStreamTimeout) {
+      Duration resumeStreamTimeout,
+      boolean cleanupOnKeepAlive) {
     this.tag = tag;
     this.resumableFramesStore = resumableFramesStore;
     this.resumeStreamTimeout = resumeStreamTimeout;
+    this.cleanupOnKeepAlive = cleanupOnKeepAlive;
 
     resumableFramesStore
         .saveFrames(resumeSaveStreamRequestListener.apply(resumeSaveFrames))
@@ -131,11 +134,8 @@ public Mono<Void> sendOne(ByteBuf frame) {
 
   @Override
   public Mono<Void> send(Publisher<ByteBuf> frames) {
-    return Mono.defer(
-        () -> {
-          upstreams.onNext(Flux.from(frames));
-          return framesSent;
-        });
+    upstreams.onNext(Flux.from(frames));
+    return framesSent;
   }
 
   @Override
@@ -164,7 +164,9 @@ public long impliedPosition() {
   @Override
   public void onImpliedPosition(long remoteImpliedPos) {
     logger.debug("Got remote position from keep-alive: {}", remoteImpliedPos);
-    releaseFramesToPosition(remoteImpliedPos);
+    if (cleanupOnKeepAlive) {
+      dispatch(new ReleaseFrames(remoteImpliedPos));
+    }
   }
 
   @Override
@@ -268,7 +270,7 @@ private void doResume(
           Mono.error(
               new ResumeStateException(
                   localPosition, localImpliedPosition,
-                  remotePosition, remoteImpliedPos));
+                  remotePosition, remoteImpliedPosition));
     }
 
     sendResumeFrame
@@ -320,7 +322,7 @@ private void onNewConnection(ResumePositionsConnection connection) {
 
   private void disconnect(DuplexConnection connection) {
     /*do not report late disconnects on old connection if new one is available*/
-    if (curConnection == connection && state.isActive()) {
+    if (curConnection == connection && state != State.DISCONNECTED) {
       Throwable err = new ClosedChannelException();
       state = State.DISCONNECTED;
       logger.debug("{} Inner connection disconnected: {}", tag, err.getClass().getSimpleName());
@@ -330,7 +332,7 @@ private void disconnect(DuplexConnection connection) {
 
   /*remove frames confirmed by implied pos,
   set current pos accordingly*/
-  private void releaseFramesToPosition(Long remoteImpliedPos) {
+  private void releaseFramesToPosition(long remoteImpliedPos) {
     resumableFramesStore.releaseFrames(remoteImpliedPos);
   }
 
@@ -351,22 +353,12 @@ static boolean isResumableFrame(ByteBuf frame) {
     }
   }
 
-  private enum State {
-    CONNECTED(true),
-    RESUME_STARTED(true),
-    RESUME(true),
-    RESUME_COMPLETED(true),
-    DISCONNECTED(false);
-
-    private final boolean active;
-
-    State(boolean active) {
-      this.active = active;
-    }
-
-    public boolean isActive() {
-      return active;
-    }
+  static class State {
+    static int CONNECTED = 0;
+    static int RESUME_STARTED = 1;
+    static int RESUME = 2;
+    static int RESUME_COMPLETED = 3;
+    static int DISCONNECTED = 4;
   }
 
   class ResumeStart implements Runnable {
@@ -407,4 +399,17 @@ public void run() {
       doResumeComplete();
     }
   }
+
+  private class ReleaseFrames implements Runnable {
+    private final long remoteImpliedPos;
+
+    public ReleaseFrames(long remoteImpliedPos) {
+      this.remoteImpliedPos = remoteImpliedPos;
+    }
+
+    @Override
+    public void run() {
+      releaseFramesToPosition(remoteImpliedPos);
+    }
+  }
 }

rsocket-core/src/main/java/io/rsocket/resume/ServerRSocketSession.java

@@ -50,13 +50,18 @@ public ServerRSocketSession(
       Duration resumeStreamTimeout,
       Function<? super ByteBuf, ? extends ResumableFramesStore> resumeStoreFactory,
       ByteBuf resumeToken,
-      KeepAliveData keepAliveData) {
+      KeepAliveData keepAliveData,
+      boolean cleanupStoreOnKeepAlive) {
     this.allocator = allocator;
     this.keepAliveData = keepAliveData;
     this.resumeToken = resumeToken;
     this.resumableConnection =
         new ResumableDuplexConnection(
-            "server", duplexConnection, resumeStoreFactory.apply(resumeToken), resumeStreamTimeout);
+            "server",
+            duplexConnection,
+            resumeStoreFactory.apply(resumeToken),
+            resumeStreamTimeout,
+            cleanupStoreOnKeepAlive);
 
     Mono<ResumePositionsConnection> timeout =
         resumableConnection

rsocket-core/src/main/java/io/rsocket/resume/SessionManager.java

@@ -22,7 +22,7 @@
 import javax.annotation.Nullable;
 
 public class SessionManager {
-  private boolean isDisposed;
+  private volatile boolean isDisposed;
   private final Map<ByteBuf, ServerRSocketSession> sessions = new ConcurrentHashMap<>();
 
   public ServerRSocketSession save(ServerRSocketSession session) {
@@ -34,14 +34,17 @@ public ServerRSocketSession save(ServerRSocketSession session) {
           .onClose()
           .doOnSuccess(
               v -> {
-                sessions.remove(token);
+                if (isDisposed || sessions.get(token) == session) {
+                  sessions.remove(token);
+                }
                 token.release();
               })
           .subscribe();
-      ServerRSocketSession prev = sessions.put(token, session);
-      if (prev != null) {
-        prev.dispose();
+      ServerRSocketSession prevSession = sessions.remove(token);
+      if (prevSession != null) {
+        prevSession.dispose();
       }
+      sessions.put(token, session);
     }
     return session;
   }

rsocket-examples/src/test/java/io/rsocket/resume/ResumeIntegrationTest.java

@@ -211,8 +211,9 @@ private static Mono<RSocket> newClientRSocket(
     return RSocketFactory.connect()
         .resume()
         .resumeSessionDuration(Duration.ofSeconds(sessionDurationSeconds))
-        .resumeStore(t -> new InMemoryResumableFramesStore("client", 10_000_000))
-        .keepAliveTickPeriod(Duration.ofSeconds(30))
+        .resumeStore(t -> new InMemoryResumableFramesStore("client", 500_000))
+        .resumeCleanupOnKeepAlive()
+        .keepAliveTickPeriod(Duration.ofSeconds(5))
         .keepAliveAckTimeout(Duration.ofMinutes(5))
         .errorConsumer(errConsumer)
         .resumeStrategy(() -> new PeriodicResumeStrategy(Duration.ofSeconds(1)))
@@ -227,8 +228,9 @@ private static Mono<CloseableChannel> newServerRSocket() {
   private static Mono<CloseableChannel> newServerRSocket(int sessionDurationSeconds) {
     return RSocketFactory.receive()
         .resume()
-        .resumeStore(t -> new InMemoryResumableFramesStore("server", 10_000_000))
+        .resumeStore(t -> new InMemoryResumableFramesStore("server", 500_000))
         .resumeSessionDuration(Duration.ofSeconds(sessionDurationSeconds))
+        .resumeCleanupOnKeepAlive()
         .acceptor((setupPayload, rSocket) -> Mono.just(new TestResponderRSocket()))
         .transport(serverTransport(SERVER_HOST, SERVER_PORT))
         .start();