Fixes possible elements' leaks occurrence (#799)

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/internal/UnboundedProcessor.java

@@ -43,28 +43,43 @@
 public final class UnboundedProcessor<T> extends FluxProcessor<T, T>
     implements Fuseable.QueueSubscription<T>, Fuseable {
 
+  final Queue<T> queue;
+  final Queue<T> priorityQueue;
+
+  volatile boolean done;
+  Throwable error;
+  // important to not loose the downstream too early and miss discard hook, while
+  // having relevant hasDownstreams()
+  boolean hasDownstream;
+  volatile CoreSubscriber<? super T> actual;
+
+  volatile boolean cancelled;
+
+  volatile int once;
+
   @SuppressWarnings("rawtypes")
   static final AtomicIntegerFieldUpdater<UnboundedProcessor> ONCE =
       AtomicIntegerFieldUpdater.newUpdater(UnboundedProcessor.class, "once");
 
+  volatile int wip;
+
   @SuppressWarnings("rawtypes")
   static final AtomicIntegerFieldUpdater<UnboundedProcessor> WIP =
       AtomicIntegerFieldUpdater.newUpdater(UnboundedProcessor.class, "wip");
 
+  volatile int discardGuard;
+
+  @SuppressWarnings("rawtypes")
+  static final AtomicIntegerFieldUpdater<UnboundedProcessor> DISCARD_GUARD =
+      AtomicIntegerFieldUpdater.newUpdater(UnboundedProcessor.class, "discardGuard");
+
+  volatile long requested;
+
   @SuppressWarnings("rawtypes")
   static final AtomicLongFieldUpdater<UnboundedProcessor> REQUESTED =
       AtomicLongFieldUpdater.newUpdater(UnboundedProcessor.class, "requested");
 
-  final Queue<T> queue;
-  final Queue<T> priorityQueue;
-  volatile boolean done;
-  Throwable error;
-  volatile CoreSubscriber<? super T> actual;
-  volatile boolean cancelled;
-  volatile int once;
-  volatile int wip;
-  volatile long requested;
-  volatile boolean outputFused;
+  boolean outputFused;
 
   public UnboundedProcessor() {
     this.queue = new MpscUnboundedArrayQueue<>(Queues.SMALL_BUFFER_SIZE);
@@ -79,6 +94,7 @@ public int getBufferSize() {
   @Override
   public Object scanUnsafe(Attr key) {
     if (Attr.BUFFERED == key) return queue.size();
+    if (Attr.PREFETCH == key) return Integer.MAX_VALUE;
     return super.scanUnsafe(key);
   }
 
@@ -143,8 +159,8 @@ void drainFused(Subscriber<? super T> a) {
     for (; ; ) {
 
       if (cancelled) {
-        clear();
-        actual = null;
+        this.clear();
+        hasDownstream = false;
         return;
       }
 
@@ -153,7 +169,7 @@ void drainFused(Subscriber<? super T> a) {
       a.onNext(null);
 
       if (d) {
-        actual = null;
+        hasDownstream = false;
 
         Throwable ex = error;
         if (ex != null) {
@@ -173,6 +189,9 @@ void drainFused(Subscriber<? super T> a) {
 
   public void drain() {
     if (WIP.getAndIncrement(this) != 0) {
+      if (cancelled) {
+        this.clear();
+      }
       return;
     }
 
@@ -199,13 +218,13 @@ public void drain() {
 
   boolean checkTerminated(boolean d, boolean empty, Subscriber<? super T> a) {
     if (cancelled) {
-      clear();
-      actual = null;
+      this.clear();
+      hasDownstream = false;
       return true;
     }
     if (d && empty) {
       Throwable e = error;
-      actual = null;
+      hasDownstream = false;
       if (e != null) {
         a.onError(e);
       } else {
@@ -226,10 +245,6 @@ public void onSubscribe(Subscription s) {
     }
   }
 
-  public long available() {
-    return requested;
-  }
-
   @Override
   public int getPrefetch() {
     return Integer.MAX_VALUE;
@@ -308,7 +323,7 @@ public void subscribe(CoreSubscriber<? super T> actual) {
       actual.onSubscribe(this);
       this.actual = actual;
       if (cancelled) {
-        this.actual = null;
+        this.hasDownstream = false;
       } else {
         drain();
       }
@@ -335,8 +350,8 @@ public void cancel() {
     cancelled = true;
 
     if (WIP.getAndIncrement(this) == 0) {
-      clear();
-      actual = null;
+      this.clear();
+      hasDownstream = false;
     }
   }
 
@@ -362,16 +377,29 @@ public boolean isEmpty() {
 
   @Override
   public void clear() {
-    while (!queue.isEmpty()) {
-      T t = queue.poll();
-      if (t != null) {
-        release(t);
-      }
+    if (DISCARD_GUARD.getAndIncrement(this) != 0) {
+      return;
     }
-    while (!priorityQueue.isEmpty()) {
-      T t = priorityQueue.poll();
-      if (t != null) {
-        release(t);
+
+    int missed = 1;
+
+    for (; ; ) {
+      while (!queue.isEmpty()) {
+        T t = queue.poll();
+        if (t != null) {
+          release(t);
+        }
+      }
+      while (!priorityQueue.isEmpty()) {
+        T t = priorityQueue.poll();
+        if (t != null) {
+          release(t);
+        }
+      }
+
+      missed = DISCARD_GUARD.addAndGet(this, -missed);
+      if (missed == 0) {
+        break;
       }
     }
   }
@@ -413,14 +441,18 @@ public long downstreamCount() {
 
   @Override
   public boolean hasDownstreams() {
-    return actual != null;
+    return hasDownstream;
   }
 
   void release(T t) {
     if (t instanceof ReferenceCounted) {
       ReferenceCounted refCounted = (ReferenceCounted) t;
       if (refCounted.refCnt() > 0) {
-        refCounted.release();
+        try {
+          refCounted.release();
+        } catch (Throwable ex) {
+          // no ops
+        }
       }
     }
   }

rsocket-core/src/main/java/io/rsocket/resume/ResumableDuplexConnection.java

@@ -223,6 +223,10 @@ public boolean isDisposed() {
   }
 
   private void sendFrame(ByteBuf f) {
+    if (disposed.get()) {
+      f.release();
+      return;
+    }
     /*resuming from store so no need to save again*/
     if (state != State.RESUME && isResumableFrame(f)) {
       resumeSaveFrames.onNext(f);