provides extra hooks to ensure we capture all discarded elements

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

Hide file tree

Showing 3 changed files with 133 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,109 @@
+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,7 @@ 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 +261,8 @@ 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,7 @@ 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 +326,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 +360,7 @@ 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 +428,9 @@ 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 +467,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,7 @@
 
 /** 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 +420,9 @@ 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 +475,13 @@ 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 +505,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