finalize GS functionality

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

Author: OlegDokuka

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

@@ -210,6 +210,11 @@ public void dispose() {
 
   @Override
   public void disposeGracefully() {
+    getDuplexConnection()
+        .sendFrame(
+            0,
+            ErrorFrameCodec.encode(
+                getAllocator(), 0, new ConnectionCloseException("Graceful Shutdown")));
     this.onGracefulShutdownStartedSink.tryEmitEmpty();
     super.terminate();
   }

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

@@ -56,7 +56,6 @@ class RSocketResponder extends RequesterResponderSupport implements RSocket {
 
   private final RSocket requestHandler;
   private final Sinks.Empty<Void> onThisSideClosedSink;
-  private final Mono<Void> onRequesterGracefullShutdownStarted;
 
   @Nullable private final ResponderLeaseTracker leaseHandler;
 
@@ -91,7 +90,6 @@ class RSocketResponder extends RequesterResponderSupport implements RSocket {
 
     this.leaseHandler = leaseHandler;
     this.onThisSideClosedSink = onThisSideClosedSink;
-    this.onRequesterGracefullShutdownStarted = onRequesterGracefulShutdownStarted;
 
     connection
         .onClose()

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

@@ -238,11 +238,11 @@ void sendFirstPayload(Payload payload, long initialRequestN) {
         return;
       }
 
-      sm.remove(streamId, this);
-
       final ByteBuf cancelFrame = CancelFrameCodec.encode(allocator, streamId);
       connection.sendFrame(streamId, cancelFrame);
 
+      sm.remove(streamId, this);
+
       if (requestInterceptor != null) {
         requestInterceptor.onCancel(streamId, FrameType.REQUEST_STREAM);
       }
@@ -276,12 +276,13 @@ public final void cancel() {
 
     if (isFirstFrameSent(previousState)) {
       final int streamId = this.streamId;
-      this.requesterResponderSupport.remove(streamId, this);
 
       ReassemblyUtils.synchronizedRelease(this, previousState);
 
       this.connection.sendFrame(streamId, CancelFrameCodec.encode(this.allocator, streamId));
 
+      this.requesterResponderSupport.remove(streamId, this);
+
       final RequestInterceptor requestInterceptor = this.requestInterceptor;
       if (requestInterceptor != null) {
         requestInterceptor.onCancel(streamId, FrameType.REQUEST_STREAM);
@@ -309,13 +310,14 @@ public final void handlePayload(Payload p) {
       }
 
       final int streamId = this.streamId;
-      this.requesterResponderSupport.remove(streamId, this);
 
       final IllegalStateException cause =
           Exceptions.failWithOverflow(
               "The number of messages received exceeds the number requested");
       this.connection.sendFrame(streamId, CancelFrameCodec.encode(this.allocator, streamId));
 
+      this.requesterResponderSupport.remove(streamId, this);
+
       final RequestInterceptor requestInterceptor = this.requestInterceptor;
       if (requestInterceptor != null) {
         requestInterceptor.onTerminate(streamId, FrameType.REQUEST_STREAM, cause);

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

@@ -144,6 +144,8 @@ public void onNext(Payload p) {
         final ByteBuf errorFrame = ErrorFrameCodec.encode(allocator, streamId, e);
         sender.sendFrame(streamId, errorFrame);
 
+        this.requesterResponderSupport.remove(streamId, this);
+
         final RequestInterceptor requestInterceptor = this.requestInterceptor;
         if (requestInterceptor != null) {
           requestInterceptor.onTerminate(streamId, FrameType.REQUEST_STREAM, e);
@@ -162,6 +164,8 @@ public void onNext(Payload p) {
               new CanceledException("Failed to validate payload. Cause" + e.getMessage()));
       sender.sendFrame(streamId, errorFrame);
 
+      this.requesterResponderSupport.remove(streamId, this);
+
       final RequestInterceptor requestInterceptor = this.requestInterceptor;
       if (requestInterceptor != null) {
         requestInterceptor.onTerminate(streamId, FrameType.REQUEST_STREAM, e);
@@ -176,6 +180,8 @@ public void onNext(Payload p) {
         return;
       }
 
+      this.requesterResponderSupport.remove(streamId, this);
+
       final RequestInterceptor requestInterceptor = this.requestInterceptor;
       if (requestInterceptor != null) {
         requestInterceptor.onTerminate(streamId, FrameType.REQUEST_STREAM, t);
@@ -195,8 +201,6 @@ boolean tryTerminateOnError() {
       return false;
     }
 
-    this.requesterResponderSupport.remove(this.streamId, this);
-
     currentSubscription.cancel();
 
     return true;
@@ -222,11 +226,12 @@ public void onError(Throwable t) {
     }
 
     final int streamId = this.streamId;
-    this.requesterResponderSupport.remove(streamId, this);
 
     final ByteBuf errorFrame = ErrorFrameCodec.encode(this.allocator, streamId, t);
     this.connection.sendFrame(streamId, errorFrame);
 
+    this.requesterResponderSupport.remove(streamId, this);
+
     final RequestInterceptor requestInterceptor = this.requestInterceptor;
     if (requestInterceptor != null) {
       requestInterceptor.onTerminate(streamId, FrameType.REQUEST_STREAM, t);
@@ -246,11 +251,12 @@ public void onComplete() {
     }
 
     final int streamId = this.streamId;
-    this.requesterResponderSupport.remove(streamId, this);
 
     final ByteBuf completeFrame = PayloadFrameCodec.encodeComplete(this.allocator, streamId);
     this.connection.sendFrame(streamId, completeFrame);
 
+    this.requesterResponderSupport.remove(streamId, this);
+
     final RequestInterceptor requestInterceptor = this.requestInterceptor;
     if (requestInterceptor != null) {
       requestInterceptor.onTerminate(streamId, FrameType.REQUEST_STREAM, null);
@@ -321,7 +327,6 @@ public void handleNext(ByteBuf followingFrame, boolean hasFollows, boolean isLas
       S.lazySet(this, Operators.cancelledSubscription());
 
       final int streamId = this.streamId;
-      this.requesterResponderSupport.remove(streamId, this);
 
       this.frames = null;
       frames.release();
@@ -334,6 +339,8 @@ public void handleNext(ByteBuf followingFrame, boolean hasFollows, boolean isLas
               new CanceledException("Failed to reassemble payload. Cause: " + e.getMessage()));
       this.connection.sendFrame(streamId, errorFrame);
 
+      this.requesterResponderSupport.remove(streamId, this);
+
       final RequestInterceptor requestInterceptor = this.requestInterceptor;
       if (requestInterceptor != null) {
         requestInterceptor.onTerminate(streamId, FrameType.REQUEST_STREAM, e);
@@ -354,7 +361,6 @@ public void handleNext(ByteBuf followingFrame, boolean hasFollows, boolean isLas
         this.done = true;
 
         final int streamId = this.streamId;
-        this.requesterResponderSupport.remove(streamId, this);
 
         ReferenceCountUtil.safeRelease(frames);
 
@@ -366,6 +372,8 @@ public void handleNext(ByteBuf followingFrame, boolean hasFollows, boolean isLas
                 new CanceledException("Failed to reassemble payload. Cause: " + t.getMessage()));
         this.connection.sendFrame(streamId, errorFrame);
 
+        this.requesterResponderSupport.remove(streamId, this);
+
         final RequestInterceptor requestInterceptor = this.requestInterceptor;
         if (requestInterceptor != null) {
           requestInterceptor.onTerminate(streamId, FrameType.REQUEST_STREAM, t);

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

@@ -82,6 +82,7 @@
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Stream;
+import org.assertj.core.api.Assertions;
 import org.assertj.core.api.Assumptions;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -96,6 +97,7 @@
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
+import reactor.core.Disposable;
 import reactor.core.Scannable;
 import reactor.core.publisher.BaseSubscriber;
 import reactor.core.publisher.Flux;
@@ -1452,6 +1454,57 @@ public void testWorkaround959(String type) {
     }
   }
 
+  @Test
+  public void testDisposeGracefully() {
+    System.out.println(
+        FrameHeaderCodec.frameType(
+            Unpooled.wrappedBuffer(ByteBufUtil.decodeHexDump("000000012400"))));
+    final RSocketRequester rSocketRequester = rule.socket;
+    final AssertSubscriber<Void> onGracefulShutdownSubscriber =
+        rule.thisGracefulShutdownSink.asMono().subscribeWith(AssertSubscriber.create());
+    final AssertSubscriber<Void> onCloseSubscriber =
+        rSocketRequester.onClose().subscribeWith(new AssertSubscriber<>());
+
+    final Disposable stream = rSocketRequester.requestStream(EmptyPayload.INSTANCE).subscribe();
+
+    FrameAssert.assertThat(rule.connection.awaitFrame())
+        .typeOf(REQUEST_STREAM)
+        .hasClientSideStreamId()
+        .hasNoLeaks();
+
+    Assertions.assertThat(rSocketRequester.isDisposed()).isFalse();
+
+    rSocketRequester.disposeGracefully();
+    Assertions.assertThat(rSocketRequester.isDisposed()).isFalse();
+    onGracefulShutdownSubscriber.assertNotTerminated();
+
+    FrameAssert.assertThat(rule.connection.awaitFrame())
+        .typeOf(FrameType.ERROR)
+        .hasStreamIdZero()
+        .hasData("Graceful Shutdown")
+        .hasNoLeaks();
+
+    stream.dispose();
+    Assertions.assertThat(rSocketRequester.isDisposed()).isFalse();
+    Assertions.assertThat(rule.connection.isDisposed()).isFalse();
+    onGracefulShutdownSubscriber.assertTerminated();
+
+    FrameAssert.assertThat(rule.connection.awaitFrame())
+        .typeOf(CANCEL)
+        .hasClientSideStreamId()
+        .hasNoLeaks();
+
+    rule.otherGracefulShutdownSink.tryEmitEmpty();
+    Assertions.assertThat(rSocketRequester.isDisposed()).isTrue();
+    Assertions.assertThat(rule.connection.isDisposed()).isTrue();
+    onCloseSubscriber.assertNotTerminated();
+
+    rule.otherClosedSink.tryEmitEmpty();
+    onCloseSubscriber.assertTerminated();
+
+    rule.assertHasNoLeaks();
+  }
+
   public static class ClientSocketRule extends AbstractSocketRule<RSocketRequester> {
 
     protected Sinks.Empty<Void> onGracefulShutdownStartedSink;

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

@@ -67,12 +67,14 @@
 import io.rsocket.util.ByteBufPayload;
 import io.rsocket.util.DefaultPayload;
 import io.rsocket.util.EmptyPayload;
+import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
+import org.assertj.core.api.Assertions;
 import org.assertj.core.api.Assumptions;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -1179,6 +1181,65 @@ public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
     rule.assertHasNoLeaks();
   }
 
+  @Test
+  void testGracefulShutdown() {
+    final AssertSubscriber<Void> onCloseSubscriber = AssertSubscriber.create();
+    final AssertSubscriber<Void> onGracefulShutdownSubscriber = AssertSubscriber.create();
+    final Sinks.Empty<Void> onDisposeGracefullySink = Sinks.unsafe().empty();
+
+    boolean[] disposed = new boolean[] {false};
+    boolean[] disposedGracefully = new boolean[] {false};
+
+    rule.setAcceptingSocket(
+        new RSocket() {
+
+          @Override
+          public Flux<Payload> requestStream(Payload payload) {
+            return Flux.interval(Duration.ofMillis(100))
+                .takeUntilOther(onDisposeGracefullySink.asMono())
+                .map(tick -> ByteBufPayload.create(String.valueOf(tick)));
+          }
+
+          @Override
+          public void dispose() {
+            disposed[0] = true;
+          }
+
+          @Override
+          public void disposeGracefully() {
+            disposedGracefully[0] = true;
+          }
+        });
+
+    rule.connection.addToReceivedBuffer(
+        RequestStreamFrameCodec.encode(
+            rule.allocator, 1, false, Long.MAX_VALUE, null, Unpooled.EMPTY_BUFFER));
+
+    rule.onCloseSink.asMono().subscribe(onCloseSubscriber);
+    rule.onGracefulShutdownSink.asMono().subscribe(onGracefulShutdownSubscriber);
+
+    rule.onGracefulShutdownStartedSink.tryEmitEmpty();
+    Assertions.assertThat(disposed[0]).isFalse();
+    Assertions.assertThat(disposedGracefully[0]).isTrue();
+    Assertions.assertThat(rule.connection.isDisposed()).isFalse();
+    onCloseSubscriber.assertNotTerminated();
+    onGracefulShutdownSubscriber.assertNotTerminated();
+
+    onDisposeGracefullySink.tryEmitEmpty();
+    Assertions.assertThat(disposed[0]).isFalse();
+    Assertions.assertThat(disposedGracefully[0]).isTrue();
+    Assertions.assertThat(rule.connection.isDisposed()).isFalse();
+    onCloseSubscriber.assertNotTerminated();
+    onGracefulShutdownSubscriber.assertTerminated();
+
+    rule.connection.dispose();
+    Assertions.assertThat(disposed[0]).isTrue();
+    Assertions.assertThat(disposedGracefully[0]).isTrue();
+    Assertions.assertThat(rule.connection.isDisposed()).isTrue();
+    onCloseSubscriber.assertTerminated();
+    onGracefulShutdownSubscriber.assertTerminated();
+  }
+
   public static class ServerSocketRule extends AbstractSocketRule<RSocketResponder> {
 
     private RSocket acceptingSocket;

rsocket-examples/build.gradle

@@ -28,7 +28,6 @@ dependencies {
     implementation "io.micrometer:micrometer-core"
     implementation "io.micrometer:micrometer-tracing"
     implementation project(":rsocket-micrometer")
-    testImplementation 'org.awaitility:awaitility'
 
     runtimeOnly 'ch.qos.logback:logback-classic'
 
@@ -37,6 +36,7 @@ dependencies {
     testImplementation 'org.mockito:mockito-core'
     testImplementation 'org.assertj:assertj-core'
     testImplementation 'io.projectreactor:reactor-test'
+    testImplementation 'org.awaitility:awaitility'
     testImplementation "io.micrometer:micrometer-test"
     testImplementation "io.micrometer:micrometer-tracing-integration-test"