Refactors the input for LoadbalanceRSocketClient (#924)

Signed-off-by: Rossen Stoyanchev <rstoyanchev@vmware.com>

Author: rstoyanchev

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

@@ -18,14 +18,16 @@
 import io.rsocket.Payload;
 import io.rsocket.RSocket;
 import io.rsocket.RSocketClient;
+import io.rsocket.core.RSocketConnector;
 import java.util.List;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.annotation.Nullable;
 
 /**
- * {@link RSocketClient} implementation that uses a pool and a {@link LoadbalanceStrategy} to select
- * the {@code RSocket} to use for each request.
+ * {@link RSocketClient} implementation that uses a {@link LoadbalanceStrategy} to select the {@code
+ * RSocket} to use for a given request from a pool of possible targets.
  *
  * @since 1.1
  */
@@ -72,45 +74,109 @@ public void dispose() {
     rSocketPool.dispose();
   }
 
+  /**
+   * Shortcut to create an {@link LoadbalanceRSocketClient} with round robin loadalancing.
+   * Effectively a shortcut for:
+   *
+   * <pre class="cdoe">
+   * LoadbalanceRSocketClient.builder(targetPublisher)
+   *    .connector(RSocketConnector.create())
+   *    .build();
+   * </pre>
+   *
+   * @param connector the {@link Builder#connector(RSocketConnector) to use
+   * @param targetPublisher publisher that periodically refreshes the list of targets to loadbalance across.
+   * @return the created client instance
+   */
   public static LoadbalanceRSocketClient create(
-      LoadbalanceStrategy loadbalanceStrategy,
-      Publisher<List<LoadbalanceRSocketSource>> rSocketSuppliersPublisher) {
-    return new LoadbalanceRSocketClient(
-        new RSocketPool(rSocketSuppliersPublisher, loadbalanceStrategy));
+      RSocketConnector connector, Publisher<List<LoadbalanceTarget>> targetPublisher) {
+    return builder(targetPublisher).connector(connector).build();
   }
 
-  public static LoadbalanceRSocketClient create(
-      Publisher<List<LoadbalanceRSocketSource>> rSocketSuppliersPublisher) {
-    return create(new RoundRobinLoadbalanceStrategy(), rSocketSuppliersPublisher);
-  }
-
-  public static Builder builder() {
-    return new Builder();
+  /**
+   * Return a builder to create an {@link LoadbalanceRSocketClient} with.
+   *
+   * @param targetPublisher publisher that periodically refreshes the list of targets to loadbalance
+   *     across.
+   * @return the builder instance
+   */
+  public static Builder builder(Publisher<List<LoadbalanceTarget>> targetPublisher) {
+    return new Builder(targetPublisher);
   }
 
+  /** Builder for creating an {@link LoadbalanceRSocketClient}. */
   public static class Builder {
 
-    LoadbalanceStrategy loadbalanceStrategy;
+    private final Publisher<List<LoadbalanceTarget>> targetPublisher;
+
+    @Nullable private RSocketConnector connector;
 
-    Builder() {}
+    @Nullable LoadbalanceStrategy loadbalanceStrategy;
 
-    public Builder withWeightedLoadbalanceStrategy() {
-      return withCustomLoadbalanceStrategy(new WeightedLoadbalanceStrategy());
+    Builder(Publisher<List<LoadbalanceTarget>> targetPublisher) {
+      this.targetPublisher = targetPublisher;
     }
 
-    public Builder withRoundRobinLoadbalanceStrategy() {
-      return withCustomLoadbalanceStrategy(new RoundRobinLoadbalanceStrategy());
+    /**
+     * The given {@link RSocketConnector} is used as a template to produce the {@code Mono<RSocket>}
+     * source for each {@link LoadbalanceTarget}. This is done by passing the {@code
+     * ClientTransport} contained in every target to the {@code connect} method of the given
+     * connector instance.
+     *
+     * <p>By default this is initialized with {@link RSocketConnector#create()}.
+     *
+     * @param connector the connector to use as a template
+     */
+    public Builder connector(RSocketConnector connector) {
+      this.connector = connector;
+      return this;
+    }
+
+    /**
+     * Switch to using a round-robin strategy for selecting a target.
+     *
+     * <p>This is the strategy used by default.
+     */
+    public Builder roundRobinLoadbalanceStrategy() {
+      this.loadbalanceStrategy = new RoundRobinLoadbalanceStrategy();
+      return this;
     }
 
-    public Builder withCustomLoadbalanceStrategy(LoadbalanceStrategy strategy) {
+    /**
+     * Switch to using a strategy that assigns a weight to each pooled {@code RSocket} based on
+     * actual usage stats, and uses that to make a choice.
+     *
+     * <p>By default this strategy is not used.
+     */
+    public Builder weightedLoadbalanceStrategy() {
+      this.loadbalanceStrategy = new WeightedLoadbalanceStrategy();
+      return this;
+    }
+
+    /**
+     * Switch to using a custom strategy for loadbalancing.
+     *
+     * @see #roundRobinLoadbalanceStrategy()
+     */
+    public Builder customLoadbalanceStrategy(LoadbalanceStrategy strategy) {
       this.loadbalanceStrategy = strategy;
       return this;
     }
 
-    public LoadbalanceRSocketClient build(
-        Publisher<List<LoadbalanceRSocketSource>> rSocketSuppliersPublisher) {
+    /** Build the {@link LoadbalanceRSocketClient} instance. */
+    public LoadbalanceRSocketClient build() {
       return new LoadbalanceRSocketClient(
-          new RSocketPool(rSocketSuppliersPublisher, this.loadbalanceStrategy));
+          new RSocketPool(initConnector(), this.targetPublisher, initLoadbalanceStrategy()));
+    }
+
+    private RSocketConnector initConnector() {
+      return (this.connector != null ? this.connector : RSocketConnector.create());
+    }
+
+    private LoadbalanceStrategy initLoadbalanceStrategy() {
+      return (this.loadbalanceStrategy != null
+          ? this.loadbalanceStrategy
+          : new RoundRobinLoadbalanceStrategy());
     }
   }
 }

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

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2002-2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.rsocket.loadbalance;
+
+import io.rsocket.transport.ClientTransport;
+
+/**
+ * Simple container for a key and a {@link ClientTransport}, representing a specific target for
+ * loadbalancing purposes. The key is used to compare previous and new targets when refreshing the
+ * list of target to use. The transport is used to connect to the target.
+ *
+ * @since 1.1
+ */
+public class LoadbalanceTarget {
+
+  final String key;
+  final ClientTransport transport;
+
+  private LoadbalanceTarget(String key, ClientTransport transport) {
+    this.key = key;
+    this.transport = transport;
+  }
+
+  /** Return the key for this target. */
+  public String getKey() {
+    return key;
+  }
+
+  /** Return the transport to use to connect to the target. */
+  public ClientTransport getTransport() {
+    return transport;
+  }
+
+  /**
+   * Create a an instance of {@link LoadbalanceTarget} with the given key and {@link
+   * ClientTransport}. The key can be anything that can be used to identify identical targets, e.g.
+   * a SocketAddress, URL, etc.
+   *
+   * @param key the key to use to identify identical targets
+   * @param transport the transport to use for connecting to the target
+   * @return the created instance
+   */
+  public static LoadbalanceTarget from(String key, ClientTransport transport) {
+    return new LoadbalanceTarget(key, transport);
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (this == other) {
+      return true;
+    }
+    if (other == null || getClass() != other.getClass()) {
+      return false;
+    }
+    LoadbalanceTarget that = (LoadbalanceTarget) other;
+    return key.equals(that.key);
+  }
+
+  @Override
+  public int hashCode() {
+    return key.hashCode();
+  }
+}

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

@@ -33,7 +33,8 @@ final class PooledWeightedRSocket extends ResolvingOperator<RSocket>
     implements CoreSubscriber<RSocket>, WeightedRSocket {
 
   final RSocketPool parent;
-  final LoadbalanceRSocketSource loadbalanceRSocketSource;
+  final Mono<RSocket> rSocketSource;
+  final LoadbalanceTarget loadbalanceTarget;
   final Stats stats;
 
   volatile Subscription s;
@@ -42,10 +43,14 @@ final class PooledWeightedRSocket extends ResolvingOperator<RSocket>
       AtomicReferenceFieldUpdater.newUpdater(PooledWeightedRSocket.class, Subscription.class, "s");
 
   PooledWeightedRSocket(
-      RSocketPool parent, LoadbalanceRSocketSource loadbalanceRSocketSource, Stats stats) {
+      RSocketPool parent,
+      Mono<RSocket> rSocketSource,
+      LoadbalanceTarget loadbalanceTarget,
+      Stats stats) {
     this.parent = parent;
+    this.rSocketSource = rSocketSource;
+    this.loadbalanceTarget = loadbalanceTarget;
     this.stats = stats;
-    this.loadbalanceRSocketSource = loadbalanceRSocketSource;
   }
 
   @Override
@@ -103,7 +108,7 @@ public void onNext(RSocket value) {
 
   @Override
   protected void doSubscribe() {
-    this.loadbalanceRSocketSource.source().subscribe(this);
+    this.rSocketSource.subscribe(this);
   }
 
   @Override
@@ -196,8 +201,8 @@ public Stats stats() {
     return stats;
   }
 
-  LoadbalanceRSocketSource source() {
-    return loadbalanceRSocketSource;
+  LoadbalanceTarget target() {
+    return loadbalanceTarget;
   }
 
   @Override