ensures resumable connection awaits termination of all component

Signed-off-by: Oleh Dokuka <odokuka@vmware.com>
Signed-off-by: Oleh Dokuka <oleh.dokuka@icloud.com>
Signed-off-by: OlegDokuka <odokuka@vmware.com>

Author: Oleh Dokuka

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

@@ -4,6 +4,7 @@
 import io.netty.buffer.Unpooled;
 import io.rsocket.DuplexConnection;
 import java.nio.channels.ClosedChannelException;
+import reactor.core.Disposable;
 import reactor.core.publisher.Mono;
 import reactor.util.function.Tuple2;
 import reactor.util.function.Tuples;
@@ -25,8 +26,24 @@ class ResumableClientSetup extends ClientSetup {
 
   @Override
   Mono<Tuple2<ByteBuf, DuplexConnection>> init(DuplexConnection connection) {
-    return Mono.<Tuple2<ByteBuf, DuplexConnection>>create(
-            sink -> sink.onRequest(__ -> new SetupHandlingDuplexConnection(connection, sink)))
-        .or(connection.onClose().then(Mono.error(ClosedChannelException::new)));
+    return Mono.create(
+        sink -> {
+          sink.onRequest(
+              __ -> {
+                new SetupHandlingDuplexConnection(connection, sink);
+              });
+
+          Disposable subscribe =
+              connection
+                  .onClose()
+                  .doFinally(__ -> sink.error(new ClosedChannelException()))
+                  .subscribe();
+          sink.onCancel(
+              () -> {
+                subscribe.dispose();
+                connection.dispose();
+                connection.receive().subscribe();
+              });
+        });
   }
 }

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

@@ -60,6 +60,7 @@ void dispose() {}
 
   void sendError(DuplexConnection duplexConnection, RSocketErrorException exception) {
     duplexConnection.sendErrorAndClose(exception);
+    duplexConnection.receive().subscribe();
   }
 
   static class DefaultServerSetup extends ServerSetup {

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

@@ -29,7 +29,6 @@
 import io.rsocket.frame.ResumeOkFrameCodec;
 import io.rsocket.keepalive.KeepAliveSupport;
 import java.time.Duration;
-import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.Function;
 import org.reactivestreams.Subscription;
@@ -79,31 +78,49 @@ public ClientRSocketSession(
     this.resumeToken = resumeToken;
     this.session = resumeToken.toString(CharsetUtil.UTF_8);
     this.connectionFactory =
-        connectionFactory.flatMap(
-            dc -> {
-              final long impliedPosition = resumableFramesStore.frameImpliedPosition();
-              final long position = resumableFramesStore.framePosition();
-              dc.sendFrame(
-                  0,
-                  ResumeFrameCodec.encode(
-                      dc.alloc(),
-                      resumeToken.retain(),
-                      // server uses this to release its cache
-                      impliedPosition, //  observed on the client side
-                      // server uses this to check whether there is no mismatch
-                      position //  sent from the client sent
-                      ));
-
-              if (logger.isDebugEnabled()) {
-                logger.debug(
-                    "Side[client]|Session[{}]. ResumeFrame[impliedPosition[{}], position[{}]] has been sent.",
-                    session,
-                    impliedPosition,
-                    position);
-              }
-
-              return connectionTransformer.apply(dc);
-            });
+        connectionFactory
+            .doOnDiscard(
+                DuplexConnection.class,
+                c -> {
+                  final ConnectionErrorException connectionErrorException =
+                      new ConnectionErrorException("resumption_server=[Session Expired]");
+                  c.sendErrorAndClose(connectionErrorException);
+                  c.receive().subscribe();
+                })
+            .flatMap(
+                dc -> {
+                  final long impliedPosition = resumableFramesStore.frameImpliedPosition();
+                  final long position = resumableFramesStore.framePosition();
+                  dc.sendFrame(
+                      0,
+                      ResumeFrameCodec.encode(
+                          dc.alloc(),
+                          resumeToken.retain(),
+                          // server uses this to release its cache
+                          impliedPosition, //  observed on the client side
+                          // server uses this to check whether there is no mismatch
+                          position //  sent from the client sent
+                          ));
+
+                  if (logger.isDebugEnabled()) {
+                    logger.debug(
+                        "Side[client]|Session[{}]. ResumeFrame[impliedPosition[{}], position[{}]] has been sent.",
+                        session,
+                        impliedPosition,
+                        position);
+                  }
+
+                  return connectionTransformer
+                      .apply(dc)
+                      .doOnDiscard(
+                          Tuple2.class,
+                          tuple2 -> {
+                            if (logger.isDebugEnabled()) {
+                              logger.debug("try to reestablish from discard");
+                            }
+                            tryReestablishSession(tuple2);
+                          });
+                });
     this.resumableFramesStore = resumableFramesStore;
     this.allocator = resumableDuplexConnection.alloc();
     this.resumeSessionDuration = resumeSessionDuration;
@@ -160,11 +177,20 @@ public void onImpliedPosition(long remoteImpliedPos) {
 
   @Override
   public void dispose() {
-    Operators.terminate(S, this);
+    if (logger.isDebugEnabled()) {
+      logger.debug("Side[client]|Session[{}]. Disposing", session);
+    }
+
+    boolean result = Operators.terminate(S, this);
+
+    if (logger.isDebugEnabled()) {
+      logger.debug("Side[client]|Session[{}]. Sessions[isDisposed={}]", session, result);
+    }
 
     reconnectDisposable.dispose();
     resumableConnection.dispose();
-    resumableFramesStore.dispose();
+    // frame store is disposed by resumable connection
+    // resumableFramesStore.dispose();
 
     if (resumeToken.refCnt() > 0) {
       resumeToken.release();
@@ -177,6 +203,9 @@ public boolean isDisposed() {
   }
 
   void tryReestablishSession(Tuple2<ByteBuf, DuplexConnection> tuple2) {
+    if (logger.isDebugEnabled()) {
+      logger.debug("Active subscription is canceled {}", s == Operators.cancelledSubscription());
+    }
     ByteBuf shouldBeResumeOKFrame = tuple2.getT1();
     DuplexConnection nextDuplexConnection = tuple2.getT2();
 
@@ -189,9 +218,9 @@ void tryReestablishSession(Tuple2<ByteBuf, DuplexConnection> tuple2) {
       }
       final ConnectionErrorException connectionErrorException =
           new ConnectionErrorException("RESUME_OK frame must be received before any others");
-      resumableConnection.dispose(connectionErrorException);
+      resumableConnection.dispose(nextDuplexConnection, connectionErrorException);
       nextDuplexConnection.sendErrorAndClose(connectionErrorException);
-      nextDuplexConnection.receive().subscribe().dispose();
+      nextDuplexConnection.receive().subscribe();
 
       throw connectionErrorException; // throw to retry connection again
     }
@@ -227,10 +256,10 @@ void tryReestablishSession(Tuple2<ByteBuf, DuplexConnection> tuple2) {
           }
           final ConnectionErrorException t = new ConnectionErrorException(e.getMessage(), e);
 
-          resumableConnection.dispose(t);
+          resumableConnection.dispose(nextDuplexConnection, t);
 
           nextDuplexConnection.sendErrorAndClose(t);
-          nextDuplexConnection.receive().subscribe().dispose();
+          nextDuplexConnection.receive().subscribe();
 
           return;
         }
@@ -244,7 +273,7 @@ void tryReestablishSession(Tuple2<ByteBuf, DuplexConnection> tuple2) {
           final ConnectionErrorException connectionErrorException =
               new ConnectionErrorException("resumption_server=[Session Expired]");
           nextDuplexConnection.sendErrorAndClose(connectionErrorException);
-          nextDuplexConnection.receive().subscribe().dispose();
+          nextDuplexConnection.receive().subscribe();
           return;
         }
 
@@ -263,7 +292,7 @@ void tryReestablishSession(Tuple2<ByteBuf, DuplexConnection> tuple2) {
           final ConnectionErrorException connectionErrorException =
               new ConnectionErrorException("resumption_server_pos=[Session Expired]");
           nextDuplexConnection.sendErrorAndClose(connectionErrorException);
-          nextDuplexConnection.receive().subscribe().dispose();
+          nextDuplexConnection.receive().subscribe();
           // no need to do anything since connection resumable connection is liklly to
           // be disposed
         }
@@ -278,10 +307,10 @@ void tryReestablishSession(Tuple2<ByteBuf, DuplexConnection> tuple2) {
         final ConnectionErrorException connectionErrorException =
             new ConnectionErrorException("resumption_server_pos=[" + remoteImpliedPos + "]");
 
-        resumableConnection.dispose(connectionErrorException);
+        resumableConnection.dispose(nextDuplexConnection, connectionErrorException);
 
         nextDuplexConnection.sendErrorAndClose(connectionErrorException);
-        nextDuplexConnection.receive().subscribe().dispose();
+        nextDuplexConnection.receive().subscribe();
       }
     } else if (frameType == FrameType.ERROR) {
       final RuntimeException exception = Exceptions.from(0, shouldBeResumeOKFrame);
@@ -292,13 +321,14 @@ void tryReestablishSession(Tuple2<ByteBuf, DuplexConnection> tuple2) {
             exception);
       }
       if (exception instanceof RejectedResumeException) {
-        resumableConnection.dispose(exception);
+        resumableConnection.dispose(nextDuplexConnection, exception);
         nextDuplexConnection.dispose();
-        nextDuplexConnection.receive().subscribe().dispose();
+        nextDuplexConnection.receive().subscribe();
         return;
       }
 
       nextDuplexConnection.dispose();
+      nextDuplexConnection.receive().subscribe();
       throw exception; // assume retryable exception
     } else {
       if (logger.isDebugEnabled()) {
@@ -309,10 +339,10 @@ void tryReestablishSession(Tuple2<ByteBuf, DuplexConnection> tuple2) {
       final ConnectionErrorException connectionErrorException =
           new ConnectionErrorException("RESUME_OK frame must be received before any others");
 
-      resumableConnection.dispose(connectionErrorException);
+      resumableConnection.dispose(nextDuplexConnection, connectionErrorException);
 
       nextDuplexConnection.sendErrorAndClose(connectionErrorException);
-      nextDuplexConnection.receive().subscribe().dispose();
+      nextDuplexConnection.receive().subscribe();
 
       // no need to do anything since remote server rejected our connection completely
     }
@@ -349,11 +379,7 @@ public void onError(Throwable t) {
       Operators.onErrorDropped(t, currentContext());
     }
 
-    if (t instanceof TimeoutException) {
-      resumableConnection.dispose();
-    } else {
-      resumableConnection.dispose(t);
-    }
+    resumableConnection.dispose();
   }
 
   @Override

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

@@ -168,9 +168,9 @@ void drain(long expectedState) {
 
       if (isConnected(expectedState)) {
         if (isTerminated(expectedState)) {
-          handleTerminal(this.terminal);
+          handleTerminated(qs, this.terminal);
         } else if (isDisposed()) {
-          handleTerminal(new CancellationException("Disposed"));
+          handleDisposed();
         } else if (hasFrames(expectedState)) {
           handlePendingFrames(qs);
         }
@@ -402,14 +402,28 @@ void handleFrame(ByteBuf frame) {
     handleConnectionFrame(frame);
   }
 
-  void handleTerminal(@Nullable Throwable t) {
+  void handleTerminated(Fuseable.QueueSubscription<ByteBuf> qs, @Nullable Throwable t) {
+    for (; ; ) {
+      final ByteBuf frame = qs.poll();
+      final boolean empty = frame == null;
+
+      if (empty) {
+        break;
+      }
+
+      handleFrame(frame);
+    }
     if (t != null) {
       this.actual.onError(t);
     } else {
       this.actual.onComplete();
     }
   }
 
+  void handleDisposed() {
+    this.actual.onError(new CancellationException("Disposed"));
+  }
+
   void handleConnectionFrame(ByteBuf frame) {
     this.actual.onNext(frame);
   }