add configs to enable fast request failure on broker and historical (#4540)

* add configs to enable fast request failure on broker

* address review comments

* fix styling error

* fix style error

* have enableRequestLimit config instead of having user specify max limit

* add comment

* fix style error

* add UT fo LimitRequestsFilter

* address review comments

* fix test

* make LimitRequestsFilterTest more robust

* fix JettyQosTest

docs/content/configuration/broker.md

@@ -36,7 +36,9 @@ Druid uses Jetty to serve HTTP requests.
 |Property|Description|Default|
 |--------|-----------|-------|
 |`druid.server.http.numThreads`|Number of threads for HTTP requests.|max(10, (Number of cores * 17) / 16 + 2) + 30|
+|`druid.server.http.queueSize`|Size of the worker queue used by Jetty server to temporarily store incoming client connections. If this value is set and a request is rejected by jetty because queue is full then client would observe request failure with TCP connection being closed immediately with a completely empty response from server.|Unbounded|
 |`druid.server.http.maxIdleTime`|The Jetty max idle time for a connection.|PT5m|
+|`druid.server.http.enableRequestLimit`|If enabled, no requests would be queued in jetty queue and "HTTP 429 Too Many Requests" error response would be sent. |false|
 |`druid.server.http.defaultQueryTimeout`|Query timeout in millis, beyond which unfinished queries will be cancelled|300000|
 |`druid.server.http.maxScatterGatherBytes`|Maximum number of bytes gathered from data nodes such as historicals and realtime processes to execute a query. This is an advance configuration that allows to protect in case broker is under heavy load and not utilizing the data gathered in memory fast enough and leading to OOMs. This limit can be further reduced at query time using `maxScatterGatherBytes` in the context. Note that having large limit is not necessarily bad if broker is never under heavy concurrent load in which case data gathered is processed quickly and freeing up the memory used.|Long.MAX_VALUE|
 |`druid.broker.http.numConnections`|Size of connection pool for the Broker to connect to historical and real-time processes. If there are more queries than this number that all need to speak to the same node, then they will queue up.|20|

docs/content/configuration/historical.md

@@ -47,7 +47,9 @@ Druid uses Jetty to serve HTTP requests.
 |Property|Description|Default|
 |--------|-----------|-------|
 |`druid.server.http.numThreads`|Number of threads for HTTP requests.|max(10, (Number of cores * 17) / 16 + 2) + 30|
+|`druid.server.http.queueSize`|Size of the worker queue used by Jetty server to temporarily store incoming client connections. If this value is set and a request is rejected by jetty because queue is full then client would observe request failure with TCP connection being closed immediately with a completely empty response from server.|Unbounded|
 |`druid.server.http.maxIdleTime`|The Jetty max idle time for a connection.|PT5m|
+|`druid.server.http.enableRequestLimit`|If enabled, no requests would be queued in jetty queue and "HTTP 429 Too Many Requests" error response would be sent. |false|
 |`druid.server.http.defaultQueryTimeout`|Query timeout in millis, beyond which unfinished queries will be cancelled|300000|
 
 #### Processing

server/src/main/java/io/druid/server/initialization/ServerConfig.java

@@ -34,6 +34,13 @@ public class ServerConfig
   @Min(1)
   private int numThreads = Math.max(10, (Runtime.getRuntime().availableProcessors() * 17) / 16 + 2) + 30;
 
+  @JsonProperty
+  @Min(1)
+  private int queueSize = Integer.MAX_VALUE;
+
+  @JsonProperty
+  private boolean enableRequestLimit = false;
+
   @JsonProperty
   @NotNull
   private Period maxIdleTime = new Period("PT5m");
@@ -51,6 +58,16 @@ public class ServerConfig
     return numThreads;
   }
 
+  public int getQueueSize()
+  {
+    return queueSize;
+  }
+
+  public boolean isEnableRequestLimit()
+  {
+    return enableRequestLimit;
+  }
+
   public Period getMaxIdleTime()
   {
     return maxIdleTime;
@@ -77,6 +94,8 @@ public class ServerConfig
     }
     ServerConfig that = (ServerConfig) o;
     return numThreads == that.numThreads &&
+           queueSize == that.queueSize &&
+           enableRequestLimit == that.enableRequestLimit &&
            defaultQueryTimeout == that.defaultQueryTimeout &&
            maxScatterGatherBytes == that.maxScatterGatherBytes &&
            Objects.equals(maxIdleTime, that.maxIdleTime);
@@ -85,17 +104,13 @@ public class ServerConfig
   @Override
   public int hashCode()
   {
-    return Objects.hash(numThreads, maxIdleTime, defaultQueryTimeout, maxScatterGatherBytes);
+    return Objects.hash(
-  }
+        numThreads,
-
+        queueSize,
-  @Override
+        enableRequestLimit,
-  public String toString()
+        maxIdleTime,
-  {
+        defaultQueryTimeout,
-    return "ServerConfig{" +
+        maxScatterGatherBytes
-           "numThreads=" + numThreads +
+    );
-           ", maxIdleTime=" + maxIdleTime +
-           ", defaultQueryTimeout=" + defaultQueryTimeout +
-           ", maxScatterGatherBytes=" + maxScatterGatherBytes +
-           '}';
   }
 }

server/src/main/java/io/druid/server/initialization/jetty/JettyServerModule.java

@@ -22,16 +22,13 @@ package io.druid.server.initialization.jetty;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.jaxrs.json.JacksonJsonProvider;
 import com.fasterxml.jackson.jaxrs.smile.JacksonSmileProvider;
-import com.google.common.collect.Iterables;
 import com.google.common.primitives.Ints;
 import com.google.inject.Binder;
 import com.google.inject.Binding;
-import com.google.inject.ConfigurationException;
 import com.google.inject.Inject;
 import com.google.inject.Injector;
 import com.google.inject.Key;
 import com.google.inject.Provides;
-import com.google.inject.ProvisionException;
 import com.google.inject.Scopes;
 import com.google.inject.Singleton;
 import com.google.inject.multibindings.Multibinder;
@@ -51,6 +48,7 @@ import io.druid.guice.annotations.JSR311Resource;
 import io.druid.guice.annotations.Json;
 import io.druid.guice.annotations.Self;
 import io.druid.guice.annotations.Smile;
+import io.druid.java.util.common.RE;
 import io.druid.java.util.common.lifecycle.Lifecycle;
 import io.druid.java.util.common.logger.Logger;
 import io.druid.server.DruidNode;
@@ -79,6 +77,7 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -183,9 +182,24 @@ public class JettyServerModule extends JerseyServletModule
       Binding<SslContextFactory> sslContextFactoryBinding
   )
   {
-    final QueuedThreadPool threadPool = new QueuedThreadPool();
+    // adjusting to make config.getNumThreads() mean, "number of threads
-    threadPool.setMinThreads(config.getNumThreads());
+    // that concurrently handle the requests".
-    threadPool.setMaxThreads(config.getNumThreads());
+    int numServerThreads = config.getNumThreads() + getMaxJettyAcceptorsSelectorsNum(node);
+
+    final QueuedThreadPool threadPool;
+    if (config.getQueueSize() == Integer.MAX_VALUE) {
+      threadPool = new QueuedThreadPool();
+      threadPool.setMinThreads(numServerThreads);
+      threadPool.setMaxThreads(numServerThreads);
+    } else {
+      threadPool = new QueuedThreadPool(
+          numServerThreads,
+          numServerThreads,
+          60000, // same default is used in other case when threadPool = new QueuedThreadPool()
+          new LinkedBlockingQueue<>(config.getQueueSize())
+      );
+    }
+
     threadPool.setDaemon(true);
 
     final Server server = new Server(threadPool);
@@ -261,8 +275,8 @@ public class JettyServerModule extends JerseyServletModule
     try {
       initializer.initialize(server, injector);
     }
-    catch (ConfigurationException e) {
+    catch (Exception e) {
-      throw new ProvisionException(Iterables.getFirst(e.getErrorMessages(), null).getMessage());
+      throw new RE(e, "server initialization exception");
     }
 
     lifecycle.addHandler(
@@ -288,6 +302,14 @@ public class JettyServerModule extends JerseyServletModule
     );
   }
 
+  private static int getMaxJettyAcceptorsSelectorsNum(DruidNode druidNode)
+  {
+    // This computation is based on Jetty v9.3.19 which uses upto 8(4 acceptors and 4 selectors) threads per
+    // ServerConnector
+    int numServerConnector = (druidNode.isEnablePlaintextPort() ? 1 : 0) + (druidNode.isEnableTlsPort() ? 1 : 0);
+    return numServerConnector * 8;
+  }
+
   @Provides
   @Singleton
   public JettyMonitor getJettyMonitor(

server/src/main/java/io/druid/server/initialization/jetty/LimitRequestsFilter.java

@@ -0,0 +1,85 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets 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 io.druid.server.initialization.jetty;
+
+import com.google.common.base.Preconditions;
+
+import javax.servlet.Filter;
+import javax.servlet.FilterChain;
+import javax.servlet.FilterConfig;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletResponse;
+import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ */
+public class LimitRequestsFilter implements Filter
+{
+  private final int maxActiveRequests;
+
+  private final AtomicInteger activeRequestsCount = new AtomicInteger();
+
+  public LimitRequestsFilter(int maxActiveRequests)
+  {
+    Preconditions.checkArgument(
+        maxActiveRequests > 0 && maxActiveRequests < Integer.MAX_VALUE,
+        "maxActiveRequests must be > 0 and < Integer.MAX_VALUE."
+    );
+    this.maxActiveRequests = maxActiveRequests;
+  }
+
+  @Override
+  public void init(FilterConfig filterConfig) throws ServletException
+  {
+  }
+
+  @Override
+  public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
+      throws IOException, ServletException
+  {
+
+    int curr = activeRequestsCount.incrementAndGet();
+    try {
+      if (curr <= maxActiveRequests) {
+        chain.doFilter(request, response);
+      } else {
+        // See https://tools.ietf.org/html/rfc6585 for status code 429 explanation.
+        ((HttpServletResponse) response).sendError(429, "Too Many Requests");
+      }
+    }
+    finally {
+      activeRequestsCount.decrementAndGet();
+    }
+  }
+
+  @Override
+  public void destroy()
+  {
+
+  }
+
+  public int getActiveRequestsCount()
+  {
+    return activeRequestsCount.get();
+  }
+}

server/src/test/java/io/druid/server/initialization/JettyQosTest.java

@@ -80,7 +80,7 @@ public class JettyQosTest extends BaseJettyTest
                 binder.bind(AuthorizerMapper.class).toInstance(AuthTestUtils.TEST_AUTHORIZER_MAPPER);
                 JettyBindings.addQosFilter(binder, "/slow/*", 2);
                 final ServerConfig serverConfig = new ObjectMapper().convertValue(
-                    ImmutableMap.of("numThreads", "10"),
+                    ImmutableMap.of("numThreads", "2"),
                     ServerConfig.class
                 );
                 binder.bind(ServerConfig.class).toInstance(serverConfig);

server/src/test/java/io/druid/server/initialization/jetty/LimitRequestsFilterTest.java

@@ -0,0 +1,149 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets 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 io.druid.server.initialization.jetty;
+
+import io.druid.java.util.common.ISE;
+import org.easymock.EasyMock;
+import org.junit.Test;
+
+import javax.servlet.FilterChain;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletResponse;
+import java.io.IOException;
+import java.util.concurrent.CountDownLatch;
+
+/**
+ */
+public class LimitRequestsFilterTest
+{
+  @Test(timeout = 5000L)
+  public void testSimple() throws Exception
+  {
+    LimitRequestsFilter filter = new LimitRequestsFilter(2);
+
+    CountDownLatch latch1 = createAndStartRequestThread(
+        filter,
+        EasyMock.createStrictMock(ServletRequest.class),
+        EasyMock.createStrictMock(HttpServletResponse.class)
+    );
+
+    CountDownLatch latch2 = createAndStartRequestThread(
+        filter,
+        EasyMock.createStrictMock(ServletRequest.class),
+        EasyMock.createStrictMock(HttpServletResponse.class)
+    );
+
+    while (filter.getActiveRequestsCount() != 2) {
+      Thread.sleep(100);
+    }
+
+    //now further requests should fail
+    HttpServletResponse resp = EasyMock.createMock(HttpServletResponse.class);
+    resp.sendError(429, "Too Many Requests");
+    EasyMock.expectLastCall().times(2);
+    EasyMock.replay(resp);
+
+    filter.doFilter(
+        EasyMock.createStrictMock(ServletRequest.class),
+        resp,
+        EasyMock.createStrictMock(FilterChain.class)
+    );
+
+    filter.doFilter(
+        EasyMock.createStrictMock(ServletRequest.class),
+        resp,
+        EasyMock.createStrictMock(FilterChain.class)
+    );
+
+    EasyMock.verify(resp);
+
+    //release one of the pending requests
+    latch1.countDown();
+
+    while (filter.getActiveRequestsCount() != 1) {
+      Thread.sleep(100);
+    }
+
+    //now requests should go through
+    FilterChain chain = EasyMock.createMock(FilterChain.class);
+    chain.doFilter(EasyMock.anyObject(), EasyMock.anyObject());
+    EasyMock.expectLastCall().times(2);
+    EasyMock.replay(chain);
+
+    filter.doFilter(
+        EasyMock.createStrictMock(ServletRequest.class),
+        EasyMock.createStrictMock(HttpServletResponse.class),
+        chain
+    );
+
+    filter.doFilter(
+        EasyMock.createStrictMock(ServletRequest.class),
+        EasyMock.createStrictMock(HttpServletResponse.class),
+        chain
+    );
+
+    EasyMock.verify(chain);
+
+    latch2.countDown();
+
+    while (filter.getActiveRequestsCount() != 0) {
+      Thread.sleep(100);
+    }
+  }
+
+  private CountDownLatch createAndStartRequestThread(LimitRequestsFilter filter, ServletRequest req, HttpServletResponse resp)
+  {
+    CountDownLatch latch = new CountDownLatch(1);
+    new Thread(
+        () -> {
+          try {
+            filter.doFilter(req, resp, new TestFilterChain(latch));
+          }
+          catch (Exception e) {
+            throw new ISE(e, "exception");
+          }
+        }
+    ).start();
+    return latch;
+  }
+
+  private static class TestFilterChain implements FilterChain
+  {
+    private final CountDownLatch latch;
+
+    TestFilterChain(CountDownLatch latch)
+    {
+      this.latch = latch;
+    }
+
+    @Override
+    public void doFilter(ServletRequest request, ServletResponse response) throws IOException, ServletException
+    {
+      try {
+        latch.await();
+      }
+      catch (InterruptedException ex) {
+        throw new ISE(ex, "exception");
+      }
+    }
+  }
+}

services/src/main/java/io/druid/cli/QueryJettyServerInitializer.java

@@ -20,6 +20,7 @@
 package io.druid.cli;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
 import com.google.inject.Injector;
@@ -27,9 +28,10 @@ import com.google.inject.Key;
 import com.google.inject.servlet.GuiceFilter;
 import io.druid.guice.annotations.Json;
 import io.druid.java.util.common.logger.Logger;
+import io.druid.server.initialization.ServerConfig;
 import io.druid.server.initialization.jetty.JettyServerInitUtils;
 import io.druid.server.initialization.jetty.JettyServerInitializer;
-import io.druid.server.security.AuthConfig;
+import io.druid.server.initialization.jetty.LimitRequestsFilter;
 import io.druid.server.security.AuthenticationUtils;
 import io.druid.server.security.Authenticator;
 import io.druid.server.security.AuthenticatorMapper;
@@ -37,6 +39,7 @@ import org.eclipse.jetty.server.Handler;
 import org.eclipse.jetty.server.Server;
 import org.eclipse.jetty.server.handler.HandlerList;
 import org.eclipse.jetty.servlet.DefaultServlet;
+import org.eclipse.jetty.servlet.FilterHolder;
 import org.eclipse.jetty.servlet.ServletContextHandler;
 import org.eclipse.jetty.servlet.ServletHolder;
 
@@ -47,14 +50,17 @@ import java.util.Set;
  */
 public class QueryJettyServerInitializer implements JettyServerInitializer
 {
-  private static Logger log = new Logger(QueryJettyServerInitializer.class);
+  private static final Logger log = new Logger(QueryJettyServerInitializer.class);
 
   private final List<Handler> extensionHandlers;
 
+  private final ServerConfig serverConfig;
+
   @Inject
-  public QueryJettyServerInitializer(Set<Handler> extensionHandlers)
+  public QueryJettyServerInitializer(Set<Handler> extensionHandlers, ServerConfig serverConfig)
   {
     this.extensionHandlers = ImmutableList.copyOf(extensionHandlers);
+    this.serverConfig = serverConfig;
   }
 
   @Override
@@ -63,7 +69,20 @@ public class QueryJettyServerInitializer implements JettyServerInitializer
     final ServletContextHandler root = new ServletContextHandler(ServletContextHandler.SESSIONS);
     root.addServlet(new ServletHolder(new DefaultServlet()), "/*");
 
-    final AuthConfig authConfig = injector.getInstance(AuthConfig.class);
+    // Add LimitRequestsFilter as first in the chain if enabled.
+    if (serverConfig.isEnableRequestLimit()) {
+      //To reject xth request, limit should be set to x-1 because (x+1)st request wouldn't reach filter
+      // but rather wait on jetty queue.
+      Preconditions.checkArgument(
+          serverConfig.getNumThreads() > 1,
+          "numThreads must be > 1 to enable Request Limit Filter."
+      );
+      log.info("Enabling Request Limit Filter with limit [%d].", serverConfig.getNumThreads() - 1);
+      root.addFilter(new FilterHolder(new LimitRequestsFilter(serverConfig.getNumThreads() - 1)),
+                     "/*", null
+      );
+    }
+
     final ObjectMapper jsonMapper = injector.getInstance(Key.get(ObjectMapper.class, Json.class));
     final AuthenticatorMapper authenticatorMapper = injector.getInstance(AuthenticatorMapper.class);