Prefer @SpringBootApplication-annotated class when finding main class

Closes gh-6496

spring-boot-tools/spring-boot-antlib/src/main/java/org/springframework/boot/ant/FindMainClass.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.
@@ -35,6 +35,8 @@ import org.springframework.util.StringUtils;
  */
 public class FindMainClass extends Task {
 
+	private static final String SPRING_BOOT_APPLICATION_CLASS_NAME = "org.springframework.boot.autoconfigure.SpringBootApplication";
+
 	private String mainClass;
 
 	private File classesRoot;
@@ -70,10 +72,11 @@ public class FindMainClass extends Task {
 		}
 		try {
 			if (this.classesRoot.isDirectory()) {
-				return MainClassFinder.findSingleMainClass(this.classesRoot);
+				return MainClassFinder.findSingleMainClass(this.classesRoot,
+						SPRING_BOOT_APPLICATION_CLASS_NAME);
 			}
-			return MainClassFinder.findSingleMainClass(new JarFile(this.classesRoot),
-					"/");
+			return MainClassFinder.findSingleMainClass(new JarFile(this.classesRoot), "/",
+					SPRING_BOOT_APPLICATION_CLASS_NAME);
 		}
 		catch (IOException ex) {
 			throw new BuildException(ex);

spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/run/FindMainClassTask.java

@@ -41,6 +41,8 @@ import org.springframework.boot.loader.tools.MainClassFinder;
  */
 public class FindMainClassTask extends DefaultTask {
 
+	private static final String SPRING_BOOT_APPLICATION_CLASS_NAME = "org.springframework.boot.autoconfigure.SpringBootApplication";
+
 	@Input
 	private SourceSetOutput mainClassSourceSetOutput;
 
@@ -105,7 +107,8 @@ public class FindMainClassTask extends DefaultTask {
 						+ this.mainClassSourceSetOutput.getClassesDir());
 				try {
 					mainClass = MainClassFinder.findSingleMainClass(
-							this.mainClassSourceSetOutput.getClassesDir());
+							this.mainClassSourceSetOutput.getClassesDir(),
+							SPRING_BOOT_APPLICATION_CLASS_NAME);
 					project.getLogger().info("Computed main class: " + mainClass);
 				}
 				catch (IOException ex) {

spring-boot-tools/spring-boot-loader-tools/src/main/java/org/springframework/boot/loader/tools/MainClassFinder.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.
@@ -29,12 +29,14 @@ import java.util.Collections;
 import java.util.Comparator;
 import java.util.Deque;
 import java.util.Enumeration;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 
+import org.springframework.asm.AnnotationVisitor;
 import org.springframework.asm.ClassReader;
 import org.springframework.asm.ClassVisitor;
 import org.springframework.asm.MethodVisitor;
@@ -46,6 +48,7 @@ import org.springframework.asm.Type;
  * search.
  *
  * @author Phillip Webb
+ * @author Andy Wilkinson
  */
 public abstract class MainClassFinder {
 
@@ -79,24 +82,39 @@ public abstract class MainClassFinder {
 	 * @throws IOException if the folder cannot be read
 	 */
 	public static String findMainClass(File rootFolder) throws IOException {
-		return doWithMainClasses(rootFolder, new ClassNameCallback<String>() {
+		return doWithMainClasses(rootFolder, new MainClassCallback<String>() {
 			@Override
-			public String doWith(String className) {
-				return className;
+			public String doWith(MainClass mainClass) {
+				return mainClass.getName();
 			}
 		});
 	}
 
 	/**
-	 * Find a single main class from a given folder.
+	 * Find a single main class from the given {@code rootFolder}.
 	 * @param rootFolder the root folder to search
 	 * @return the main class or {@code null}
 	 * @throws IOException if the folder cannot be read
 	 */
 	public static String findSingleMainClass(File rootFolder) throws IOException {
-		MainClassesCallback callback = new MainClassesCallback();
+		return findSingleMainClass(rootFolder, null);
+	}
+
+	/**
+	 * Find a single main class from the given {@code rootFolder}. A main class annotated
+	 * with an annotation with the given {@code annotationName} will be preferred over a
+	 * main class with no such annotation.
+	 * @param rootFolder the root folder to search
+	 * @param annotationName the name of the annotation that may be present on the main
+	 * class
+	 * @return the main class or {@code null}
+	 * @throws IOException if the folder cannot be read
+	 */
+	public static String findSingleMainClass(File rootFolder, String annotationName)
+			throws IOException {
+		SingleMainClassCallback callback = new SingleMainClassCallback(annotationName);
 		MainClassFinder.doWithMainClasses(rootFolder, callback);
-		return callback.getMainClass();
+		return callback.getMainClassName();
 	}
 
 	/**
@@ -108,7 +126,7 @@ public abstract class MainClassFinder {
 	 * @return the first callback result or {@code null}
 	 * @throws IOException in case of I/O errors
 	 */
-	static <T> T doWithMainClasses(File rootFolder, ClassNameCallback<T> callback)
+	static <T> T doWithMainClasses(File rootFolder, MainClassCallback<T> callback)
 			throws IOException {
 		if (!rootFolder.exists()) {
 			return null; // nothing to do
@@ -125,10 +143,12 @@ public abstract class MainClassFinder {
 			if (file.isFile()) {
 				InputStream inputStream = new FileInputStream(file);
 				try {
-					if (isMainClass(inputStream)) {
+					ClassDescriptor classDescriptor = createClassDescriptor(inputStream);
+					if (classDescriptor != null && classDescriptor.isMainMethodFound()) {
 						String className = convertToClassName(file.getAbsolutePath(),
 								prefix);
-						T result = callback.doWith(className);
+						T result = callback.doWith(new MainClass(className,
+								classDescriptor.getAnnotationNames()));
 						if (result != null) {
 							return result;
 						}
@@ -168,10 +188,10 @@ public abstract class MainClassFinder {
 	public static String findMainClass(JarFile jarFile, String classesLocation)
 			throws IOException {
 		return doWithMainClasses(jarFile, classesLocation,
-				new ClassNameCallback<String>() {
+				new MainClassCallback<String>() {
 					@Override
-					public String doWith(String className) {
-						return className;
+					public String doWith(MainClass mainClass) {
+						return mainClass.getName();
 					}
 				});
 	}
@@ -185,9 +205,25 @@ public abstract class MainClassFinder {
 	 */
 	public static String findSingleMainClass(JarFile jarFile, String classesLocation)
 			throws IOException {
-		MainClassesCallback callback = new MainClassesCallback();
+		return findSingleMainClass(jarFile, classesLocation, null);
+	}
+
+	/**
+	 * Find a single main class in a given jar file. A main class annotated with an
+	 * annotation with the given {@code annotationName} will be preferred over a main
+	 * class with no such annotation.
+	 * @param jarFile the jar file to search
+	 * @param classesLocation the location within the jar containing classes
+	 * @param annotationName the name of the annotation that may be present on the main
+	 * class
+	 * @return the main class or {@code null}
+	 * @throws IOException if the jar file cannot be read
+	 */
+	public static String findSingleMainClass(JarFile jarFile, String classesLocation,
+			String annotationName) throws IOException {
+		SingleMainClassCallback callback = new SingleMainClassCallback(annotationName);
 		MainClassFinder.doWithMainClasses(jarFile, classesLocation, callback);
-		return callback.getMainClass();
+		return callback.getMainClassName();
 	}
 
 	/**
@@ -200,17 +236,19 @@ public abstract class MainClassFinder {
 	 * @throws IOException in case of I/O errors
 	 */
 	static <T> T doWithMainClasses(JarFile jarFile, String classesLocation,
-			ClassNameCallback<T> callback) throws IOException {
+			MainClassCallback<T> callback) throws IOException {
 		List<JarEntry> classEntries = getClassEntries(jarFile, classesLocation);
 		Collections.sort(classEntries, new ClassEntryComparator());
 		for (JarEntry entry : classEntries) {
 			InputStream inputStream = new BufferedInputStream(
 					jarFile.getInputStream(entry));
 			try {
-				if (isMainClass(inputStream)) {
+				ClassDescriptor classDescriptor = createClassDescriptor(inputStream);
+				if (classDescriptor != null && classDescriptor.isMainMethodFound()) {
 					String className = convertToClassName(entry.getName(),
 							classesLocation);
-					T result = callback.doWith(className);
+					T result = callback.doWith(new MainClass(className,
+							classDescriptor.getAnnotationNames()));
 					if (result != null) {
 						return result;
 					}
@@ -248,15 +286,15 @@ public abstract class MainClassFinder {
 		return classEntries;
 	}
 
-	private static boolean isMainClass(InputStream inputStream) {
+	private static ClassDescriptor createClassDescriptor(InputStream inputStream) {
 		try {
 			ClassReader classReader = new ClassReader(inputStream);
-			MainMethodFinder mainMethodFinder = new MainMethodFinder();
-			classReader.accept(mainMethodFinder, ClassReader.SKIP_CODE);
-			return mainMethodFinder.isFound();
+			ClassDescriptor classDescriptor = new ClassDescriptor();
+			classReader.accept(classDescriptor, ClassReader.SKIP_CODE);
+			return classDescriptor;
 		}
 		catch (IOException ex) {
-			return false;
+			return null;
 		}
 	}
 
@@ -279,21 +317,29 @@ public abstract class MainClassFinder {
 
 	}
 
-	private static class MainMethodFinder extends ClassVisitor {
+	private static class ClassDescriptor extends ClassVisitor {
 
-		private boolean found;
+		private final Set<String> annotationNames = new LinkedHashSet<String>();
 
-		MainMethodFinder() {
+		private boolean mainMethodFound;
+
+		ClassDescriptor() {
 			super(Opcodes.ASM4);
 		}
 
+		@Override
+		public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
+			this.annotationNames.add(Type.getType(desc).getClassName());
+			return null;
+		}
+
 		@Override
 		public MethodVisitor visitMethod(int access, String name, String desc,
 				String signature, String[] exceptions) {
 			if (isAccess(access, Opcodes.ACC_PUBLIC, Opcodes.ACC_STATIC)
 					&& MAIN_METHOD_NAME.equals(name)
 					&& MAIN_METHOD_TYPE.getDescriptor().equals(desc)) {
-				this.found = true;
+				this.mainMethodFound = true;
 			}
 			return null;
 		}
@@ -307,24 +353,90 @@ public abstract class MainClassFinder {
 			return true;
 		}
 
-		public boolean isFound() {
-			return this.found;
+		boolean isMainMethodFound() {
+			return this.mainMethodFound;
+		}
+
+		Set<String> getAnnotationNames() {
+			return this.annotationNames;
 		}
 
 	}
 
 	/**
-	 * Callback interface used to receive class names.
-	 * @param <T> the result type
+	 * Callback for handling {@link MainClass MainClasses}.
+	 *
+	 * @param <T> the callback's return type
 	 */
-	public interface ClassNameCallback<T> {
+	interface MainClassCallback<T> {
 
 		/**
-		 * Handle the specified class name.
-		 * @param className the class name
+		 * Handle the specified main class.
+		 * @param mainClass the mainClass
 		 * @return a non-null value if processing should end or {@code null} to continue
 		 */
-		T doWith(String className);
+		T doWith(MainClass mainClass);
+
+	}
+
+	/**
+	 * A class with a {@code main} method.
+	 */
+	static final class MainClass {
+
+		private final String name;
+
+		private final Set<String> annotationNames;
+
+		/**
+		 * Creates a new {@code MainClass} rather represents the main class with the given
+		 * {@code name}. The class is annotated with the annotations with the given
+		 * {@code annotationNames}.
+		 *
+		 * @param name the name of the class
+		 * @param annotationNames the names of the annotations on the class
+		 */
+		MainClass(String name, Set<String> annotationNames) {
+			this.name = name;
+			this.annotationNames = Collections
+					.unmodifiableSet(new HashSet<String>(annotationNames));
+		}
+
+		String getName() {
+			return this.name;
+		}
+
+		Set<String> getAnnotationNames() {
+			return this.annotationNames;
+		}
+
+		@Override
+		public String toString() {
+			return this.name;
+		}
+
+		@Override
+		public int hashCode() {
+			return this.name.hashCode();
+		}
+
+		@Override
+		public boolean equals(Object obj) {
+			if (this == obj) {
+				return true;
+			}
+			if (obj == null) {
+				return false;
+			}
+			if (getClass() != obj.getClass()) {
+				return false;
+			}
+			MainClass other = (MainClass) obj;
+			if (!this.name.equals(other.name)) {
+				return false;
+			}
+			return true;
+		}
 
 	}
 
@@ -332,23 +444,42 @@ public abstract class MainClassFinder {
 	 * Find a single main class, throwing an {@link IllegalStateException} if multiple
 	 * candidates exist.
 	 */
-	private static class MainClassesCallback implements ClassNameCallback<Object> {
+	private static final class SingleMainClassCallback
+			implements MainClassCallback<Object> {
 
-		private final Set<String> classNames = new LinkedHashSet<String>();
+		private final Set<MainClass> mainClasses = new LinkedHashSet<MainClass>();
+
+		private final String annotationName;
+
+		private SingleMainClassCallback(String annotationName) {
+			this.annotationName = annotationName;
+		}
 
 		@Override
-		public Object doWith(String className) {
-			this.classNames.add(className);
+		public Object doWith(MainClass mainClass) {
+			this.mainClasses.add(mainClass);
 			return null;
 		}
 
-		public String getMainClass() {
-			if (this.classNames.size() > 1) {
+		private String getMainClassName() {
+			Set<MainClass> matchingMainClasses = new LinkedHashSet<MainClass>();
+			if (this.annotationName != null) {
+				for (MainClass mainClass : this.mainClasses) {
+					if (mainClass.getAnnotationNames().contains(this.annotationName)) {
+						matchingMainClasses.add(mainClass);
+					}
+				}
+			}
+			if (matchingMainClasses.isEmpty()) {
+				matchingMainClasses.addAll(this.mainClasses);
+			}
+			if (matchingMainClasses.size() > 1) {
 				throw new IllegalStateException(
 						"Unable to find a single main class from the following candidates "
-								+ this.classNames);
+								+ matchingMainClasses);
 			}
-			return this.classNames.isEmpty() ? null : this.classNames.iterator().next();
+			return matchingMainClasses.isEmpty() ? null
+					: matchingMainClasses.iterator().next().getName();
 		}
 
 	}

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

@@ -59,6 +59,8 @@ public class Repackager {
 
 	private static final long FIND_WARNING_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
 
+	private static final String SPRING_BOOT_APPLICATION_CLASS_NAME = "org.springframework.boot.autoconfigure.SpringBootApplication";
+
 	private List<MainClassTimeoutWarningListener> mainClassTimeoutListeners = new ArrayList<MainClassTimeoutWarningListener>();
 
 	private String mainClass;
@@ -383,7 +385,7 @@ public class Repackager {
 
 	protected String findMainMethod(JarFile source) throws IOException {
 		return MainClassFinder.findSingleMainClass(source,
-				this.layout.getClassesLocation());
+				this.layout.getClassesLocation(), SPRING_BOOT_APPLICATION_CLASS_NAME);
 	}
 
 	private void renameFile(File file, File dest) {

spring-boot-tools/spring-boot-loader-tools/src/test/java/org/springframework/boot/loader/tools/MainClassFinderTests.java

@@ -26,7 +26,9 @@ import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.junit.rules.TemporaryFolder;
 
-import org.springframework.boot.loader.tools.MainClassFinder.ClassNameCallback;
+import org.springframework.boot.loader.tools.MainClassFinder.MainClass;
+import org.springframework.boot.loader.tools.MainClassFinder.MainClassCallback;
+import org.springframework.boot.loader.tools.sample.AnnotatedClassWithMainMethod;
 import org.springframework.boot.loader.tools.sample.ClassWithMainMethod;
 import org.springframework.boot.loader.tools.sample.ClassWithoutMainMethod;
 
@@ -87,6 +89,16 @@ public class MainClassFinderTests {
 		MainClassFinder.findSingleMainClass(this.testJarFile.getJarFile(), "");
 	}
 
+	@Test
+	public void findSingleJarSearchPrefersAnnotatedMainClass() throws Exception {
+		this.testJarFile.addClass("a/B.class", ClassWithMainMethod.class);
+		this.testJarFile.addClass("a/b/c/E.class", AnnotatedClassWithMainMethod.class);
+		String mainClass = MainClassFinder.findSingleMainClass(
+				this.testJarFile.getJarFile(), "",
+				"org.springframework.boot.loader.tools.sample.SomeApplication");
+		assertThat(mainClass).isEqualTo("a.b.c.E");
+	}
+
 	@Test
 	public void findMainClassInJarSubLocation() throws Exception {
 		this.testJarFile.addClass("a/B.class", ClassWithMainMethod.class);
@@ -132,6 +144,16 @@ public class MainClassFinderTests {
 		MainClassFinder.findSingleMainClass(this.testJarFile.getJarSource());
 	}
 
+	@Test
+	public void findSingleFolderSearchPrefersAnnotatedMainClass() throws Exception {
+		this.testJarFile.addClass("a/B.class", ClassWithMainMethod.class);
+		this.testJarFile.addClass("a/b/c/E.class", AnnotatedClassWithMainMethod.class);
+		String mainClass = MainClassFinder.findSingleMainClass(
+				this.testJarFile.getJarSource(),
+				"org.springframework.boot.loader.tools.sample.SomeApplication");
+		assertThat(mainClass).isEqualTo("a.b.c.E");
+	}
+
 	@Test
 	public void doWithFolderMainMethods() throws Exception {
 		this.testJarFile.addClass("a/b/c/D.class", ClassWithMainMethod.class);
@@ -150,17 +172,17 @@ public class MainClassFinderTests {
 		this.testJarFile.addClass("a/b/F.class", ClassWithoutMainMethod.class);
 		this.testJarFile.addClass("a/b/G.class", ClassWithMainMethod.class);
 		ClassNameCollector callback = new ClassNameCollector();
-		MainClassFinder.doWithMainClasses(this.testJarFile.getJarFile(), "", callback);
+		MainClassFinder.doWithMainClasses(this.testJarFile.getJarFile(), null, callback);
 		assertThat(callback.getClassNames().toString()).isEqualTo("[a.b.G, a.b.c.D]");
 	}
 
-	private static class ClassNameCollector implements ClassNameCallback<Object> {
+	private static class ClassNameCollector implements MainClassCallback<Object> {
 
 		private final List<String> classNames = new ArrayList<String>();
 
 		@Override
-		public Object doWith(String className) {
-			this.classNames.add(className);
+		public Object doWith(MainClass mainClass) {
+			this.classNames.add(mainClass.getName());
 			return null;
 		}
 

spring-boot-tools/spring-boot-loader-tools/src/test/java/org/springframework/boot/loader/tools/sample/AnnotatedClassWithMainMethod.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.loader.tools.sample;
+
+/**
+ * Sample annotated class with a main method.
+ *
+ * @author Andy Wilkinson
+ */
+@SomeApplication
+public class AnnotatedClassWithMainMethod {
+
+	public void run() {
+		System.out.println("Hello World");
+	}
+
+	public static void main(String[] args) {
+		new AnnotatedClassWithMainMethod().run();
+	}
+
+}

spring-boot-tools/spring-boot-loader-tools/src/test/java/org/springframework/boot/loader/tools/sample/SomeApplication.java

@@ -0,0 +1,33 @@
+/*
+ * 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.loader.tools.sample;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Test annotation for a main application class.
+ *
+ * @author Andy Wilkinson
+ */
+@Target(ElementType.TYPE)
+@Retention(RetentionPolicy.RUNTIME)
+@interface SomeApplication {
+
+}

spring-boot-tools/spring-boot-maven-plugin/src/main/java/org/springframework/boot/maven/AbstractRunMojo.java

@@ -52,6 +52,8 @@ import org.springframework.boot.loader.tools.MainClassFinder;
  */
 public abstract class AbstractRunMojo extends AbstractDependencyFilterMojo {
 
+	private static final String SPRING_BOOT_APPLICATION_CLASS_NAME = "org.springframework.boot.autoconfigure.SpringBootApplication";
+
 	private static final String SPRING_LOADED_AGENT_CLASS_NAME = "org.springsource.loaded.agent.SpringLoadedAgent";
 
 	/**
@@ -374,7 +376,8 @@ public abstract class AbstractRunMojo extends AbstractDependencyFilterMojo {
 		String mainClass = this.mainClass;
 		if (mainClass == null) {
 			try {
-				mainClass = MainClassFinder.findSingleMainClass(this.classesDirectory);
+				mainClass = MainClassFinder.findSingleMainClass(this.classesDirectory,
+						SPRING_BOOT_APPLICATION_CLASS_NAME);
 			}
 			catch (IOException ex) {
 				throw new MojoExecutionException(ex.getMessage(), ex);