Introduce DepJarPlugin based on AssemblyPlugin and use it in SparkBuild.

3 files changed, 116 insertions, 36 deletions

diff --git a/project/DepJar.scala b/project/DepJar.scala
new file mode 100644
index 0000000000..1d54005690
--- /dev/null
+++ b/project/DepJar.scala
@@ -0,0 +1,108 @@
+import sbt._
+import Keys._
+import java.io.PrintWriter
+import scala.collection.mutable
+import scala.io.Source
+import Project.Initialize
+
+/*
+ * This is based on the AssemblyPlugin. For now it was easier to copy and modify than to wait for
+ * the required changes needed for us to customise it so that it does what we want. We may revisit
+ * this in the future.
+ */
+object DepJarPlugin extends Plugin {
+  val DepJar = config("dep-jar") extend(Runtime)
+  val depJar = TaskKey[File]("dep-jar", "Builds a single-file jar of all dependencies.")
+
+  val jarName           = SettingKey[String]("jar-name")
+  val outputPath        = SettingKey[File]("output-path")
+  val excludedFiles     = SettingKey[Seq[File] => Seq[File]]("excluded-files")
+  val conflictingFiles  = SettingKey[Seq[File] => Seq[File]]("conflicting-files")
+
+  private def assemblyTask: Initialize[Task[File]] =
+    (test, packageOptions, cacheDirectory, outputPath,
+        fullClasspath, excludedFiles, conflictingFiles, streams) map {
+      (test, options, cacheDir, jarPath, cp, exclude, conflicting, s) =>
+        IO.withTemporaryDirectory { tempDir =>
+          val srcs = assemblyPaths(tempDir, cp, exclude, conflicting, s.log)
+          val config = new Package.Configuration(srcs, jarPath, options)
+          Package(config, cacheDir, s.log)
+          jarPath
+        }
+    }
+
+  private def assemblyPackageOptionsTask: Initialize[Task[Seq[PackageOption]]] =
+    (packageOptions in Compile, mainClass in DepJar) map { (os, mainClass) =>
+      mainClass map { s =>
+        os find { o => o.isInstanceOf[Package.MainClass] } map { _ => os
+        } getOrElse { Package.MainClass(s) +: os }
+      } getOrElse {os}
+    }
+
+  private def assemblyExcludedFiles(base: Seq[File]): Seq[File] = {
+    ((base / "scala" ** "*") +++ // exclude scala library
+      (base / "spark" ** "*") +++ // exclude Spark classes
+      ((base / "META-INF" ** "*") --- // generally ignore the hell out of META-INF
+        (base / "META-INF" / "services" ** "*") --- // include all service providers
+        (base / "META-INF" / "maven" ** "*"))).get // include all Maven POMs and such
+  }
+
+  private def assemblyPaths(tempDir: File, classpath: Classpath,
+      exclude: Seq[File] => Seq[File], conflicting: Seq[File] => Seq[File], log: Logger) = {
+    import sbt.classpath.ClasspathUtilities
+
+    val (libs, directories) = classpath.map(_.data).partition(ClasspathUtilities.isArchive)
+    val services = mutable.Map[String, mutable.ArrayBuffer[String]]()
+    for(jar <- libs) {
+      val jarName = jar.asFile.getName
+      log.info("Including %s".format(jarName))
+      IO.unzip(jar, tempDir)
+      IO.delete(conflicting(Seq(tempDir)))
+      val servicesDir = tempDir / "META-INF" / "services"
+      if (servicesDir.asFile.exists) {
+       for (service <- (servicesDir ** "*").get) {
+         val serviceFile = service.asFile
+         if (serviceFile.exists && serviceFile.isFile) {
+           val entries = services.getOrElseUpdate(serviceFile.getName, new mutable.ArrayBuffer[String]())
+           for (provider <- Source.fromFile(serviceFile).getLines) {
+             if (!entries.contains(provider)) {
+               entries += provider
+             }
+           }
+         }
+       }
+     }
+    }
+
+    for ((service, providers) <- services) {
+      log.debug("Merging providers for %s".format(service))
+      val serviceFile = (tempDir / "META-INF" / "services" / service).asFile
+      val writer = new PrintWriter(serviceFile)
+      for (provider <- providers.map { _.trim }.filter { !_.isEmpty }) {
+        log.debug("-  %s".format(provider))
+        writer.println(provider)
+      }
+      writer.close()
+    }
+
+    val base = tempDir +: directories
+    val descendants = ((base ** (-DirectoryFilter)) --- exclude(base)).get
+    descendants x relativeTo(base)
+  }
+
+  lazy val depJarSettings = inConfig(DepJar)(Seq(
+    depJar <<= packageBin.identity,
+    packageBin <<= assemblyTask,
+    jarName <<= (name, version) { (name, version) => name + "-dep-" + version + ".jar" },
+    outputPath <<= (target, jarName) { (t, s) => t / s },
+    test <<= (test in Test).identity,
+    mainClass <<= (mainClass in Runtime).identity,
+    fullClasspath <<= (fullClasspath in Runtime).identity,
+    packageOptions <<= assemblyPackageOptionsTask,
+    excludedFiles := assemblyExcludedFiles _,
+    conflictingFiles := assemblyExcludedFiles _
+  )) ++
+  Seq(
+    depJar <<= (depJar in DepJar).identity
+  )
+}
\ No newline at end of file
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index cd2e3211fe..22279b7919 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -30,7 +30,7 @@ object SparkBuild extends Build {
 
   val slf4jVersion = "1.6.1"
 
-  //FIXME DepJar and XmlTestReport
+  //FIXME XmlTestReport
   def coreSettings = sharedSettings ++ Seq(libraryDependencies ++= Seq(
     "com.google.guava" % "guava" % "r09",
     "log4j" % "log4j" % "1.2.16",
@@ -39,15 +39,17 @@ object SparkBuild extends Build {
     "com.ning" % "compress-lzf" % "0.7.0",
     "org.apache.hadoop" % "hadoop-core" % "0.20.2",
     "asm" % "asm-all" % "3.3.1"
-  ))
+  )) ++ DepJarPlugin.depJarSettings
 
-  //FIXME DepJar and XmlTestReport
-  def replSettings = sharedSettings ++ Seq(libraryDependencies <+= scalaVersion("org.scala-lang" % "scala-compiler" % _))
+  //FIXME XmlTestReport
+  def replSettings = sharedSettings ++
+      Seq(libraryDependencies <+= scalaVersion("org.scala-lang" % "scala-compiler" % _)) ++
+      DepJarPlugin.depJarSettings
 
   def examplesSettings = sharedSettings ++ Seq(libraryDependencies += "colt" % "colt" % "1.2.0")
 
-  //FIXME DepJar and XmlTestReport
-  def bagelSettings = sharedSettings
+  //FIXME XmlTestReport
+  def bagelSettings = sharedSettings ++ DepJarPlugin.depJarSettings
 }
 
 // Project mixin for an XML-based ScalaTest report. Unfortunately
@@ -71,31 +73,3 @@ object SparkBuild extends Build {
 //    None
 //  }.dependsOn(compile, testCompile).describedAs("Generate XML test report.")
 //}
-
-// Project mixin for creating a JAR with  a project's dependencies. This is based
-// on the AssemblyBuilder plugin, but because this plugin attempts to package Scala
-// and our project too, we leave that out using our own exclude filter (depJarExclude).
-//trait DepJar extends AssemblyBuilder {
-//  def depJarExclude(base: PathFinder) = {
-//    (base / "scala" ** "*") +++ // exclude scala library
-//    (base / "spark" ** "*") +++ // exclude Spark classes
-//    ((base / "META-INF" ** "*") --- // generally ignore the hell out of META-INF
-//     (base / "META-INF" / "services" ** "*") --- // include all service providers
-//     (base / "META-INF" / "maven" ** "*")) // include all Maven POMs and such
-//  }
-//
-//  def depJarTempDir = outputPath / "dep-classes"
-//
-//  def depJarOutputPath =
-//    outputPath / (name.toLowerCase.replace(" ", "-") + "-dep-" + version.toString + ".jar")
-//
-//  lazy val depJar = {
-//    packageTask(
-//      Path.lazyPathFinder(assemblyPaths(depJarTempDir,
-//                                        assemblyClasspath,
-//                                        assemblyExtraJars,
-//                                        depJarExclude)),
-//      depJarOutputPath,
-//      packageOptions)
-//  }.dependsOn(compile).describedAs("Bundle project's dependencies into a JAR.")
-//}
diff --git a/project/plugins/build.sbt b/project/plugins/build.sbt
index 8585d2c6a1..91c6cb4df1 100644
--- a/project/plugins/build.sbt
+++ b/project/plugins/build.sbt
@@ -11,5 +11,3 @@ libraryDependencies += "com.github.mpeltonen" %% "sbt-idea" % "0.10.0"
 libraryDependencies <<= (libraryDependencies, sbtVersion) { (deps, version) =>
   deps :+ ("com.typesafe.sbteclipse" %% "sbteclipse" % "1.2" extra("sbtversion" -> version))
 }
-
-libraryDependencies <+= (sbtVersion) { sv => "com.eed3si9n" %% "sbt-assembly" % ("sbt" + sv + "_0.3") }
\ No newline at end of file