Basic support for profiling JMH benchmarks (#187)

* Add poc of JMH support

* update readme

* update workflow

.github/workflows/site.yml

@@ -22,7 +22,7 @@ jobs:
       with:
         fetch-depth: '0'
     - name: Setup Scala
-      uses: actions/setup-java@v3.11.0
+      uses: actions/setup-java@v3.9.0
       with:
         distribution: temurin
         java-version: 17
@@ -45,7 +45,7 @@ jobs:
       with:
         fetch-depth: '0'
     - name: Setup Scala
-      uses: actions/setup-java@v3.11.0
+      uses: actions/setup-java@v3.9.0
       with:
         distribution: temurin
         java-version: 17
@@ -70,7 +70,7 @@ jobs:
         ref: ${{ github.head_ref }}
         fetch-depth: '0'
     - name: Setup Scala
-      uses: actions/setup-java@v3.11.0
+      uses: actions/setup-java@v3.9.0
       with:
         distribution: temurin
         java-version: 17
@@ -84,7 +84,7 @@ jobs:
         git add README.md
         git commit -m "Update README.md" || echo "No changes to commit"
     - name: Create Pull Request
-      uses: peter-evans/create-pull-request@v5.0.0
+      uses: peter-evans/create-pull-request@v4.2.3
       with:
         body: |-
           Autogenerated changes after running the `sbt docs/generateReadme` command of the [zio-sbt-website](https://zio.dev/zio-sbt) plugin.

README.md

@@ -128,6 +128,69 @@ val testEffect = ZIO.unit
 val testEffect = CostCenter.withChildCostCenter("foo.Foo.testEffect(Foo.scala:12)")(ZIO.unit)
 ```
 
+## Jmh Support
+
+ZIO Profiling offers an integration with the Java Microbenchmark Harness (JMH). In order to profile a jmh benchmark, first ensure that the sources are properly tagged using the tagging plugin. Next, add a dependency to the jmh module to your benchmarking module:
+```scala
+libraryDependencies += "dev.zio" %% "zio-profiling-jmh" % "0.1.2"
+```
+
+In your actual benchmarks, ensure that you are running ZIO effects using the methods in `zio.profiling.jmh.BenchmarkUtils`. A possible benchmark might look like this
+```scala
+package zio.redis.benchmarks.lists
+
+import org.openjdk.jmh.annotations._
+import zio.profiling.jmh.BenchmarkUtils
+import zio.redis._
+import zio.redis.benchmarks._
+import zio.{Scope => _, _}
+
+import java.util.concurrent.TimeUnit
+
+@State(Scope.Thread)
+@BenchmarkMode(Array(Mode.Throughput))
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Measurement(iterations = 15)
+@Warmup(iterations = 15)
+@Fork(2)
+class BlMoveBenchmarks extends BenchmarkRuntime {
+
+  @Param(Array("500"))
+  var count: Int = _
+
+  private var items: List[String] = _
+
+  private val key = "test-list"
+
+  private def execute(query: ZIO[Redis, RedisError, Unit]): Unit =
+    BenchmarkUtils.unsafeRun(query.provideLayer(BenchmarkRuntime.Layer))
+
+  @Setup(Level.Trial)
+  def setup(): Unit = {
+    items = (0 to count).toList.map(_.toString)
+    execute(ZIO.serviceWithZIO[Redis](_.rPush(key, items.head, items.tail: _*).unit))
+  }
+
+  @TearDown(Level.Trial)
+  def tearDown(): Unit =
+    execute(ZIO.serviceWithZIO[Redis](_.del(key).unit))
+
+  @Benchmark
+  def zio(): Unit = execute(
+    ZIO.foreachDiscard(items)(_ =>
+      ZIO.serviceWithZIO[Redis](_.blMove(key, key, Side.Left, Side.Right, 1.second).returning[String])
+    )
+  )
+}
+```
+
+Once the benchmark is set up properly, you can specify the profiler from the jmh command line. Using sbt-jmh, it might look like this:
+```
+Jmh/run -i 3 -wi 3 -f1 -t1 -prof zio.profiling.jmh.JmhZioProfiler zio.redis.benchmarks.lists.BlMoveBenchmarks.zio
+```
+
+The profiler output will be written to a file in the directory the JVM has been invoked from.
+
 ## Documentation
 
 Learn more on the [ZIO Profiling homepage](https://zio.dev/zio-profiling/)!

build.sbt

@@ -28,7 +28,7 @@ addCommandAlias("prepare", "fix; fmt")
 lazy val root = project
   .in(file("."))
   .settings(publish / skip := true)
-  .aggregate(core, taggingPlugin, examples, benchmarks, docs)
+  .aggregate(core, jmh, taggingPlugin, examples, benchmarks, docs)
 
 lazy val core = project
   .in(file("zio-profiling"))
@@ -43,6 +43,16 @@ lazy val core = project
     testFrameworks += new TestFramework("zio.test.sbt.ZTestFramework")
   )
 
+lazy val jmh = project
+  .in(file("zio-profiling-jmh"))
+  .dependsOn(core)
+  .settings(
+    stdSettings("zio-profiling-jmh"),
+    libraryDependencies ++= Seq(
+      "org.openjdk.jmh" % "jmh-core" % jmhVersion
+    )
+  )
+
 lazy val taggingPlugin = project
   .in(file("zio-profiling-tagging-plugin"))
   .settings(
@@ -51,15 +61,13 @@ lazy val taggingPlugin = project
     pluginDefinitionSettings
   )
 
-lazy val taggingPluginJar = taggingPlugin / Compile / packageTask
-
 lazy val examples = project
   .in(file("examples"))
   .dependsOn(core, taggingPlugin % "plugin")
   .settings(
     stdSettings("examples"),
     publish / skip := true,
-    scalacOptions += s"-Xplugin:${taggingPluginJar.value.getAbsolutePath}"
+    scalacOptions += s"-Xplugin:${(taggingPlugin / Compile / packageTask).value.getAbsolutePath}"
   )
 
 lazy val benchmarks = project

docs/index.md

@@ -127,3 +127,66 @@ val testEffect = ZIO.unit
 
 val testEffect = CostCenter.withChildCostCenter("foo.Foo.testEffect(Foo.scala:12)")(ZIO.unit)
 ```
+
+## Jmh Support
+
+ZIO Profiling offers an integration with the Java Microbenchmark Harness (JMH). In order to profile a jmh benchmark, first ensure that the sources are properly tagged using the tagging plugin. Next, add a dependency to the jmh module to your benchmarking module:
+```scala
+libraryDependencies += "dev.zio" %% "zio-profiling-jmh" % "@VERSION@"
+```
+
+In your actual benchmarks, ensure that you are running ZIO effects using the methods in `zio.profiling.jmh.BenchmarkUtils`. A possible benchmark might look like this
+```scala
+package zio.redis.benchmarks.lists
+
+import org.openjdk.jmh.annotations._
+import zio.profiling.jmh.BenchmarkUtils
+import zio.redis._
+import zio.redis.benchmarks._
+import zio.{Scope => _, _}
+
+import java.util.concurrent.TimeUnit
+
+@State(Scope.Thread)
+@BenchmarkMode(Array(Mode.Throughput))
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Measurement(iterations = 15)
+@Warmup(iterations = 15)
+@Fork(2)
+class BlMoveBenchmarks extends BenchmarkRuntime {
+
+  @Param(Array("500"))
+  var count: Int = _
+
+  private var items: List[String] = _
+
+  private val key = "test-list"
+
+  private def execute(query: ZIO[Redis, RedisError, Unit]): Unit =
+    BenchmarkUtils.unsafeRun(query.provideLayer(BenchmarkRuntime.Layer))
+
+  @Setup(Level.Trial)
+  def setup(): Unit = {
+    items = (0 to count).toList.map(_.toString)
+    execute(ZIO.serviceWithZIO[Redis](_.rPush(key, items.head, items.tail: _*).unit))
+  }
+
+  @TearDown(Level.Trial)
+  def tearDown(): Unit =
+    execute(ZIO.serviceWithZIO[Redis](_.del(key).unit))
+
+  @Benchmark
+  def zio(): Unit = execute(
+    ZIO.foreachDiscard(items)(_ =>
+      ZIO.serviceWithZIO[Redis](_.blMove(key, key, Side.Left, Side.Right, 1.second).returning[String])
+    )
+  )
+}
+```
+
+Once the benchmark is set up properly, you can specify the profiler from the jmh command line. Using sbt-jmh, it might look like this:
+```
+Jmh/run -i 3 -wi 3 -f1 -t1 -prof zio.profiling.jmh.JmhZioProfiler zio.redis.benchmarks.lists.BlMoveBenchmarks.zio
+```
+
+The profiler output will be written to a file in the directory the JVM has been invoked from.

project/BuildHelper.scala

@@ -51,7 +51,7 @@ object BuildHelper {
           )
       },
       semanticdbEnabled := scalaVersion.value == defaulScalaVersion,
-      semanticdbOptions += "-P:semanticdb:synthetics:on",
+      semanticdbOptions ++= (if (scalaVersion.value != Scala3) List("-P:semanticdb:synthetics:on") else Nil),
       semanticdbVersion                                          := scalafixSemanticdb.revision,
       ThisBuild / scalafixScalaBinaryVersion                     := CrossVersion.binaryScalaVersion(scalaVersion.value),
       ThisBuild / scalafixDependencies += "com.github.liancheng" %% "organize-imports" % organizeImportsVersion,

project/Dependencies.scala

@@ -1,7 +1,7 @@
 object Dependencies {
   val collectionCompatVersion = "2.8.1"
+  val jmhVersion              = "1.36"
   val organizeImportsVersion  = "0.6.0"
   val silencerVersion         = "1.7.12"
-
-  val zioVersion = "2.0.10"
+  val zioVersion              = "2.0.10"
 }

zio-profiling-jmh/src/main/resources/META-INF/services/org.openjdk.jmh.profile.Profiler

@@ -0,0 +1 @@
+zio.profiling.jmh.JmhZioProfiler

zio-profiling-jmh/src/main/scala/zio/profiling/jmh/BenchmarkUtils.scala

@@ -0,0 +1,22 @@
+package zio.profiling.jmh
+
+import zio._
+import zio.profiling.sampling.SamplingProfilerSupervisor
+
+import java.util.concurrent.atomic.AtomicReference
+
+object BenchmarkUtils {
+
+  private[jmh] val runtimeRef: AtomicReference[Runtime.Scoped[SamplingProfilerSupervisor]] = new AtomicReference()
+
+  def getRuntime(): Runtime[Any] = {
+    val customRt = runtimeRef.get()
+    if (customRt ne null) customRt else Runtime.default
+  }
+
+  def unsafeRun[E, A](zio: ZIO[Any, E, A]): A =
+    Unsafe.unsafe { implicit unsafe =>
+      getRuntime().unsafe.run(zio).getOrThrowFiberFailure()
+    }
+
+}

zio-profiling-jmh/src/main/scala/zio/profiling/jmh/JmhZioProfiler.scala

@@ -0,0 +1,107 @@
+package zio.profiling.jmh
+
+import org.openjdk.jmh.infra.{BenchmarkParams, IterationParams}
+import org.openjdk.jmh.profile.InternalProfiler
+import org.openjdk.jmh.results.{IterationResult, Result, TextResult}
+import org.openjdk.jmh.runner.IterationType
+import zio._
+import zio.profiling.sampling.SamplingProfiler
+
+import java.io.{File, PrintWriter, StringWriter}
+import java.lang.System
+import java.nio.file.Files
+import java.util.Collections
+import java.{util => ju}
+import scala.annotation.unused
+
+class JmhZioProfiler(@unused initLine: String) extends InternalProfiler {
+  private val outDir: File                   = new File(System.getProperty("user.dir"))
+  private var trialOutDir: Option[File]      = None
+  private var warmupStarted: Boolean         = false
+  private var measurementStarted: Boolean    = false
+  private var measurementIterationCount: Int = 0
+
+  def this() = this("")
+
+  def getDescription(): String =
+    "zio-profiling profiler provider."
+
+  def beforeIteration(benchmarkParams: BenchmarkParams, iterationParams: IterationParams): Unit = {
+    if (trialOutDir.isEmpty) {
+      createTrialOutDir(benchmarkParams);
+    }
+
+    if (iterationParams.getType() == IterationType.WARMUP) {
+      if (!warmupStarted) {
+        start()
+        warmupStarted = true
+      }
+    }
+
+    if (iterationParams.getType() == IterationType.MEASUREMENT) {
+      if (!measurementStarted) {
+        if (warmupStarted) {
+          reset()
+        }
+        measurementStarted = true
+      }
+    }
+  }
+
+  def afterIteration(
+    benchmarkParams: BenchmarkParams,
+    iterationParams: IterationParams,
+    result: IterationResult
+  ): ju.Collection[_ <: Result[_]] = {
+    if (iterationParams.getType() == IterationType.MEASUREMENT) {
+      measurementIterationCount += 1
+      if (measurementIterationCount == iterationParams.getCount()) {
+        Collections.singletonList(stopAndDump());
+      }
+    }
+
+    Collections.emptyList();
+  }
+
+  private def createTrialOutDir(benchmarkParams: BenchmarkParams): Unit = {
+    val fileName = benchmarkParams.id().replace("%", "_")
+    trialOutDir = Some(new File(outDir, fileName))
+    trialOutDir.foreach(_.mkdirs())
+  }
+
+  private def start(): Unit = Unsafe.unsafe { implicit unsafe =>
+    BenchmarkUtils.runtimeRef.set(SamplingProfiler().supervisedRuntime)
+  }
+
+  private def stopAndDump(): TextResult = {
+    val sw = new StringWriter()
+    val pw = new PrintWriter(sw)
+    val rt = BenchmarkUtils.runtimeRef.getAndSet(null)
+
+    if (rt ne null) {
+      val supervisor = rt.environment.get
+      val resultPath = writeFile("profile.folded", supervisor.unsafeValue().stackCollapse.mkString("\n"))
+      Unsafe.unsafe(implicit u => rt.unsafe.shutdown())
+
+      pw.println("zio-profiler results:")
+      pw.println(s"  ${resultPath}")
+    }
+
+    pw.flush();
+    pw.close();
+    new TextResult(sw.toString(), "zio-profiling")
+  }
+
+  private def reset(): Unit = Unsafe.unsafe { implicit unsafe =>
+    val runtime = BenchmarkUtils.runtimeRef.get()
+    if (runtime ne null) {
+      runtime.environment.get.reset()
+    }
+  }
+
+  private def writeFile(name: String, content: String) = {
+    val out = new File(trialOutDir.get, name)
+    Files.write(out.toPath(), content.getBytes())
+  }
+
+}

zio-profiling-tagging-plugin/src/main/scala-2/zio/profiling/plugins/TaggingPlugin.scala

@@ -27,12 +27,12 @@ class TaggingPlugin(val global: Global) extends Plugin {
 
     class TaggingTransformer(unit: CompilationUnit) extends TypingTransformer(unit) {
       override def transform(tree: Tree): Tree = tree match {
-        case valDef @ ValDef(_, _, ZioTypeTree(t1, t2, t3), rhs) =>
+        case valDef @ ValDef(_, _, ZioTypeTree(t1, t2, t3), rhs) if isNonAbstract(valDef) =>
           val transformedRhs = tagEffectTree(descriptiveName(tree), rhs, t1, t2, t3)
           val typedRhs       = localTyper.typed(transformedRhs)
           val updated        = treeCopy.ValDef(tree, valDef.mods, valDef.name, valDef.tpt, rhs = typedRhs)
           super.transform(updated)
-        case defDef @ DefDef(_, _, _, _, ZioTypeTree(t1, t2, t3), rhs) =>
+        case defDef @ DefDef(_, _, _, _, ZioTypeTree(t1, t2, t3), rhs) if isNonAbstract(defDef) =>
           val transformedRhs = tagEffectTree(descriptiveName(tree), rhs, t1, t2, t3)
           val typedRhs       = localTyper.typed(transformedRhs)
           val updated =
@@ -42,6 +42,9 @@ class TaggingPlugin(val global: Global) extends Plugin {
           super.transform(tree)
       }
 
+      private def isNonAbstract(tree: ValOrDefDef): Boolean =
+        !tree.mods.isDeferred
+
       private def descriptiveName(tree: Tree): String = {
         val fullName   = tree.symbol.fullNameString
         val sourceFile = tree.pos.source.file.name