ensures LocalDuplexConnection#onClose notifies about termination correctly

At the moment, the onClose hook has no "wait until cleaned" logic, which leads to unpredicted behaviors when used with resumability or others scenarios where we need to wait until all the queues are cleaned and there are no other resources in use (e.g. ByteBufs). For that porpuse, this commit adds onFinalizeHook to the UnboundedProcessor so we can now listen when the UnboundedProcessor is finalized and only after that send the onClose signal

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

Author: OlegDokuka

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

@@ -45,6 +45,7 @@ public final class UnboundedProcessor extends FluxProcessor<ByteBuf, ByteBuf>
 
   final Queue<ByteBuf> queue;
   final Queue<ByteBuf> priorityQueue;
+  final Runnable onFinalizedHook;
 
   boolean cancelled;
   boolean done;
@@ -88,6 +89,11 @@ public final class UnboundedProcessor extends FluxProcessor<ByteBuf, ByteBuf>
   boolean outputFused;
 
   public UnboundedProcessor() {
+    this(() -> {});
+  }
+
+  public UnboundedProcessor(Runnable onFinalizedHook) {
+    this.onFinalizedHook = onFinalizedHook;
     this.queue = new MpscUnboundedArrayQueue<>(Queues.SMALL_BUFFER_SIZE);
     this.priorityQueue = new MpscUnboundedArrayQueue<>(Queues.SMALL_BUFFER_SIZE);
   }
@@ -793,6 +799,9 @@ static long markTerminatedOrFinalized(UnboundedProcessor instance) {
       }
 
       if (STATE.compareAndSet(instance, state, nextState | FLAG_TERMINATED)) {
+        if (isFinalized(nextState)) {
+          instance.onFinalizedHook.run();
+        }
         return state;
       }
     }
@@ -906,6 +915,7 @@ static void clearAndFinalize(UnboundedProcessor instance) {
 
       if (STATE.compareAndSet(
           instance, state, (state & ~MAX_WIP_VALUE & ~FLAG_HAS_VALUE) | FLAG_FINALIZED)) {
+        instance.onFinalizedHook.run();
         break;
       }
     }

rsocket-transport-local/src/main/java/io/rsocket/transport/local/LocalClientTransport.java

@@ -77,14 +77,17 @@ public Mono<DuplexConnection> connect() {
             return Mono.error(new IllegalArgumentException("Could not find server: " + name));
           }
 
-          UnboundedProcessor in = new UnboundedProcessor();
-          UnboundedProcessor out = new UnboundedProcessor();
-          Sinks.Empty<Void> closeSink = Sinks.empty();
+          Sinks.One<Object> inSink = Sinks.one();
+          Sinks.One<Object> outSink = Sinks.one();
+          UnboundedProcessor in = new UnboundedProcessor(() -> inSink.tryEmitValue(inSink));
+          UnboundedProcessor out = new UnboundedProcessor(() -> outSink.tryEmitValue(outSink));
 
-          server.apply(new LocalDuplexConnection(name, allocator, out, in, closeSink)).subscribe();
+          Mono<Void> onClose = inSink.asMono().zipWith(outSink.asMono()).then();
 
-          return Mono.just(
-              (DuplexConnection) new LocalDuplexConnection(name, allocator, in, out, closeSink));
+          server.apply(new LocalDuplexConnection(name, allocator, out, in, onClose)).subscribe();
+
+          return Mono.<DuplexConnection>just(
+              new LocalDuplexConnection(name, allocator, in, out, onClose));
         });
   }
 }

rsocket-transport-local/src/main/java/io/rsocket/transport/local/LocalDuplexConnection.java

@@ -27,11 +27,9 @@
 import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
 import reactor.core.Fuseable;
-import reactor.core.Scannable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Operators;
-import reactor.core.publisher.Sinks;
 
 /** An implementation of {@link DuplexConnection} that connects inside the same JVM. */
 final class LocalDuplexConnection implements DuplexConnection {
@@ -40,7 +38,7 @@ final class LocalDuplexConnection implements DuplexConnection {
   private final ByteBufAllocator allocator;
   private final Flux<ByteBuf> in;
 
-  private final Sinks.Empty<Void> onClose;
+  private final Mono<Void> onClose;
 
   private final UnboundedProcessor out;
 
@@ -58,7 +56,7 @@ final class LocalDuplexConnection implements DuplexConnection {
       ByteBufAllocator allocator,
       Flux<ByteBuf> in,
       UnboundedProcessor out,
-      Sinks.Empty<Void> onClose) {
+      Mono<Void> onClose) {
     this.address = new LocalSocketAddress(name);
     this.allocator = Objects.requireNonNull(allocator, "allocator must not be null");
     this.in = Objects.requireNonNull(in, "in must not be null");
@@ -69,24 +67,23 @@ final class LocalDuplexConnection implements DuplexConnection {
   @Override
   public void dispose() {
     out.onComplete();
-    onClose.tryEmitEmpty();
   }
 
   @Override
-  @SuppressWarnings("ConstantConditions")
   public boolean isDisposed() {
-    return onClose.scan(Scannable.Attr.TERMINATED) || onClose.scan(Scannable.Attr.CANCELLED);
+    return out.isDisposed();
   }
 
   @Override
   public Mono<Void> onClose() {
-    return onClose.asMono();
+    return onClose;
   }
 
   @Override
   public Flux<ByteBuf> receive() {
     return in.transform(
-        Operators.<ByteBuf, ByteBuf>lift((__, actual) -> new ByteBufReleaserOperator(actual)));
+        Operators.<ByteBuf, ByteBuf>lift(
+            (__, actual) -> new ByteBufReleaserOperator(actual, this)));
   }
 
   @Override
@@ -119,11 +116,14 @@ static class ByteBufReleaserOperator
       implements CoreSubscriber<ByteBuf>, Subscription, Fuseable.QueueSubscription<ByteBuf> {
 
     final CoreSubscriber<? super ByteBuf> actual;
+    final LocalDuplexConnection parent;
 
     Subscription s;
 
-    public ByteBufReleaserOperator(CoreSubscriber<? super ByteBuf> actual) {
+    public ByteBufReleaserOperator(
+        CoreSubscriber<? super ByteBuf> actual, LocalDuplexConnection parent) {
       this.actual = actual;
+      this.parent = parent;
     }
 
     @Override
@@ -142,11 +142,13 @@ public void onNext(ByteBuf buf) {
 
     @Override
     public void onError(Throwable t) {
+      parent.out.onError(t);
       actual.onError(t);
     }
 
     @Override
     public void onComplete() {
+      parent.out.onComplete();
       actual.onComplete();
     }
 
@@ -158,6 +160,7 @@ public void request(long n) {
     @Override
     public void cancel() {
       s.cancel();
+      parent.out.onComplete();
     }
 
     @Override