Add ExitCodeExceptionMapper support

Add ExitCodeExceptionMapper strategy interface which can be used to map
exceptions to exit codes.

Closes gh-4803

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

@@ -0,0 +1,35 @@
+/*
+ * 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 org.springframework.boot;
+
+/**
+ * Strategy interface that can be used to provide a mapping between exceptions and exit
+ * codes.
+ *
+ * @author Phillip Webb
+ * @since 1.3.2
+ */
+public interface ExitCodeExceptionMapper {
+
+	/**
+	 * Returns the exit code that should be returned from the application.
+	 * @param exception the exception causing the application to exit
+	 * @return the exit code or {@code 0}.
+	 */
+	int getExitCode(Throwable exception);
+
+}

spring-boot/src/main/java/org/springframework/boot/ExitCodeGenerators.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.
@@ -36,6 +36,27 @@ class ExitCodeGenerators implements Iterable<ExitCodeGenerator> {
 
 	private List<ExitCodeGenerator> generators = new ArrayList<ExitCodeGenerator>();
 
+	public void addAll(Throwable exception, ExitCodeExceptionMapper... mappers) {
+		Assert.notNull(exception, "Exception must not be null");
+		Assert.notNull(mappers, "Mappers must not be null");
+		addAll(exception, Arrays.asList(mappers));
+	}
+
+	public void addAll(Throwable exception,
+			Iterable<? extends ExitCodeExceptionMapper> mappers) {
+		Assert.notNull(exception, "Exception must not be null");
+		Assert.notNull(mappers, "Mappers must not be null");
+		for (ExitCodeExceptionMapper mapper : mappers) {
+			add(exception, mapper);
+		}
+	}
+
+	public void add(Throwable exception, ExitCodeExceptionMapper mapper) {
+		Assert.notNull(exception, "Exception must not be null");
+		Assert.notNull(mapper, "Mapper must not be null");
+		add(new MappedExitCodeGenerator(exception, mapper));
+	}
+
 	public void addAll(ExitCodeGenerator... generators) {
 		Assert.notNull(generators, "Generators must not be null");
 		addAll(Arrays.asList(generators));
@@ -79,4 +100,25 @@ class ExitCodeGenerators implements Iterable<ExitCodeGenerator> {
 		return exitCode;
 	}
 
+	/**
+	 * Adapts an {@link ExitCodeExceptionMapper} to an {@link ExitCodeGenerator}.
+	 */
+	private static class MappedExitCodeGenerator implements ExitCodeGenerator {
+
+		private final Throwable exception;
+
+		private final ExitCodeExceptionMapper mapper;
+
+		MappedExitCodeGenerator(Throwable exception, ExitCodeExceptionMapper mapper) {
+			this.exception = exception;
+			this.mapper = mapper;
+		}
+
+		@Override
+		public int getExitCode() {
+			return this.mapper.getExitCode(this.exception);
+		}
+
+	}
+
 }

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

@@ -858,7 +858,7 @@ public class SpringApplication {
 
 	private void handeExitCode(ConfigurableApplicationContext context,
 			Throwable exception) {
-		int exitCode = getExitCodeFromException(exception);
+		int exitCode = getExitCodeFromException(context, exception);
 		if (exitCode != 0) {
 			if (context != null) {
 				context.publishEvent(new ExitCodeEvent(context, exitCode));
@@ -870,14 +870,32 @@ public class SpringApplication {
 		}
 	}
 
-	private int getExitCodeFromException(Throwable exception) {
+	private int getExitCodeFromException(ConfigurableApplicationContext context,
+			Throwable exception) {
+		int exitCode = getExitCodeFromMappedException(context, exception);
+		if (exitCode == 0) {
+			exitCode = getExitCodeFromExitCodeGeneratorException(exception);
+		}
+		return exitCode;
+	}
+
+	private int getExitCodeFromMappedException(ConfigurableApplicationContext context,
+			Throwable exception) {
+		ExitCodeGenerators generators = new ExitCodeGenerators();
+		Collection<ExitCodeExceptionMapper> beans = context
+				.getBeansOfType(ExitCodeExceptionMapper.class).values();
+		generators.addAll(exception, beans);
+		return generators.getExitCode();
+	}
+
+	private int getExitCodeFromExitCodeGeneratorException(Throwable exception) {
 		if (exception == null) {
 			return 0;
 		}
 		if (exception instanceof ExitCodeGenerator) {
 			return ((ExitCodeGenerator) exception).getExitCode();
 		}
-		return getExitCodeFromException(exception.getCause());
+		return getExitCodeFromExitCodeGeneratorException(exception.getCause());
 	}
 
 	SpringBootExceptionHandler getSpringBootExceptionHandler() {

spring-boot/src/test/java/org/springframework/boot/ExitCodeGeneratorsTests.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.
@@ -16,6 +16,7 @@
 
 package org.springframework.boot;
 
+import java.io.IOException;
 import java.util.List;
 
 import org.junit.Rule;
@@ -84,10 +85,36 @@ public class ExitCodeGeneratorsTests {
 		assertThat(generators.getExitCode(), equalTo(3));
 	}
 
+	@Test
+	public void getExitCodeWhenUsingExitCodeExceptionMapperShouldCallMapper()
+			throws Exception {
+		ExitCodeGenerators generators = new ExitCodeGenerators();
+		Exception e = new IOException();
+		generators.add(e, mockMapper(IllegalStateException.class, 1));
+		generators.add(e, mockMapper(IOException.class, 2));
+		generators.add(e, mockMapper(UnsupportedOperationException.class, 3));
+		assertThat(generators.getExitCode(), equalTo(2));
+	}
+
 	private ExitCodeGenerator mockGenerator(int exitCode) {
 		ExitCodeGenerator generator = mock(ExitCodeGenerator.class);
 		given(generator.getExitCode()).willReturn(exitCode);
 		return generator;
 	}
 
+	private ExitCodeExceptionMapper mockMapper(final Class<?> exceptionType,
+			final int exitCode) {
+		return new ExitCodeExceptionMapper() {
+
+			@Override
+			public int getExitCode(Throwable exception) {
+				if (exceptionType.isInstance(exception)) {
+					return exitCode;
+				}
+				return 0;
+			}
+
+		};
+	}
+
 }

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

@@ -590,6 +590,31 @@ public class SpringApplicationTests {
 		assertThat(listener.getExitCode(), equalTo(11));
 	}
 
+	@Test
+	public void exitWithExplicitCodeFromMappedException() throws Exception {
+		final SpringBootExceptionHandler handler = mock(SpringBootExceptionHandler.class);
+		SpringApplication application = new SpringApplication(
+				MappedExitCodeCommandLineRunConfig.class) {
+
+			@Override
+			SpringBootExceptionHandler getSpringBootExceptionHandler() {
+				return handler;
+			}
+
+		};
+		ExitCodeListener listener = new ExitCodeListener();
+		application.addListeners(listener);
+		application.setWebEnvironment(false);
+		try {
+			application.run();
+			fail("Did not throw");
+		}
+		catch (IllegalStateException ex) {
+		}
+		verify(handler).registerExitCode(11);
+		assertThat(listener.getExitCode(), equalTo(11));
+	}
+
 	@Test
 	public void defaultCommandLineArgs() throws Exception {
 		SpringApplication application = new SpringApplication(ExampleConfig.class);
@@ -909,6 +934,38 @@ public class SpringApplicationTests {
 
 	}
 
+	@Configuration
+	static class MappedExitCodeCommandLineRunConfig {
+
+		@Bean
+		public CommandLineRunner runner() {
+			return new CommandLineRunner() {
+
+				@Override
+				public void run(String... args) throws Exception {
+					throw new IllegalStateException();
+				}
+
+			};
+		}
+
+		@Bean
+		public ExitCodeExceptionMapper exceptionMapper() {
+			return new ExitCodeExceptionMapper() {
+
+				@Override
+				public int getExitCode(Throwable exception) {
+					if (exception instanceof IllegalStateException) {
+						return 11;
+					}
+					return 0;
+				}
+
+			};
+		}
+
+	}
+
 	static class ExitStatusException extends RuntimeException
 			implements ExitCodeGenerator {