Address review comments

bazelbuild · Aug 9, 2023 · ab31c6b · ab31c6b
1 parent e19b541
commit ab31c6b
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 3 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/BUILD b/src/main/java/com/google/devtools/build/lib/rules/cpp/BUILD
@@ -38,6 +38,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
         "//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
+        "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
         "//src/main/java/com/google/devtools/build/lib/actions:localhost_capacity",
         "//src/main/java/com/google/devtools/build/lib/actions:middleman_type",
         "//src/main/java/com/google/devtools/build/lib/actions:runfiles_supplier",

diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java
@@ -40,6 +40,7 @@
 import com.google.devtools.build.lib.actions.CommandLines.CommandLineAndParamFileInfo;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionRequirements;
+import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.ResourceSet;
 import com.google.devtools.build.lib.actions.SimpleSpawn;
 import com.google.devtools.build.lib.actions.SimpleSpawn.LocalResourcesSupplier;
@@ -335,6 +336,7 @@ public ActionResult execute(ActionExecutionContext actionExecutionContext)
       throws ActionExecutionException, InterruptedException {
     LocalResourcesEstimator localResourcesEstimator =
         new LocalResourcesEstimator(
+            actionExecutionContext,
             OS.getCurrent(),
             linkCommandLine.getLinkerInputArtifacts());
 
@@ -491,6 +493,7 @@ protected String getRawProgressMessage() {
    */
   @VisibleForTesting
   static class LocalResourcesEstimator implements LocalResourcesSupplier {
+    private final ActionExecutionContext actionExecutionContext;
     private final OS os;
     private final NestedSet<Artifact> inputs;
 
@@ -509,7 +512,9 @@ private static class LazyData {
 
     private LazyData lazyData = null;
 
-    public LocalResourcesEstimator(OS os, NestedSet<Artifact> inputs) {
+    public LocalResourcesEstimator(
+        ActionExecutionContext actionExecutionContext, OS os, NestedSet<Artifact> inputs) {
+      this.actionExecutionContext = actionExecutionContext;
       this.os = os;
       this.inputs = inputs;
     }
@@ -521,7 +526,14 @@ private LazyData doCostlyEstimation() {
       for (Artifact input : inputs.toList()) {
         inputsCount += 1;
         try {
-          inputsBytes += input.getPath().getFileSize();
+          FileArtifactValue value =
+              actionExecutionContext.getInputMetadataProvider().getInputMetadata(input);
+          if (value != null) {
+            inputsBytes += value.getSize();
+          } else {
+            logger.atWarning().log(
+                "Linker metrics: failed to get size of %s: no metadata", input.getExecPath());
+          }
         } catch (Exception e) {
           logger.atWarning().log(
               "Linker metrics: failed to get size of %s: %s (ignored)", input.getExecPath(), e);

diff --git a/src/test/java/com/google/devtools/build/lib/rules/cpp/BUILD b/src/test/java/com/google/devtools/build/lib/rules/cpp/BUILD
@@ -451,6 +451,7 @@ java_test(
         "//src/test/java/com/google/devtools/build/lib/testutil:TestUtils",
         "//third_party:guava",
         "//third_party:junit4",
+        "//third_party:mockito",
         "//third_party:truth",
     ],
 )

diff --git a/src/test/java/com/google/devtools/build/lib/rules/cpp/CppLinkActionTest.java b/src/test/java/com/google/devtools/build/lib/rules/cpp/CppLinkActionTest.java
@@ -16,6 +16,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.mock;
 
 import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
@@ -25,6 +26,7 @@
 import com.google.common.primitives.Ints;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
+import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.ArtifactExpander;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
@@ -785,9 +787,12 @@ private static NestedSet<Artifact> createInputs(RuleContext ruleContext, int cou
 
   private ResourceSet estimateResourceConsumptionLocal(
       RuleContext ruleContext, OS os, int inputsCount) throws Exception {
+    ActionExecutionContext actionExecutionContext = mock(ActionExecutionContext.class);
+
     NestedSet<Artifact> inputs = createInputs(ruleContext, inputsCount);
     try {
-      LocalResourcesEstimator estimator = new LocalResourcesEstimator(os, inputs);
+      LocalResourcesEstimator estimator =
+          new LocalResourcesEstimator(actionExecutionContext, os, inputs);
       return estimator.get();
     } finally {
       for (Artifact input : inputs.toList()) {