ensures behaviour of the first payload sending and racing with cancel

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

Author: OlegDokuka

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

@@ -54,7 +54,7 @@
 import io.rsocket.util.MonoLifecycleHandler;
 import java.nio.channels.ClosedChannelException;
 import java.util.concurrent.CancellationException;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.Consumer;
 import java.util.function.LongConsumer;
@@ -242,9 +242,14 @@ private Mono<Payload> handleRequestResponse(final Payload payload) {
             new MonoLifecycleHandler<Payload>() {
               @Override
               public void doOnSubscribe() {
-                final ByteBuf requestFrame =
-                    RequestResponseFrameFlyweight.encodeReleasingPayload(
-                        allocator, streamId, payload);
+                final ByteBuf requestFrame;
+                try {
+                  requestFrame =
+                      RequestResponseFrameFlyweight.encodeReleasingPayload(
+                          allocator, streamId, payload);
+                } catch (IllegalReferenceCountException e) {
+                  return;
+                }
 
                 sendProcessor.onNext(requestFrame);
               }
@@ -260,6 +265,7 @@ public void doOnTerminal(
                 removeStreamReceiver(streamId);
               }
             });
+
     receivers.put(streamId, receiver);
 
     return receiver.doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
@@ -281,7 +287,7 @@ private Flux<Payload> handleRequestStream(final Payload payload) {
 
     final UnboundedProcessor<ByteBuf> sendProcessor = this.sendProcessor;
     final UnicastProcessor<Payload> receiver = UnicastProcessor.create();
-    final AtomicBoolean payloadReleasedFlag = new AtomicBoolean(false);
+    final AtomicInteger wip = new AtomicInteger(0);
 
     receivers.put(streamId, receiver);
 
@@ -295,30 +301,56 @@ private Flux<Payload> handleRequestStream(final Payload payload) {
               public void accept(long n) {
                 if (firstRequest && !receiver.isDisposed()) {
                   firstRequest = false;
-                  if (!payloadReleasedFlag.getAndSet(true)) {
-                    sendProcessor.onNext(
-                        RequestStreamFrameFlyweight.encodeReleasingPayload(
-                            allocator, streamId, n, payload));
+                  if (wip.getAndIncrement() != 0) {
+                    // no need to do anything.
+                    // stream was canceled and fist payload has already been discarded
+                    return;
+                  }
+                  int missed = 1;
+                  boolean firstHasBeenSent = false;
+                  for (; ; ) {
+                    if (!firstHasBeenSent) {
+                      ByteBuf frame;
+                      try {
+                        frame =
+                            RequestStreamFrameFlyweight.encodeReleasingPayload(
+                                allocator, streamId, n, payload);
+                      } catch (IllegalReferenceCountException e) {
+                        return;
+                      }
+
+                      sendProcessor.onNext(frame);
+                      firstHasBeenSent = true;
+                    } else {
+                      // if first frame was sent but we cycling again, it means that wip was
+                      // incremented at doOnCancel
+                      sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
+                      return;
+                    }
+
+                    missed = wip.addAndGet(-missed);
+                    if (missed == 0) {
+                      return;
+                    }
                   }
-                } else if (contains(streamId) && !receiver.isDisposed()) {
+                } else {
                   sendProcessor.onNext(RequestNFrameFlyweight.encode(allocator, streamId, n));
                 }
               }
             })
-        .doOnError(
-            t -> {
-              if (contains(streamId) && !receiver.isDisposed()) {
-                sendProcessor.onNext(ErrorFrameFlyweight.encode(allocator, streamId, t));
-              }
-            })
         .doOnCancel(
             () -> {
-              if (!payloadReleasedFlag.getAndSet(true)) {
-                payload.release();
+              if (wip.getAndIncrement() != 0) {
+                return;
               }
-              if (contains(streamId) && !receiver.isDisposed()) {
-                sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
+
+              // check if we need to release payload
+              // only applicable if the cancel appears earlier than actual request
+              if (payload.refCnt() > 0) {
+                payload.release();
               }
+
+              sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
             })
         .doFinally(s -> removeStreamReceiver(streamId))
         .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
@@ -330,30 +362,32 @@ private Flux<Payload> handleChannel(Flux<Payload> request) {
       return Flux.error(err);
     }
 
-    return request.switchOnFirst(
-        (s, flux) -> {
-          Payload payload = s.get();
-          if (payload != null) {
-            if (!PayloadValidationUtils.isValid(mtu, payload)) {
-              payload.release();
-              final IllegalArgumentException t =
-                  new IllegalArgumentException(INVALID_PAYLOAD_ERROR_MESSAGE);
-              errorConsumer.accept(t);
-              return Mono.error(t);
-            }
-            return handleChannel(payload, flux);
-          } else {
-            return flux;
-          }
-        },
-        false);
+    return request
+        .switchOnFirst(
+            (s, flux) -> {
+              Payload payload = s.get();
+              if (payload != null) {
+                if (!PayloadValidationUtils.isValid(mtu, payload)) {
+                  payload.release();
+                  final IllegalArgumentException t =
+                      new IllegalArgumentException(INVALID_PAYLOAD_ERROR_MESSAGE);
+                  errorConsumer.accept(t);
+                  return Mono.error(t);
+                }
+                return handleChannel(payload, flux);
+              } else {
+                return flux;
+              }
+            },
+            false)
+        .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
   }
 
   private Flux<? extends Payload> handleChannel(Payload initialPayload, Flux<Payload> inboundFlux) {
     final UnboundedProcessor<ByteBuf> sendProcessor = this.sendProcessor;
-    final AtomicBoolean payloadReleasedFlag = new AtomicBoolean(false);
     final int streamId = streamIdSupplier.nextStreamId(receivers);
 
+    final AtomicInteger wip = new AtomicInteger(0);
     final UnicastProcessor<Payload> receiver = UnicastProcessor.create();
     final BaseSubscriber<Payload> upstreamSubscriber =
         new BaseSubscriber<Payload>() {
@@ -421,43 +455,65 @@ protected void hookFinally(SignalType type) {
               public void accept(long n) {
                 if (firstRequest) {
                   firstRequest = false;
-                  senders.put(streamId, upstreamSubscriber);
-                  receivers.put(streamId, receiver);
-
-                  inboundFlux
-                      .limitRate(Queues.SMALL_BUFFER_SIZE)
-                      .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER)
-                      .subscribe(upstreamSubscriber);
-                  if (!payloadReleasedFlag.getAndSet(true)) {
-                    ByteBuf frame =
-                        RequestChannelFrameFlyweight.encodeReleasingPayload(
-                            allocator, streamId, false, n, initialPayload);
-
-                    sendProcessor.onNext(frame);
+                  if (wip.getAndIncrement() != 0) {
+                    // no need to do anything.
+                    // stream was canceled and fist payload has already been discarded
+                    return;
+                  }
+                  int missed = 1;
+                  boolean firstHasBeenSent = false;
+                  for (; ; ) {
+                    if (!firstHasBeenSent) {
+                      ByteBuf frame;
+                      try {
+                        frame =
+                            RequestChannelFrameFlyweight.encodeReleasingPayload(
+                                allocator, streamId, false, n, initialPayload);
+                      } catch (IllegalReferenceCountException e) {
+                        return;
+                      }
+
+                      senders.put(streamId, upstreamSubscriber);
+                      receivers.put(streamId, receiver);
+
+                      inboundFlux
+                          .limitRate(Queues.SMALL_BUFFER_SIZE)
+                          .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER)
+                          .subscribe(upstreamSubscriber);
+
+                      sendProcessor.onNext(frame);
+                      firstHasBeenSent = true;
+                    } else {
+                      // if first frame was sent but we cycling again, it means that wip was
+                      // incremented at doOnCancel
+                      senders.remove(streamId, upstreamSubscriber);
+                      receivers.remove(streamId, receiver);
+                      sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
+                      return;
+                    }
+
+                    missed = wip.addAndGet(-missed);
+                    if (missed == 0) {
+                      return;
+                    }
                   }
                 } else {
                   sendProcessor.onNext(RequestNFrameFlyweight.encode(allocator, streamId, n));
                 }
               }
             })
-        .doOnError(
-            t -> {
-              if (receivers.remove(streamId, receiver)) {
-                upstreamSubscriber.cancel();
-              }
-            })
-        .doOnComplete(() -> receivers.remove(streamId, receiver))
+        .doOnError(t -> upstreamSubscriber.cancel())
         .doOnCancel(
             () -> {
-              if (!payloadReleasedFlag.getAndSet(true)) {
-                initialPayload.release();
-              }
-              if (receivers.remove(streamId, receiver)) {
-                sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
-                upstreamSubscriber.cancel();
+              upstreamSubscriber.cancel();
+              if (wip.getAndIncrement() != 0) {
+                return;
               }
+
+              // need to send frame only if RequestChannelFrame was sent
+              sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
             })
-        .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
+        .doFinally(__ -> receivers.remove(streamId, receiver));
   }
 
   private Mono<Void> handleMetadataPush(Payload payload) {