1

Author: dkropachev

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

@@ -29,7 +29,6 @@
 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;
@@ -53,17 +52,15 @@
 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.Iterator;
 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;
@@ -193,28 +190,44 @@ CompletionStage<DriverChannel> connect(
       isNegotiating = true;
     }
 
+    PortIterator portIterator = null;
+
+    if (shardId != null && shardingInfo != null) {
+      int lowestPort =
+          context
+              .getConfig()
+              .getDefaultProfile()
+              .getInt(DefaultDriverOption.ADVANCED_SHARD_AWARENESS_PORT_LOW);
+      int highestPort =
+          context
+              .getConfig()
+              .getDefaultProfile()
+              .getInt(DefaultDriverOption.ADVANCED_SHARD_AWARENESS_PORT_HIGH);
+      portIterator =
+          new PortIterator(lowestPort, highestPort, shardingInfo.getShardsCount(), shardId)
+              .iterator();
+    }
+
     connect(
         endPoint,
-        shardingInfo,
-        shardId,
         options,
         nodeMetricUpdater,
         currentVersion,
         isNegotiating,
         attemptedVersions,
+        portIterator,
         resultFuture);
     return resultFuture;
   }
 
   private void connect(
       EndPoint endPoint,
-      NodeShardingInfo shardingInfo,
-      Integer shardId,
       DriverChannelOptions options,
       NodeMetricUpdater nodeMetricUpdater,
       ProtocolVersion currentVersion,
       boolean isNegotiating,
       List<ProtocolVersion> attemptedVersions,
+      PortIterator portIterator,
       CompletableFuture<DriverChannel> resultFuture) {
 
     NettyOptions nettyOptions = context.getNettyOptions();
@@ -230,26 +243,14 @@ private void connect(
     nettyOptions.afterBootstrapInitialized(bootstrap);
 
     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);
-      }
+    if (portIterator == null) {
       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));
+      if (!portIterator.hasNext()) {
+        portIterator.reset();
       }
+      connectFuture =
+          bootstrap.connect(endPoint.resolve(), new InetSocketAddress(portIterator.next()));
     }
 
     connectFuture.addListener(
@@ -299,19 +300,29 @@ private void connect(
                     downgraded.get());
                 connect(
                     endPoint,
-                    shardingInfo,
-                    shardId,
                     options,
                     nodeMetricUpdater,
                     downgraded.get(),
                     true,
                     attemptedVersions,
+                    portIterator,
                     resultFuture);
               } else {
                 resultFuture.completeExceptionally(
                     UnsupportedProtocolVersionException.forNegotiation(
                         endPoint, attemptedVersions));
               }
+            } else if (error instanceof java.net.BindException
+                || error.getCause() instanceof java.net.BindException) {
+              connect(
+                  endPoint,
+                  options,
+                  nodeMetricUpdater,
+                  currentVersion,
+                  true,
+                  attemptedVersions,
+                  portIterator,
+                  resultFuture);
             } else {
               // Note: might be completed already if the failure happened in initializer(), this is
               // fine
@@ -330,7 +341,7 @@ ChannelInitializer<Channel> initializer(
       CompletableFuture<DriverChannel> resultFuture) {
     return new ChannelFactoryInitializer(
         endPoint, protocolVersion, options, nodeMetricUpdater, resultFuture);
-  };
+  }
 
   class ChannelFactoryInitializer extends ChannelInitializer<Channel> {
 
@@ -440,87 +451,40 @@ 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);
+  static class PortIterator implements Iterator<Integer> {
+    private final int highestPort;
+    private final int startPort;
+    private int currentPort;
+    private final int shardCount;
+
+    PortIterator(int lowestPort, int highestPort, int shardCount, int shardId) {
+      this.highestPort = highestPort;
+      this.currentPort = lowestPort + (shardCount - lowestPort % shardCount) + shardId;
+      this.startPort = currentPort;
+      this.shardCount = shardCount;
+    }
 
-    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);
+    @Override
+    public Integer next() {
+      try {
+        return currentPort;
+      } finally {
+        this.currentPort += 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)) {
-            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)) {
-              foundPort = port;
-              break;
-            }
-          }
-        }
+    }
 
-        // No luck! All ports taken!
-        if (foundPort == -1) {
-          return -1;
-        }
-      } while (!lastPort.compareAndSet(lastPortValue, foundPort));
+    @Override
+    public boolean hasNext() {
+      return currentPort <= highestPort;
+    }
 
-      return foundPort;
+    public void reset() {
+      this.currentPort = startPort;
     }
 
-    public static boolean isTcpPortAvailable(int port) {
-      try {
-        ServerSocket serverSocket = new ServerSocket();
-        try {
-          serverSocket.setReuseAddress(false);
-          serverSocket.bind(new InetSocketAddress(port), 1);
-          return true;
-        } finally {
-          serverSocket.close();
-        }
-      } catch (IOException ex) {
-        return false;
-      }
+    public PortIterator iterator() {
+      this.reset();
+      return this;
     }
   }
 }