A basic vertx plugin.

VertxPlugin.scala

/*
 * Copyright (c) 2013 Plausible Labs Cooperative, Inc.
 * All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */

import java.net.URLClassLoader

import java.util.zip.{ZipEntry, ZipOutputStream}
import org.vertx.java.core.json.JsonObject
import org.vertx.java.platform.impl.ModuleIdentifier
import org.vertx.java.platform.{PlatformManager, PlatformLocator}
import sbt._
import sbt.classpath.ClasspathUtilities
import sbt.Keys._
import complete.DefaultParsers._

import org.vertx.java.core.AsyncResult
import org.vertx.java.core.Handler

import java.nio.file._

import sbt.Task
import scala.collection.JavaConverters._

import scala.concurrent.duration.Duration
import scala.concurrent.{Await, promise}

object VertxPlugin extends Plugin {
  val vertxConfig = config("vertx")

  object VertxKeys {
    val vertxAssembly         = TaskKey[Unit]("vertx-assembly", "Builds a single-file deployable vert.x jar.")
    val vertxRun              = InputKey[Unit]("vertx-run", "Run a vert.x module in place.")

    val vertxModuleName       = TaskKey[String]("vertx-module-name", "The vert.x module name.")
    val vertxJarName          = TaskKey[String]("vertx-jar-name", "The jar name to be used when assemblying a self-contained executable jar")
    val vertxOutputPath       = TaskKey[File]("vertx-output-path", "The output path to be used when assembling vert.x artifacts.")

    val vertxModuleDirectory  = SettingKey[File]("vertx-module-directory", "The vert.x module runtime directory.")
  }

  import VertxKeys._

  lazy val vertxSettings: Seq[Setting[_]] = Seq[Setting[_]](
    vertxModuleName := organization.value + "~" + name.value + "~" + version.value,

    /* Assembly */
    vertxJarName := name.value + "-fat-" + version.value + ".jar",
    vertxOutputPath := (crossTarget in Compile).value / vertxJarName.value,
    test in vertxConfig := (test in Test).value,

    /* Runtime configuration */
    vertxModuleDirectory <<= (crossTarget in Compile) { _ / "vertx-mods" },

    /* Tasks */
    vertxRun <<= runModuleTask(fullClasspath in Runtime, resourceDirectories in Compile),
    vertxAssembly <<= packageModuleTask(fullClasspath in Runtime)
  )

  /**
   * Execute (and wait on) an asynchronous vert.x task.
   *
   * @param moduleDirectory The directory to use for the vertx.mods system property.
   * @param classpath The target process' classpath.
   * @param log The logger to use for this task.
   * @param action The action to execute.
   */
  private def vertxExecute[T] (moduleDirectory:File, classpath:Seq[File], log:Logger, running:(()=>Unit), action:(PlatformManager, Handler[AsyncResult[T]])=>Unit): Boolean = {
    /* Sadly we can only do this globally */
    System.setProperty("vertx.mods", moduleDirectory.getAbsolutePath)

    /*
     * Carve out a class loader distinct from sbt's.
     * This fixes a number of issues related to ServiceLoader interaction
     * with sbt's class loader, and also allows us to directly utilize the target's classpath.
     */
    val classpathArray = classpath.map(_.toURI.toURL).toArray
    val originalClassLoader = getClass.getClassLoader
    val newClassLoader = new URLClassLoader(classpathArray, originalClassLoader)

    Thread.currentThread().setContextClassLoader(newClassLoader);
    try {
      val pm = PlatformLocator.factory.createPlatformManager()
      val result = promise[Either[Throwable, Boolean]]
      action(pm, new Handler[AsyncResult[T]] {
        override def handle (event: AsyncResult[T]) = {
          if (event.succeeded()) {
            /* We inform the caller that the module started, but we leave our
             * future hanging. A future implementation could instead trigger the future
             * upon module termination. */
            running()
          } else {
            result.success(Left(event.cause()))
          }
        }
      })

      Await.result(result.future, Duration.Inf) match {
        case Right(didStart) => didStart
        case Left(failure) => {
          log.error(s"Failed to execute vert.x task: $failure (cause: ${failure.getCause})")
          false
        }
      }
    } finally {
      /* Restore the SBT class loader */
      Thread.currentThread().setContextClassLoader(originalClassLoader);
    }
  }

  /**
   * Create, use, and clean up a resource.
   *
   * @param resource A function that returns the resource to be closed.
   * @param cleanup The fuction to use to close the resource
   * @param action The function that will make use of the resource.
   * @tparam A The resource type.
   * @tparam B The result type.
   * @return
   */
  private def cleanly[A,B](resource: => A)(cleanup: A => Unit)(action: A => B): B = {
    /* Fetch a single instance of the resource */
    val r = resource
    try {
      action(r)
    } finally {
      cleanup(r)
    }
  }

  /**
   * Implementation of the 'vertx-assembly' task
   * @param classpath The full class path for the project.
   */
  private def packageModuleTask (classpath: Def.Initialize[Task[Classpath]]) = Def.task {
    val log = streams.value.log

    /* The full class path */
    val classpaths:Seq[File] = classpath.value.files

    /* The vert.x module directory */
    val modDir = vertxModuleDirectory.value

    /* Our module name */
    val modName = vertxModuleName.value

    /* The directory to which we'll write the module */
    val modOutputDir = modDir / modName

    /* The generated jar file */
    val jarPath:File = vertxOutputPath.value

    /* Clear out the existing module directory */
    IO.delete(modOutputDir)
    IO.createDirectory(modOutputDir)

    /* Populate the vert.x module directory */
    val (libs, dirs) = classpaths.partition(ClasspathUtilities.isArchive)
    for (dir <- dirs) {
      IO.copyDirectory(dir, modOutputDir, preserveLastModified = true)
    }

    for (lib <- libs) {
      IO.copyFile(lib, modOutputDir / "lib" / lib.getName, preserveLastModified = true)
    }

    /* Pull in all nested modules */
    vertxExecute(modDir, classpaths, log, () => {
      /* Dependency fetching is running */
      // do nothing
    }, (pm, handler:Handler[AsyncResult[Void]]) => {
      pm.pullInDependencies(modName, handler)
    }) match {
      case true => log.info("Fetched module dependencies.")
      case false => throw new RuntimeException("Failed to fetch module dependencies.")
    }

    /* Let vert.x perform the final jar packaging */
    vertxExecute(modDir, classpaths, log, () => {
      /* Packaging is running */
      // do nothing
    }, (pm, handler:Handler[AsyncResult[Void]]) => {
        pm.makeFatJar(modName, streams.value.cacheDirectory.getAbsolutePath, handler)
    }) match {
      case true => {
        /* XXX: the vert.x implementation of makeFatJar() hard-codes the jar name based on the
         * module name. We have to manually move the result back to where it belongs. */
        val modID = new ModuleIdentifier(modName)
        val hardcodedJarPath = streams.value.cacheDirectory / s"${modID.getName}-${modID.getVersion}-fat.jar"
        IO.move(hardcodedJarPath, jarPath)

        /*
         * vert.x 2.1M1's fat loader uses the platform class loader to find lang.properties,
         * which means we have to insert the file as a jar visible to said class loader.
         *
         * This has been fixed in vert.x master by allowing the language settings to be specified directly
         * in a module's own mod.json file.
         */
        val langsFile:File = modOutputDir / "langs.properties"
        if (langsFile.exists) {
          val fsEnv = Map[String, Object]("create" -> "true").asJava
          val jarURI = new URI("jar:file", null, jarPath.getAbsolutePath, null)

          cleanly (FileSystems.newFileSystem(jarURI, fsEnv)) (_.close) (jarFS => {
            val libPath = jarFS.getPath("lib/vertx-lang-fat-configuration-injection.jar")
            cleanly (Files.newOutputStream(libPath, StandardOpenOption.CREATE)) (_.close) (os => {
              val zipper = new ZipOutputStream(os)
              zipper.putNextEntry(new ZipEntry("langs.properties"))
              Files.copy(langsFile.toPath, zipper)
              zipper.closeEntry()
              zipper.finish()
            })
          })
        }

        log.info(s"Wrote assembly to $jarPath.")
      }
      case false => log.error("Failed to write assembly.")
    }

  }

  private val vertxModuleNameParser = OptSpace ~> StringBasic.examples("<arg>").?

  /**
   * Implementation of the 'runmod' task.
   */
  private def runModuleTask(classpath: Def.Initialize[Task[Classpath]], resourceDirectories:SettingKey[Seq[File]]): Def.Initialize[InputTask[Unit]] = Def.inputTask {
    val mod = vertxModuleNameParser.parsed
    val instances = 1 // TODO - make configurable
    val config = new JsonObject() // TODO - make configurable
    val log = streams.value.log

    val classpaths:Seq[File] = resourceDirectories.value ++ classpath.value.files

    vertxExecute(vertxModuleDirectory.value, classpaths, log, () => {
      log.info("CTRL-C to stop vert.x server")
    }, (pm, handler:Handler[AsyncResult[String]]) => {
      pm.deployModuleFromClasspath(mod.getOrElse(vertxModuleName.value), config, instances, classpaths.map(_.toURI.toURL).toArray, handler)
    })
  }
}