fixes AssertSubscriber

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

Author: OlegDokuka

rsocket-core/src/test/java/io/rsocket/internal/UnboundedProcessorTest.java

@@ -25,18 +25,29 @@
 import io.netty.util.ReferenceCountUtil;
 import io.rsocket.buffer.LeaksTrackingByteBufAllocator;
 import io.rsocket.internal.subscriber.AssertSubscriber;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 import reactor.core.Fuseable;
+import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Operators;
 import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 import reactor.test.util.RaceTestUtils;
 
 public class UnboundedProcessorTest {
 
+  @BeforeAll
+  public static void setup() {
+    Hooks.onErrorDropped(__ -> {});
+  }
+
+  public static void teardown() {
+    Hooks.resetOnErrorDropped();
+  }
+
   @ParameterizedTest(
       name =
           "Test that emitting {0} onNext before subscribe and requestN should deliver all the signals once the subscriber is available")
@@ -96,7 +107,7 @@ public void testPrioritizedSending(boolean fusedCase) {
   public void ensureUnboundedProcessorDisposesQueueProperly(boolean withFusionEnabled) {
     final LeaksTrackingByteBufAllocator allocator =
         LeaksTrackingByteBufAllocator.instrument(ByteBufAllocator.DEFAULT);
-    for (int i = 0; i < 10000000; i++) {
+    for (int i = 0; i < 100000; i++) {
       final UnboundedProcessor<ByteBuf> unboundedProcessor = new UnboundedProcessor<>();
 
       final ByteBuf buffer1 = allocator.buffer(1);
@@ -135,7 +146,7 @@ public void ensureUnboundedProcessorDisposesQueueProperly(boolean withFusionEnab
 
   @RepeatedTest(
       name =
-          "Ensures that racing between onNext | dispose | cancel | request(n) will not cause any issues and leaks",
+          "Ensures that racing between onNext | dispose | cancel | request(n) will not cause any issues and leaks in async backFused mode",
       value = 100000)
   @Timeout(10)
   public void ensureUnboundedProcessorDisposesQueueProperlyAsyncMode() {

rsocket-core/src/test/java/io/rsocket/internal/subscriber/AssertSubscriber.java

@@ -105,6 +105,8 @@ public class AssertSubscriber<T> implements CoreSubscriber<T>, Subscription {
 
   private final CountDownLatch cdl = new CountDownLatch(1);
 
+  volatile boolean done;
+
   volatile Subscription s;
 
   volatile long requested;
@@ -879,53 +881,77 @@ public final boolean isTerminated() {
 
   @Override
   public void onComplete() {
+    done = true;
     completionCount++;
+
+    if (establishedFusionMode == Fuseable.ASYNC) {
+      drain();
+      return;
+    }
+
     cdl.countDown();
   }
 
   @Override
   public void onError(Throwable t) {
+    done = true;
     errors.add(t);
+
+    if (establishedFusionMode == Fuseable.ASYNC) {
+      drain();
+      return;
+    }
+
     cdl.countDown();
   }
 
   @Override
   public void onNext(T t) {
     if (establishedFusionMode == Fuseable.ASYNC) {
-      if (this.wip != 0 || WIP.getAndIncrement(this) != 0) {
-        if (isCancelled()) {
-          qs.clear();
+      drain();
+    } else {
+      valueCount++;
+      if (valuesStorage) {
+        List<T> nextValuesSnapshot;
+        for (; ; ) {
+          nextValuesSnapshot = values;
+          nextValuesSnapshot.add(t);
+          if (NEXT_VALUES.compareAndSet(this, nextValuesSnapshot, nextValuesSnapshot)) {
+            break;
+          }
         }
-        return;
       }
+    }
+  }
 
-      int m = 0;
-      for (; ; ) {
-        if (isCancelled()) {
-          qs.clear();
-          break;
-        }
-        t = qs.poll();
-        if (t == null) {
-          m = WIP.addAndGet(this, -m);
-          if (m == 0) {
-            break;
-          }
-          continue;
+  void drain() {
+    if (this.wip != 0 || WIP.getAndIncrement(this) != 0) {
+      if (isCancelled()) {
+        qs.clear();
+      }
+      return;
+    }
+
+    T t;
+    int m = 0;
+    for (; ; ) {
+      if (isCancelled()) {
+        qs.clear();
+        break;
+      }
+      boolean done = this.done;
+      t = qs.poll();
+      if (t == null) {
+        if (done) {
+          cdl.countDown();
+          return;
         }
-        valueCount++;
-        if (valuesStorage) {
-          List<T> nextValuesSnapshot;
-          for (; ; ) {
-            nextValuesSnapshot = values;
-            nextValuesSnapshot.add(t);
-            if (NEXT_VALUES.compareAndSet(this, nextValuesSnapshot, nextValuesSnapshot)) {
-              break;
-            }
-          }
+        m = WIP.addAndGet(this, -m);
+        if (m == 0) {
+          break;
         }
+        continue;
       }
-    } else {
       valueCount++;
       if (valuesStorage) {
         List<T> nextValuesSnapshot;