feat: add jmh

timothyklim · Jul 29, 2023 · a53ec24 · a53ec24
1 parent 4a70bee
commit a53ec24
Show file tree

Hide file tree

Showing 13 changed files with 174 additions and 36 deletions.
diff --git a/deps/Vars.scala b/deps/Vars.scala
@@ -2,7 +2,7 @@ package rules_scala3.deps
 
 import java.io.File
 
-case class Vars(
+final case class Vars(
     projectRoot: File,
     depsDirName: String,
     bazelExtFileName: String,

diff --git a/jmh/ArrayBench.scala b/jmh/ArrayBench.scala
@@ -0,0 +1,38 @@
+package test
+
+import java.util.concurrent.TimeUnit.*
+
+import org.openjdk.jmh.annotations.*
+@State(Scope.Benchmark)
+class ArrayState:
+  val array = (1 to 100).toArray
+
+  val h = 123
+
+class ArrayBench:
+  @Benchmark
+  @BenchmarkMode(Array(Mode.SingleShotTime, Mode.Throughput))
+  @OutputTimeUnit(MILLISECONDS)
+  def testFold(state: ArrayState): Unit =
+    val builder = Vector.newBuilder[Int]
+    state.array.foldLeft(builder)(_ += _)
+    ()
+
+  @Benchmark
+  @BenchmarkMode(Array(Mode.SingleShotTime, Mode.Throughput))
+  @OutputTimeUnit(MILLISECONDS)
+  def testArray(state: ArrayState): Unit =
+    val builder = Vector.newBuilder[Int]
+    var _i = 0
+    while _i < state.array.length do
+      builder += state.array(_i)
+      _i += 1
+    ()
+
+  @Benchmark
+  @BenchmarkMode(Array(Mode.SingleShotTime, Mode.Throughput))
+  @OutputTimeUnit(MILLISECONDS)
+  def testForeach(state: ArrayState): Unit =
+    val builder = Vector.newBuilder[Int]
+    state.array.foreach(builder += _)
+    ()
diff --git a/jmh/BUILD b/jmh/BUILD
@@ -1,12 +1,18 @@
-load("//rules:scala.bzl", "scala_binary")
+load("//jmh:jmh.bzl", "scala_jmh_benchmark")
+load("@rules_java//java:defs.bzl", "java_binary")
 
-scala_binary(
-    name = "jmh",
-    srcs = glob(["*.scala"]),
-    main_class = "rules_scala3.jmh.JmhRunner",
-    scala = "//scala:bootstrap_3_3",
-    visibility = ["//visibility:public"],
-    deps = [
+java_binary(
+    name = "generator",
+    main_class = "org.openjdk.jmh.generators.bytecode.JmhBytecodeGenerator",
+    runtime_deps = [
         "@//3rdparty/jvm/org/openjdk/jmh:jmh_core",
+        "@//3rdparty/jvm/org/openjdk/jmh:jmh_generator_asm",
+        "@//3rdparty/jvm/org/openjdk/jmh:jmh_generator_bytecode",
+        "@//3rdparty/jvm/org/openjdk/jmh:jmh_generator_reflection",
     ],
 )
+
+scala_jmh_benchmark(
+  name = "test",
+  srcs = ["ArrayBench.scala"],
+)
diff --git a/jmh/JmhRunner.scala b/jmh/JmhRunner.scala
diff --git a/jmh/jmh.bzl b/jmh/jmh.bzl
@@ -0,0 +1,102 @@
+load("@rules_java//java:defs.bzl", "java_binary", "java_library")
+load("//rules:scala.bzl", "scala_library")
+load("//rules/common:private/utils.bzl", _safe_name = "safe_name")
+
+def _jmh_benchmark(ctx):
+    info = ctx.attr.src[JavaInfo]
+
+    input_bytecode_dir = ctx.actions.declare_directory("{}_bytecode".format(_safe_name(ctx.attr.name)))
+    output_source_dir = ctx.actions.declare_directory("{}_src".format(_safe_name(ctx.attr.name)))
+    output_resource_dir = ctx.actions.declare_directory("{}_resources".format(_safe_name(ctx.attr.name)))
+
+    class_jar = info.outputs.jars[0].class_jar
+    ctx.actions.run_shell(
+        inputs = [class_jar],
+        outputs = [input_bytecode_dir],
+        arguments = [class_jar.path, input_bytecode_dir.path],
+        command = "unzip $1 -d $2",
+        progress_message = "Unzip classes from jar",
+        use_default_shell_env = True,
+    )
+    ctx.actions.run(
+        outputs = [output_source_dir, output_resource_dir],
+        inputs = [input_bytecode_dir],
+        executable = ctx.executable._generator,
+        arguments = [input_bytecode_dir.path, output_source_dir.path, output_resource_dir.path, ctx.attr.generator_type],
+        progress_message = "Generating benchmark code for %s" % ctx.label,
+    )
+    ctx.actions.run_shell(
+        inputs = [output_source_dir],
+        outputs = [ctx.outputs.srcjar],
+        arguments = [ctx.executable._zipper.path, output_source_dir.path, output_source_dir.short_path, ctx.outputs.srcjar.path],
+        command = """$1 c $4 META-INF/= $(find -L $2 -type f | while read v; do echo ${v#"$2/"}=$v; done)""",
+        progress_message = "Bundling generated Java sources into srcjar",
+        tools = [ctx.executable._zipper],
+        use_default_shell_env = True,
+    )
+    ctx.actions.run_shell(
+        inputs = [output_resource_dir],
+        outputs = [ctx.outputs.resources_jar],
+        arguments = [ctx.executable._zipper.path, output_resource_dir.path, output_resource_dir.short_path, ctx.outputs.resources_jar.path],
+        command = """$1 c $4 META-INF/= $(find -L $2 -type f | while read v; do echo ${v#"$2/"}=$v; done)""",
+        progress_message = "Bundling generated resources into jar",
+        tools = [ctx.executable._zipper],
+        use_default_shell_env = True,
+    )
+
+generate_jmh_benchmark = rule(
+    implementation = _jmh_benchmark,
+    attrs = {
+        "src": attr.label(mandatory = True, providers = [[JavaInfo]]),
+        "generator_type": attr.string(
+            default = "asm",
+            mandatory = False,
+        ),
+        "_generator": attr.label(executable = True, cfg = "exec", default = "@//jmh:generator"),
+        "_zipper": attr.label(cfg = "exec", default = "@bazel_tools//tools/zip:zipper", executable = True),
+    },
+    outputs = {
+        "srcjar": "%{name}.srcjar",
+        "resources_jar": "%{name}_resources.jar",
+    },
+)
+
+def scala_jmh_benchmark(**kwargs):
+  name = kwargs["name"]
+  deps = kwargs.get("deps", [])
+  lib = "%s_generator" % name
+
+  scala_library(
+      name = lib,
+      srcs = kwargs["srcs"],
+      deps = deps + [
+          "@//3rdparty/jvm/org/openjdk/jmh:jmh_core",
+      ],
+      runtime_deps = kwargs.get("runtime_deps", []),
+      resources = kwargs.get("resources", []),
+      resource_jars = kwargs.get("resource_jars", []),
+      scala = "//scala:bootstrap_3_3",
+  )
+
+  codegen = name + "_codegen"
+  generate_jmh_benchmark(
+      name = codegen,
+      src = lib,
+      generator_type = kwargs.get("generator_type", "asm"),
+  )
+  compiled_lib = name + "_compiled_jmh_lib"
+  java_library(
+      name = compiled_lib,
+      srcs = ["%s.srcjar" % codegen],
+      deps = deps + [
+          "@//3rdparty/jvm/org/openjdk/jmh:jmh_core",
+          lib,
+      ],
+      runtime_deps = ["%s_resources.jar" % codegen],
+  )
+  java_binary(
+      name = name,
+      runtime_deps = [compiled_lib],
+      data = kwargs.get("data", []),
+      main_class = "org.openjdk.jmh.Main",
+  )
diff --git a/rules/private/phases/phase_bootstrap_compile.bzl b/rules/private/phases/phase_bootstrap_compile.bzl
@@ -17,7 +17,7 @@ def phase_bootstrap_compile(ctx, g):
     if g.classpaths.plugin:
         fail("plugins aren't supported for bootstrap_scala rules")
     if g.classpaths.src_jars:
-        fail("source jars supported for bootstrap_scala rules")
+        fail("source jars aren't supported for bootstrap_scala rules")
 
     inputs = depset(
         ctx.files.srcs,
@@ -38,9 +38,6 @@ def phase_bootstrap_compile(ctx, g):
         command = _strip_margin(
             """
             |set -eo pipefail
-            |
-            |mkdir -p tmp/classes
-            |
             |{java} \\
             |  -Dscala.usejavacp=true \\
             |  -classpath {compiler_classpath} \\

diff --git a/rules/providers.bzl b/rules/providers.bzl
@@ -149,7 +149,7 @@ declare_zinc_configuration = rule(
             default = "@rules_scala3//scala/workers/zinc/compile",
             allow_files = True,
             executable = True,
-            cfg = "host",
+            cfg = "exec",
         ),
     },
     doc = "Creates a `ZincConfiguration`.",

diff --git a/rules/scala.bzl b/rules/scala.bzl
@@ -59,10 +59,10 @@ _compile_private_attributes = {
     ),
     "_host_javabase": attr.label(
         default = Label("@bazel_tools//tools/jdk:current_java_runtime"),
-        cfg = "host",
+        cfg = "exec",
     ),
     "_singlejar": attr.label(
-        cfg = "host",
+        cfg = "exec",
         default = "@remote_java_tools//:singlejar_cc_bin",
         executable = True,
     ),
@@ -72,12 +72,12 @@ _compile_private_attributes = {
     "_jdk": attr.label(
         default = Label("@bazel_tools//tools/jdk:current_java_runtime"),
         providers = [java_common.JavaRuntimeInfo],
-        cfg = "host",
+        cfg = "exec",
     ),
     "_jar_creator": attr.label(
         default = Label("@remote_java_tools//:ijar_cc_binary"),
         executable = True,
-        cfg = "host",
+        cfg = "exec",
     ),
 }
 
@@ -188,11 +188,11 @@ _testing_private_attributes = {
     # in https://github.com/bazelbuild/bazel/blob/0.22.0/src/main/java/com/google/devtools/build/lib/bazel/rules/java/BazelJavaTestRule.java#L69-L76
     "_jacocorunner": attr.label(
         default = Label("@bazel_tools//tools/jdk:JacocoCoverage"),
-        cfg = "host",
+        cfg = "exec",
     ),
     "_lcov_merger": attr.label(
         default = Label("@bazel_tools//tools/test/CoverageOutputGenerator/java/com/google/devtools/coverageoutputgenerator:Main"),
-        cfg = "host",
+        cfg = "exec",
     ),
 }
 
@@ -376,7 +376,7 @@ _scala_repl_private_attributes = _dicts.add(
     _runtime_private_attributes,
     {
         "_runner": attr.label(
-            cfg = "host",
+            cfg = "exec",
             executable = True,
             default = "@rules_scala3//scala/workers/zinc/repl",
         ),
@@ -538,19 +538,19 @@ _configure_zinc_scala = rule(
             default = "@rules_scala3//scala/workers/zinc/compile",
             allow_files = True,
             executable = True,
-            cfg = "host",
+            cfg = "exec",
         ),
         "_deps_worker": attr.label(
             default = "@rules_scala3//scala/workers/deps",
             allow_files = True,
             executable = True,
-            cfg = "host",
+            cfg = "exec",
         ),
         "_code_coverage_instrumentation_worker": attr.label(
             default = "@rules_scala3//scala/workers/jacoco/instrumenter",
             allow_files = True,
             executable = True,
-            cfg = "host",
+            cfg = "exec",
         ),
     },
     implementation = _configure_zinc_scala_implementation,

diff --git a/rules/scala/private/doc.bzl b/rules/scala/private/doc.bzl
@@ -14,7 +14,7 @@ load(
 
 scaladoc_private_attributes = {
     "_runner": attr.label(
-        cfg = "host",
+        cfg = "exec",
         executable = True,
         default = "//scala/workers/zinc/doc",
     ),

diff --git a/rules/scala/private/import.bzl b/rules/scala/private/import.bzl
@@ -10,7 +10,7 @@ scala_import_private_attributes = {
     ),
     "_host_javabase": attr.label(
         default = Label("@bazel_tools//tools/jdk:current_java_runtime"),
-        cfg = "host",
+        cfg = "exec",
     ),
 }
 

diff --git a/rules/scala/private/scalajs.bzl b/rules/scala/private/scalajs.bzl
@@ -92,7 +92,7 @@ scalajs_link = rule(
             default = "@rules_scala3//scala/workers/scalajs:scalajs_linker",
             allow_files = True,
             executable = True,
-            cfg = "host",
+            cfg = "exec",
         ),
     },
     implementation = _scalajs_link_impl,

diff --git a/rules/scala_proto.bzl b/rules/scala_proto.bzl
@@ -15,15 +15,15 @@ scala_proto_library = rule(
             "protoc": attr.label(
                 doc = "protoc compiler",
                 default = "@com_google_protobuf//:protoc",
-                cfg = "host",
+                cfg = "exec",
                 allow_single_file = True,
                 executable = True,
             ),
             "scalapb_options": attr.string_list(
                 doc = "ScalaPB gen options: flat_package, java_conversions, grpc, single_line_to_proto_string, ascii_format_to_string, no_lenses",
             ),
             "grpc_web": attr.bool(default = False),
-            "_zipper": attr.label(cfg = "host", default = "@bazel_tools//tools/zip:zipper", executable = True),
+            "_zipper": attr.label(cfg = "exec", default = "@bazel_tools//tools/zip:zipper", executable = True),
         },
     ),
     doc = """
@@ -52,7 +52,7 @@ scala_proto_toolchain = rule(
             doc = "The compiler to use to generate Scala form proto sources",
             allow_files = True,
             executable = True,
-            cfg = "host",
+            cfg = "exec",
         ),
         "compiler_supports_workers": attr.bool(default = False),
     },

diff --git a/rules/scalafmt/private/test.bzl b/rules/scalafmt/private/test.bzl
@@ -10,7 +10,7 @@ scala_format_attributes = {
         doc = "The Scalafmt configuration file.",
     ),
     "_fmt": attr.label(
-        cfg = "host",
+        cfg = "exec",
         default = "@rules_scala3//rules/scalafmt",
         executable = True,
     ),