cleanup: resume token is released in ClientRSocketSession instead of

         RSocketFactory

         KeepAliveConnection: dispose KeepAliveHandler in onClose()

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

Author: mostroverkhov

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

@@ -278,8 +278,6 @@ public Mono<RSocket> start() {
                           errorConsumer,
                           StreamIdSupplier.clientSupplier());
 
-                  rSocketClient.onClose().doFinally(s -> resumeToken.release()).subscribe();
-
                   RSocket wrappedRSocketClient = plugins.applyClient(rSocketClient);
 
                   RSocket unwrappedServerSocket = acceptor.get().apply(wrappedRSocketClient);

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

@@ -31,7 +31,9 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import org.reactivestreams.Publisher;
-import reactor.core.publisher.*;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.MonoProcessor;
 
 public class KeepAliveConnection extends DuplexConnectionProxy
     implements ResumePositionsConnection {
@@ -126,12 +128,15 @@ public Flux<ByteBuf> receive() {
   }
 
   @Override
-  public void dispose() {
-    KeepAliveHandler keepAliveHandler = keepAliveHandlerReady.peek();
-    if (keepAliveHandler != null) {
-      keepAliveHandler.dispose();
-    }
-    super.dispose();
+  public Mono<Void> onClose() {
+    return super.onClose()
+        .doFinally(
+            s -> {
+              KeepAliveHandler keepAliveHandler = keepAliveHandlerReady.peek();
+              if (keepAliveHandler != null) {
+                keepAliveHandler.dispose();
+              }
+            });
   }
 
   @Override

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

@@ -48,6 +48,9 @@ public ClientRSocketSession(
         new ResumableDuplexConnection(
             "client", duplexConnection, config.resumeStore(), config.resumeStreamTimeout());
 
+    /*session completed: release token initially retained in resumeToken(ByteBuf)*/
+    onClose().doFinally(s -> resumeToken.release()).subscribe();
+
     resumableConnection
         .connectionErrors()
         .flatMap(