Add profiling for Bzlmod operations

The resolution of Bazel modules and evaluation of module extensions is now represented in the profile.

Profile from a run of `bazel build //src:bazel-dev --enable_bzlmod --nobuild`:
![Bzlmod profile](https://github.com/bazelbuild/bazel/assets/4312191/22e03604-f832-4f64-be0e-a3281b94da3e)

Closes #19291.

PiperOrigin-RevId: 559466489
Change-Id: If0989a4a6728f8dbf659197a7acbd3a0fcef9471

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -80,6 +80,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
+        "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//third_party:caffeine",
@@ -198,6 +199,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/packages",
+        "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/repository_directory_value",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/repository_function",
         "//src/main/java/com/google/devtools/build/lib/skyframe:bzl_load_value",

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileFunction.java

@@ -22,6 +22,9 @@
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
+import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalDeps.Code;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue.Precomputed;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
@@ -80,7 +83,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       return null;
     }
 
-    try {
+    try (SilentCloseable c = Profiler.instance().profile(ProfilerTask.BZLMOD, "parse lockfile")) {
       return getLockfileValue(lockfilePath);
     } catch (IOException | JsonSyntaxException | NullPointerException e) {
       throw new BazelLockfileFunctionException(

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionFunction.java

@@ -29,6 +29,9 @@
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
+import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalDeps.Code;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue.Precomputed;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -61,34 +64,52 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     if (root == null) {
       return null;
     }
-    ImmutableMap<ModuleKey, InterimModule> initialDepGraph = Discovery.run(env, root);
-    if (initialDepGraph == null) {
-      return null;
+    ImmutableMap<ModuleKey, InterimModule> initialDepGraph;
+    try (SilentCloseable c = Profiler.instance().profile(ProfilerTask.BZLMOD, "discovery")) {
+      initialDepGraph = Discovery.run(env, root);
+      if (initialDepGraph == null) {
+        return null;
+      }
     }
 
     Selection.Result selectionResult;
-    try {
+    try (SilentCloseable c = Profiler.instance().profile(ProfilerTask.BZLMOD, "selection")) {
       selectionResult = Selection.run(initialDepGraph, root.getOverrides());
     } catch (ExternalDepsException e) {
       throw new BazelModuleResolutionFunctionException(e, Transience.PERSISTENT);
     }
     ImmutableMap<ModuleKey, InterimModule> resolvedDepGraph = selectionResult.getResolvedDepGraph();
 
-    verifyRootModuleDirectDepsAreAccurate(
-        initialDepGraph.get(ModuleKey.ROOT),
-        resolvedDepGraph.get(ModuleKey.ROOT),
-        Objects.requireNonNull(CHECK_DIRECT_DEPENDENCIES.get(env)),
-        env.getListener());
+    try (SilentCloseable c =
+        Profiler.instance().profile(ProfilerTask.BZLMOD, "verify root module direct deps")) {
+      verifyRootModuleDirectDepsAreAccurate(
+          initialDepGraph.get(ModuleKey.ROOT),
+          resolvedDepGraph.get(ModuleKey.ROOT),
+          Objects.requireNonNull(CHECK_DIRECT_DEPENDENCIES.get(env)),
+          env.getListener());
+    }
 
-    checkBazelCompatibility(
-        resolvedDepGraph.values(),
-        Objects.requireNonNull(BAZEL_COMPATIBILITY_MODE.get(env)),
-        env.getListener());
+    try (SilentCloseable c =
+        Profiler.instance().profile(ProfilerTask.BZLMOD, "check bazel compatibility")) {
+      checkBazelCompatibility(
+          resolvedDepGraph.values(),
+          Objects.requireNonNull(BAZEL_COMPATIBILITY_MODE.get(env)),
+          env.getListener());
+    }
 
-    checkNoYankedVersions(resolvedDepGraph);
+    try (SilentCloseable c =
+        Profiler.instance().profile(ProfilerTask.BZLMOD, "check no yanked versions")) {
+      checkNoYankedVersions(resolvedDepGraph);
+    }
+
+    ImmutableMap<ModuleKey, Module> finalDepGraph;
+    try (SilentCloseable c =
+        Profiler.instance().profile(ProfilerTask.BZLMOD, "compute final dep graph")) {
+      finalDepGraph =
+          computeFinalDepGraph(resolvedDepGraph, root.getOverrides(), env.getListener());
+    }
 
-    ImmutableMap<ModuleKey, Module> finalDepGraph =
-        computeFinalDepGraph(resolvedDepGraph, root.getOverrides(), env.getListener());
+    Profiler.instance().profile(ProfilerTask.BZLMOD, "module resolution completed").close();
 
     return BazelModuleResolutionValue.create(finalDepGraph, selectionResult.getUnprunedDepGraph());
   }
@@ -241,6 +262,12 @@ private static RepoSpec computeRepoSpec(
   static Module moduleFromInterimModule(
       InterimModule interim, ModuleOverride override, ExtendedEventHandler eventHandler)
       throws BazelModuleResolutionFunctionException, InterruptedException {
+    RepoSpec repoSpec;
+    try (SilentCloseable c =
+        Profiler.instance()
+            .profile(ProfilerTask.BZLMOD, () -> "compute repo spec: " + interim.getKey())) {
+      repoSpec = computeRepoSpec(interim, override, eventHandler);
+    }
     return Module.builder()
         .setName(interim.getName())
         .setVersion(interim.getVersion())
@@ -249,7 +276,7 @@ static Module moduleFromInterimModule(
         .setExecutionPlatformsToRegister(interim.getExecutionPlatformsToRegister())
         .setToolchainsToRegister(interim.getToolchainsToRegister())
         .setDeps(ImmutableMap.copyOf(Maps.transformValues(interim.getDeps(), DepSpec::toModuleKey)))
-        .setRepoSpec(computeRepoSpec(interim, override, eventHandler))
+        .setRepoSpec(repoSpec)
         .setExtensionUsages(interim.getExtensionUsages())
         .build();
   }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistry.java

@@ -24,6 +24,9 @@
 import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
+import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.gson.FieldNamingPolicy;
@@ -81,7 +84,8 @@ private String constructUrl(String base, String... segments) {
   /** Grabs a file from the given URL. Returns {@link Optional#empty} if the file doesn't exist. */
   private Optional<byte[]> grabFile(String url, ExtendedEventHandler eventHandler)
       throws IOException, InterruptedException {
-    try {
+    try (SilentCloseable c =
+        Profiler.instance().profile(ProfilerTask.BZLMOD, () -> "download file: " + url)) {
       return Optional.of(
           downloadManager.downloadAndReadOneUrl(new URL(url), eventHandler, clientEnv));
     } catch (FileNotFoundException e) {

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionId.java

@@ -76,7 +76,9 @@ public static ModuleExtensionId create(
   }
 
   public String asTargetString() {
+    String isolationKeyPart = getIsolationKey().map(key -> "%" + key).orElse("");
     return String.format(
-        "%s%%%s", getBzlFileLabel().getUnambiguousCanonicalForm(), getExtensionName());
+        "%s%%%s%s",
+        getBzlFileLabel().getUnambiguousCanonicalForm(), getExtensionName(), isolationKeyPart);
   }
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java

@@ -31,6 +31,9 @@
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.packages.DotBazelFileSyntaxChecker;
 import com.google.devtools.build.lib.packages.StarlarkExportable;
+import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.rules.repository.RepositoryDirectoryValue;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalDeps.Code;
 import com.google.devtools.build.lib.skyframe.ClientEnvironmentFunction;
@@ -131,8 +134,12 @@ public SkyValue compute(SkyKey skyKey, Environment env)
 
     ModuleFileValue.Key moduleFileKey = (ModuleFileValue.Key) skyKey;
     ModuleKey moduleKey = moduleFileKey.getModuleKey();
-    GetModuleFileResult getModuleFileResult =
-        getModuleFile(moduleKey, moduleFileKey.getOverride(), allowedYankedVersions, env);
+    GetModuleFileResult getModuleFileResult;
+    try (SilentCloseable c =
+        Profiler.instance().profile(ProfilerTask.BZLMOD, () -> "fetch module file: " + moduleKey)) {
+      getModuleFileResult =
+          getModuleFile(moduleKey, moduleFileKey.getOverride(), allowedYankedVersions, env);
+    }
     if (getModuleFileResult == null) {
       return null;
     }
@@ -277,7 +284,10 @@ private ModuleFileGlobals execModuleFile(
 
     ModuleFileGlobals moduleFileGlobals =
         new ModuleFileGlobals(builtinModules, moduleKey, registry, ignoreDevDeps);
-    try (Mutability mu = Mutability.create("module file", moduleKey)) {
+    try (SilentCloseable c =
+            Profiler.instance()
+                .profile(ProfilerTask.BZLMOD, () -> "evaluate module file: " + moduleKey);
+        Mutability mu = Mutability.create("module file", moduleKey)) {
       new DotBazelFileSyntaxChecker("MODULE.bazel files", /* canLoadBzl= */ false)
           .check(starlarkFile);
       net.starlark.java.eval.Module predeclaredEnv =

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java

@@ -35,6 +35,9 @@
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.rules.repository.NeedsSkyframeRestartException;
 import com.google.devtools.build.lib.rules.repository.RepositoryFunction;
 import com.google.devtools.build.lib.runtime.ProcessWrapper;
@@ -487,7 +490,11 @@ private RunModuleExtensionResult runModuleExtension(
       thread.setPrintHandler(Event.makeDebugPrintHandler(env.getListener()));
       moduleContext = createContext(env, usagesValue, starlarkSemantics, extensionId, extension);
       threadContext.storeInThread(thread);
-      try {
+      try (SilentCloseable c =
+          Profiler.instance()
+              .profile(
+                  ProfilerTask.BZLMOD,
+                  () -> "evaluate module extension: " + extensionId.asTargetString())) {
         Object returnValue =
             Starlark.fastcall(
                 thread, extension.getImplementation(), new Object[] {moduleContext}, new Object[0]);

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

@@ -28,6 +28,7 @@ public enum ProfilerTask {
   ACTION_RELEASE("action resource release", Threshold.TEN_MILLIS),
   ACTION_UPDATE("update action information", Threshold.TEN_MILLIS),
   ACTION_COMPLETE("complete action execution"),
+  BZLMOD("bazel module processing"),
   INFO("general information"),
   CREATE_PACKAGE("package creation"),
   REMOTE_EXECUTION("remote action execution"),

src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandDispatcher.java

@@ -46,6 +46,7 @@
 import com.google.devtools.build.lib.exec.ExecutionOptions;
 import com.google.devtools.build.lib.profiler.MemoryProfiler;
 import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.runtime.proto.InvocationPolicyOuterClass.InvocationPolicy;
 import com.google.devtools.build.lib.server.FailureDetails;
@@ -585,7 +586,8 @@ private BlazeCommandResult execExclusively(
         // Compute the repo mapping of the main repo and re-parse options so that we get correct
         // values for label-typed options.
         env.getEventBus().post(new MainRepoMappingComputationStartingEvent());
-        try {
+        try (SilentCloseable c =
+            Profiler.instance().profile(ProfilerTask.BZLMOD, "compute main repo mapping")) {
           RepositoryMapping mainRepoMapping =
               env.getSkyframeExecutor().getMainRepoMapping(reporter);
           optionsParser = optionsParser.toBuilder().withConversionContext(mainRepoMapping).build();
@@ -606,10 +608,14 @@ private BlazeCommandResult execExclusively(
           result = BlazeCommandResult.detailedExitCode(earlyExitCode);
           return result;
         }
-        optionHandler =
-            new BlazeOptionHandler(
-                runtime, workspace, command, commandAnnotation, optionsParser, invocationPolicy);
-        earlyExitCode = optionHandler.parseOptions(args, reporter);
+        try (SilentCloseable c =
+            Profiler.instance()
+                .profile(ProfilerTask.BZLMOD, "reparse options with main repo mapping")) {
+          optionHandler =
+              new BlazeOptionHandler(
+                  runtime, workspace, command, commandAnnotation, optionsParser, invocationPolicy);
+          earlyExitCode = optionHandler.parseOptions(args, reporter);
+        }
         if (!earlyExitCode.isSuccess()) {
           reporter.post(
               new NoBuildEvent(
@@ -620,7 +626,10 @@ private BlazeCommandResult execExclusively(
       }
 
       // Parse starlark options.
-      earlyExitCode = optionHandler.parseStarlarkOptions(env);
+      try (SilentCloseable c =
+          Profiler.instance().profile(ProfilerTask.BZLMOD, "parse starlark options")) {
+        earlyExitCode = optionHandler.parseStarlarkOptions(env);
+      }
       if (!earlyExitCode.isSuccess()) {
         reporter.post(
             new NoBuildEvent(