provides better naming; cleanups docs; minors fixes

Co-authored-by: Rossen Stoyanchev <[email protected]>
Signed-off-by: Oleh Dokuka <[email protected]>

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/loadbalance/ResolvingPooledRSocket.java renamed to rsocket-core/src/main/java/io/rsocket/loadbalance/DefaultPooledRSocket.java

@@ -28,21 +28,24 @@
 import reactor.core.publisher.Operators;
 import reactor.util.context.Context;
 
-/** Specific interface for all RSocket store in {@link RSocketPool} */
-final class ResolvingPooledRSocket extends ResolvingOperator<RSocket>
+/** Default implementation of {@link PooledRSocket} stored in {@link RSocketPool} */
+final class DefaultPooledRSocket extends ResolvingOperator<RSocket>
     implements CoreSubscriber<RSocket>, PooledRSocket {
 
-  final LoadbalanceTarget loadbalanceTarget;
+  final RSocketPool parent;
+  final LoadbalanceRSocketSource loadbalanceRSocketSource;
   final Stats stats;
 
   volatile Subscription s;
 
-  static final AtomicReferenceFieldUpdater<ResolvingPooledRSocket, Subscription> S =
-      AtomicReferenceFieldUpdater.newUpdater(ResolvingPooledRSocket.class, Subscription.class, "s");
+  static final AtomicReferenceFieldUpdater<DefaultPooledRSocket, Subscription> S =
+      AtomicReferenceFieldUpdater.newUpdater(DefaultPooledRSocket.class, Subscription.class, "s");
 
-  ResolvingPooledRSocket(LoadbalanceTarget loadbalanceTarget, Stats stats) {
+  DefaultPooledRSocket(
+      RSocketPool parent, LoadbalanceRSocketSource loadbalanceRSocketSource, Stats stats) {
+    this.parent = parent;
     this.stats = stats;
-    this.loadbalanceTarget = loadbalanceTarget;
+    this.loadbalanceRSocketSource = loadbalanceRSocketSource;
   }
 
   @Override
@@ -100,21 +103,61 @@ public void onNext(RSocket value) {
 
   @Override
   protected void doSubscribe() {
-    this.loadbalanceTarget.source().subscribe(this);
+    this.loadbalanceRSocketSource.source().subscribe(this);
   }
 
   @Override
   protected void doOnValueResolved(RSocket value) {
+    stats.setAvailability(1.0);
     value.onClose().subscribe(null, t -> this.invalidate(), this::invalidate);
   }
 
   @Override
   protected void doOnValueExpired(RSocket value) {
+    stats.setAvailability(0.0);
     value.dispose();
+    this.dispose();
+  }
+
+  @Override
+  public void dispose() {
+    super.dispose();
   }
 
   @Override
   protected void doOnDispose() {
+    final RSocketPool parent = this.parent;
+    for (; ; ) {
+      final PooledRSocket[] sockets = parent.activeSockets;
+      final int activeSocketsCount = sockets.length;
+
+      int index = -1;
+      for (int i = 0; i < activeSocketsCount; i++) {
+        if (sockets[i] == this) {
+          index = i;
+          break;
+        }
+      }
+
+      if (index == -1) {
+        break;
+      }
+
+      final int lastIndex = activeSocketsCount - 1;
+      final PooledRSocket[] 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 (RSocketPool.ACTIVE_SOCKETS.compareAndSet(parent, sockets, newSockets)) {
+        break;
+      }
+    }
+    stats.setAvailability(0.0);
     Operators.terminate(S, this);
   }
 
@@ -154,47 +197,21 @@ public Stats stats() {
   }
 
   @Override
-  public LoadbalanceTarget supplier() {
-    return loadbalanceTarget;
+  public LoadbalanceRSocketSource source() {
+    return loadbalanceRSocketSource;
   }
 
   @Override
   public double availability() {
     return stats.availability();
   }
 
-  /**
-   * Try to dispose this instance if possible. Otherwise, if there is ongoing requests, mark this as
-   * pending for removal and dispose once all the requests are terminated.<br>
-   * This operation may be cancelled if {@link #markActive()} is invoked prior this instance has
-   * been disposed
-   *
-   * @return {@code true} if this instance was disposed
-   */
-  @Override
-  public boolean markForRemoval() {
-    // FIXME: provide real logic here
-    this.dispose();
-    return true;
-  }
-
-  /**
-   * Try to restore state of this RSocket to be active after marking as pending removal again.
-   *
-   * @return {@code true} if marked as active. Otherwise, should be treated as it was disposed.
-   */
-  @Override
-  public boolean markActive() {
-    return false;
-  }
-
   static final class RequestTrackingMonoInner<RESULT>
       extends MonoDeferredResolution<RESULT, RSocket> {
 
     long startTime;
 
-    RequestTrackingMonoInner(
-        ResolvingPooledRSocket parent, Payload payload, FrameType requestType) {
+    RequestTrackingMonoInner(DefaultPooledRSocket parent, Payload payload, FrameType requestType) {
       super(parent, payload, requestType);
     }
 
@@ -228,7 +245,7 @@ public void accept(RSocket rSocket, Throwable t) {
             return;
         }
 
-        startTime = ((ResolvingPooledRSocket) parent).stats.startRequest();
+        startTime = ((DefaultPooledRSocket) parent).stats.startRequest();
 
         source.subscribe((CoreSubscriber) this);
       } else {
@@ -240,7 +257,7 @@ public void accept(RSocket rSocket, Throwable t) {
     public void onComplete() {
       final long state = this.requested;
       if (state != TERMINATED_STATE && REQUESTED.compareAndSet(this, state, TERMINATED_STATE)) {
-        final Stats stats = ((ResolvingPooledRSocket) parent).stats;
+        final Stats stats = ((DefaultPooledRSocket) parent).stats;
         final long now = stats.stopRequest(startTime);
         stats.record(now - startTime);
         super.onComplete();
@@ -251,7 +268,7 @@ public void onComplete() {
     public void onError(Throwable t) {
       final long state = this.requested;
       if (state != TERMINATED_STATE && REQUESTED.compareAndSet(this, state, TERMINATED_STATE)) {
-        Stats stats = ((ResolvingPooledRSocket) parent).stats;
+        Stats stats = ((DefaultPooledRSocket) parent).stats;
         stats.stopRequest(startTime);
         stats.recordError(0.0);
         super.onError(t);
@@ -267,7 +284,7 @@ public void cancel() {
 
       if (state == STATE_SUBSCRIBED) {
         this.s.cancel();
-        ((ResolvingPooledRSocket) parent).stats.stopRequest(startTime);
+        ((DefaultPooledRSocket) parent).stats.stopRequest(startTime);
       } else {
         this.parent.remove(this);
         ReferenceCountUtil.safeRelease(this.payload);
@@ -279,7 +296,7 @@ static final class RequestTrackingFluxInner<INPUT>
       extends FluxDeferredResolution<INPUT, RSocket> {
 
     RequestTrackingFluxInner(
-        ResolvingPooledRSocket parent, INPUT fluxOrPayload, FrameType requestType) {
+        DefaultPooledRSocket parent, INPUT fluxOrPayload, FrameType requestType) {
       super(parent, fluxOrPayload, requestType);
     }
 
@@ -312,7 +329,7 @@ public void accept(RSocket rSocket, Throwable t) {
             return;
         }
 
-        ((ResolvingPooledRSocket) parent).stats.startStream();
+        ((DefaultPooledRSocket) parent).stats.startStream();
 
         source.subscribe(this);
       } else {
@@ -324,7 +341,7 @@ public void accept(RSocket rSocket, Throwable t) {
     public void onComplete() {
       final long state = this.requested;
       if (state != TERMINATED_STATE && REQUESTED.compareAndSet(this, state, TERMINATED_STATE)) {
-        ((ResolvingPooledRSocket) parent).stats.stopStream();
+        ((DefaultPooledRSocket) parent).stats.stopStream();
         super.onComplete();
       }
     }
@@ -333,7 +350,7 @@ public void onComplete() {
     public void onError(Throwable t) {
       final long state = this.requested;
       if (state != TERMINATED_STATE && REQUESTED.compareAndSet(this, state, TERMINATED_STATE)) {
-        ((ResolvingPooledRSocket) parent).stats.stopStream();
+        ((DefaultPooledRSocket) parent).stats.stopStream();
         super.onError(t);
       }
     }
@@ -347,7 +364,7 @@ public void cancel() {
 
       if (state == STATE_SUBSCRIBED) {
         this.s.cancel();
-        ((ResolvingPooledRSocket) parent).stats.stopStream();
+        ((DefaultPooledRSocket) parent).stats.stopStream();
       } else {
         this.parent.remove(this);
         if (requestType == FrameType.REQUEST_STREAM) {

rsocket-core/src/main/java/io/rsocket/loadbalance/LoadBalancedRSocketClient.java renamed to rsocket-core/src/main/java/io/rsocket/loadbalance/LoadbalanceRSocketClient.java

@@ -23,11 +23,17 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-public class LoadBalancedRSocketClient implements RSocketClient {
+/**
+ * {@link RSocketClient} implementation that uses a pool and a {@link LoadbalanceStrategy} to select
+ * the {@code RSocket} to use for each request.
+ *
+ * @since 1.1
+ */
+public class LoadbalanceRSocketClient implements RSocketClient {
 
   private final RSocketPool rSocketPool;
 
-  LoadBalancedRSocketClient(RSocketPool rSocketPool) {
+  private LoadbalanceRSocketClient(RSocketPool rSocketPool) {
     this.rSocketPool = rSocketPool;
   }
 
@@ -66,14 +72,15 @@ public void dispose() {
     rSocketPool.dispose();
   }
 
-  public static LoadBalancedRSocketClient create(
+  public static LoadbalanceRSocketClient create(
       LoadbalanceStrategy loadbalanceStrategy,
-      Publisher<List<LoadbalanceTarget>> rSocketSuppliersPublisher) {
-    return new LoadBalancedRSocketClient(
+      Publisher<List<LoadbalanceRSocketSource>> rSocketSuppliersPublisher) {
+    return new LoadbalanceRSocketClient(
         new RSocketPool(rSocketSuppliersPublisher, loadbalanceStrategy));
   }
 
-  public static RSocketClient create(Publisher<List<LoadbalanceTarget>> rSocketSuppliersPublisher) {
+  public static RSocketClient create(
+      Publisher<List<LoadbalanceRSocketSource>> rSocketSuppliersPublisher) {
     return create(new RoundRobinLoadbalanceStrategy(), rSocketSuppliersPublisher);
   }
 }

rsocket-core/src/main/java/io/rsocket/loadbalance/LoadbalanceTarget.java renamed to rsocket-core/src/main/java/io/rsocket/loadbalance/LoadbalanceRSocketSource.java

@@ -19,12 +19,12 @@
 import io.rsocket.RSocket;
 import reactor.core.publisher.Mono;
 
-public class LoadbalanceTarget {
+public class LoadbalanceRSocketSource {
 
   final String serverKey;
   final Mono<RSocket> source;
 
-  public LoadbalanceTarget(String serverKey, Mono<RSocket> source) {
+  private LoadbalanceRSocketSource(String serverKey, Mono<RSocket> source) {
     this.serverKey = serverKey;
     this.source = source;
   }
@@ -38,7 +38,7 @@ public boolean equals(Object o) {
     if (this == o) return true;
     if (o == null || getClass() != o.getClass()) return false;
 
-    LoadbalanceTarget that = (LoadbalanceTarget) o;
+    LoadbalanceRSocketSource that = (LoadbalanceRSocketSource) o;
 
     return serverKey.equals(that.serverKey);
   }
@@ -47,4 +47,8 @@ public boolean equals(Object o) {
   public int hashCode() {
     return serverKey.hashCode();
   }
+
+  public static LoadbalanceRSocketSource from(String serverKey, Mono<RSocket> source) {
+    return new LoadbalanceRSocketSource(serverKey, source);
+  }
 }

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

@@ -15,7 +15,14 @@
  */
 package io.rsocket.loadbalance;
 
+import java.util.function.Supplier;
+
+@FunctionalInterface
 public interface LoadbalanceStrategy {
 
   PooledRSocket select(PooledRSocket[] availableRSockets);
+
+  default Supplier<Stats> statsSupplier() {
+    return Stats::noOps;
+  }
 }

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

@@ -21,9 +21,5 @@ public interface PooledRSocket extends RSocket {
 
   Stats stats();
 
-  LoadbalanceTarget supplier();
-
-  boolean markForRemoval();
-
-  boolean markActive();
+  LoadbalanceRSocketSource source();
 }