improves loadbalance test coverage and provides fixes

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

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/loadbalance/MonoDeferredResolution.java

@@ -59,7 +59,7 @@ abstract class MonoDeferredResolution<RESULT, R> extends Mono<RESULT>
   }
 
   @Override
-  public void subscribe(CoreSubscriber<? super RESULT> actual) {
+  public final void subscribe(CoreSubscriber<? super RESULT> actual) {
     if (this.requested == STATE_UNSUBSCRIBED
         && REQUESTED.compareAndSet(this, STATE_UNSUBSCRIBED, STATE_SUBSCRIBER_SET)) {
 
@@ -145,7 +145,7 @@ public final void onNext(RESULT payload) {
   }
 
   @Override
-  public void onError(Throwable t) {
+  public final void onError(Throwable t) {
     if (this.done) {
       Operators.onErrorDropped(t, this.actual.currentContext());
       return;
@@ -156,7 +156,7 @@ public void onError(Throwable t) {
   }
 
   @Override
-  public void onComplete() {
+  public final void onComplete() {
     if (this.done) {
       return;
     }
@@ -206,7 +206,7 @@ public final void request(long n) {
     }
   }
 
-  public void cancel() {
+  public final void cancel() {
     long state = REQUESTED.getAndSet(this, STATE_TERMINATED);
     if (state == STATE_TERMINATED) {
       return;

rsocket-core/src/main/java/io/rsocket/loadbalance/PooledRSocket.java

@@ -108,22 +108,16 @@ protected void doSubscribe() {
 
   @Override
   protected void doOnValueResolved(RSocket value) {
-    value.onClose().subscribe(null, t -> this.invalidate(), this::invalidate);
+    value.onClose().subscribe(null, t -> this.doCleanup(), this::doCleanup);
   }
 
-  @Override
-  protected void doOnValueExpired(RSocket value) {
-    value.dispose();
-    this.dispose();
-  }
+  void doCleanup() {
+    if (isDisposed()) {
+      return;
+    }
 
-  @Override
-  public void dispose() {
-    super.dispose();
-  }
+    this.dispose();
 
-  @Override
-  protected void doOnDispose() {
     final RSocketPool parent = this.parent;
     for (; ; ) {
       final PooledRSocket[] sockets = parent.activeSockets;
@@ -141,20 +135,35 @@ protected void doOnDispose() {
         break;
       }
 
-      final int lastIndex = activeSocketsCount - 1;
-      final PooledRSocket[] newSockets = new PooledRSocket[lastIndex];
-      if (index != 0) {
-        System.arraycopy(sockets, 0, newSockets, 0, index);
-      }
+      final PooledRSocket[] newSockets;
+      if (activeSocketsCount == 1) {
+        newSockets = RSocketPool.EMPTY;
+      } else {
+        final int lastIndex = activeSocketsCount - 1;
+
+        newSockets = new PooledRSocket[lastIndex];
+        if (index != 0) {
+          System.arraycopy(sockets, 0, newSockets, 0, index);
+        }
 
-      if (index != lastIndex) {
-        System.arraycopy(sockets, index + 1, newSockets, index, lastIndex - index);
+        if (index != lastIndex) {
+          System.arraycopy(sockets, index + 1, newSockets, index, lastIndex - index);
+        }
       }
 
       if (RSocketPool.ACTIVE_SOCKETS.compareAndSet(parent, sockets, newSockets)) {
         break;
       }
     }
+  }
+
+  @Override
+  protected void doOnValueExpired(RSocket value) {
+    value.dispose();
+  }
+
+  @Override
+  protected void doOnDispose() {
     Operators.terminate(S, this);
   }
 
@@ -231,7 +240,7 @@ public void accept(RSocket rSocket, Throwable t) {
 
         source.subscribe((CoreSubscriber) this);
       } else {
-        parent.add(this);
+        parent.observe(this);
       }
     }
   }
@@ -273,7 +282,7 @@ public void accept(RSocket rSocket, Throwable t) {
 
         source.subscribe(this);
       } else {
-        parent.add(this);
+        parent.observe(this);
       }
     }
   }