provides extra hooks to ensure we capture all discarded elements

Signed-off-by: Oleh Dokuka <shadowgun@i.ua>
rsocket · Apr 10, 2020 · 626dd81 · 626dd81
1 parent c2475c2
commit 626dd81
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 10 deletions.
diff --git a/rsocket-core/src/main/java/io/rsocket/core/CleanOnClearQueueDecorator.java b/rsocket-core/src/main/java/io/rsocket/core/CleanOnClearQueueDecorator.java
@@ -0,0 +1,108 @@
+package io.rsocket.core;
+
+import io.netty.util.ReferenceCountUtil;
+import io.rsocket.Payload;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.Queue;
+
+final class CleanOnClearQueueDecorator implements Queue<Payload> {
+  final Queue<Payload> delegate;
+
+  CleanOnClearQueueDecorator(Queue<Payload> delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public void clear() {
+    Payload p;
+    while ((p = delegate.poll()) != null) {
+      ReferenceCountUtil.safeRelease(p);
+    }
+  }
+
+  @Override
+  public int size() {
+    return delegate.size();
+  }
+
+  @Override
+  public boolean isEmpty() {
+    return delegate.isEmpty();
+  }
+
+  @Override
+  public boolean contains(Object o) {
+    return delegate.contains(o);
+  }
+
+  @Override
+  public Iterator<Payload> iterator() {
+    return delegate.iterator();
+  }
+
+  @Override
+  public Object[] toArray() {
+    return delegate.toArray();
+  }
+
+  @Override
+  public <T> T[] toArray(T[] a) {
+    return delegate.toArray(a);
+  }
+
+  @Override
+  public boolean add(Payload payload) {
+    return delegate.add(payload);
+  }
+
+  @Override
+  public boolean remove(Object o) {
+    return delegate.remove(o);
+  }
+
+  @Override
+  public boolean containsAll(Collection<?> c) {
+    return delegate.containsAll(c);
+  }
+
+  @Override
+  public boolean addAll(Collection<? extends Payload> c) {
+    return delegate.addAll(c);
+  }
+
+  @Override
+  public boolean removeAll(Collection<?> c) {
+    return delegate.retainAll(c);
+  }
+
+  @Override
+  public boolean retainAll(Collection<?> c) {
+    return delegate.retainAll(c);
+  }
+
+  @Override
+  public boolean offer(Payload payload) {
+    return delegate.offer(payload);
+  }
+
+  @Override
+  public Payload remove() {
+    return delegate.remove();
+  }
+
+  @Override
+  public Payload poll() {
+    return delegate.poll();
+  }
+
+  @Override
+  public Payload element() {
+    return delegate.element();
+  }
+
+  @Override
+  public Payload peek() {
+    return delegate.peek();
+  }
+}
diff --git a/rsocket-core/src/main/java/io/rsocket/core/RSocketRequester.java b/rsocket-core/src/main/java/io/rsocket/core/RSocketRequester.java
@@ -23,6 +23,7 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
 import io.netty.util.collection.IntObjectMap;
 import io.rsocket.DuplexConnection;
 import io.rsocket.Payload;
@@ -77,6 +78,8 @@ class RSocketRequester implements RSocket {
       AtomicReferenceFieldUpdater.newUpdater(
           RSocketRequester.class, Throwable.class, "terminationError");
   private static final Exception CLOSED_CHANNEL_EXCEPTION = new ClosedChannelException();
+  private static final Consumer<ReferenceCounted> DROPPED_ELEMENTS_CONSUMER =
+      ReferenceCountUtil::safeRelease;
 
   static {
     CLOSED_CHANNEL_EXCEPTION.setStackTrace(new StackTraceElement[0]);
@@ -259,7 +262,7 @@ public void doOnTerminal(
             });
     receivers.put(streamId, receiver);
 
-    return receiver;
+    return receiver.doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
   }
 
   private Flux<Payload> handleRequestStream(final Payload payload) {
@@ -277,7 +280,8 @@ private Flux<Payload> handleRequestStream(final Payload payload) {
     int streamId = streamIdSupplier.nextStreamId(receivers);
 
     final UnboundedProcessor<ByteBuf> sendProcessor = this.sendProcessor;
-    final UnicastProcessor<Payload> receiver = UnicastProcessor.create();
+    final UnicastProcessor<Payload> receiver =
+        UnicastProcessor.create(new CleanOnClearQueueDecorator(Queues.<Payload>unbounded().get()));
     final AtomicBoolean payloadReleasedFlag = new AtomicBoolean(false);
 
     receivers.put(streamId, receiver);
@@ -323,7 +327,8 @@ public void accept(long n) {
                 sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
               }
             })
-        .doFinally(s -> removeStreamReceiver(streamId));
+        .doFinally(s -> removeStreamReceiver(streamId))
+        .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
   }
 
   private Flux<Payload> handleChannel(Flux<Payload> request) {
@@ -356,7 +361,8 @@ private Flux<? extends Payload> handleChannel(Payload initialPayload, Flux<Paylo
     final AtomicBoolean payloadReleasedFlag = new AtomicBoolean(false);
     final int streamId = streamIdSupplier.nextStreamId(receivers);
 
-    final UnicastProcessor<Payload> receiver = UnicastProcessor.create();
+    final UnicastProcessor<Payload> receiver =
+        UnicastProcessor.create(new CleanOnClearQueueDecorator(Queues.<Payload>unbounded().get()));
     final BaseSubscriber<Payload> upstreamSubscriber =
         new BaseSubscriber<Payload>() {
 
@@ -424,7 +430,10 @@ public void accept(long n) {
                   senders.put(streamId, upstreamSubscriber);
                   receivers.put(streamId, receiver);
 
-                  inboundFlux.limitRate(Queues.SMALL_BUFFER_SIZE).subscribe(upstreamSubscriber);
+                  inboundFlux
+                      .limitRate(Queues.SMALL_BUFFER_SIZE)
+                      .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER)
+                      .subscribe(upstreamSubscriber);
                   if (!payloadReleasedFlag.getAndSet(true)) {
                     ByteBuf frame =
                         RequestChannelFrameFlyweight.encode(
@@ -461,7 +470,8 @@ public void accept(long n) {
                 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/core/RSocketResponder.java b/rsocket-core/src/main/java/io/rsocket/core/RSocketResponder.java
@@ -21,6 +21,7 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
 import io.netty.util.collection.IntObjectMap;
 import io.rsocket.DuplexConnection;
 import io.rsocket.Payload;
@@ -45,6 +46,8 @@
 
 /** Responder side of RSocket. Receives {@link ByteBuf}s from a peer's {@link RSocketRequester} */
 class RSocketResponder implements ResponderRSocket {
+  private static final Consumer<ReferenceCounted> DROPPED_ELEMENTS_CONSUMER =
+      ReferenceCountUtil::safeRelease;
 
   private final DuplexConnection connection;
   private final RSocket requestHandler;
@@ -418,7 +421,7 @@ protected void hookFinally(SignalType type) {
         };
 
     sendingSubscriptions.put(streamId, subscriber);
-    response.subscribe(subscriber);
+    response.doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER).subscribe(subscriber);
   }
 
   private void handleStream(int streamId, Flux<Payload> response, int initialRequestN) {
@@ -471,11 +474,15 @@ protected void hookFinally(SignalType type) {
         };
 
     sendingSubscriptions.put(streamId, subscriber);
-    response.limitRate(Queues.SMALL_BUFFER_SIZE).subscribe(subscriber);
+    response
+        .limitRate(Queues.SMALL_BUFFER_SIZE)
+        .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER)
+        .subscribe(subscriber);
   }
 
   private void handleChannel(int streamId, Payload payload, int initialRequestN) {
-    UnicastProcessor<Payload> frames = UnicastProcessor.create();
+    UnicastProcessor<Payload> frames =
+        UnicastProcessor.create(new CleanOnClearQueueDecorator(Queues.<Payload>unbounded().get()));
     channelProcessors.put(streamId, frames);
 
     Flux<Payload> payloads =
@@ -499,7 +506,8 @@ public void accept(long l) {
                     sendProcessor.onNext(RequestNFrameFlyweight.encode(allocator, streamId, n));
                   }
                 })
-            .doFinally(signalType -> channelProcessors.remove(streamId));
+            .doFinally(signalType -> channelProcessors.remove(streamId))
+            .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER);
 
     // not chained, as the payload should be enqueued in the Unicast processor before this method
     // returns