4.x: Add advanced shard awareness (#517)

* Add config options for advanced shard awareness

Adds config options for toggling the feature on and defining
port ranges to use.
By default the feature will be enabled.

* Copy 3.x PortAllocator

* Extend ChannelFactory methods with shardId params

With advanced shard awareness target shard is a necessary parameter
for connecting. Extends non public methods with new parameters and
provides one additional public method.
Passing `null` as shard id or ShardingInfo will lead to previous non-shard
aware behavior.

* Enable advanced shard awareness

Makes `addMissingChannels` use advanced shard awareness.
It will now specify target shard when adding missing channels for
specific shards.
In case returned channels do not match requested shards warnings are logged.
Initial connection to the node works on previous rules, meaning it uses
arbitrary local port for connection to arbitrary shard.

Adds AdvancedShardAwarenessIT that has several methods displaying
the difference between establishing connections with option enabled
and disabled.

Adds ChannelPoolShardAwareInitTest

Author: Bouncheck

core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java

@@ -141,6 +141,17 @@ public enum DefaultDriverOption implements DriverOption {
    * <p>Value-type: boolean
    */
   CONNECTION_WARN_INIT_ERROR("advanced.connection.warn-on-init-error"),
+  /**
+   * Whether to use advanced shard awareness.
+   *
+   * <p>Value-type: boolean
+   */
+  CONNECTION_ADVANCED_SHARD_AWARENESS_ENABLED(
+      "advanced.connection.advanced-shard-awareness.enabled"),
+  /** Inclusive lower bound of port range to use in advanced shard awareness */
+  ADVANCED_SHARD_AWARENESS_PORT_LOW("advanced.connection.advanced-shard-awareness.port-low"),
+  /** Inclusive upper bound of port range to use in advanced shard awareness */
+  ADVANCED_SHARD_AWARENESS_PORT_HIGH("advanced.connection.advanced-shard-awareness.port-high"),
   /**
    * The number of connections in the LOCAL pool.
    *

core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java

@@ -276,6 +276,9 @@ protected static void fillWithDriverDefaults(OptionsMap map) {
     map.put(TypedDriverOption.CONNECTION_MAX_REQUESTS, 1024);
     map.put(TypedDriverOption.CONNECTION_MAX_ORPHAN_REQUESTS, 256);
     map.put(TypedDriverOption.CONNECTION_WARN_INIT_ERROR, true);
+    map.put(TypedDriverOption.CONNECTION_ADVANCED_SHARD_AWARENESS_ENABLED, true);
+    map.put(TypedDriverOption.ADVANCED_SHARD_AWARENESS_PORT_LOW, 10000);
+    map.put(TypedDriverOption.ADVANCED_SHARD_AWARENESS_PORT_HIGH, 65535);
     map.put(TypedDriverOption.RECONNECT_ON_INIT, false);
     map.put(TypedDriverOption.RECONNECTION_POLICY_CLASS, "ExponentialReconnectionPolicy");
     map.put(TypedDriverOption.RECONNECTION_BASE_DELAY, Duration.ofSeconds(1));

core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java

@@ -175,6 +175,18 @@ public String toString() {
   /** Whether to log non-fatal errors when the driver tries to open a new connection. */
   public static final TypedDriverOption<Boolean> CONNECTION_WARN_INIT_ERROR =
       new TypedDriverOption<>(DefaultDriverOption.CONNECTION_WARN_INIT_ERROR, GenericType.BOOLEAN);
+  /** Whether to use advanced shard awareness */
+  public static final TypedDriverOption<Boolean> CONNECTION_ADVANCED_SHARD_AWARENESS_ENABLED =
+      new TypedDriverOption<>(
+          DefaultDriverOption.CONNECTION_ADVANCED_SHARD_AWARENESS_ENABLED, GenericType.BOOLEAN);
+  /** Inclusive lower bound of port range to use in advanced shard awareness */
+  public static final TypedDriverOption<Integer> ADVANCED_SHARD_AWARENESS_PORT_LOW =
+      new TypedDriverOption<>(
+          DefaultDriverOption.ADVANCED_SHARD_AWARENESS_PORT_LOW, GenericType.INTEGER);
+  /** Inclusive upper bound of port range to use in advanced shard awareness */
+  public static final TypedDriverOption<Integer> ADVANCED_SHARD_AWARENESS_PORT_HIGH =
+      new TypedDriverOption<>(
+          DefaultDriverOption.ADVANCED_SHARD_AWARENESS_PORT_HIGH, GenericType.INTEGER);
   /** The number of connections in the LOCAL pool. */
   public static final TypedDriverOption<Integer> CONNECTION_POOL_LOCAL_SIZE =
       new TypedDriverOption<>(DefaultDriverOption.CONNECTION_POOL_LOCAL_SIZE, GenericType.INTEGER);

core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelFactory.java

@@ -29,8 +29,10 @@
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.config.DriverConfig;
 import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
+import com.datastax.oss.driver.api.core.context.DriverContext;
 import com.datastax.oss.driver.api.core.metadata.EndPoint;
 import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.metadata.NodeShardingInfo;
 import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric;
 import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric;
 import com.datastax.oss.driver.internal.core.config.typesafe.TypesafeDriverConfig;
@@ -51,13 +53,17 @@
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelOption;
 import io.netty.channel.ChannelPipeline;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.ServerSocket;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import net.jcip.annotations.ThreadSafe;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -153,12 +159,27 @@ public CompletionStage<DriverChannel> connect(Node node, DriverChannelOptions op
     } else {
       nodeMetricUpdater = NoopNodeMetricUpdater.INSTANCE;
     }
-    return connect(node.getEndPoint(), options, nodeMetricUpdater);
+    return connect(node.getEndPoint(), null, null, options, nodeMetricUpdater);
+  }
+
+  public CompletionStage<DriverChannel> connect(
+      Node node, Integer shardId, DriverChannelOptions options) {
+    NodeMetricUpdater nodeMetricUpdater;
+    if (node instanceof DefaultNode) {
+      nodeMetricUpdater = ((DefaultNode) node).getMetricUpdater();
+    } else {
+      nodeMetricUpdater = NoopNodeMetricUpdater.INSTANCE;
+    }
+    return connect(node.getEndPoint(), node.getShardingInfo(), shardId, options, nodeMetricUpdater);
   }
 
   @VisibleForTesting
   CompletionStage<DriverChannel> connect(
-      EndPoint endPoint, DriverChannelOptions options, NodeMetricUpdater nodeMetricUpdater) {
+      EndPoint endPoint,
+      NodeShardingInfo shardingInfo,
+      Integer shardId,
+      DriverChannelOptions options,
+      NodeMetricUpdater nodeMetricUpdater) {
     CompletableFuture<DriverChannel> resultFuture = new CompletableFuture<>();
 
     ProtocolVersion currentVersion;
@@ -174,6 +195,8 @@ CompletionStage<DriverChannel> connect(
 
     connect(
         endPoint,
+        shardingInfo,
+        shardId,
         options,
         nodeMetricUpdater,
         currentVersion,
@@ -185,6 +208,8 @@ CompletionStage<DriverChannel> connect(
 
   private void connect(
       EndPoint endPoint,
+      NodeShardingInfo shardingInfo,
+      Integer shardId,
       DriverChannelOptions options,
       NodeMetricUpdater nodeMetricUpdater,
       ProtocolVersion currentVersion,
@@ -204,7 +229,28 @@ private void connect(
 
     nettyOptions.afterBootstrapInitialized(bootstrap);
 
-    ChannelFuture connectFuture = bootstrap.connect(endPoint.resolve());
+    ChannelFuture connectFuture;
+    if (shardId == null || shardingInfo == null) {
+      if (shardId != null) {
+        LOG.debug(
+            "Requested connection to shard {} but shardingInfo is currently missing for Node at endpoint {}. Falling back to arbitrary local port.",
+            shardId,
+            endPoint);
+      }
+      connectFuture = bootstrap.connect(endPoint.resolve());
+    } else {
+      int localPort =
+          PortAllocator.getNextAvailablePort(shardingInfo.getShardsCount(), shardId, context);
+      if (localPort == -1) {
+        LOG.warn(
+            "Could not find free port for shard {} at {}. Falling back to arbitrary local port.",
+            shardId,
+            endPoint);
+        connectFuture = bootstrap.connect(endPoint.resolve());
+      } else {
+        connectFuture = bootstrap.connect(endPoint.resolve(), new InetSocketAddress(localPort));
+      }
+    }
 
     connectFuture.addListener(
         cf -> {
@@ -253,6 +299,8 @@ private void connect(
                     downgraded.get());
                 connect(
                     endPoint,
+                    shardingInfo,
+                    shardId,
                     options,
                     nodeMetricUpdater,
                     downgraded.get(),
@@ -391,4 +439,92 @@ protected void initChannel(Channel channel) {
       }
     }
   }
+
+  static class PortAllocator {
+    private static final AtomicInteger lastPort = new AtomicInteger(-1);
+    private static final Logger LOG = LoggerFactory.getLogger(PortAllocator.class);
+
+    public static int getNextAvailablePort(int shardCount, int shardId, DriverContext context) {
+      int lowPort =
+          context
+              .getConfig()
+              .getDefaultProfile()
+              .getInt(DefaultDriverOption.ADVANCED_SHARD_AWARENESS_PORT_LOW);
+      int highPort =
+          context
+              .getConfig()
+              .getDefaultProfile()
+              .getInt(DefaultDriverOption.ADVANCED_SHARD_AWARENESS_PORT_HIGH);
+      if (highPort - lowPort < shardCount) {
+        LOG.error(
+            "There is not enough ports in range [{},{}] for {} shards. Update your configuration.",
+            lowPort,
+            highPort,
+            shardCount);
+      }
+      int lastPortValue, foundPort = -1;
+      do {
+        lastPortValue = lastPort.get();
+
+        // We will scan from lastPortValue
+        // (or lowPort is there was no lastPort or lastPort is too low)
+        int scanStart = lastPortValue == -1 ? lowPort : lastPortValue;
+        if (scanStart < lowPort) {
+          scanStart = lowPort;
+        }
+
+        // Round it up to "% shardCount == shardId"
+        scanStart += (shardCount - scanStart % shardCount) + shardId;
+
+        // Scan from scanStart upwards to highPort.
+        for (int port = scanStart; port <= highPort; port += shardCount) {
+          if (isTcpPortAvailable(port, context)) {
+            foundPort = port;
+            break;
+          }
+        }
+
+        // If we started scanning from a high scanStart port
+        // there might have been not enough ports left that are
+        // smaller than highPort. Scan from the beginning
+        // from the lowPort.
+        if (foundPort == -1) {
+          scanStart = lowPort + (shardCount - lowPort % shardCount) + shardId;
+
+          for (int port = scanStart; port <= highPort; port += shardCount) {
+            if (isTcpPortAvailable(port, context)) {
+              foundPort = port;
+              break;
+            }
+          }
+        }
+
+        // No luck! All ports taken!
+        if (foundPort == -1) {
+          return -1;
+        }
+      } while (!lastPort.compareAndSet(lastPortValue, foundPort));
+
+      return foundPort;
+    }
+
+    public static boolean isTcpPortAvailable(int port, DriverContext context) {
+      try {
+        ServerSocket serverSocket = new ServerSocket();
+        try {
+          serverSocket.setReuseAddress(
+              context
+                  .getConfig()
+                  .getDefaultProfile()
+                  .getBoolean(DefaultDriverOption.SOCKET_REUSE_ADDRESS, false));
+          serverSocket.bind(new InetSocketAddress(port), 1);
+          return true;
+        } finally {
+          serverSocket.close();
+        }
+      } catch (IOException ex) {
+        return false;
+      }
+    }
+  }
 }

core/src/main/java/com/datastax/oss/driver/internal/core/pool/ChannelPool.java

@@ -58,6 +58,7 @@
 import io.netty.util.concurrent.EventExecutor;
 import io.netty.util.concurrent.Future;
 import io.netty.util.concurrent.GenericFutureListener;
+import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
@@ -489,9 +490,23 @@ private CompletionStage<Boolean> addMissingChannels() {
           channels.length * wantedCount - Arrays.stream(channels).mapToInt(ChannelSet::size).sum();
       LOG.debug("[{}] Trying to create {} missing channels", logPrefix, missing);
       DriverChannelOptions options = buildDriverOptions();
-      for (int i = 0; i < missing; i++) {
-        CompletionStage<DriverChannel> channelFuture = channelFactory.connect(node, options);
-        pendingChannels.add(channelFuture);
+      for (int shard = 0; shard < channels.length; shard++) {
+        LOG.trace(
+            "[{}] Missing {} channels for shard {}",
+            logPrefix,
+            wantedCount - channels[shard].size(),
+            shard);
+        for (int p = channels[shard].size(); p < wantedCount; p++) {
+          CompletionStage<DriverChannel> channelFuture;
+          if (config
+              .getDefaultProfile()
+              .getBoolean(DefaultDriverOption.CONNECTION_ADVANCED_SHARD_AWARENESS_ENABLED)) {
+            channelFuture = channelFactory.connect(node, shard, options);
+          } else {
+            channelFuture = channelFactory.connect(node, options);
+          }
+          pendingChannels.add(channelFuture);
+        }
       }
       return CompletableFutures.allDone(pendingChannels)
           .thenApplyAsync(this::onAllConnected, adminExecutor);
@@ -551,6 +566,23 @@ private boolean onAllConnected(@SuppressWarnings("unused") Void v) {
                 channel);
             channel.forceClose();
           } else {
+            if (config
+                    .getDefaultProfile()
+                    .getBoolean(DefaultDriverOption.CONNECTION_ADVANCED_SHARD_AWARENESS_ENABLED)
+                && channel.localAddress() instanceof InetSocketAddress
+                && channel.getShardingInfo() != null) {
+              int port = ((InetSocketAddress) channel.localAddress()).getPort();
+              int actualShard = channel.getShardId();
+              int targetShard = port % channel.getShardingInfo().getShardsCount();
+              if (actualShard != targetShard) {
+                LOG.warn(
+                    "[{}] New channel {} connected to shard {}, but shard {} was requested. If this is not transient check your driver AND cluster configuration of shard aware port.",
+                    logPrefix,
+                    channel,
+                    actualShard,
+                    targetShard);
+              }
+            }
             LOG.debug("[{}] New channel added {}", logPrefix, channel);
             if (channels[channel.getShardId()].size() < wantedCount) {
               addChannel(channel);

core/src/main/resources/reference.conf

@@ -535,6 +535,52 @@ datastax-java-driver {
     #   change.
     # Overridable in a profile: no
     warn-on-init-error = true
+
+
+    advanced-shard-awareness {
+      # Whether to use advanced shard awareness when trying to open new connections.
+      #
+      # Requires passing shard-aware port as contact point (usually 19042 or 19142(ssl)).
+      # Having this enabled makes sense only for ScyllaDB clusters.
+      #
+      # Reduces number of reconnections driver needs to fully initialize connection pool.
+      # In short it's a feature that allows targeting particular shard when connecting to a node by using specific
+      # local port number.
+      # For context see https://www.scylladb.com/2021/04/27/connect-faster-to-scylla-with-a-shard-aware-port/
+      #
+      # If set to false the driver will not attempt to use this feature. This means connection's local port
+      # will be random according to system rules and driver will keep opening connections until it gets right shards.
+      # In such case non-shard aware port is recommended (by default 9042 or 9142).
+      # If set to true the driver will attempt to use it and will log warnings each time something
+      # makes it not possible.
+      #
+      # If the node for some reason does not report it's sharding info the driver
+      # will log a warning and create connection the same way as if this feature was disabled.
+      # If the cluster ignores the request for specific shard warning will also be logged,
+      # although the local port will already be chosen according to advanced shard awareness rules.
+      #
+      # Required: yes
+      # Modifiable at runtime: yes, the new value will be used for connections created after the
+      # change.
+      # Overridable in a profile: no
+      enabled = true
+
+      # Inclusive lower bound of port range to use in advanced shard awareness
+      # The driver will attempt to reserve ports for connection only within the range.
+      # Required: yes
+      # Modifiable at runtime: yes, the new value will be used for calls after the
+      # change.
+      # Overridable in a profile: no
+      port-low = 10000
+
+      # Inclusive upper bound of port range to use in advanced shard awareness.
+      # The driver will attempt to reserve ports for connection only within the range.
+      # Required: yes
+      # Modifiable at runtime: yes, the new value will be used for calls after the
+      # change.
+      # Overridable in a profile: no
+      port-high = 65535
+    }
   }
 
   # Advanced options for the built-in load-balancing policies.

core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryAvailableIdsTest.java

@@ -61,7 +61,11 @@ public void should_report_available_ids() {
     // When
     CompletionStage<DriverChannel> channelFuture =
         factory.connect(
-            SERVER_ADDRESS, DriverChannelOptions.builder().build(), NoopNodeMetricUpdater.INSTANCE);
+            SERVER_ADDRESS,
+            null,
+            null,
+            DriverChannelOptions.builder().build(),
+            NoopNodeMetricUpdater.INSTANCE);
     completeSimpleChannelInit();
 
     // Then