Add support for connection throttling in JDisc

4 files changed, 167 insertions, 0 deletions

diff --git a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ConnectionThrottler.java b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ConnectionThrottler.java
new file mode 100644
index 00000000000..370ac0aa788
--- /dev/null
+++ b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/ConnectionThrottler.java
@@ -0,0 +1,128 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.jdisc.http.server.jetty;
+
+import com.yahoo.jdisc.http.ConnectorConfig;
+import org.eclipse.jetty.server.AcceptRateLimit;
+import org.eclipse.jetty.server.ConnectionLimit;
+import org.eclipse.jetty.server.Connector;
+import org.eclipse.jetty.server.LowResourceMonitor;
+import org.eclipse.jetty.util.component.LifeCycle;
+
+import java.time.Duration;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.TimeUnit;
+
+import static java.util.Collections.singleton;
+
+/**
+ * Throttles new connections using {@link LowResourceMonitor}, {@link AcceptRateLimit} and {@link ConnectionLimit}.
+ *
+ * @author bjorncs
+ */
+class ConnectionThrottler {
+
+    private final Object monitor = new Object();
+    private final Queue<Runnable> throttleResetters = new ArrayDeque<>();
+    private final Collection<LifeCycle> beans = new ArrayList<>();
+    private final Connector connector;
+    private int throttlersCount;
+
+    ConnectionThrottler(Connector connector, ConnectorConfig.Throttling config) {
+        this.connector = connector;
+        Duration idleTimeout = fromSeconds(config.idleTimeout());
+        if (config.maxAcceptRate() != -1) {
+            beans.add(new CoordinatedAcceptRateLimit(config.maxAcceptRate(), fromSeconds(config.maxAcceptRatePeriod())));
+        }
+        if (config.maxConnections() != -1) {
+            beans.add(new CoordinatedConnectionLimit(config.maxConnections(), idleTimeout));
+        }
+        if (config.maxMemoryUsage() != -1) {
+            beans.add(new CoordinatedLowResourcesLimit(config.maxMemoryUsage(), idleTimeout));
+        }
+    }
+
+    void registerBeans() {
+        beans.forEach(bean -> connector.getServer().addBean(connector));
+    }
+
+    private static Duration fromSeconds(double seconds) {
+        return Duration.ofMillis((long) (seconds * 1000));
+    }
+
+    private void onThrottle(Runnable throttleResetter) {
+        synchronized (monitor) {
+            ++throttlersCount;
+            throttleResetters.offer(throttleResetter);
+        }
+    }
+
+    private void onReset() {
+        List<Runnable> resetters = new ArrayList<>();
+        synchronized (monitor) {
+            if (--throttlersCount == 0) {
+                resetters.addAll(throttleResetters);
+                throttleResetters.clear();
+            }
+        }
+        resetters.forEach(Runnable::run);
+    }
+
+    private class CoordinatedLowResourcesLimit extends LowResourceMonitor {
+        CoordinatedLowResourcesLimit(int maxMemoryUsageMegaBytes, Duration idleTimeout) {
+            super(connector.getServer());
+            super.setMonitoredConnectors(singleton(connector));
+            super.setMaxMemory(maxMemoryUsageMegaBytes * 1024 * 1024L);
+            super.setLowResourcesIdleTimeout((int)idleTimeout.toMillis());
+        }
+
+        @Override
+        protected void setLowResources() {
+            super.setLowResources();
+            ConnectionThrottler.this.onThrottle(() -> super.clearLowResources());
+        }
+
+        @Override
+        protected void clearLowResources() {
+            ConnectionThrottler.this.onReset();
+        }
+    }
+
+    private class CoordinatedConnectionLimit extends ConnectionLimit {
+        CoordinatedConnectionLimit(int maxConnections, Duration idleTimeout) {
+            super(maxConnections, connector);
+            super.setIdleTimeout(idleTimeout.toMillis());
+        }
+
+        @Override
+        protected void limit() {
+            super.limit();
+            ConnectionThrottler.this.onThrottle(() -> super.unlimit());
+        }
+
+        @Override
+        protected void unlimit() {
+            ConnectionThrottler.this.onReset();
+        }
+    }
+
+    private class CoordinatedAcceptRateLimit extends AcceptRateLimit {
+        CoordinatedAcceptRateLimit(int limit, Duration period) {
+            super(limit, period.toMillis(), TimeUnit.MILLISECONDS, connector);
+        }
+
+        @Override
+        protected void limit() {
+            super.limit();
+            ConnectionThrottler.this.onThrottle(() -> super.unlimit());
+        }
+
+        @Override
+        protected void unlimit() {
+            ConnectionThrottler.this.onReset();
+        }
+    }
+}
diff --git a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/JDiscServerConnector.java b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/JDiscServerConnector.java
index f7d6e1717af..a80e694ed30 100644
--- a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/JDiscServerConnector.java
+++ b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/JDiscServerConnector.java
@@ -50,6 +50,10 @@ class JDiscServerConnector extends ServerConnector {
 
         this.statistics = new ServerConnectionStatistics();
         addBean(statistics);
+        ConnectorConfig.Throttling throttlingConfig = config.throttling();
+        if (throttlingConfig.enabled()) {
+            new ConnectionThrottler(this, throttlingConfig).registerBeans();
+        }
     }
 
     @Override
diff --git a/jdisc_http_service/src/main/resources/configdefinitions/jdisc.http.connector.def b/jdisc_http_service/src/main/resources/configdefinitions/jdisc.http.connector.def
index 157ffabdd63..ae4e66a236b 100644
--- a/jdisc_http_service/src/main/resources/configdefinitions/jdisc.http.connector.def
+++ b/jdisc_http_service/src/main/resources/configdefinitions/jdisc.http.connector.def
@@ -41,6 +41,24 @@ tcpKeepAliveEnabled                 bool     default=false
 # Enable/disable TCP_NODELAY (disable/enable Nagle's algorithm).
 tcpNoDelay                          bool     default=true
 
+# Whether to enable connection throttling. New connections will be dropped when a threshold is exceeded.
+throttling.enabled                  bool     default=false
+
+# Max number of connections.
+throttling.maxConnections           int      default=-1
+
+# Max memory usage in megabytes (totalMemory - freeMemory).
+throttling.maxMemoryUsage           int      default=-1
+
+# Max connection accept rate.
+throttling.maxAcceptRate            int      default=-1
+
+# Accept rate sample period in seconds. Used in conjunction with throttling.maxAcceptRate.
+throttling.maxAcceptRatePeriod      double   default=1.0
+
+# Idle timeout in seconds applied to endpoints when a threshold is exceeded (except accept rate threshold).
+throttling.idleTimeout              double   default=1.0
+
 # Whether to enable SSL for this connector.
 ssl.enabled                         bool     default=false
 
diff --git a/jdisc_http_service/src/test/java/com/yahoo/jdisc/http/server/jetty/HttpServerTest.java b/jdisc_http_service/src/test/java/com/yahoo/jdisc/http/server/jetty/HttpServerTest.java
index 30d5f9e657a..479cf514e30 100644
--- a/jdisc_http_service/src/test/java/com/yahoo/jdisc/http/server/jetty/HttpServerTest.java
+++ b/jdisc_http_service/src/test/java/com/yahoo/jdisc/http/server/jetty/HttpServerTest.java
@@ -16,6 +16,7 @@ import com.yahoo.jdisc.handler.RequestHandler;
 import com.yahoo.jdisc.handler.ResponseDispatch;
 import com.yahoo.jdisc.handler.ResponseHandler;
 import com.yahoo.jdisc.http.ConnectorConfig;
+import com.yahoo.jdisc.http.ConnectorConfig.Throttling;
 import com.yahoo.jdisc.http.Cookie;
 import com.yahoo.jdisc.http.HttpRequest;
 import com.yahoo.jdisc.http.HttpResponse;
@@ -480,6 +481,22 @@ public class HttpServerTest {
         assertThat(driver.close(), is(true));
     }
 
+    @Test
+    public void requireThatConnectionThrottleDoesNotBlockConnectionsBelowThreshold() throws Exception {
+        final TestDriver driver = TestDrivers.newConfiguredInstance(
+                new EchoRequestHandler(),
+                new ServerConfig.Builder(),
+                new ConnectorConfig.Builder()
+                        .throttling(new Throttling.Builder()
+                                            .enabled(true)
+                                            .maxAcceptRate(10)
+                                            .maxMemoryUsage(100*1024)
+                                            .maxConnections(10)));
+        driver.client().get("/status.html")
+                .expectStatusCode(is(OK));
+        assertThat(driver.close(), is(true));
+    }
+
     private static RequestHandler mockRequestHandler() {
         final RequestHandler mockRequestHandler = mock(RequestHandler.class);
         when(mockRequestHandler.refer()).thenReturn(References.NOOP_REFERENCE);