provides workaround for FluxPublishOn to ensure that all elements are released in case of racing

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

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/core/CleanOnClearQueueDecorator.java

@@ -7,6 +7,12 @@
 import java.util.Queue;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+/**
+ * This queue decorator is temporary solution in order to workaround leaks in UnicastProcessor and
+ * FluxPublishOn. <br>
+ * For more information about the root-cause of leaks, please see
+ * https://github.com/reactor/reactor-core/pull/2114
+ */
 final class CleanOnClearQueueDecorator extends AtomicBoolean implements Queue<Payload> {
   final Queue<Payload> delegate;
 
@@ -23,6 +29,23 @@ public void clear() {
     }
   }
 
+  @Override
+  public boolean offer(Payload payload) {
+    if (get()) {
+      ReferenceCountUtil.safeRelease(payload);
+      return true;
+    }
+
+    boolean result = delegate.offer(payload);
+
+    // ensures in case of racing offered element is released for sure
+    if (get()) {
+      ReferenceCountUtil.safeRelease(payload);
+    }
+
+    return result;
+  }
+
   @Override
   public int size() {
     return delegate.size();
@@ -83,15 +106,6 @@ public boolean retainAll(Collection<?> c) {
     return delegate.retainAll(c);
   }
 
-  @Override
-  public boolean offer(Payload payload) {
-    if (get()) {
-      ReferenceCountUtil.safeRelease(payload);
-      return true;
-    }
-    return delegate.offer(payload);
-  }
-
   @Override
   public Payload remove() {
     return delegate.remove();

rsocket-core/src/main/java/io/rsocket/core/LimitRateOperator.java

@@ -0,0 +1,46 @@
+package io.rsocket.core;
+
+import io.rsocket.Payload;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
+import java.util.Queue;
+import java.util.function.Supplier;
+import reactor.core.publisher.Flux;
+import reactor.core.scheduler.Schedulers;
+import reactor.util.concurrent.Queues;
+
+/**
+ * This class is a temporary solution to workaround memory leaks related to bug in reactor. <br>
+ * For more information, pleas see https://github.com/reactor/reactor-core/pull/2114
+ */
+@SuppressWarnings("unchecked")
+final class LimitRateOperator {
+
+  static final Constructor<Flux<Payload>> fluxPublishOnCtr;
+
+  static {
+    try {
+      Class<?> fluxPublishOnClass = Class.forName("reactor.core.publisher.FluxPublishOn");
+      fluxPublishOnCtr =
+          (Constructor<Flux<Payload>>) fluxPublishOnClass.getDeclaredConstructors()[0];
+      fluxPublishOnCtr.setAccessible(true);
+    } catch (ClassNotFoundException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  static Flux<Payload> applyLimitRateOperator(Flux<Payload> source, int limit) {
+    Supplier<Queue<Payload>> queueSupplier =
+        () -> {
+          Queue<Payload> queue = Queues.<Payload>get(limit).get();
+          return new CleanOnClearQueueDecorator(queue);
+        };
+
+    try {
+      return fluxPublishOnCtr.newInstance(
+          source, Schedulers.immediate(), true, limit, limit, queueSupplier);
+    } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) {
+      throw new RuntimeException(e);
+    }
+  }
+}

rsocket-core/src/main/java/io/rsocket/core/RSocketRequester.java

@@ -16,6 +16,7 @@
 
 package io.rsocket.core;
 
+import static io.rsocket.core.LimitRateOperator.applyLimitRateOperator;
 import static io.rsocket.core.PayloadValidationUtils.INVALID_PAYLOAD_ERROR_MESSAGE;
 import static io.rsocket.keepalive.KeepAliveSupport.ClientKeepAliveSupport;
 import static io.rsocket.keepalive.KeepAliveSupport.KeepAlive;
@@ -431,7 +432,7 @@ public void accept(long n) {
                   receivers.put(streamId, receiver);
 
                   inboundFlux
-                      .limitRate(Queues.SMALL_BUFFER_SIZE)
+                      .transform(f -> applyLimitRateOperator(f, Queues.SMALL_BUFFER_SIZE))
                       .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER)
                       .subscribe(upstreamSubscriber);
                   if (!payloadReleasedFlag.getAndSet(true)) {

rsocket-core/src/main/java/io/rsocket/core/RSocketResponder.java

@@ -16,6 +16,7 @@
 
 package io.rsocket.core;
 
+import static io.rsocket.core.LimitRateOperator.applyLimitRateOperator;
 import static io.rsocket.core.PayloadValidationUtils.INVALID_PAYLOAD_ERROR_MESSAGE;
 
 import io.netty.buffer.ByteBuf;
@@ -475,7 +476,7 @@ protected void hookFinally(SignalType type) {
 
     sendingSubscriptions.put(streamId, subscriber);
     response
-        .limitRate(Queues.SMALL_BUFFER_SIZE)
+        .transform(f -> applyLimitRateOperator(f, Queues.SMALL_BUFFER_SIZE))
         .doOnDiscard(ReferenceCounted.class, DROPPED_ELEMENTS_CONSUMER)
         .subscribe(subscriber);
   }