From 5f7683b5b81e32e7265a944c1e867e252023d7b8 Mon Sep 17 00:00:00 2001
From: Nicholas Chammas <nicholas.chammas@gmail.com>
Date: Tue, 2 Dec 2014 18:39:56 -0500
Subject: [PATCH] parallelize tests

---
 project/SparkBuild.scala | 71 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 68 insertions(+), 3 deletions(-)

diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index e4b1b96527fbd..abc4ad8b8b48d 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -23,6 +23,7 @@ import scala.collection.JavaConversions._
 import sbt._
 import sbt.Classpaths.publishTask
 import sbt.Keys._
+import sbt.Tests._
 import sbtunidoc.Plugin.genjavadocSettings
 import sbtunidoc.Plugin.UnidocKeys.unidocGenjavadocVersion
 import com.typesafe.sbt.pom.{loadEffectivePom, PomBuild, SbtPomKeys}
@@ -416,13 +417,44 @@ object Unidoc {
 object TestSettings {
   import BuildCommons._
 
+  def singleTests(tests: Seq[TestDefinition]) =
+    tests map { test =>
+      new Group(
+        name = test.name,
+        tests = Seq(test),
+        runPolicy = SubProcess(javaOptions = Seq.empty[String]))
+    }
+
+  def singleForkedTests(tests: Seq[TestDefinition], javaOptions: Seq[String]) =
+    tests map { test =>
+      new Group(
+        name = test.name,
+        tests = Seq(test),
+        runPolicy = SubProcess(javaOptions = javaOptions))
+    }
+
+  def groupBySuite(tests: Seq[TestDefinition], javaOptions: Seq[String]) = {
+    // println(javaOptions.mkString("\n"))
+    tests groupBy (_.name.split('.').slice(0,5).mkString(".")) map {
+      case (suite, tests) =>
+        new Group(
+          name = suite,
+          tests = tests,
+          // runPolicy = Tests.InProcess)
+          runPolicy = SubProcess(javaOptions = javaOptions))
+    } toSeq
+  }
+
   lazy val settings = Seq (
     // Fork new JVMs for tests and set Java options for those
     fork := true,
+    fork in Test := true,
+    
     // Setting SPARK_DIST_CLASSPATH is a simple way to make sure any child processes
     // launched by the tests have access to the correct test-time classpath.
     envVars in Test += ("SPARK_DIST_CLASSPATH" ->
       (fullClasspath in Test).value.files.map(_.getAbsolutePath).mkString(":").stripSuffix(":")),
+    
     javaOptions in Test += "-Dspark.test.home=" + sparkHome,
     javaOptions in Test += "-Dspark.testing=1",
     javaOptions in Test += "-Dspark.port.maxRetries=100",
@@ -436,14 +468,47 @@ object TestSettings {
     javaOptions in Test ++= "-Xmx3g -XX:PermSize=128M -XX:MaxNewSize=256m -XX:MaxPermSize=1g"
       .split(" ").toSeq,
     javaOptions += "-Xmx3g",
+
     // Show full stack trace and duration in test cases.
     testOptions in Test += Tests.Argument("-oDF"),
     testOptions += Tests.Argument(TestFrameworks.JUnit, "-v", "-a"),
     // Enable Junit testing.
     libraryDependencies += "com.novocode" % "junit-interface" % "0.9" % "test",
-    // Only allow one test at a time, even across projects, since they run in the same JVM
-    parallelExecution in Test := false,
-    concurrentRestrictions in Global += Tags.limit(Tags.Test, 1),
+
+    // Parallelize execution of tests
+    parallelExecution := true,
+    parallelExecution in Test := true,
+    testForkedParallel in Test := false,
+
+    concurrentRestrictions in Global := Seq(
+      Tags.limit(Tags.Compile, 8),
+      Tags.limit(Tags.ForkedTestGroup, 8)
+    ),
+
+    testGrouping in Test := ((definedTests in Test, javaOptions in Test) map groupBySuite).value,
+    testGrouping in Test := {
+      val original: Seq[Tests.Group] = (testGrouping in Test).value
+
+      original.map { group =>
+        val forkOptions = ForkOptions(
+          bootJars = Nil,
+          javaHome = javaHome.value,
+          connectInput = connectInput.value,
+          // outputStrategy = outputStrategy.value,
+          runJVMOptions = (javaOptions in Test).value,
+          // workingDirectory = Some(new File(System.getProperty("user.dir"))),
+          workingDirectory = Some(baseDirectory.value),
+          envVars = envVars.value
+        )
+
+        group.copy(runPolicy = Tests.SubProcess(forkOptions))
+      }
+    },
+    // testGrouping in Test <<= (definedTests in Test, javaOptions in Test) map groupBySuite,
+    // testGrouping in Test <<= definedTests in Test map singleTests,
+    // testGrouping in Test <<= (definedTests in Test, javaOptions in Test) map singleForkedTests,
+    logBuffered in Test := true,
+
     // Remove certain packages from Scaladoc
     scalacOptions in (Compile, doc) := Seq(
       "-groups",