[SCB-1023] limit max connection by vertx metrics

Author: wujimin

foundations/foundation-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/server/TcpServer.java

@@ -26,12 +26,15 @@
 import org.apache.servicecomb.foundation.ssl.SSLOptionFactory;
 import org.apache.servicecomb.foundation.vertx.AsyncResultCallback;
 import org.apache.servicecomb.foundation.vertx.VertxTLSBuilder;
+import org.apache.servicecomb.foundation.vertx.metrics.DefaultTcpServerMetrics;
+import org.apache.servicecomb.foundation.vertx.metrics.metric.DefaultServerEndpointMetric;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import io.vertx.core.Vertx;
 import io.vertx.core.net.NetServer;
 import io.vertx.core.net.NetServerOptions;
+import io.vertx.core.net.impl.NetSocketImpl;
 
 public class TcpServer {
   private static final Logger LOGGER = LoggerFactory.getLogger(TcpServer.class);
@@ -62,6 +65,16 @@ public void init(Vertx vertx, String sslKey, AsyncResultCallback<InetSocketAddre
     }
 
     netServer.connectHandler(netSocket -> {
+      DefaultTcpServerMetrics serverMetrics = (DefaultTcpServerMetrics) ((NetSocketImpl) netSocket).metrics();
+      DefaultServerEndpointMetric endpointMetric = serverMetrics.getEndpointMetric();
+      long connectedCount = endpointMetric.getCurrentConnectionCount();
+      int connectionLimit = getConnectionLimit();
+      if (connectedCount > connectionLimit) {
+        netSocket.close();
+        endpointMetric.onRejectByConnectionLimit();
+        return;
+      }
+
       TcpServerConnection connection = createTcpServerConnection();
       connection.init(netSocket);
     });
@@ -81,6 +94,10 @@ public void init(Vertx vertx, String sslKey, AsyncResultCallback<InetSocketAddre
     });
   }
 
+  protected int getConnectionLimit() {
+    return Integer.MAX_VALUE;
+  }
+
   protected TcpServerConnection createTcpServerConnection() {
     return new TcpServerConnection();
   }

foundations/foundation-vertx/src/test/java/org/apache/servicecomb/foundation/vertx/server/TestTcpServer.java

@@ -21,14 +21,20 @@
 
 import org.apache.servicecomb.foundation.common.net.URIEndpointObject;
 import org.apache.servicecomb.foundation.vertx.AsyncResultCallback;
+import org.apache.servicecomb.foundation.vertx.metrics.DefaultTcpServerMetrics;
+import org.apache.servicecomb.foundation.vertx.metrics.metric.DefaultServerEndpointMetric;
+import org.junit.Assert;
 import org.junit.Test;
 
 import io.vertx.core.Handler;
 import io.vertx.core.Vertx;
 import io.vertx.core.net.NetServer;
 import io.vertx.core.net.NetServerOptions;
 import io.vertx.core.net.NetSocket;
+import io.vertx.core.net.impl.NetSocketImpl;
 import mockit.Expectations;
+import mockit.Mock;
+import mockit.MockUp;
 import mockit.Mocked;
 
 public class TestTcpServer {
@@ -83,4 +89,59 @@ public void testTcpServerSSL(@Mocked Vertx vertx, @Mocked AsyncResultCallback<In
     // assert done in Expectations
     server.init(vertx, "", callback);
   }
+
+  Handler<NetSocket> connectHandler;
+
+  boolean netSocketClosed;
+
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  @Test
+  public void testConnectionLimit(@Mocked Vertx vertx, @Mocked AsyncResultCallback<InetSocketAddress> callback,
+      @Mocked NetServer netServer, @Mocked NetSocketImpl netSocket) {
+    DefaultServerEndpointMetric endpointMetric = new DefaultServerEndpointMetric(null);
+    DefaultTcpServerMetrics tcpServerMetrics = new DefaultTcpServerMetrics(endpointMetric);
+
+    new MockUp<NetServer>(netServer) {
+      @Mock
+      NetServer connectHandler(Handler<NetSocket> handler) {
+        connectHandler = handler;
+        return netServer;
+      }
+    };
+    new MockUp<NetSocketImpl>(netSocket) {
+      @Mock
+      void close() {
+        netSocketClosed = true;
+      }
+    };
+    new Expectations() {
+      {
+        vertx.createNetServer((NetServerOptions) any);
+        result = netServer;
+        netServer.listen(anyInt, anyString, (Handler) any);
+        netSocket.metrics();
+        result = tcpServerMetrics;
+      }
+    };
+    URIEndpointObject endpointObject = new URIEndpointObject("highway://127.0.0.1:6663?sslEnabled=true");
+    TcpServer server = new TcpServerForTest(endpointObject) {
+      @Override
+      protected int getConnectionLimit() {
+        return 2;
+      }
+    };
+    // assert done in Expectations
+    server.init(vertx, "", callback);
+
+    // no problem
+    endpointMetric.onConnect();
+    endpointMetric.onConnect();
+    connectHandler.handle(netSocket);
+
+    // reject
+    endpointMetric.onConnect();
+    connectHandler.handle(netSocket);
+    Assert.assertTrue(netSocketClosed);
+    Assert.assertEquals(1, endpointMetric.getRejectByConnectionLimitCount());
+  }
 }

transports/transport-highway/src/main/java/org/apache/servicecomb/transport/highway/HighwayServer.java

@@ -22,6 +22,8 @@
 import org.apache.servicecomb.foundation.vertx.server.TcpServer;
 import org.apache.servicecomb.foundation.vertx.server.TcpServerConnection;
 
+import com.netflix.config.DynamicPropertyFactory;
+
 public class HighwayServer extends TcpServer {
   private Endpoint endpoint;
 
@@ -34,4 +36,10 @@ public HighwayServer(Endpoint endpoint) {
   protected TcpServerConnection createTcpServerConnection() {
     return new HighwayServerConnection(endpoint);
   }
+
+  @Override
+  protected int getConnectionLimit() {
+    return DynamicPropertyFactory.getInstance()
+        .getIntProperty("servicecomb.highway.server.connection-limit", Integer.MAX_VALUE).get();
+  }
 }

transports/transport-rest/transport-rest-vertx/src/main/java/org/apache/servicecomb/transport/rest/vertx/RestServerVerticle.java

@@ -32,6 +32,8 @@
 import org.apache.servicecomb.foundation.ssl.SSLOption;
 import org.apache.servicecomb.foundation.ssl.SSLOptionFactory;
 import org.apache.servicecomb.foundation.vertx.VertxTLSBuilder;
+import org.apache.servicecomb.foundation.vertx.metrics.DefaultHttpServerMetrics;
+import org.apache.servicecomb.foundation.vertx.metrics.metric.DefaultServerEndpointMetric;
 import org.apache.servicecomb.swagger.invocation.exception.CommonExceptionData;
 import org.apache.servicecomb.swagger.invocation.exception.InvocationException;
 import org.apache.servicecomb.transport.rest.vertx.accesslog.AccessLogConfiguration;
@@ -52,6 +54,7 @@
 import io.vertx.core.http.HttpServer;
 import io.vertx.core.http.HttpServerOptions;
 import io.vertx.core.http.HttpServerResponse;
+import io.vertx.core.net.impl.ConnectionBase;
 import io.vertx.ext.web.Router;
 import io.vertx.ext.web.RoutingContext;
 import io.vertx.ext.web.handler.CorsHandler;
@@ -89,6 +92,17 @@ public void start(Future<Void> startFuture) throws Exception {
       initDispatcher(mainRouter);
       HttpServer httpServer = createHttpServer();
       httpServer.requestHandler(mainRouter::accept);
+      httpServer.connectionHandler(connection -> {
+        DefaultHttpServerMetrics serverMetrics = (DefaultHttpServerMetrics) ((ConnectionBase) connection).metrics();
+        DefaultServerEndpointMetric endpointMetric = serverMetrics.getEndpointMetric();
+        long connectedCount = endpointMetric.getCurrentConnectionCount();
+        int connectionLimit = DynamicPropertyFactory.getInstance()
+            .getIntProperty("servicecomb.rest.server.connection-limit", Integer.MAX_VALUE).get();
+        if (connectedCount > connectionLimit) {
+          connection.close();
+          endpointMetric.onRejectByConnectionLimit();
+        }
+      });
       httpServer.exceptionHandler(e -> {
         LOGGER.error("Unexpected error in server.{}", ExceptionUtils.getExceptionMessageWithoutTrace(e));
       });