Add support for configuring embedded Jetty's max queue capacity

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/ServerProperties.java

@@ -29,6 +29,7 @@
 import java.util.Map;
 
 import io.undertow.UndertowOptions;
+import io.undertow.server.handlers.RequestLimitingHandler;
 
 import org.springframework.boot.context.properties.ConfigurationProperties;
 import org.springframework.boot.context.properties.DeprecatedConfigurationProperty;
@@ -1307,6 +1308,21 @@ public static class Undertow {
 		 */
 		private Integer workerThreads;
 
+		/**
+		 * Number of max concurrent requests. When specified will put a
+		 * {@link RequestLimitingHandler} in place. The default is null. Can only be
+		 * specified in Servlet environments.
+		 */
+		private Integer maxRequests;
+
+		/**
+		 * Maximum number of requests to queue up when running with a
+		 * {@link RequestLimitingHandler} and the {@link #maxRequests} is reached. The
+		 * default is -1 (unbounded). Can only be specified in Servlet environments and is
+		 * only used when {@link #maxRequests} is specified.
+		 */
+		private Integer maxQueueCapacity;
+
 		/**
 		 * Whether to allocate buffers outside the Java heap. The default is derived from
 		 * the maximum amount of memory that is available to the JVM.
@@ -1403,6 +1419,22 @@ public void setWorkerThreads(Integer workerThreads) {
 			this.workerThreads = workerThreads;
 		}
 
+		public Integer getMaxRequests() {
+			return this.maxRequests;
+		}
+
+		public void setMaxRequests(Integer maxRequests) {
+			this.maxRequests = maxRequests;
+		}
+
+		public Integer getMaxQueueCapacity() {
+			return this.maxQueueCapacity;
+		}
+
+		public void setMaxQueueCapacity(Integer maxQueueCapacity) {
+			this.maxQueueCapacity = maxQueueCapacity;
+		}
+
 		public Boolean getDirectBuffers() {
 			return this.directBuffers;
 		}

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/servlet/ServletWebServerFactoryAutoConfiguration.java

@@ -78,6 +78,13 @@ public TomcatServletWebServerFactoryCustomizer tomcatServletWebServerFactoryCust
 		return new TomcatServletWebServerFactoryCustomizer(serverProperties);
 	}
 
+	@Bean
+	@ConditionalOnClass(name = "io.undertow.Undertow")
+	public UndertowServletWebServerFactoryCustomizer undertowServletWebServerFactoryCustomizer(
+			ServerProperties serverProperties) {
+		return new UndertowServletWebServerFactoryCustomizer(serverProperties);
+	}
+
 	@Bean
 	@ConditionalOnMissingFilterBean(ForwardedHeaderFilter.class)
 	@ConditionalOnProperty(value = "server.forward-headers-strategy", havingValue = "framework")

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/servlet/UndertowServletWebServerFactoryCustomizer.java

@@ -16,6 +16,11 @@
 
 package org.springframework.boot.autoconfigure.web.servlet;
 
+import io.undertow.Handlers;
+import io.undertow.server.HandlerWrapper;
+import io.undertow.server.HttpHandler;
+import io.undertow.server.handlers.RequestLimitingHandler;
+
 import org.springframework.boot.autoconfigure.web.ServerProperties;
 import org.springframework.boot.web.embedded.undertow.UndertowServletWebServerFactory;
 import org.springframework.boot.web.server.WebServerFactoryCustomizer;
@@ -25,6 +30,7 @@
  * Servlet web servers.
  *
  * @author Andy Wilkinson
+ * @author Chris Bono
  * @since 2.1.7
  */
 public class UndertowServletWebServerFactoryCustomizer
@@ -36,10 +42,55 @@ public UndertowServletWebServerFactoryCustomizer(ServerProperties serverProperti
 		this.serverProperties = serverProperties;
 	}
 
+	// TODO why was this not being used before now?
+
 	@Override
 	public void customize(UndertowServletWebServerFactory factory) {
+
 		factory.addDeploymentInfoCustomizers((deploymentInfo) -> deploymentInfo
 				.setEagerFilterInit(this.serverProperties.getUndertow().isEagerFilterInit()));
+
+		addRateLimiterIfMaxRequestsSet(factory);
+	}
+
+	/**
+	 * Installs a {@link RequestLimitingHandler} in the initial handler chain if
+	 * {@link ServerProperties.Undertow#maxRequests} is set. If installed, it will use
+	 * {@link ServerProperties.Undertow#maxQueueCapacity} to control the size of the queue
+	 * in the rate limting handler.
+	 * @param factory the factory to add the deployment info customizer on
+	 */
+	private void addRateLimiterIfMaxRequestsSet(UndertowServletWebServerFactory factory) {
+		Integer maxRequests = this.serverProperties.getUndertow().getMaxRequests();
+		if (maxRequests == null) {
+			return;
+		}
+		Integer maxQueueCapacity = this.serverProperties.getUndertow().getMaxQueueCapacity();
+		factory.addDeploymentInfoCustomizers((deploymentInfo) -> deploymentInfo.addInitialHandlerChainWrapper(
+				new RequestLimiterHandlerWrapper(maxRequests, (maxQueueCapacity != null) ? maxQueueCapacity : -1)));
+	}
+
+	/**
+	 * Explicit implementation of HandlerWrapper rather than Lambda to make it possible to
+	 * verify the handler wrappers at runtime during tests. This could be instead written
+	 * as a Lambda but then tests can not see what type of handler wrapper it is.
+	 */
+	static class RequestLimiterHandlerWrapper implements HandlerWrapper {
+
+		Integer maxRequests;
+
+		Integer maxQueueCapacity;
+
+		RequestLimiterHandlerWrapper(Integer maxRequests, Integer maxQueueCapacity) {
+			this.maxRequests = maxRequests;
+			this.maxQueueCapacity = maxQueueCapacity;
+		}
+
+		@Override
+		public HttpHandler wrap(HttpHandler handler) {
+			return Handlers.requestLimitingHandler(this.maxRequests, this.maxQueueCapacity, handler);
+		}
+
 	}
 
 }

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/ServerPropertiesTests.java

@@ -259,6 +259,30 @@ void testCustomizeUndertowSocketOption() {
 				"true");
 	}
 
+	@Test
+	void testCustomizeUndertowWorkerThreads() {
+		bind("server.undertow.worker-threads", "150");
+		assertThat(this.properties.getUndertow().getWorkerThreads()).isEqualTo(150);
+	}
+
+	@Test
+	void testCustomizeUndertowIoThreads() {
+		bind("server.undertow.io-threads", "10");
+		assertThat(this.properties.getUndertow().getIoThreads()).isEqualTo(10);
+	}
+
+	@Test
+	void testCustomizeUndertowMaxRequests() {
+		bind("server.undertow.max-requests", "200");
+		assertThat(this.properties.getUndertow().getMaxRequests()).isEqualTo(200);
+	}
+
+	@Test
+	void testCustomizeUndertowMaxQueueCapacity() {
+		bind("server.undertow.max-queue-capacity", "100");
+		assertThat(this.properties.getUndertow().getMaxQueueCapacity()).isEqualTo(100);
+	}
+
 	@Test
 	void testCustomizeJettyAccessLog() {
 		Map<String, String> map = new HashMap<>();

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/servlet/ServletWebServerFactoryAutoConfigurationTests.java

@@ -231,7 +231,7 @@ void undertowDeploymentInfoCustomizerBeanIsAddedToFactory() {
 						.withPropertyValues("server.port:0");
 		runner.run((context) -> {
 			UndertowServletWebServerFactory factory = context.getBean(UndertowServletWebServerFactory.class);
-			assertThat(factory.getDeploymentInfoCustomizers()).hasSize(1);
+			assertThat(factory.getDeploymentInfoCustomizers()).hasSize(2);
 		});
 	}
 

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/servlet/UndertowServletWebServerFactoryCustomizerTests.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2012-2019 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.boot.autoconfigure.web.servlet;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.boot.autoconfigure.web.ServerProperties;
+import org.springframework.boot.autoconfigure.web.servlet.UndertowServletWebServerFactoryCustomizer.RequestLimiterHandlerWrapper;
+import org.springframework.boot.context.properties.bind.Bindable;
+import org.springframework.boot.context.properties.bind.Binder;
+import org.springframework.boot.context.properties.source.ConfigurationPropertySources;
+import org.springframework.boot.web.embedded.undertow.UndertowServletWebServer;
+import org.springframework.boot.web.embedded.undertow.UndertowServletWebServerFactory;
+import org.springframework.mock.env.MockEnvironment;
+import org.springframework.test.context.support.TestPropertySourceUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link UndertowServletWebServerFactoryCustomizer}.
+ *
+ * @author Chris Bono
+ */
+class UndertowServletWebServerFactoryCustomizerTests {
+
+	private UndertowServletWebServerFactoryCustomizer customizer;
+
+	private MockEnvironment environment;
+
+	private ServerProperties serverProperties;
+
+	@BeforeEach
+	void setup() {
+		this.environment = new MockEnvironment();
+		this.serverProperties = new ServerProperties();
+		ConfigurationPropertySources.attach(this.environment);
+		this.customizer = new UndertowServletWebServerFactoryCustomizer(this.serverProperties);
+	}
+
+	@Test
+	void eagerFilterInitEnabled() {
+		bind("server.undertow.eager-filter-init=true");
+		UndertowServletWebServer server = customizeAndGetServer();
+		assertThat(server.getDeploymentManager().getDeployment().getDeploymentInfo().isEagerFilterInit()).isTrue();
+	}
+
+	@Test
+	void eagerFilterInitDisabled() {
+		bind("server.undertow.eager-filter-init=false");
+		UndertowServletWebServer server = customizeAndGetServer();
+		assertThat(server.getDeploymentManager().getDeployment().getDeploymentInfo().isEagerFilterInit()).isFalse();
+	}
+
+	@Test
+	void limiterNotAddedWhenMaxRequestsNotSet() {
+		bind("server.undertow.max-requests=");
+		UndertowServletWebServer server = customizeAndGetServer();
+		assertThat(server.getDeploymentManager().getDeployment().getDeploymentInfo().getInitialHandlerChainWrappers()
+				.stream().anyMatch(RequestLimiterHandlerWrapper.class::isInstance)).isFalse();
+	}
+
+	@Test
+	void limiterAddedWhenMaxRequestSetWithDefaultMaxQueueCapacity() {
+		bind("server.undertow.max-requests=200");
+		UndertowServletWebServer server = customizeAndGetServer();
+		RequestLimiterHandlerWrapper requestLimiter = server.getDeploymentManager().getDeployment().getDeploymentInfo()
+				.getInitialHandlerChainWrappers().stream().filter(RequestLimiterHandlerWrapper.class::isInstance)
+				.findFirst().map(RequestLimiterHandlerWrapper.class::cast).orElse(null);
+		assertThat(requestLimiter).isNotNull();
+		assertThat(requestLimiter.maxRequests).isEqualTo(200);
+		assertThat(requestLimiter.maxQueueCapacity).isEqualTo(-1);
+	}
+
+	@Test
+	void limiterAddedWhenMaxRequestSetWithCustomMaxQueueCapacity() {
+		bind("server.undertow.max-requests=200", "server.undertow.max-queue-capacity=100");
+		UndertowServletWebServer server = customizeAndGetServer();
+		RequestLimiterHandlerWrapper requestLimiter = server.getDeploymentManager().getDeployment().getDeploymentInfo()
+				.getInitialHandlerChainWrappers().stream().filter(RequestLimiterHandlerWrapper.class::isInstance)
+				.findFirst().map(RequestLimiterHandlerWrapper.class::cast).orElse(null);
+		assertThat(requestLimiter).isNotNull();
+		assertThat(requestLimiter.maxRequests).isEqualTo(200);
+		assertThat(requestLimiter.maxQueueCapacity).isEqualTo(100);
+	}
+
+	private void bind(String... inlinedProperties) {
+		TestPropertySourceUtils.addInlinedPropertiesToEnvironment(this.environment, inlinedProperties);
+		new Binder(ConfigurationPropertySources.get(this.environment)).bind("server",
+				Bindable.ofInstance(this.serverProperties));
+	}
+
+	private UndertowServletWebServer customizeAndGetServer() {
+		UndertowServletWebServerFactory factory = customizeAndGetFactory();
+		return (UndertowServletWebServer) factory.getWebServer();
+	}
+
+	private UndertowServletWebServerFactory customizeAndGetFactory() {
+		UndertowServletWebServerFactory factory = new UndertowServletWebServerFactory(0);
+		this.customizer.customize(factory);
+		return factory;
+	}
+
+}