ensures RRRSubscriber doesn't cancel subscription on onNext (#918)

rsocket · Aug 22, 2020 · ff9b02a · ff9b02a
1 parent 6368fae
commit ff9b02a
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 5 deletions.
diff --git a/rsocket-core/src/main/java/io/rsocket/core/RequestResponseResponderSubscriber.java b/rsocket-core/src/main/java/io/rsocket/core/RequestResponseResponderSubscriber.java
@@ -58,6 +58,7 @@ final class RequestResponseResponderSubscriber
 
   final RSocket handler;
 
+  boolean done;
   CompositeByteBuf frames;
 
   volatile Subscription s;
@@ -109,13 +110,24 @@ public void onSubscribe(Subscription subscription) {
 
   @Override
   public void onNext(@Nullable Payload p) {
-    if (!Operators.terminate(S, this)) {
+    if (this.done) {
       if (p != null) {
         p.release();
       }
       return;
     }
 
+    final Subscription currentSubscription = this.s;
+    if (currentSubscription == Operators.cancelledSubscription()
+        || !S.compareAndSet(this, currentSubscription, Operators.cancelledSubscription())) {
+      if (p != null) {
+        p.release();
+      }
+      return;
+    }
+
+    this.done = true;
+
     final int streamId = this.streamId;
     final UnboundedProcessor<ByteBuf> sender = this.sendProcessor;
     final ByteBufAllocator allocator = this.allocator;
@@ -131,6 +143,8 @@ public void onNext(@Nullable Payload p) {
     final int mtu = this.mtu;
     try {
       if (!isValid(mtu, this.maxFrameLength, p, false)) {
+        currentSubscription.cancel();
+
         p.release();
 
         final ByteBuf errorFrame =
@@ -143,6 +157,8 @@ public void onNext(@Nullable Payload p) {
         return;
       }
     } catch (IllegalReferenceCountException e) {
+      currentSubscription.cancel();
+
       final ByteBuf errorFrame =
           ErrorFrameCodec.encode(
               allocator,
@@ -155,16 +171,26 @@ public void onNext(@Nullable Payload p) {
     try {
       sendReleasingPayload(streamId, FrameType.NEXT_COMPLETE, mtu, p, sender, allocator, false);
     } catch (Throwable ignored) {
+      currentSubscription.cancel();
     }
   }
 
   @Override
   public void onError(Throwable t) {
-    if (S.getAndSet(this, Operators.cancelledSubscription()) == Operators.cancelledSubscription()) {
+    if (this.done) {
+      logger.debug("Dropped error", t);
+      return;
+    }
+
+    final Subscription currentSubscription = this.s;
+    if (currentSubscription == Operators.cancelledSubscription()
+        || !S.compareAndSet(this, currentSubscription, Operators.cancelledSubscription())) {
       logger.debug("Dropped error", t);
       return;
     }
 
+    this.done = true;
+
     final int streamId = this.streamId;
 
     this.requesterResponderSupport.remove(streamId, this);

diff --git a/rsocket-core/src/test/java/io/rsocket/core/RSocketResponderTest.java b/rsocket-core/src/test/java/io/rsocket/core/RSocketResponderTest.java
@@ -143,17 +143,24 @@ public void testHandleKeepAlive() throws Exception {
 
   @Test
   @Timeout(2_000)
-  @Disabled
   public void testHandleResponseFrameNoError() throws Exception {
     final int streamId = 4;
     rule.connection.clearSendReceiveBuffers();
-
+    final TestPublisher<Payload> testPublisher = TestPublisher.create();
+    rule.setAcceptingSocket(
+        new RSocket() {
+          @Override
+          public Mono<Payload> requestResponse(Payload payload) {
+            return testPublisher.mono();
+          }
+        });
     rule.sendRequest(streamId, FrameType.REQUEST_RESPONSE);
-
+    testPublisher.complete();
     assertThat(
         "Unexpected frame sent.",
         frameType(rule.connection.awaitSend()),
         anyOf(is(FrameType.COMPLETE), is(FrameType.NEXT_COMPLETE)));
+    testPublisher.assertWasNotCancelled();
   }
 
   @Test