Perform frontier serialization after the build command completes, if …

…enabled.

This moves the callsite from `DumpCommand` into an `afterCommand` hook for the serialization module, which triggers with a non-empty `--serialize_frontier_profile` (now) build option.

PiperOrigin-RevId: 671227197
Change-Id: I8665d18069393102e59db13316f27d468f48075d

src/main/java/com/google/devtools/build/lib/runtime/commands/BUILD

@@ -95,7 +95,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe:skyframe_cluster",
         "//src/main/java/com/google/devtools/build/lib/skyframe:skyframe_stats",
         "//src/main/java/com/google/devtools/build/lib/skyframe/config",
-        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis:frontier_serializer",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis:options",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",
         "//src/main/java/com/google/devtools/build/lib/util:command",

src/main/java/com/google/devtools/build/lib/runtime/commands/BuildCommand.java

@@ -35,6 +35,7 @@
 import com.google.devtools.build.lib.runtime.KeepGoingOption;
 import com.google.devtools.build.lib.runtime.LoadingPhaseThreadsOption;
 import com.google.devtools.build.lib.skyframe.SkyfocusOptions;
+import com.google.devtools.build.lib.skyframe.serialization.analysis.RemoteAnalysisCachingOptions;
 import com.google.devtools.build.lib.util.DetailedExitCode;
 import com.google.devtools.common.options.OptionsParsingResult;
 import java.util.List;
@@ -57,6 +58,7 @@
       LoadingPhaseThreadsOption.class,
       BuildEventProtocolOptions.class,
       SkyfocusOptions.class,
+      RemoteAnalysisCachingOptions.class,
     },
     usesConfigurationOptions = true,
     shortDescription = "Builds the specified targets.",

src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java

@@ -17,7 +17,6 @@
 import static com.google.devtools.build.lib.runtime.Command.BuildPhase.NONE;
 
 import com.google.common.base.Splitter;
-import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.buildtool.SkyframeMemoryDumper;
@@ -50,7 +49,6 @@
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
 import com.google.devtools.build.lib.skyframe.SkyframeStats;
 import com.google.devtools.build.lib.skyframe.config.BuildConfigurationKey;
-import com.google.devtools.build.lib.skyframe.serialization.analysis.FrontierSerializer;
 import com.google.devtools.build.lib.util.MemoryAccountant.Stats;
 import com.google.devtools.build.lib.util.RegexFilter;
 import com.google.devtools.build.lib.util.RegexFilter.RegexFilterConverter;
@@ -298,14 +296,6 @@ public static class DumpOptions extends OptionsBase {
         effectTags = {OptionEffectTag.BAZEL_MONITORING},
         help = "Dump the memory use of the given Skyframe node.")
     public MemoryMode memory;
-
-    @Option(
-        name = "serialized_frontier_profile",
-        defaultValue = "",
-        documentationCategory = OptionDocumentationCategory.OUTPUT_SELECTION,
-        effectTags = {OptionEffectTag.BAZEL_MONITORING},
-        help = "Dump a profile of serialized frontier bytes. Specifies the output path.")
-    public String serializedFrontierProfile;
   }
 
   /** Different ways to dump information about Skyframe. */
@@ -344,8 +334,7 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
             || dumpOptions.dumpRules
             || dumpOptions.starlarkMemory != null
             || dumpOptions.dumpSkyframe != SkyframeDumpOption.OFF
-            || dumpOptions.memory != null
-            || !Strings.isNullOrEmpty(dumpOptions.serializedFrontierProfile);
+            || dumpOptions.memory != null;
     if (!anyOutput) {
       Collection<Class<? extends OptionsBase>> optionList = new ArrayList<>();
       optionList.add(DumpOptions.class);
@@ -421,12 +410,6 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
         case OFF -> {}
       }
 
-      if (!Strings.isNullOrEmpty(dumpOptions.serializedFrontierProfile) && failure.isEmpty()) {
-        failure =
-            FrontierSerializer.dumpFrontierSerializationProfile(
-                out, env, dumpOptions.serializedFrontierProfile);
-      }
-
       return failure.orElse(BlazeCommandResult.success());
     } catch (InterruptedException e) {
       env.getReporter().error(null, "Interrupted", e);

src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD

@@ -53,6 +53,13 @@ java_library(
         ":serialization_registry_setup_helpers",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis:frontier_serializer",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis:options",
+        "//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",
+        "//src/main/java/com/google/devtools/build/lib/util:detailed_exit_code",
+        "//src/main/protobuf:failure_details_java_proto",
+        "//third_party:guava",
+        "//third_party:jsr305",
     ],
 )
 

src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationModule.java

@@ -13,14 +13,31 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import static com.google.common.base.Strings.isNullOrEmpty;
+
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.runtime.BlazeModule;
 import com.google.devtools.build.lib.runtime.BlazeRuntime;
+import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.runtime.WorkspaceBuilder;
+import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
+import com.google.devtools.build.lib.server.FailureDetails.RemoteAnalysisCaching.Code;
+import com.google.devtools.build.lib.skyframe.serialization.analysis.FrontierSerializer;
+import com.google.devtools.build.lib.skyframe.serialization.analysis.RemoteAnalysisCachingOptions;
+import com.google.devtools.build.lib.util.AbruptExitException;
+import com.google.devtools.build.lib.util.DetailedExitCode;
+import java.util.Optional;
+import javax.annotation.Nullable;
 
 /** A {@link BlazeModule} to store Skyframe serialization lifecycle hooks. */
 public class SerializationModule extends BlazeModule {
 
+  @Nullable private RemoteAnalysisCachingOptions options;
+
+  // Retained only for analysis caching operations in afterCommand. Must be garbage collected once
+  // it's no longer needed. See the javadoc for CommandEnvironment for more information.
+  @Nullable private CommandEnvironment env;
+
   @Override
   public void workspaceInit(
       BlazeRuntime runtime, BlazeDirectories directories, WorkspaceBuilder builder) {
@@ -30,7 +47,6 @@ public void workspaceInit(
       // documentation HTML.
       return;
     }
-
     // This is injected as a callback instead of evaluated eagerly to avoid forcing the somewhat
     // expensive AutoRegistry.get call on clients that don't require it.
     runtime.initAnalysisCodecRegistry(
@@ -42,4 +58,47 @@ public void workspaceInit(
                 directories.getWorkspace().getBaseName())));
   }
 
+  @Override
+  public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
+    RemoteAnalysisCachingOptions options =
+        env.getOptions().getOptions(RemoteAnalysisCachingOptions.class);
+    if (options == null) {
+      // not a supported command
+      return;
+    }
+
+    if (!isNullOrEmpty(options.serializedFrontierProfile)) {
+      this.options = options;
+      this.env = env;
+    }
+  }
+
+  @Override
+  public void afterCommand() throws AbruptExitException {
+    if (options == null) {
+      return;
+    }
+
+    try {
+      if (!isNullOrEmpty(options.serializedFrontierProfile)) {
+        Optional<FailureDetail> failureDetail =
+            FrontierSerializer.dumpFrontierSerializationProfile(
+                env, options.serializedFrontierProfile);
+        if (failureDetail.isPresent()) {
+          throw new AbruptExitException(DetailedExitCode.of(failureDetail.get()));
+        }
+      }
+    } catch (InterruptedException e) {
+      throw new AbruptExitException(
+          DetailedExitCode.of(
+              FrontierSerializer.createFailureDetail(
+                  "frontier serializer was interrupted: " + e.getMessage(),
+                  Code.SERIALIZED_FRONTIER_PROFILE_FAILED)));
+    } finally {
+      // Do not retain these objects between invocations.
+      this.options = null;
+      this.env = null;
+      System.gc();
+    }
+  }
 }

src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis/BUILD

@@ -16,7 +16,6 @@ java_library(
     srcs = ["FrontierSerializer.java"],
     deps = [
         "//src/main/java/com/google/devtools/build/lib:runtime",
-        "//src/main/java/com/google/devtools/build/lib:runtime/blaze_command_result",
         "//src/main/java/com/google/devtools/build/lib/actions:action_lookup_key",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_options",
@@ -41,3 +40,11 @@ java_library(
         "//third_party/protobuf:protobuf_java",
     ],
 )
+
+java_library(
+    name = "options",
+    srcs = ["RemoteAnalysisCachingOptions.java"],
+    deps = [
+        "//src/main/java/com/google/devtools/common/options",
+    ],
+)

src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis/FrontierSerializer.java

@@ -35,10 +35,10 @@
 import com.google.devtools.build.lib.buildtool.BuildTool;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.PathFragmentPrefixTrie;
+import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.packages.RuleClassProvider;
 import com.google.devtools.build.lib.pkgcache.LoadingFailedException;
-import com.google.devtools.build.lib.runtime.BlazeCommandResult;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.server.FailureDetails;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
@@ -65,7 +65,6 @@
 import java.io.BufferedOutputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
-import java.io.PrintStream;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Optional;
@@ -89,8 +88,8 @@ private FrontierSerializer() {}
    *
    * @return empty if successful, otherwise a result containing the appropriate error
    */
-  public static Optional<BlazeCommandResult> dumpFrontierSerializationProfile(
-      PrintStream out, CommandEnvironment env, String path) throws InterruptedException {
+  public static Optional<FailureDetail> dumpFrontierSerializationProfile(
+      CommandEnvironment env, String path) throws InterruptedException {
     // Starts initializing ObjectCodecs in a background thread as it can take some time.
     var futureCodecs = new FutureTask<>(() -> initObjectCodecs(env));
     commonPool().execute(futureCodecs);
@@ -100,17 +99,22 @@ public static Optional<BlazeCommandResult> dumpFrontierSerializationProfile(
     switch (computeDirectoryMatcher(env.getSkyframeExecutor(), env.getReporter())) {
       case DirectoryMatcherError error:
         return Optional.of(
-            createFailureResult(error.message(), Code.SERIALIZED_FRONTIER_PROFILE_FAILED));
+            createFailureDetail(error.message(), Code.SERIALIZED_FRONTIER_PROFILE_FAILED));
       case ActiveDirectoryMatcher matcher:
         directoryMatcher = matcher;
         break;
     }
-    out.format("Determined active directories in %s\n", stopwatch);
+    env.getReporter()
+        .handle(Event.info(String.format("Determined active directories in %s\n", stopwatch)));
 
     InMemoryGraph graph = env.getSkyframeExecutor().getEvaluator().getInMemoryGraph();
     ConcurrentHashMap<ActionLookupKey, SelectionMarking> selection =
         computeSelection(graph, directoryMatcher);
-    out.format("Found %d active or frontier keys in %s\n", selection.size(), stopwatch);
+    env.getReporter()
+        .handle(
+            Event.info(
+                String.format(
+                    "Found %d active or frontier keys in %s", selection.size(), stopwatch)));
 
     var profileCollector = new ProfileCollector();
     var fingerprintValueService =
@@ -130,9 +134,10 @@ public static Optional<BlazeCommandResult> dumpFrontierSerializationProfile(
     if (codecs == null) {
       String message = "serialization not supported";
       env.getReporter().error(null, message);
-      return Optional.of(createFailureResult(message, Code.SERIALIZED_FRONTIER_PROFILE_FAILED));
+      return Optional.of(createFailureDetail(message, Code.SERIALIZED_FRONTIER_PROFILE_FAILED));
     }
-    out.format("Initializing codecs took %s\n", stopwatch);
+
+    env.getReporter().handle(Event.info(String.format("Initializing codecs took %s\n", stopwatch)));
 
     var writeStatuses = Collections.synchronizedList(new ArrayList<ListenableFuture<Void>>());
     AtomicInteger frontierValueCount = new AtomicInteger();
@@ -164,7 +169,12 @@ public static Optional<BlazeCommandResult> dumpFrontierSerializationProfile(
             writeStatuses.add(immediateFailedFuture(e));
           }
         });
-    out.format("Serialized %s frontier entries in %s\n", frontierValueCount, stopwatch);
+
+    env.getReporter()
+        .handle(
+            Event.info(
+                String.format(
+                    "Serialized %s frontier entries in %s\n", frontierValueCount, stopwatch)));
 
     try {
       var unusedNull =
@@ -176,17 +186,20 @@ public static Optional<BlazeCommandResult> dumpFrontierSerializationProfile(
         message = "with unexpected exception type " + cause.getClass().getName() + ": " + message;
       }
       env.getReporter().error(/* location= */ null, message, cause);
-      return Optional.of(createFailureResult(message, Code.SERIALIZED_FRONTIER_PROFILE_FAILED));
+      return Optional.of(createFailureDetail(message, Code.SERIALIZED_FRONTIER_PROFILE_FAILED));
     }
-    out.format("Waiting for write futures took an additional %s\n", stopwatch);
+    env.getReporter()
+        .handle(
+            Event.info(
+                String.format("Waiting for write futures took an additional %s\n", stopwatch)));
 
     try (var fileOutput = new FileOutputStream(path);
         var bufferedOutput = new BufferedOutputStream(fileOutput)) {
       profileCollector.toProto().writeTo(bufferedOutput);
     } catch (IOException e) {
       String message = "Error writing serialization profile to file: " + e.getMessage();
       env.getReporter().error(null, message, e);
-      return Optional.of(createFailureResult(message, Code.SERIALIZED_FRONTIER_PROFILE_FAILED));
+      return Optional.of(createFailureDetail(message, Code.SERIALIZED_FRONTIER_PROFILE_FAILED));
     }
     return Optional.empty();
   }
@@ -376,12 +389,11 @@ public String toString() {
     }
   }
 
-  static BlazeCommandResult createFailureResult(String message, Code detailedCode) {
-    return BlazeCommandResult.failureDetail(
-        FailureDetail.newBuilder()
-            .setMessage(message)
-            .setRemoteAnalysisCaching(
-                FailureDetails.RemoteAnalysisCaching.newBuilder().setCode(detailedCode))
-            .build());
+  public static FailureDetail createFailureDetail(String message, Code detailedCode) {
+    return FailureDetail.newBuilder()
+        .setMessage(message)
+        .setRemoteAnalysisCaching(
+            FailureDetails.RemoteAnalysisCaching.newBuilder().setCode(detailedCode))
+        .build();
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis/RemoteAnalysisCachingOptions.java

@@ -0,0 +1,31 @@
+// 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.
+package com.google.devtools.build.lib.skyframe.serialization.analysis;
+
+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;
+
+/** Options for caching analysis results remotely. */
+public class RemoteAnalysisCachingOptions extends OptionsBase {
+
+  @Option(
+      name = "serialized_frontier_profile",
+      defaultValue = "",
+      documentationCategory = OptionDocumentationCategory.OUTPUT_SELECTION,
+      effectTags = {OptionEffectTag.BAZEL_MONITORING},
+      help = "Dump a profile of serialized frontier bytes. Specifies the output path.")
+  public String serializedFrontierProfile;
+}

src/test/java/com/google/devtools/build/lib/skyframe/serialization/analysis/BUILD

@@ -10,7 +10,6 @@ java_library(
     testonly = True,
     srcs = ["FrontierSerializerTestBase.java"],
     deps = [
-        "//src/main/java/com/google/devtools/build/lib:runtime/blaze_command_result",
         "//src/main/java/com/google/devtools/build/lib/actions:action_lookup_key",
         "//src/main/java/com/google/devtools/build/lib/collect",
         "//src/main/java/com/google/devtools/build/lib/skyframe:aspect_key_creator",