ZOOKEEPER-3146: Limit the maximum client connections per IP in NettyServerCnxnFactory

Add similar max cnxn per ip throttling logic to Netty implementation.

Author: Fangmin Lyu <allenlyu@fb.com>

Reviewers: hanm@apache.org

Closes #623 from lvfangmin/ZOOKEEPER-3146

src/java/main/org/apache/zookeeper/server/NettyServerCnxnFactory.java

@@ -105,7 +105,17 @@ public class NettyServerCnxnFactory extends ServerCnxnFactory {
                 LOG.trace("Channel connected " + e);
             }
 
-            NettyServerCnxn cnxn = new NettyServerCnxn(ctx.getChannel(),
+            Channel channel = ctx.getChannel();
+            InetAddress addr = ((InetSocketAddress) channel.getRemoteAddress())
+                    .getAddress();
+            if (maxClientCnxns > 0 && getClientCnxnCount(addr) >= maxClientCnxns) {
+                LOG.warn("Too many connections from {} - max is {}", addr,
+                        maxClientCnxns);
+                channel.close();
+                return;
+            }
+
+            NettyServerCnxn cnxn = new NettyServerCnxn(channel,
                     zkServer, NettyServerCnxnFactory.this);
             ctx.setAttachment(cnxn);
 
@@ -537,6 +547,12 @@ public class NettyServerCnxnFactory extends ServerCnxnFactory {
         }
     }
 
+    private int getClientCnxnCount(InetAddress addr) {
+        Set<NettyServerCnxn> s = ipMap.get(addr);
+        if (s == null) return 0;
+        return s.size();
+    }
+
     @Override
     public void resetAllConnectionStats() {
         // No need to synchronize since cnxns is backed by a ConcurrentHashMap

src/java/test/org/apache/zookeeper/test/SessionTest.java

@@ -390,4 +390,23 @@ public class SessionTest extends ZKTestCase {
         zk.close();
         LOG.info(zk.toString());
     }
+
+    @Test
+    public void testMaximumCnxnPerIP() throws Exception {
+        final int maxClientCnxnsPerIP = 3;
+        serverFactory.setMaxClientCnxnsPerHost(maxClientCnxnsPerIP);
+        ZooKeeper[] clients = new ZooKeeper[maxClientCnxnsPerIP + 1];
+        for (int i = 0; i < clients.length; i++) {
+            CountdownWatcher watcher = new CountdownWatcher();
+            // wait for 3s
+            int timeout = 3000;
+            clients[i] = new DisconnectableZooKeeper(HOSTPORT, timeout, watcher);
+            boolean result = watcher.clientConnected.await(timeout, TimeUnit.MILLISECONDS);
+            if (i >= maxClientCnxnsPerIP) {
+                Assert.assertFalse(result);
+            } else {
+                Assert.assertTrue(result);
+            }
+        }
+    }
 }