[7.1.0] Add support for additional command profiler event types.

src/main/java/com/google/devtools/build/lib/profiler/CommandProfilerModule.java

@@ -21,53 +21,81 @@
 import com.google.devtools.build.lib.runtime.Command;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.common.options.EnumConverter;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionDocumentationCategory;
 import com.google.devtools.common.options.OptionEffectTag;
 import com.google.devtools.common.options.OptionsBase;
 import java.time.Duration;
+import java.util.Locale;
 import javax.annotation.Nullable;
 import one.profiler.AsyncProfiler;
 
-/** Bazel module to record pprof-compatible profiles for single invocations. */
+/** Bazel module to record a Java Flight Recorder profile for a single command. */
 public class CommandProfilerModule extends BlazeModule {
 
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
   private static final Duration PROFILING_INTERVAL = Duration.ofMillis(10);
 
+  /** The type of profile to capture. */
+  enum ProfileType {
+    CPU,
+    WALL,
+    ALLOC,
+    LOCK;
+
+    @Override
+    public String toString() {
+      return name().toLowerCase(Locale.US);
+    }
+  }
+
+  /** Options converter for --experimental_command_profile. */
+  public static final class ProfileTypeConverter extends EnumConverter<ProfileType> {
+
+    public ProfileTypeConverter() {
+      super(ProfileType.class, "--experimental_command_profile setting");
+    }
+  }
+
   /** CommandProfilerModule options. */
   public static final class Options extends OptionsBase {
+
     @Option(
         name = "experimental_command_profile",
-        defaultValue = "false",
+        defaultValue = "null",
+        converter = ProfileTypeConverter.class,
         documentationCategory = OptionDocumentationCategory.LOGGING,
         effectTags = {OptionEffectTag.UNKNOWN},
         help =
-            "Records a Java Flight Recorder CPU profile into a profile.jfr file in the output base"
-                + " directory. The syntax and semantics of this flag might change in the future to"
-                + " support different profile types or output formats; use at your own risk.")
-    public boolean captureCommandProfile;
+            "Records a Java Flight Recorder profile for the duration of the command. One of the"
+                + " supported profiling event types (cpu, wall, alloc or lock) must be given as an"
+                + " argument. The profile is written to a file named after the event type under the"
+                + " output base directory."
+                + " The syntax and semantics of this flag might change in the future to support"
+                + " additional profile types or output formats; use at your own risk.")
+    public ProfileType profileType;
   }
 
   @Override
   public Iterable<Class<? extends OptionsBase>> getCommandOptions(Command command) {
     return ImmutableList.of(Options.class);
   }
 
-  private boolean captureCommandProfile;
+  @Nullable private ProfileType profileType;
   @Nullable private Reporter reporter;
   @Nullable private Path outputBase;
   @Nullable private Path outputPath;
 
   @Override
   public void beforeCommand(CommandEnvironment env) {
     Options options = env.getOptions().getOptions(Options.class);
-    captureCommandProfile = options.captureCommandProfile;
+    profileType = options.profileType;
     outputBase = env.getBlazeWorkspace().getOutputBase();
     reporter = env.getReporter();
 
-    if (!captureCommandProfile) {
+    if (profileType == null) {
       // Early exit so we don't attempt to load the JNI unless necessary.
       return;
     }
@@ -77,24 +105,26 @@ public void beforeCommand(CommandEnvironment env) {
       return;
     }
 
-    outputPath = outputBase.getRelative("profile.jfr");
+    outputPath = getProfilerOutputPath(profileType);
 
     try {
-      profiler.execute(getProfilerCommand(outputPath));
+      profiler.execute(getProfilerCommand(profileType, outputPath));
     } catch (Exception e) {
       // This may occur if the user has insufficient privileges to capture performance events.
-      reporter.handle(Event.error("Starting JFR CPU profile failed: " + e));
-      captureCommandProfile = false;
+      reporter.handle(
+          Event.error(String.format("Starting JFR %s profile failed: %s", profileType, e)));
+      profileType = null;
     }
 
-    if (captureCommandProfile) {
-      reporter.handle(Event.info("Writing JFR CPU profile to " + outputPath));
+    if (profileType != null) {
+      reporter.handle(
+          Event.info(String.format("Writing JFR %s profile to %s", profileType, outputPath)));
     }
   }
 
   @Override
   public void afterCommand() {
-    if (!captureCommandProfile) {
+    if (profileType == null) {
       // Early exit so we don't attempt to load the JNI unless necessary.
       return;
     }
@@ -106,7 +136,7 @@ public void afterCommand() {
 
     profiler.stop();
 
-    captureCommandProfile = false;
+    profileType = null;
     outputBase = null;
     reporter = null;
     outputPath = null;
@@ -123,9 +153,14 @@ private static AsyncProfiler getProfiler() {
     return null;
   }
 
-  private static String getProfilerCommand(Path outputPath) {
+  private Path getProfilerOutputPath(ProfileType profileType) {
+    return outputBase.getChild(profileType + ".jfr");
+  }
+
+  private static String getProfilerCommand(ProfileType profileType, Path outputPath) {
     // See https://github.com/async-profiler/async-profiler/blob/master/src/arguments.cpp.
     return String.format(
-        "start,event=cpu,interval=%s,file=%s,jfr", PROFILING_INTERVAL.toNanos(), outputPath);
+        "start,event=%s,interval=%s,file=%s,jfr",
+        profileType, PROFILING_INTERVAL.toNanos(), outputPath);
   }
 }

src/test/java/com/google/devtools/build/lib/profiler/BUILD

@@ -17,10 +17,7 @@ filegroup(
 
 java_test(
     name = "ProfilerTests",
-    srcs = glob(
-        ["*.java"],
-        exclude = ["CommandProfilerModuleTest.java"],
-    ),
+    srcs = glob(["*.java"]),
     test_class = "com.google.devtools.build.lib.AllTests",
     runtime_deps = [
         "//src/test/java/com/google/devtools/build/lib:test_runner",
@@ -42,20 +39,3 @@ java_test(
         "//third_party:truth",
     ],
 )
-
-java_test(
-    name = "CommandProfilerModuleTest",
-    srcs = ["CommandProfilerModuleTest.java"],
-    tags = [
-        # Bazel-specific tests
-        "manual",
-    ],
-    deps = [
-        "//src/main/java/com/google/devtools/build/lib:runtime",
-        "//src/main/java/com/google/devtools/build/lib/profiler:command_profiler_module",
-        "//src/main/java/com/google/devtools/build/lib/util:os",
-        "//src/test/java/com/google/devtools/build/lib/buildtool/util",
-        "//third_party:junit4",
-        "//third_party:truth",
-    ],
-)

src/test/shell/bazel/BUILD

@@ -1079,6 +1079,14 @@ sh_test(
     tags = ["no_windows"],
 )
 
+sh_test(
+    name = "command_profiler_test",
+    size = "medium",
+    srcs = ["command_profiler_test.sh"],
+    data = [":test-deps"],
+    tags = ["no_windows"],
+)
+
 sh_test(
     name = "execroot_test",
     size = "medium",

src/test/shell/bazel/command_profiler_test.sh

@@ -0,0 +1,76 @@
+#!/bin/bash
+#
+# Copyright 2024 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -eu
+
+# Load the test setup defined in the parent directory
+CURRENT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "${CURRENT_DIR}/../integration_test_setup.sh" \
+  || { echo "integration_test_setup.sh not found!" >&2; exit 1; }
+
+
+function test_profiler_disabled() {
+  cat > BUILD <<'EOF'
+genrule(
+  name = "gen",
+  outs = ["out"],
+  cmd = "touch $@",
+)
+EOF
+
+  bazel build //:gen || fail "Expected build to succeed"
+
+  if [[ "$(ls "$(bazel info output_base)/*.jfr")" ]]; then
+    fail "Expected no profiler outputs"
+  fi
+}
+
+function do_test_profiler_enabled() {
+  local -r type="$1"
+
+  cat > BUILD <<'EOF'
+genrule(
+  name = "gen",
+  outs = ["out"],
+  cmd = "touch $@",
+)
+EOF
+
+ bazel build --experimental_command_profile="${type}" //:gen \
+     || fail "Expected build to succeed"
+
+ if ! [[ -f "$(bazel info output_base)/${type}.jfr" ]]; then
+     fail "Expected profiler output"
+ fi
+}
+
+function test_cpu_profiler_enabled() {
+  do_test_profiler_enabled cpu
+}
+
+function test_wall_profiler_enabled() {
+  do_test_profiler_enabled wall
+}
+
+function test_alloc_profiler_enabled() {
+  do_test_profiler_enabled alloc
+}
+
+function test_lock_profiler_enabled() {
+  do_test_profiler_enabled lock
+}
+
+run_suite "command profiler tests"