Merge branch 'cassandra-2.1' into trunk

CHANGES.txt

@@ -93,6 +93,7 @@
    marked are in the live set (CASSANDRA-8689)
  * cassandra-stress support for varint (CASSANDRA-8882)
 Merged from 2.0:
+ * Add ability to limit number of native connections (CASSANDRA-8086)
  * Add offline tool to relevel sstables (CASSANDRA-8301)
  * Preserve stream ID for more protocol errors (CASSANDRA-8848)
  * Fix combining token() function with multi-column relations on

conf/cassandra.yaml

@@ -436,6 +436,14 @@ native_transport_port: 9042
 # be rejected as invalid. The default is 256MB.
 # native_transport_max_frame_size_in_mb: 256
 
+# The maximum number of concurrent client connections.
+# The default is -1, which means unlimited.
+# native_transport_max_concurrent_connections: -1
+
+# The maximum number of concurrent client connections per source ip.
+# The default is -1, which means unlimited.
+# native_transport_max_concurrent_connections_per_ip: -1
+
 # Whether to start the thrift rpc server.
 start_rpc: true
 

src/java/org/apache/cassandra/config/Config.java

@@ -126,6 +126,8 @@ public class Config
     public Integer native_transport_port = 9042;
     public Integer native_transport_max_threads = 128;
     public Integer native_transport_max_frame_size_in_mb = 256;
+    public volatile Long native_transport_max_concurrent_connections = -1L;
+    public volatile Long native_transport_max_concurrent_connections_per_ip = -1L;
 
     @Deprecated
     public Integer thrift_max_message_length_in_mb = 16;

src/java/org/apache/cassandra/config/DatabaseDescriptor.java

@@ -1173,6 +1173,25 @@ public class DatabaseDescriptor
         return conf.native_transport_max_frame_size_in_mb * 1024 * 1024;
     }
 
+    public static Long getNativeTransportMaxConcurrentConnections()
+    {
+        return conf.native_transport_max_concurrent_connections;
+    }
+
+    public static void setNativeTransportMaxConcurrentConnections(long nativeTransportMaxConcurrentConnections)
+    {
+        conf.native_transport_max_concurrent_connections = nativeTransportMaxConcurrentConnections;
+    }
+
+    public static Long getNativeTransportMaxConcurrentConnectionsPerIp() {
+        return conf.native_transport_max_concurrent_connections_per_ip;
+    }
+
+    public static void setNativeTransportMaxConcurrentConnectionsPerIp(long native_transport_max_concurrent_connections_per_ip)
+    {
+        conf.native_transport_max_concurrent_connections_per_ip = native_transport_max_concurrent_connections_per_ip;
+    }
+
     public static double getCommitLogSyncBatchWindow()
     {
         return conf.commitlog_sync_batch_window_in_ms;

src/java/org/apache/cassandra/service/StorageProxy.java

@@ -2249,9 +2249,15 @@ public class StorageProxy implements StorageProxyMBean
 
     public Long getTruncateRpcTimeout() { return DatabaseDescriptor.getTruncateRpcTimeout(); }
     public void setTruncateRpcTimeout(Long timeoutInMillis) { DatabaseDescriptor.setTruncateRpcTimeout(timeoutInMillis); }
+
+    public Long getNativeTransportMaxConcurrentConnections() { return DatabaseDescriptor.getNativeTransportMaxConcurrentConnections(); }
+    public void setNativeTransportMaxConcurrentConnections(Long nativeTransportMaxConcurrentConnections) { DatabaseDescriptor.setNativeTransportMaxConcurrentConnections(nativeTransportMaxConcurrentConnections); }
+
+    public Long getNativeTransportMaxConcurrentConnectionsPerIp() { return DatabaseDescriptor.getNativeTransportMaxConcurrentConnectionsPerIp(); }
+    public void setNativeTransportMaxConcurrentConnectionsPerIp(Long nativeTransportMaxConcurrentConnections) { DatabaseDescriptor.setNativeTransportMaxConcurrentConnectionsPerIp(nativeTransportMaxConcurrentConnections); }
+
     public void reloadTriggerClasses() { TriggerExecutor.instance.reloadClasses(); }
 
-    
     public long getReadRepairAttempted() {
         return ReadRepairMetrics.attempted.getCount();
     }

src/java/org/apache/cassandra/service/StorageProxyMBean.java

@@ -49,6 +49,9 @@ public interface StorageProxyMBean
     public Long getTruncateRpcTimeout();
     public void setTruncateRpcTimeout(Long timeoutInMillis);
 
+    public void setNativeTransportMaxConcurrentConnections(Long nativeTransportMaxConcurrentConnections);
+    public Long getNativeTransportMaxConcurrentConnections();
+
     public void reloadTriggerClasses();
 
     public long getReadRepairAttempted();

src/java/org/apache/cassandra/transport/ConnectionLimitHandler.java

@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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
+ *
+ *     http://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 org.apache.cassandra.transport;
+
+
+import io.netty.channel.ChannelHandler;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+import org.apache.cassandra.config.DatabaseDescriptor;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicLong;
+
+
+/**
+ * {@link ChannelInboundHandlerAdapter} implementation which allows to limit the number of concurrent
+ * connections to the Server. Be aware this <strong>MUST</strong> be shared between all child channels.
+ */
+@ChannelHandler.Sharable
+final class ConnectionLimitHandler extends ChannelInboundHandlerAdapter
+{
+    private static final Logger logger = LoggerFactory.getLogger(ConnectionLimitHandler.class);
+    private final ConcurrentMap<InetAddress, AtomicLong> connectionsPerClient = new ConcurrentHashMap<>();
+    private final AtomicLong counter = new AtomicLong(0);
+
+    @Override
+    public void channelActive(ChannelHandlerContext ctx) throws Exception
+    {
+        final long count = counter.incrementAndGet();
+        long limit = DatabaseDescriptor.getNativeTransportMaxConcurrentConnections();
+        // Setting the limit to -1 disables it.
+        if(limit < 0)
+        {
+            limit = Long.MAX_VALUE;
+        }
+        if (count > limit)
+        {
+            // The decrement will be done in channelClosed(...)
+            logger.warn("Exceeded maximum native connection limit of {} by using {} connections", limit, count);
+            ctx.close();
+        }
+        else
+        {
+            long perIpLimit = DatabaseDescriptor.getNativeTransportMaxConcurrentConnectionsPerIp();
+            if (perIpLimit > 0)
+            {
+                InetAddress address = ((InetSocketAddress) ctx.channel().remoteAddress()).getAddress();
+
+                AtomicLong perIpCount = connectionsPerClient.get(address);
+                if (perIpCount == null)
+                {
+                    perIpCount = new AtomicLong(0);
+
+                    AtomicLong old = connectionsPerClient.putIfAbsent(address, perIpCount);
+                    if (old != null)
+                    {
+                        perIpCount = old;
+                    }
+                }
+                if (perIpCount.incrementAndGet() > perIpLimit)
+                {
+                    // The decrement will be done in channelClosed(...)
+                    logger.warn("Exceeded maximum native connection limit per ip of {} by using {} connections", perIpLimit, perIpCount);
+                    ctx.close();
+                    return;
+                }
+            }
+            ctx.fireChannelActive();
+        }
+    }
+
+    @Override
+    public void channelInactive(ChannelHandlerContext ctx) throws Exception
+    {
+        counter.decrementAndGet();
+        InetAddress address = ((InetSocketAddress) ctx.channel().remoteAddress()).getAddress();
+
+        AtomicLong count = connectionsPerClient.get(address);
+        if (count != null)
+        {
+            if (count.decrementAndGet() <= 0)
+            {
+                connectionsPerClient.remove(address);
+            }
+        }
+        ctx.fireChannelInactive();
+    }
+}

src/java/org/apache/cassandra/transport/Server.java

@@ -262,6 +262,7 @@ public class Server implements CassandraDaemon.Server
         private static final Frame.Compressor frameCompressor = new Frame.Compressor();
         private static final Frame.Encoder frameEncoder = new Frame.Encoder();
         private static final Message.Dispatcher dispatcher = new Message.Dispatcher();
+        private static final ConnectionLimitHandler connectionLimitHandler = new ConnectionLimitHandler();
 
         private final Server server;
 
@@ -274,6 +275,14 @@ public class Server implements CassandraDaemon.Server
         {
             ChannelPipeline pipeline = channel.pipeline();
 
+            // Add the ConnectionLimitHandler to the pipeline if configured to do so.
+            if (DatabaseDescriptor.getNativeTransportMaxConcurrentConnections() > 0
+                    || DatabaseDescriptor.getNativeTransportMaxConcurrentConnectionsPerIp() > 0)
+            {
+                // Add as first to the pipeline so the limit is enforced as first action.
+                pipeline.addFirst("connectionLimitHandler", connectionLimitHandler);
+            }
+
             //pipeline.addLast("debug", new LoggingHandler());
 
             pipeline.addLast("frameDecoder", new Frame.Decoder(server.connectionFactory));