Prevent early initialization of Servlets

Defer loading of Servlets (with a loadOnStartup priority) until
the `EmbeddedServletContainer.start()` method is called. This prevents
issues with the DispatcherServlet being initialized before the embedded
ApplicationContext is fully initialized.

spring-boot/src/main/java/org/springframework/boot/context/embedded/EmbeddedServletContainerFactory.java

 /*
- * Copyright 2012-2013 the original author or authors.
+ * Copyright 2012-2014 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.
@@ -22,7 +22,7 @@ import org.springframework.boot.context.embedded.tomcat.TomcatEmbeddedServletCon
 
 /**
  * Factory interface that can be used to create {@link EmbeddedServletContainer}s.
- * Implementations are encouraged to extends
+ * Implementations are encouraged to extend
  * {@link AbstractEmbeddedServletContainerFactory} when possible.
  * 
  * @author Phillip Webb

spring-boot/src/main/java/org/springframework/boot/context/embedded/jetty/JettyEmbeddedServletContainer.java

 /*
- * Copyright 2012-2013 the original author or authors.
+ * Copyright 2012-2014 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.
@@ -19,6 +19,7 @@ package org.springframework.boot.context.embedded.jetty;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.eclipse.jetty.server.Connector;
+import org.eclipse.jetty.server.Handler;
 import org.eclipse.jetty.server.Server;
 import org.springframework.boot.context.embedded.EmbeddedServletContainer;
 import org.springframework.boot.context.embedded.EmbeddedServletContainerException;
@@ -84,6 +85,11 @@ public class JettyEmbeddedServletContainer implements EmbeddedServletContainer {
 		}
 		try {
 			this.server.start();
+			for (Handler handler : this.server.getHandlers()) {
+				if (handler instanceof JettyEmbeddedWebAppContext) {
+					((JettyEmbeddedWebAppContext) handler).deferredInitialize();
+				}
+			}
 			Connector[] connectors = this.server.getConnectors();
 			for (Connector connector : connectors) {
 				connector.start();

spring-boot/src/main/java/org/springframework/boot/context/embedded/jetty/JettyEmbeddedServletContainerFactory.java

 /*
- * Copyright 2012-2013 the original author or authors.
+ * Copyright 2012-2014 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.
@@ -96,7 +96,7 @@ public class JettyEmbeddedServletContainerFactory extends
 	@Override
 	public EmbeddedServletContainer getEmbeddedServletContainer(
 			ServletContextInitializer... initializers) {
-		WebAppContext context = new WebAppContext();
+		JettyEmbeddedWebAppContext context = new JettyEmbeddedWebAppContext();
 		int port = (getPort() >= 0 ? getPort() : 0);
 		Server server = new Server(new InetSocketAddress(getAddress(), port));
 

spring-boot/src/main/java/org/springframework/boot/context/embedded/jetty/JettyEmbeddedWebAppContext.java

+/*
+ * Copyright 2012-2014 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
+ *
+ *      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 org.springframework.boot.context.embedded.jetty;
+
+import org.eclipse.jetty.servlet.ServletHandler;
+import org.eclipse.jetty.webapp.WebAppContext;
+
+/**
+ * Jetty {@link WebAppContext} used by {@link JettyEmbeddedServletContainer} to support
+ * deferred initialization.
+ * 
+ * @author Phillip Webb
+ */
+class JettyEmbeddedWebAppContext extends WebAppContext {
+
+	@Override
+	protected ServletHandler newServletHandler() {
+		return new JettyEmbeddedServletHandler();
+	}
+
+	public void deferredInitialize() throws Exception {
+		((JettyEmbeddedServletHandler) getServletHandler()).deferredInitialize();
+	}
+
+	private static class JettyEmbeddedServletHandler extends ServletHandler {
+
+		@Override
+		public void initialize() throws Exception {
+		}
+
+		public void deferredInitialize() throws Exception {
+			super.initialize();
+		}
+
+	}
+
+}

spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedContext.java

+/*
+ * Copyright 2012-2014 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
+ *
+ *      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 org.springframework.boot.context.embedded.tomcat;
+
+import org.apache.catalina.Container;
+import org.apache.catalina.core.StandardContext;
+
+/**
+ * Tomcat {@link StandardContext} used by {@link TomcatEmbeddedServletContainer} to
+ * support deferred initialization.
+ * 
+ * @author Phillip Webb
+ */
+class TomcatEmbeddedContext extends StandardContext {
+
+	@Override
+	public void loadOnStartup(Container[] children) {
+	}
+
+	public void deferredLoadOnStartup() {
+		super.loadOnStartup(findChildren());
+	}
+
+}

spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainer.java

 /*
- * Copyright 2012-2013 the original author or authors.
+ * Copyright 2012-2014 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.
@@ -18,6 +18,7 @@ package org.springframework.boot.context.embedded.tomcat;
 
 import java.util.concurrent.atomic.AtomicInteger;
 
+import org.apache.catalina.Container;
 import org.apache.catalina.Engine;
 import org.apache.catalina.LifecycleException;
 import org.apache.catalina.LifecycleState;
@@ -111,6 +112,11 @@ public class TomcatEmbeddedServletContainer implements EmbeddedServletContainer
 		Connector connector = this.tomcat.getConnector();
 		if (connector != null && this.autoStart) {
 			try {
+				for (Container child : this.tomcat.getHost().findChildren()) {
+					if (child instanceof TomcatEmbeddedContext) {
+						((TomcatEmbeddedContext) child).deferredLoadOnStartup();
+					}
+				}
 				connector.getProtocolHandler().start();
 				this.logger.info("Tomcat started on port: " + connector.getLocalPort());
 			}

spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactory.java

 /*
- * Copyright 2012-2013 the original author or authors.
+ * Copyright 2012-2014 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.
@@ -36,7 +36,6 @@ import org.apache.catalina.LifecycleListener;
 import org.apache.catalina.Valve;
 import org.apache.catalina.Wrapper;
 import org.apache.catalina.connector.Connector;
-import org.apache.catalina.core.StandardContext;
 import org.apache.catalina.loader.WebappLoader;
 import org.apache.catalina.startup.Tomcat;
 import org.apache.catalina.startup.Tomcat.FixContextListener;
@@ -119,14 +118,11 @@ public class TomcatEmbeddedServletContainerFactory extends
 	@Override
 	public EmbeddedServletContainer getEmbeddedServletContainer(
 			ServletContextInitializer... initializers) {
-
-		Connector connector;
 		Tomcat tomcat = new Tomcat();
-
 		File baseDir = (this.baseDirectory != null ? this.baseDirectory
 				: createTempDir("tomcat"));
 		tomcat.setBaseDir(baseDir.getAbsolutePath());
-		connector = new Connector(this.protocol);
+		Connector connector = new Connector(this.protocol);
 		tomcat.getService().addConnector(connector);
 		customizeConnector(connector);
 		tomcat.setConnector(connector);
@@ -141,7 +137,7 @@ public class TomcatEmbeddedServletContainerFactory extends
 	protected void prepareContext(Host host, ServletContextInitializer[] initializers) {
 		File docBase = getValidDocumentRoot();
 		docBase = (docBase != null ? docBase : createTempDir("tomcat-docbase"));
-		Context context = new StandardContext();
+		TomcatEmbeddedContext context = new TomcatEmbeddedContext();
 		context.setName(getContextPath());
 		context.setPath(getContextPath());
 		context.setDocBase(docBase.getAbsolutePath());

spring-boot/src/test/java/org/springframework/boot/context/embedded/AbstractEmbeddedServletContainerFactoryTests.java

 /*
- * Copyright 2012-2013 the original author or authors.
+ * Copyright 2012-2014 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.
@@ -26,6 +26,7 @@ import java.util.Arrays;
 import java.util.Date;
 import java.util.concurrent.TimeUnit;
 
+import javax.servlet.GenericServlet;
 import javax.servlet.ServletContext;
 import javax.servlet.ServletException;
 import javax.servlet.ServletRequest;
@@ -168,6 +169,23 @@ public abstract class AbstractEmbeddedServletContainerFactoryTests {
 		assertThat(date[0], notNullValue());
 	}
 
+	@Test
+	public void loadOnStartAfterContextIsInitialized() throws Exception {
+		ConfigurableEmbeddedServletContainerFactory factory = getFactory();
+		final InitCountingServlet servlet = new InitCountingServlet();
+		this.container = factory
+				.getEmbeddedServletContainer(new ServletContextInitializer() {
+					@Override
+					public void onStartup(ServletContext servletContext)
+							throws ServletException {
+						servletContext.addServlet("test", servlet).setLoadOnStartup(1);
+					}
+				});
+		assertThat(servlet.getInitCount(), equalTo(0));
+		this.container.start();
+		assertThat(servlet.getInitCount(), equalTo(1));
+	}
+
 	@Test
 	public void specificPort() throws Exception {
 		ConfigurableEmbeddedServletContainerFactory factory = getFactory();
@@ -313,4 +331,23 @@ public abstract class AbstractEmbeddedServletContainerFactoryTests {
 		bean.setName("error");
 		return bean;
 	}
+
+	private static class InitCountingServlet extends GenericServlet {
+
+		private int initCount;
+
+		@Override
+		public void init() throws ServletException {
+			this.initCount++;
+		}
+
+		@Override
+		public void service(ServletRequest req, ServletResponse res)
+				throws ServletException, IOException {
+		}
+
+		public int getInitCount() {
+			return this.initCount;
+		}
+	};
 }