Cherry-pick Java execution info improvements (#21703)

* [Avoid making a copy in mergeMaps if one of the input maps is empty.](3e8ee27) * [Add the path for in-memory jdeps files to execution info on demand instead of storing it in `JavaCompileAction` and `JavaHeaderCompileAction`.](c830f21) Fixes #21661 --------- Co-authored-by: Googler <jhorvitz@google.com>
bazelbuild · Mar 14, 2024 · 7c29e52 · 7c29e52
1 parent 010ac99
commit 7c29e52
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 21 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionAnalysisMetadata.java b/src/main/java/com/google/devtools/build/lib/actions/ActionAnalysisMetadata.java
@@ -253,6 +253,12 @@ default ImmutableMap<String, String> getExecutionInfo() {
 
   static ImmutableMap<String, String> mergeMaps(
       ImmutableMap<String, String> first, ImmutableMap<String, String> second) {
+    if (first.isEmpty()) {
+      return second;
+    }
+    if (second.isEmpty()) {
+      return first;
+    }
     return ImmutableMap.<String, String>builderWithExpectedSize(first.size() + second.size())
         .putAll(first)
         .putAll(second)

diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java
@@ -51,6 +51,7 @@
 import com.google.devtools.build.lib.actions.EmptyRunfilesSupplier;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
+import com.google.devtools.build.lib.actions.ExecutionRequirements;
 import com.google.devtools.build.lib.actions.ParamFileInfo;
 import com.google.devtools.build.lib.actions.ParameterFile;
 import com.google.devtools.build.lib.actions.PathMapper;
@@ -661,7 +662,16 @@ public Sequence<String> getStarlarkArgv() throws EvalException, InterruptedExcep
   /** Returns the out-of-band execution data for this action. */
   @Override
   public ImmutableMap<String, String> getExecutionInfo() {
-    return mergeMaps(super.getExecutionInfo(), executionInfo);
+    var result = mergeMaps(super.getExecutionInfo(), executionInfo);
+    if (outputDepsProto == null
+        || !configuration.getFragment(JavaConfiguration.class).inmemoryJdepsFiles()) {
+      return result;
+    }
+    return mergeMaps(
+        result,
+        ImmutableMap.of(
+            ExecutionRequirements.REMOTE_EXECUTION_INLINE_OUTPUTS,
+            outputDepsProto.getExecPathString()));
   }
 
   @Override

diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileActionBuilder.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileActionBuilder.java
@@ -23,7 +23,6 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.Artifact;
-import com.google.devtools.build.lib.actions.ExecutionRequirements;
 import com.google.devtools.build.lib.actions.extra.ExtraActionInfo;
 import com.google.devtools.build.lib.actions.extra.JavaCompileInfo;
 import com.google.devtools.build.lib.analysis.RuleContext;
@@ -230,20 +229,6 @@ public JavaCompileAction build() throws RuleErrorException {
       classpathMode = JavaClasspathMode.OFF;
     }
 
-    if (outputs.depsProto() != null) {
-      JavaConfiguration javaConfiguration =
-          ruleContext.getConfiguration().getFragment(JavaConfiguration.class);
-      if (javaConfiguration.inmemoryJdepsFiles()) {
-        executionInfo =
-            ImmutableMap.<String, String>builderWithExpectedSize(this.executionInfo.size() + 1)
-                .putAll(this.executionInfo)
-                .put(
-                    ExecutionRequirements.REMOTE_EXECUTION_INLINE_OUTPUTS,
-                    outputs.depsProto().getExecPathString())
-                .buildOrThrow();
-      }
-    }
-
     NestedSet<Artifact> tools = toolsBuilder.build();
     mandatoryInputsBuilder.addTransitive(tools);
     NestedSet<Artifact> mandatoryInputs = mandatoryInputsBuilder.build();

diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileAction.java
@@ -16,6 +16,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.devtools.build.lib.actions.ActionAnalysisMetadata.mergeMaps;
 import static com.google.devtools.build.lib.actions.ParameterFile.ParameterFileType.UNQUOTED;
 import static com.google.devtools.build.lib.rules.java.JavaCompileActionBuilder.UTF8_ENVIRONMENT;
 import static java.nio.charset.StandardCharsets.ISO_8859_1;
@@ -80,6 +81,7 @@ public final class JavaHeaderCompileAction extends SpawnAction {
   private static final String DIRECT_CLASSPATH_MNEMONIC = "Turbine";
 
   private final boolean insertDependencies;
+  private final boolean inMemoryJdeps;
   private final NestedSet<Artifact> additionalArtifactsForPathMapping;
 
   private JavaHeaderCompileAction(
@@ -96,6 +98,7 @@ private JavaHeaderCompileAction(
       String mnemonic,
       OutputPathsMode outputPathsMode,
       boolean insertDependencies,
+      boolean inMemoryJdeps,
       NestedSet<Artifact> additionalArtifactsForPathMapping) {
     super(
         owner,
@@ -111,9 +114,24 @@ private JavaHeaderCompileAction(
         mnemonic,
         outputPathsMode);
     this.insertDependencies = insertDependencies;
+    this.inMemoryJdeps = inMemoryJdeps;
     this.additionalArtifactsForPathMapping = additionalArtifactsForPathMapping;
   }
 
+  @Override
+  public ImmutableMap<String, String> getExecutionInfo() {
+    var result = super.getExecutionInfo();
+    if (!inMemoryJdeps) {
+      return result;
+    }
+    Artifact outputDepsProto = Iterables.get(getOutputs(), 1);
+    return mergeMaps(
+        result,
+        ImmutableMap.of(
+            ExecutionRequirements.REMOTE_EXECUTION_INLINE_OUTPUTS,
+            outputDepsProto.getExecPathString()));
+  }
+
   @Override
   public NestedSet<Artifact> getAdditionalArtifactsForPathMapping() {
     return additionalArtifactsForPathMapping;
@@ -478,11 +496,6 @@ public void build(JavaToolchainProvider javaToolchain) throws RuleErrorException
       executionInfo.putAll(
           TargetUtils.getExecutionInfo(
               ruleContext.getRule(), ruleContext.isAllowTagsPropagation()));
-      if (javaConfiguration.inmemoryJdepsFiles()) {
-        executionInfo.put(
-            ExecutionRequirements.REMOTE_EXECUTION_INLINE_OUTPUTS,
-            outputDepsProto.getExecPathString());
-      }
 
       if (useDirectClasspath) {
         NestedSet<Artifact> classpath;
@@ -535,6 +548,7 @@ public void build(JavaToolchainProvider javaToolchain) throws RuleErrorException
                 // available to downstream actions. Else just do enough work to locally create the
                 // full .jdeps from the .stripped .jdeps produced on the executor.
                 /* insertDependencies= */ classpathMode == JavaClasspathMode.BAZEL,
+                javaConfiguration.inmemoryJdepsFiles(),
                 additionalArtifactsForPathMapping));
         return;
       }