Add support for connection throttling in JDisc

2 files changed, 95 insertions, 0 deletions

diff --git a/jdisc_http_service/src/test/java/com/yahoo/jdisc/http/server/jetty/ConnectionThrottlerTest.java b/jdisc_http_service/src/test/java/com/yahoo/jdisc/http/server/jetty/ConnectionThrottlerTest.java
new file mode 100644
index 00000000000..9e66418c382
--- /dev/null
+++ b/jdisc_http_service/src/test/java/com/yahoo/jdisc/http/server/jetty/ConnectionThrottlerTest.java
@@ -0,0 +1,78 @@
+// Copyright 2019 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.AbstractConnector;
+import org.eclipse.jetty.util.component.AbstractLifeCycle;
+import org.eclipse.jetty.util.statistic.RateStatistic;
+import org.eclipse.jetty.util.thread.Scheduler;
+import org.testng.annotations.Test;
+
+import java.util.concurrent.TimeUnit;
+
+import static org.mockito.Matchers.anyBoolean;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.mockito.internal.verification.VerificationModeFactory.times;
+import static org.testng.Assert.assertNotNull;
+
+/**
+ * @author bjorncs
+ */
+public class ConnectionThrottlerTest {
+
+    @Test
+    public void throttles_when_any_resource_check_exceeds_configured_threshold() {
+        Runtime runtime = mock(Runtime.class);
+        when(runtime.maxMemory()).thenReturn(100l);
+        RateStatistic rateStatistic = new RateStatistic(1, TimeUnit.HOURS);
+        MockScheduler scheduler = new MockScheduler();
+        ConnectorConfig.Throttling config = new ConnectorConfig.Throttling(new ConnectorConfig.Throttling.Builder()
+                                                                                   .maxHeapUtilization(0.8)
+                                                                                   .maxAcceptRate(1));
+
+        AbstractConnector connector = mock(AbstractConnector.class);
+
+        ConnectionThrottler throttler = new ConnectionThrottler(runtime, rateStatistic, scheduler, connector, config);
+
+        // Heap utilization above configured threshold, but connection rate below threshold.
+        when(runtime.freeMemory()).thenReturn(10l);
+        when(connector.isAccepting()).thenReturn(true);
+        throttler.onAccepting(null);
+        assertNotNull(scheduler.task);
+        verify(connector).setAccepting(false);
+
+        // Heap utilization below threshold, but connection rate above threshold.
+        when(runtime.freeMemory()).thenReturn(80l);
+        rateStatistic.record();
+        rateStatistic.record(); // above accept rate limit (2 > 1)
+        scheduler.task.run(); // run unthrottleIfBelowThresholds()
+        verify(connector, times(1)).setAccepting(anyBoolean()); // verify setAccepting has not been called any mores times
+
+        // Both heap utilization and accept rate below threshold
+        when(runtime.freeMemory()).thenReturn(80l);
+        when(connector.isAccepting()).thenReturn(false);
+        rateStatistic.reset();
+        scheduler.task.run(); // run unthrottleIfBelowThresholds()
+        verify(connector).setAccepting(true);
+
+        // Both heap utilization and accept rate below threshold
+        when(connector.isAccepting()).thenReturn(true);
+        when(runtime.freeMemory()).thenReturn(80l);
+        rateStatistic.record();
+        throttler.onAccepting(null);
+        verify(connector, times(2)).setAccepting(anyBoolean()); // verify setAccepting has not been called any mores times
+    }
+
+    private static class MockScheduler extends AbstractLifeCycle implements Scheduler {
+        Runnable task;
+
+        @Override
+        public Task schedule(Runnable task, long delay, TimeUnit units) {
+            this.task = task;
+            return () -> false;
+        }
+    }
+
+}
\ No newline at end of file
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..db0746baf2e 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)
+                                            .maxHeapUtilization(1.0)
+                                            .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);