Add JMH contrib module (#2100)

This adds the script from #182 (with minor modifications) as a module under contrib, along with tests and documentation.

Pull request: #2100

build.sc

@@ -903,6 +903,15 @@ object contrib extends MillModule {
     )
   }
 
+  object jmh extends MillInternalModule with MillAutoTestSetup with WithMillCompiler {
+    override def compileModuleDeps = Seq(scalalib)
+    override def testArgs = T {
+      Seq(
+        "-DMILL_SCALA_LIB=" + scalalib.runClasspath().map(_.path).mkString(",")
+      ) ++ scalalib.worker.testArgs()
+    }
+    override def testModuleDeps: Seq[JavaModule] = super.testModuleDeps ++ Seq(scalalib)
+  }
 }
 
 object scalanativelib extends MillModule {

contrib/jmh/src/mill/contrib/jmh/JmhModule.scala

@@ -0,0 +1,101 @@
+package mill.contrib.jmh
+
+import mill._, scalalib._, modules._
+
+/**
+ * This module provides an easy way to integrate <a href="https://openjdk.org/projects/code-tools/jmh/">JMH</a> benchmarking with Mill.
+ *
+ * Example configuration:
+ * {{{
+ * import mill._, scalalib._
+ *
+ * import $ivy.`com.lihaoyi::mill-contrib-jmh:$MILL_VERSION`
+ * import contrib.jmh.JmhModule
+ *
+ * object foo extends ScalaModule with JmhModule {
+ *   def scalaVersion = "2.13.8"
+ *   def jmhCoreVersion = "1.35"
+ * }
+ * }}}
+ *
+ * Here are some sample commands:
+ * - mill foo.runJmh             # Runs all detected jmh benchmarks
+ * - mill foo.listJmhBenchmarks  # List detected jmh benchmarks
+ * - mill foo.runJmh -h          # List available arguments to runJmh
+ * - mill foo.runJmh regexp      # Run all benchmarks matching `regexp`
+ *
+ * For Scala JMH samples see:
+ * [[https://github.com/sbt/sbt-jmh/tree/main/plugin/src/sbt-test/sbt-jmh/run/src/main/scala/org/openjdk/jmh/samples]].
+ */
+trait JmhModule extends JavaModule {
+
+  def jmhCoreVersion: T[String]
+  def jmhGeneratorByteCodeVersion: T[String] = jmhCoreVersion
+
+  def ivyDeps = super.ivyDeps() ++ Agg(ivy"org.openjdk.jmh:jmh-core:${jmhCoreVersion()}")
+
+  def runJmh(args: String*) =
+    T.command {
+      val (_, resources) = generateBenchmarkSources()
+      Jvm.runSubprocess(
+        "org.openjdk.jmh.Main",
+        classPath = (runClasspath() ++ generatorDeps()).map(_.path) ++
+          Seq(compileGeneratedSources().path, resources),
+        mainArgs = args,
+        workingDir = T.ctx().dest
+      )
+    }
+
+  def listJmhBenchmarks(args: String*) = runJmh(("-l" +: args): _*)
+
+  def compileGeneratedSources =
+    T {
+      val dest = T.ctx().dest
+      val (sourcesDir, _) = generateBenchmarkSources()
+      val sources = os.walk(sourcesDir).filter(os.isFile)
+
+      os.proc(
+        Jvm.jdkTool("javac"),
+        sources.map(_.toString),
+        "-cp",
+        (runClasspath() ++ generatorDeps()).map(_.path.toString).mkString(
+          java.io.File.pathSeparator
+        ),
+        "-d",
+        dest
+      ).call(dest)
+      PathRef(dest)
+    }
+
+  // returns sources and resources directories
+  def generateBenchmarkSources =
+    T {
+      val dest = T.ctx().dest
+
+      val sourcesDir = dest / "jmh_sources"
+      val resourcesDir = dest / "jmh_resources"
+
+      os.remove.all(sourcesDir)
+      os.makeDir.all(sourcesDir)
+      os.remove.all(resourcesDir)
+      os.makeDir.all(resourcesDir)
+
+      Jvm.runSubprocess(
+        "org.openjdk.jmh.generators.bytecode.JmhBytecodeGenerator",
+        (runClasspath() ++ generatorDeps()).map(_.path),
+        mainArgs = Seq(
+          compile().classes.path.toString,
+          sourcesDir.toString,
+          resourcesDir.toString,
+          "default"
+        )
+      )
+
+      (sourcesDir, resourcesDir)
+    }
+
+  def generatorDeps =
+    resolveDeps(
+      T { Agg(ivy"org.openjdk.jmh:jmh-generator-bytecode:${jmhGeneratorByteCodeVersion()}") }
+    )
+}

contrib/jmh/test/resources/jmh/src/Bench1.scala

@@ -0,0 +1,34 @@
+package mill.contrib.jmh
+
+import org.openjdk.jmh.annotations._
+
+object Bench1States {
+
+  @State(Scope.Benchmark)
+  class BenchmarkState {
+    @volatile
+    var x = Math.PI
+  }
+
+  @State(Scope.Thread)
+  class ThreadState {
+    @volatile
+    var x = Math.PI
+  }
+}
+
+@BenchmarkMode(Array(Mode.All))
+class Bench1 {
+
+  import Bench1States._
+
+  @Benchmark
+  def measureShared(state: BenchmarkState) = {
+    state.x += 1
+  }
+
+  @Benchmark
+  def measureUnshared(state: ThreadState) = {
+    state.x += 1
+  }
+}

contrib/jmh/test/resources/jmh/src/Bench2.scala

@@ -0,0 +1,18 @@
+package mill.contrib.jmh
+
+import org.openjdk.jmh.annotations._
+import java.util.concurrent.TimeUnit
+
+@BenchmarkMode(Array(Mode.AverageTime))
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@State(Scope.Thread)
+class Bench2 {
+
+  val x = Math.PI
+
+  @Benchmark
+  def sqrt: Double = Math.sqrt(x)
+
+  @Benchmark
+  def log: Double = Math.log(x)
+}

contrib/jmh/test/src/JmhModuleTest.scala

@@ -0,0 +1,51 @@
+package mill
+package contrib.jmh
+
+import mill.api.PathRef
+import mill.eval.EvaluatorPaths
+import mill.scalalib.ScalaModule
+import mill.util.{TestEvaluator, TestUtil}
+import os.Path
+import utest._
+import utest.framework.TestPath
+
+object JmhModuleTest extends TestSuite {
+
+  object jmh extends TestUtil.BaseModule with ScalaModule with JmhModule {
+
+    override def scalaVersion = sys.props.getOrElse("TEST_SCALA_2_13_VERSION", ???)
+    override def jmhCoreVersion = "1.35"
+    override def millSourcePath = TestUtil.getSrcPathBase() / millOuterCtx.enclosing.split('.')
+  }
+
+  val testModuleSourcesPath: Path =
+    os.pwd / "contrib" / "jmh" / "test" / "resources" / "jmh"
+
+  private def workspaceTest(m: TestUtil.BaseModule)(t: TestEvaluator => Unit)(
+      implicit tp: TestPath
+  ): Unit = {
+    val eval = new TestEvaluator(m)
+    os.remove.all(m.millSourcePath)
+    os.remove.all(eval.outPath)
+    os.makeDir.all(m.millSourcePath / os.up)
+    os.copy(testModuleSourcesPath, m.millSourcePath)
+    t(eval)
+  }
+
+  def tests = Tests {
+    test("jmh") {
+      "listJmhBenchmarks" - workspaceTest(jmh) { eval =>
+        val paths = EvaluatorPaths.resolveDestPaths(eval.outPath, jmh.listJmhBenchmarks())
+        val outFile = paths.dest / "benchmarks.out"
+        val Right((result, _)) = eval(jmh.listJmhBenchmarks("-o", outFile.toString))
+        val expected = """Benchmarks:
+                         |mill.contrib.jmh.Bench2.log
+                         |mill.contrib.jmh.Bench2.sqrt
+                         |mill.contrib.jmh.Bench1.measureShared
+                         |mill.contrib.jmh.Bench1.measureUnshared""".stripMargin
+        val out = os.read.lines(outFile).map(_.trim).mkString(System.lineSeparator())
+        assert(out == expected)
+      }
+    }
+  }
+}

docs/antora/modules/ROOT/pages/Contrib_Plugins.adoc

@@ -38,6 +38,7 @@ import $ivy.`com.lihaoyi::mill-contrib-bloop:`
 * xref:Plugin_Docker.adoc[]
 * xref:Plugin_Flyway.adoc[]
 * xref:Plugin_Gitlab.adoc[]
+* xref:Plugin_Jmh.adoc[]
 * xref:Plugin_Play.adoc[]
 * xref:Plugin_Proguard.adoc[]
 * xref:Plugin_ScalaPB.adoc[]

docs/antora/modules/ROOT/pages/Plugin_Jmh.adoc

@@ -0,0 +1,31 @@
+= JMH
+
+You can use `JmhModule` to integrate JMH testing with Mill.
+
+Example configuration:
+
+.`build.sc`
+[source,scala]
+----
+import mill._, scalalib._
+
+import $ivy.`com.lihaoyi::mill-contrib-jmh:$MILL_VERSION`
+import contrib.jmh.JmhModule
+
+object foo extends ScalaModule with JmhModule {
+  def scalaVersion = "2.13.8"
+  def jmhCoreVersion = "1.35"
+}
+----
+
+Here are some sample commands:
+
+[source,bash]
+----
+mill foo.runJmh             # Runs all detected jmh benchmarks
+mill foo.listJmhBenchmarks  # List detected jmh benchmarks
+mill foo.runJmh -h          # List available arguments to runJmh
+mill foo.runJmh regexp      # Run all benchmarks matching `regexp`
+----
+
+For Scala JMH samples see https://github.com/sbt/sbt-jmh/tree/main/plugin/src/sbt-test/sbt-jmh/run/src/main/scala/org/openjdk/jmh/samples[sbt-jmh].