Allow ExitCodeGenerator to be used on Exceptions

Update exit code support to allow the ExitCodeGenerator interface to
be placed on an Exception. Any uncaught exception implementing the
interface and returning a non `0` status will now trigger a System.exit
with the code.

Fixes gh-4803

spring-boot-cli/src/it/java/org/springframework/boot/cli/WarCommandIT.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -53,11 +53,11 @@ public class WarCommandIT {
 				.start();
 		invocation = new Invocation(process);
 		invocation.await();
-		assertThat(invocation.getErrorOutput(), containsString("onStart error"));
-		assertThat(invocation.getStandardOutput(), containsString("Tomcat started"));
-		assertThat(invocation.getStandardOutput(),
+		assertThat(invocation.getOutput(), containsString("onStart error"));
+		assertThat(invocation.getOutput(), containsString("Tomcat started"));
+		assertThat(invocation.getOutput(),
 				containsString("/WEB-INF/lib-provided/tomcat-embed-core"));
-		assertThat(invocation.getStandardOutput(),
+		assertThat(invocation.getOutput(),
 				containsString("/WEB-INF/lib-provided/tomcat-embed-core"));
 		process.destroy();
 	}

spring-boot-cli/src/it/java/org/springframework/boot/cli/infrastructure/CommandLineInvoker.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -93,21 +93,27 @@ public final class CommandLineInvoker {
 
 		private final StringBuffer out = new StringBuffer();
 
+		private final StringBuffer combined = new StringBuffer();
+
 		private final Process process;
 
 		private final List<Thread> streamReaders = new ArrayList<Thread>();
 
 		public Invocation(Process process) {
 			this.process = process;
-			this.streamReaders.add(new Thread(
-					new StreamReadingRunnable(this.process.getErrorStream(), this.err)));
-			this.streamReaders.add(new Thread(
-					new StreamReadingRunnable(this.process.getInputStream(), this.out)));
+			this.streamReaders.add(new Thread(new StreamReadingRunnable(
+					this.process.getErrorStream(), this.err, this.combined)));
+			this.streamReaders.add(new Thread(new StreamReadingRunnable(
+					this.process.getInputStream(), this.out, this.combined)));
 			for (Thread streamReader : this.streamReaders) {
 				streamReader.start();
 			}
 		}
 
+		public String getOutput() {
+			return postProcessLines(getLines(this.combined));
+		}
+
 		public String getErrorOutput() {
 			return postProcessLines(getLines(this.err));
 		}
@@ -161,13 +167,13 @@ public final class CommandLineInvoker {
 
 			private final InputStream stream;
 
-			private final StringBuffer output;
+			private final StringBuffer[] outputs;
 
 			private final byte[] buffer = new byte[4096];
 
-			private StreamReadingRunnable(InputStream stream, StringBuffer buffer) {
+			private StreamReadingRunnable(InputStream stream, StringBuffer... outputs) {
 				this.stream = stream;
-				this.output = buffer;
+				this.outputs = outputs;
 			}
 
 			@Override
@@ -175,7 +181,9 @@ public final class CommandLineInvoker {
 				int read;
 				try {
 					while ((read = this.stream.read(this.buffer)) > 0) {
-						this.output.append(new String(this.buffer, 0, read));
+						for (StringBuffer output : this.outputs) {
+							output.append(new String(this.buffer, 0, read));
+						}
 					}
 				}
 				catch (IOException ex) {

spring-boot-samples/spring-boot-sample-batch/src/main/java/sample/batch/SampleBatchApplication.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -42,6 +42,7 @@ public class SampleBatchApplication {
 
 	@Bean
 	protected Tasklet tasklet() {
+
 		return new Tasklet() {
 			@Override
 			public RepeatStatus execute(StepContribution contribution,
@@ -49,6 +50,7 @@ public class SampleBatchApplication {
 				return RepeatStatus.FINISHED;
 			}
 		};
+
 	}
 
 	@Bean

spring-boot-samples/spring-boot-sample-simple/src/main/java/sample/simple/ExitException.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2012-2016 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 sample.simple;
+
+import org.springframework.boot.ExitCodeGenerator;
+
+public class ExitException extends RuntimeException implements ExitCodeGenerator {
+
+	@Override
+	public int getExitCode() {
+		return 10;
+	}
+
+}

spring-boot-samples/spring-boot-sample-simple/src/main/java/sample/simple/SampleSimpleApplication.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -37,6 +37,9 @@ public class SampleSimpleApplication implements CommandLineRunner {
 	@Override
 	public void run(String... args) {
 		System.out.println(this.helloWorldService.getHelloMessage());
+		if (args.length > 0 && args[0].equals("exitcode")) {
+			throw new ExitException();
+		}
 	}
 
 	public static void main(String[] args) throws Exception {

spring-boot-tools/spring-boot-loader/src/main/java/org/springframework/boot/loader/MainMethodRunner.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2013 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -16,6 +16,7 @@
 
 package org.springframework.boot.loader;
 
+import java.lang.Thread.UncaughtExceptionHandler;
 import java.lang.reflect.Method;
 
 /**
@@ -53,7 +54,14 @@ public class MainMethodRunner implements Runnable {
 			mainMethod.invoke(null, new Object[] { this.args });
 		}
 		catch (Exception ex) {
+			UncaughtExceptionHandler handler = Thread.currentThread()
+					.getUncaughtExceptionHandler();
+			if (handler != null) {
+				handler.uncaughtException(Thread.currentThread(), ex);
+			}
+			else {
 				ex.printStackTrace();
+			}
 			System.exit(1);
 		}
 	}

spring-boot/src/main/java/org/springframework/boot/ExitCodeGenerator.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2013 the original author or authors.
+ * Copyright 2012-2016 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,7 +18,8 @@ package org.springframework.boot;
 
 /**
  * Interface used to generate an 'exit code' from a running command line
- * {@link SpringApplication}.
+ * {@link SpringApplication}. Since 1.3.2 this interface can be used on exceptions as well
+ * as directly on beans.
  *
  * @author Dave Syer
  * @see SpringApplication#exit(org.springframework.context.ApplicationContext,

spring-boot/src/main/java/org/springframework/boot/SpringApplication.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -827,6 +827,7 @@ public class SpringApplication {
 		}
 		try {
 			try {
+				handeExitCode(context, exception);
 				listeners.finished(context, exception);
 			}
 			finally {
@@ -847,14 +848,46 @@ public class SpringApplication {
 	 * @param exception the exception that was logged
 	 */
 	protected void registerLoggedException(Throwable exception) {
-		Thread currentThread = Thread.currentThread();
-		if (("main".equals(currentThread.getName())
-				|| "restartedMain".equals(currentThread.getName()))
-				&& "main".equals(currentThread.getThreadGroup().getName())) {
-			LoggedExceptionHandler.forCurrentThread().register(exception);
+		SpringBootExceptionHandler handler = getSpringBootExceptionHandler();
+		if (handler != null) {
+			handler.registerLoggedException(exception);
 		}
 	}
 
+	private void handeExitCode(ConfigurableApplicationContext context,
+			Throwable exception) {
+		int exitCode = getExitCodeFromException(exception);
+		if (exitCode != 0) {
+			SpringBootExceptionHandler handler = getSpringBootExceptionHandler();
+			if (handler != null) {
+				handler.registerExitCode(exitCode);
+			}
+		}
+	}
+
+	private int getExitCodeFromException(Throwable exception) {
+		if (exception == null) {
+			return 0;
+		}
+		if (exception instanceof ExitCodeGenerator) {
+			return ((ExitCodeGenerator) exception).getExitCode();
+		}
+		return getExitCodeFromException(exception.getCause());
+	}
+
+	SpringBootExceptionHandler getSpringBootExceptionHandler() {
+		if (isMainThread(Thread.currentThread())) {
+			return SpringBootExceptionHandler.forCurrentThread();
+		}
+		return null;
+	}
+
+	private boolean isMainThread(Thread currentThread) {
+		return ("main".equals(currentThread.getName())
+				|| "restartedMain".equals(currentThread.getName()))
+				&& "main".equals(currentThread.getThreadGroup().getName());
+	}
+
 	/**
 	 * Set a specific main application class that will be used as a log source and to
 	 * obtain version information. By default the main application class will be deduced.

spring-boot/src/main/java/org/springframework/boot/SpringBootExceptionHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -25,11 +25,12 @@ import java.util.List;
 import java.util.Set;
 
 /**
- * {@link UncaughtExceptionHandler} to suppress handling already logged exceptions.
+ * {@link UncaughtExceptionHandler} to suppress handling already logged exceptions and
+ * dealing with system exit.
  *
  * @author Phillip Webb
  */
-class LoggedExceptionHandler implements UncaughtExceptionHandler {
+class SpringBootExceptionHandler implements UncaughtExceptionHandler {
 
 	private static Set<String> LOG_CONFIGURATION_MESSAGES;
 
@@ -43,14 +44,20 @@ class LoggedExceptionHandler implements UncaughtExceptionHandler {
 
 	private final UncaughtExceptionHandler parent;
 
-	private final List<Throwable> exceptions = new ArrayList<Throwable>();
+	private final List<Throwable> loggedExceptions = new ArrayList<Throwable>();
 
-	LoggedExceptionHandler(UncaughtExceptionHandler parent) {
+	private int exitCode = 0;
+
+	SpringBootExceptionHandler(UncaughtExceptionHandler parent) {
 		this.parent = parent;
 	}
 
-	public void register(Throwable exception) {
-		this.exceptions.add(exception);
+	public void registerLoggedException(Throwable exception) {
+		this.loggedExceptions.add(exception);
+	}
+
+	public void registerExitCode(int exitCode) {
+		this.exitCode = exitCode;
 	}
 
 	@Override
@@ -61,7 +68,10 @@ class LoggedExceptionHandler implements UncaughtExceptionHandler {
 			}
 		}
 		finally {
-			this.exceptions.clear();
+			this.loggedExceptions.clear();
+			if (this.exitCode != 0) {
+				System.exit(this.exitCode);
+			}
 		}
 	}
 
@@ -88,7 +98,7 @@ class LoggedExceptionHandler implements UncaughtExceptionHandler {
 	}
 
 	private boolean isRegistered(Throwable ex) {
-		if (this.exceptions.contains(ex)) {
+		if (this.loggedExceptions.contains(ex)) {
 			return true;
 		}
 		if (ex instanceof InvocationTargetException) {
@@ -97,7 +107,7 @@ class LoggedExceptionHandler implements UncaughtExceptionHandler {
 		return false;
 	}
 
-	static LoggedExceptionHandler forCurrentThread() {
+	static SpringBootExceptionHandler forCurrentThread() {
 		return handler.get();
 	}
 
@@ -105,11 +115,11 @@ class LoggedExceptionHandler implements UncaughtExceptionHandler {
 	 * Thread local used to attach and track handlers.
 	 */
 	private static class LoggedExceptionHandlerThreadLocal
-			extends ThreadLocal<LoggedExceptionHandler> {
+			extends ThreadLocal<SpringBootExceptionHandler> {
 
 		@Override
-		protected LoggedExceptionHandler initialValue() {
-			LoggedExceptionHandler handler = new LoggedExceptionHandler(
+		protected SpringBootExceptionHandler initialValue() {
+			SpringBootExceptionHandler handler = new SpringBootExceptionHandler(
 					Thread.currentThread().getUncaughtExceptionHandler());
 			Thread.currentThread().setUncaughtExceptionHandler(handler);
 			return handler;

spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -82,6 +82,7 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.mockito.Matchers.anyObject;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
@@ -552,13 +553,37 @@ public class SpringApplicationTests {
 		this.context = application.run();
 		assertNotNull(this.context);
 		assertEquals(2, SpringApplication.exit(this.context, new ExitCodeGenerator() {
+
 			@Override
 			public int getExitCode() {
 				return 2;
 			}
+
 		}));
 	}
 
+	@Test
+	public void exitWithExplicitCodeFromException() throws Exception {
+		final SpringBootExceptionHandler handler = mock(SpringBootExceptionHandler.class);
+		SpringApplication application = new SpringApplication(
+				ExitCodeCommandLineRunConfig.class) {
+
+			@Override
+			SpringBootExceptionHandler getSpringBootExceptionHandler() {
+				return handler;
+			}
+
+		};
+		application.setWebEnvironment(false);
+		try {
+			application.run();
+			fail("Did not throw");
+		}
+		catch (IllegalStateException ex) {
+		}
+		verify(handler).registerExitCode(11);
+	}
+
 	@Test
 	public void defaultCommandLineArgs() throws Exception {
 		SpringApplication application = new SpringApplication(ExampleConfig.class);
@@ -858,6 +883,34 @@ public class SpringApplicationTests {
 		public TestCommandLineRunner runnerA() {
 			return new TestCommandLineRunner(Ordered.HIGHEST_PRECEDENCE);
 		}
+
+	}
+
+	@Configuration
+	static class ExitCodeCommandLineRunConfig {
+
+		@Bean
+		public CommandLineRunner runner() {
+			return new CommandLineRunner() {
+
+				@Override
+				public void run(String... args) throws Exception {
+					throw new IllegalStateException(new ExitStatusException());
+				}
+
+			};
+		}
+
+	}
+
+	static class ExitStatusException extends RuntimeException
+			implements ExitCodeGenerator {
+
+		@Override
+		public int getExitCode() {
+			return 11;
+		}
+
 	}
 
 	static class AbstractTestRunner implements ApplicationContextAware, Ordered {