diff --git a/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/run/FindMainClassTask.java b/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/run/FindMainClassTask.java
index c3d2470a92..ce35c01c41 100644
--- a/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/run/FindMainClassTask.java
+++ b/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/run/FindMainClassTask.java
@@ -16,11 +16,16 @@
 
 package org.springframework.boot.gradle.run;
 
+import java.io.File;
 import java.io.IOException;
+import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.Collection;
 
 import org.gradle.api.DefaultTask;
 import org.gradle.api.Project;
 import org.gradle.api.Task;
+import org.gradle.api.file.FileCollection;
 import org.gradle.api.plugins.ApplicationPluginConvention;
 import org.gradle.api.plugins.ExtraPropertiesExtension;
 import org.gradle.api.tasks.Input;
@@ -30,6 +35,7 @@ import org.gradle.api.tasks.TaskAction;
 
 import org.springframework.boot.gradle.SpringBootPluginExtension;
 import org.springframework.boot.loader.tools.MainClassFinder;
+import org.springframework.util.ReflectionUtils;
 
 /**
  * Task to find and set the 'mainClassName' convention when it's missing by searching the
@@ -103,13 +109,13 @@ public class FindMainClassTask extends DefaultTask {
 		if (mainClass == null) {
 			// Search
 			if (this.mainClassSourceSetOutput != null) {
-				project.getLogger().debug("Looking for main in: "
-						+ this.mainClassSourceSetOutput.getClassesDir());
+				Collection<File> classesDirs = getClassesDirs(
+						this.mainClassSourceSetOutput);
+				getProject().getLogger().debug("Looking for main in: " + classesDirs);
 				try {
-					mainClass = MainClassFinder.findSingleMainClass(
-							this.mainClassSourceSetOutput.getClassesDir(),
+					mainClass = MainClassFinder.findSingleMainClass(classesDirs,
 							SPRING_BOOT_APPLICATION_CLASS_NAME);
-					project.getLogger().info("Computed main class: " + mainClass);
+					getProject().getLogger().info("Computed main class: " + mainClass);
 				}
 				catch (IOException ex) {
 					throw new IllegalStateException("Cannot find main class", ex);
@@ -132,6 +138,17 @@ public class FindMainClassTask extends DefaultTask {
 		return mainClass;
 	}
 
+	private Collection<File> getClassesDirs(SourceSetOutput sourceSetOutput) {
+		Method getClassesDirs = ReflectionUtils.findMethod(SourceSetOutput.class,
+				"getClassesDirs");
+		if (getClassesDirs == null) {
+			return Arrays.asList(sourceSetOutput.getClassesDir());
+		}
+		FileCollection classesDirs = (FileCollection) ReflectionUtils
+				.invokeMethod(getClassesDirs, sourceSetOutput);
+		return classesDirs.getFiles();
+	}
+
 	private JavaExec findRunTask(Project project) {
 		Task runTask = project.getTasks().findByName("run");
 		if (runTask instanceof JavaExec) {
diff --git a/spring-boot-tools/spring-boot-loader-tools/src/main/java/org/springframework/boot/loader/tools/MainClassFinder.java b/spring-boot-tools/spring-boot-loader-tools/src/main/java/org/springframework/boot/loader/tools/MainClassFinder.java
index f46257618f..429caa3df8 100644
--- a/spring-boot-tools/spring-boot-loader-tools/src/main/java/org/springframework/boot/loader/tools/MainClassFinder.java
+++ b/spring-boot-tools/spring-boot-loader-tools/src/main/java/org/springframework/boot/loader/tools/MainClassFinder.java
@@ -25,6 +25,7 @@ import java.io.InputStream;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.Deque;
@@ -117,6 +118,44 @@ public abstract class MainClassFinder {
 		return callback.getMainClassName();
 	}
 
+	/**
+	 * Find a single main class from the given {@code rootFolders}. A main class annotated
+	 * with an annotation with the given {@code annotationName} will be preferred over a
+	 * main class with no such annotation.
+	 * @param rootFolders the root folders 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 a root folder cannot be read
+	 * @since 1.5.5
+	 */
+	public static String findSingleMainClass(Collection<File> rootFolders,
+			String annotationName) throws IOException {
+		SingleMainClassCallback callback = new SingleMainClassCallback(annotationName);
+		doWithMainClasses(rootFolders, callback);
+		return callback.getMainClassName();
+	}
+
+	/**
+	 * Perform the given callback operation on all main classes from the given root
+	 * folder.
+	 * @param <T> the result type
+	 * @param rootFolders the root folders
+	 * @param callback the callback
+	 * @return the first callback result or {@code null}
+	 * @throws IOException in case of I/O errors
+	 */
+	static <T> T doWithMainClasses(Collection<File> rootFolders,
+			MainClassCallback<T> callback) throws IOException {
+		for (File rootFolder : rootFolders) {
+			T result = doWithMainClasses(rootFolder, callback);
+			if (result != null) {
+				return result;
+			}
+		}
+		return null;
+	}
+
 	/**
 	 * Perform the given callback operation on all main classes from the given root
 	 * folder.