provides more ByteBuf leaks fixes (rsocket#803)

1. Ensures if there goes something during `Payload` to the frame we will not get any memory leaks in the end 2. Fixes `onDiscard` leak which did not work correctly because `actual` subscriber was nulled to early 3. Ensures that we will not have the wrong frame order in case of racing the first `Payload` sending and cancel. Also, it moves `onDiscard` hook to the very bottom in case of `requestChannel` to ensure the first payload is not leaked 4. Enables a set of tests that ensures that rsocket#757 and rsocket#733 are fully or partially fixed.
gzlicanyi · Apr 27, 2020 · a5706bf · a5706bf
1 parent 1c3927b
commit a5706bf
Show file tree

Hide file tree

Showing 14 changed files with 437 additions and 180 deletions.
diff --git a/rsocket-core/src/main/java/io/rsocket/core/RSocketRequester.java b/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;
@@ -260,6  260,7 @@ public void doOnTerminal(
                 removeStreamReceiver(streamId);
               }
             });
 
     receivers.put(streamId, receiver);
 
     return receiver.doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
@@ -281,7  282,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);
 
@@ -293,30  294,49 @@ private Flux<Payload> handleRequestStream(final Payload payload) {
 
               @Override
               public void accept(long n) {
-                if (firstRequest && !receiver.isDisposed()) {
                 if (firstRequest) {
                   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;
                   }
-                } else if (contains(streamId) && !receiver.isDisposed()) {
                   int missed = 1;
                   boolean firstHasBeenSent = false;
                   for (; ; ) {
                     if (!firstHasBeenSent) {
                       sendProcessor.onNext(
                           RequestStreamFrameFlyweight.encodeReleasingPayload(
                               allocator, streamId, n, payload));
                       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 (!receiver.isDisposed()) {
                   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()) {
 
               // check if we need to release payload
               // only applicable if the cancel appears earlier than actual request
               if (payload.refCnt() > 0) {
                 payload.release();
               } else {
                 sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
               }
             })
@@ -330,30  350,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,19  443,47 @@ 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));
@@ -442,22  492,22 @@ public void accept(long n) {
             })
         .doOnError(
             t -> {
-              if (receivers.remove(streamId, receiver)) {
-                upstreamSubscriber.cancel();
-              }
               upstreamSubscriber.cancel();
               receivers.remove(streamId, receiver);
             })
         .doOnComplete(() -> receivers.remove(streamId, receiver))
         .doOnCancel(
             () -> {
-              if (!payloadReleasedFlag.getAndSet(true)) {
-                initialPayload.release();
               upstreamSubscriber.cancel();
               if (wip.getAndIncrement() != 0) {
                 return;
               }
 
               // need to send frame only if RequestChannelFrame was sent
               if (receivers.remove(streamId, receiver)) {
                 sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
-                upstreamSubscriber.cancel();
               }
-            })
-        .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
             });
   }
 
   private Mono<Void> handleMetadataPush(Payload payload) {

diff --git a/rsocket-core/src/main/java/io/rsocket/frame/DataAndMetadataFlyweight.java b/rsocket-core/src/main/java/io/rsocket/frame/DataAndMetadataFlyweight.java
@@ -37,8  37,34 @@ static ByteBuf encode(
       boolean hasMetadata,
       ByteBuf data) {
 
-    final boolean addData = data != null && data.isReadable();
-    final boolean addMetadata = hasMetadata && metadata.isReadable();
     final boolean addData;
     if (data != null) {
       if (data.isReadable()) {
         addData = true;
       } else {
         // even though there is nothing to read, we still have to release here since nobody else
         // going to do soo
         data.release();
         addData = false;
       }
     } else {
       addData = false;
     }
 
     final boolean addMetadata;
     if (hasMetadata) {
       if (metadata.isReadable()) {
         addMetadata = true;
       } else {
         // even though there is nothing to read, we still have to release here since nobody else
         // going to do soo
         metadata.release();
         addMetadata = false;
       }
     } else {
       // has no metadata means it is null, thus no need to release anything
       addMetadata = false;
     }
 
     if (hasMetadata) {
       int length = metadata.readableBytes();

diff --git a/rsocket-core/src/main/java/io/rsocket/frame/LeaseFrameFlyweight.java b/rsocket-core/src/main/java/io/rsocket/frame/LeaseFrameFlyweight.java
@@ -14,7  14,6 @@ public static ByteBuf encode(
       @Nullable final ByteBuf metadata) {
 
     final boolean hasMetadata = metadata != null;
-    final boolean addMetadata = hasMetadata && metadata.isReadable();
 
     int flags = 0;
 
@@ -27,6  26,21 @@ public static ByteBuf encode(
             .writeInt(ttl)
             .writeInt(numRequests);
 
     final boolean addMetadata;
     if (hasMetadata) {
       if (metadata.isReadable()) {
         addMetadata = true;
       } else {
         // even though there is nothing to read, we still have to release here since nobody else
         // going to do soo
         metadata.release();
         addMetadata = false;
       }
     } else {
       // has no metadata means it is null, thus no need to release anything
       addMetadata = false;
     }
 
     if (addMetadata) {
       return allocator.compositeBuffer(2).addComponents(true, header, metadata);
     } else {

diff --git a/rsocket-core/src/main/java/io/rsocket/frame/MetadataPushFrameFlyweight.java b/rsocket-core/src/main/java/io/rsocket/frame/MetadataPushFrameFlyweight.java
@@ -2,13  2,21 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.util.IllegalReferenceCountException;
 import io.rsocket.Payload;
 
 public class MetadataPushFrameFlyweight {
 
   public static ByteBuf encodeReleasingPayload(ByteBufAllocator allocator, Payload payload) {
     final ByteBuf metadata = payload.metadata().retain();
-    payload.release();
     // releasing payload safely since it can be already released wheres we have to release retained
     // data and metadata as well
     try {
       payload.release();
     } catch (IllegalReferenceCountException e) {
       metadata.release();
       throw e;
     }
     return encode(allocator, metadata);
   }
 

diff --git a/rsocket-core/src/main/java/io/rsocket/frame/PayloadFrameFlyweight.java b/rsocket-core/src/main/java/io/rsocket/frame/PayloadFrameFlyweight.java
@@ -2,6  2,7 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.util.IllegalReferenceCountException;
 import io.rsocket.Payload;
 
 public class PayloadFrameFlyweight {
@@ -23,11  24,31 @@ public static ByteBuf encodeNextCompleteReleasingPayload(
   static ByteBuf encodeReleasingPayload(
       ByteBufAllocator allocator, int streamId, boolean complete, Payload payload) {
 
-    final boolean hasMetadata = payload.hasMetadata();
     // if refCnt exceptions throws here it is safe to do no-op
     boolean hasMetadata = payload.hasMetadata();
     // if refCnt exceptions throws here it is safe to do no-op still
     final ByteBuf metadata = hasMetadata ? payload.metadata().retain() : null;
-    final ByteBuf data = payload.data().retain();
-
-    payload.release();
     final ByteBuf data;
     // retaining data safely. May throw either NPE or RefCntE
     try {
       data = payload.data().retain();
     } catch (IllegalReferenceCountException | NullPointerException e) {
       if (hasMetadata) {
         metadata.release();
       }
       throw e;
     }
     // releasing payload safely since it can be already released wheres we have to release retained
     // data and metadata as well
     try {
       payload.release();
     } catch (IllegalReferenceCountException e) {
       data.release();
       if (hasMetadata) {
         metadata.release();
       }
       throw e;
     }
 
     return encode(allocator, streamId, false, complete, true, metadata, data);
   }