Mark SkyKeys having CPUHeavy dependencies CPUHeavy.

* This may improve walltime in high CPU settings. * This is a precursor to SkyKey prioritization, which requires it. PiperOrigin-RevId: 506714840 Change-Id: Ia2d92a39de178d5b6220232cbe5a81becec7f1fb
bazelbuild · Feb 2, 2023 · 6224c87 · 6224c87
1 parent 49f98ab
commit 6224c87
Show file tree

Hide file tree

Showing 20 changed files with 85 additions and 31 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionLookupKey.java b/src/main/java/com/google/devtools/build/lib/actions/ActionLookupKey.java
@@ -30,14 +30,14 @@
  * are subclasses of {@link ActionLookupKey}. This allows callers to easily find the value key,
  * while remaining agnostic to what action lookup values actually exist.
  */
-public interface ActionLookupKey extends ArtifactOwner, CPUHeavySkyKey {
+public abstract class ActionLookupKey extends CPUHeavySkyKey implements ArtifactOwner {
 
   /**
    * Returns the {@link BuildConfigurationKey} for the configuration associated with this key, or
    * {@code null} if this key has no associated configuration.
    */
   @Nullable
-  BuildConfigurationKey getConfigurationKey();
+  public abstract BuildConfigurationKey getConfigurationKey();
 
   /**
    * Returns {@code true} if this key <em>may</em> own shareable actions, as determined by {@link
@@ -50,7 +50,7 @@ public interface ActionLookupKey extends ArtifactOwner, CPUHeavySkyKey {
    * to determine whether the individual action can be shared - notably, for a test target,
    * compilation actions are shareable, but test actions are not.
    */
-  default boolean mayOwnShareableActions() {
+  public boolean mayOwnShareableActions() {
     return getLabel() != null;
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionValue.java
@@ -36,7 +36,7 @@ public static ActionTemplateExpansionKey key(ActionLookupKey actionLookupKey, in
 
   /** Key for {@link ActionTemplateExpansionValue} nodes. */
   @AutoCodec
-  public static final class ActionTemplateExpansionKey implements ActionLookupKey {
+  public static final class ActionTemplateExpansionKey extends ActionLookupKey {
     private static final Interner<ActionTemplateExpansionKey> interner =
         BlazeInterners.newWeakInterner();
 

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/AspectKeyCreator.java b/src/main/java/com/google/devtools/build/lib/skyframe/AspectKeyCreator.java
@@ -60,7 +60,7 @@ public static TopLevelAspectsKey createTopLevelAspectsKey(
   }
 
   /** Common superclass for {@link AspectKey} and {@link TopLevelAspectsKey}. */
-  public abstract static class AspectBaseKey implements ActionLookupKey {
+  public abstract static class AspectBaseKey extends ActionLookupKey {
     private final ConfiguredTargetKey baseConfiguredTargetKey;
     private final int hashCode;
 

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/BUILD b/src/main/java/com/google/devtools/build/lib/skyframe/BUILD
@@ -346,6 +346,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/vfs:output_service",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/skyframe",
+        "//src/main/java/com/google/devtools/build/skyframe:cpu_heavy_skykey",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/java/net/starlark/java/eval",
@@ -2241,6 +2242,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
         "//src/main/java/com/google/devtools/build/lib/util:detailed_exit_code",
         "//src/main/java/com/google/devtools/build/lib/util:exit_code",
+        "//src/main/java/com/google/devtools/build/skyframe:cpu_heavy_skykey",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//third_party:auto_value",
         "//third_party:guava",
@@ -2414,6 +2416,7 @@ java_library(
         ":sky_functions",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:serialization-constant",
+        "//src/main/java/com/google/devtools/build/skyframe:cpu_heavy_skykey",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//src/main/java/net/starlark/java/eval",
         "//third_party:guava",
@@ -2432,6 +2435,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
         "//src/main/java/com/google/devtools/build/lib/vfs",
+        "//src/main/java/com/google/devtools/build/skyframe:cpu_heavy_skykey",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//src/main/java/net/starlark/java/eval",
         "//third_party:guava",
@@ -2633,6 +2637,7 @@ java_library(
         ":sky_functions",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/toolchain_type_requirement",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
+        "//src/main/java/com/google/devtools/build/skyframe:cpu_heavy_skykey",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//third_party:auto_value",
         "//third_party:guava",

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/BuildDriverKey.java b/src/main/java/com/google/devtools/build/lib/skyframe/BuildDriverKey.java
@@ -23,7 +23,7 @@
  * Wraps an {@link ActionLookupKey}. The evaluation of this SkyKey is the entry point of analyzing
  * the {@link ActionLookupKey} and executing the associated actions.
  */
-public final class BuildDriverKey implements CPUHeavySkyKey {
+public final class BuildDriverKey extends CPUHeavySkyKey {
   private final ActionLookupKey actionLookupKey;
   private final TopLevelArtifactContext topLevelArtifactContext;
   private final TestType testType;

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/BuildInfoCollectionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/BuildInfoCollectionValue.java
@@ -57,7 +57,7 @@ public static BuildInfoKeyAndConfig key(
 
   /** Key for BuildInfoCollectionValues. */
   @AutoCodec
-  public static final class BuildInfoKeyAndConfig implements ActionLookupKey {
+  public static final class BuildInfoKeyAndConfig extends ActionLookupKey {
     private static final Interner<BuildInfoKeyAndConfig> keyInterner =
         BlazeInterners.newWeakInterner();
 

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/BzlLoadValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/BzlLoadValue.java
@@ -25,8 +25,8 @@
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.lib.vfs.RootedPath;
+import com.google.devtools.build.skyframe.CPUHeavySkyKey;
 import com.google.devtools.build.skyframe.SkyFunctionName;
-import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.util.Objects;
 import net.starlark.java.eval.Module;
@@ -73,7 +73,7 @@ public BzlVisibility getBzlVisibility() {
   private static final Interner<Key> keyInterner = BlazeInterners.newWeakInterner();
 
   /** SkyKey for a Starlark load. */
-  public abstract static class Key implements SkyKey {
+  public abstract static class Key extends CPUHeavySkyKey {
     // Closed, for class-based equals()/hashCode().
     private Key() {}
 

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetKey.java b/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetKey.java
@@ -58,7 +58,7 @@
  * BuildConfigurationKey.
  */
 @AutoCodec
-public class ConfiguredTargetKey implements ActionLookupKey {
+public class ConfiguredTargetKey extends ActionLookupKey {
   /**
    * Cache so that the number of ConfiguredTargetKey instances is {@code O(configured targets)} and
    * not {@code O(edges between configured targets)}.
@@ -101,6 +101,7 @@ public final SkyFunctionName functionName() {
   }
 
   @Nullable
+  @Override
   public final BuildConfigurationKey getConfigurationKey() {
     return configurationKey;
   }

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CoverageReportValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/CoverageReportValue.java
@@ -32,7 +32,7 @@ public final class CoverageReportValue extends BasicActionLookupValue {
     super(generatingActions);
   }
 
-  private static final class CoverageReportKey implements ActionLookupKey {
+  private static final class CoverageReportKey extends ActionLookupKey {
     private CoverageReportKey() {}
 
     @Override

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PackageValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/PackageValue.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Interner;
@@ -22,7 +23,6 @@
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
-import com.google.devtools.build.skyframe.AbstractSkyKey;
 import com.google.devtools.build.skyframe.CPUHeavySkyKey;
 import com.google.devtools.build.skyframe.NotComparableSkyValue;
 import com.google.devtools.build.skyframe.SkyFunctionName;
@@ -63,11 +63,18 @@ public static Key key(PackageIdentifier pkgIdentifier) {
   /** Skyframe key for packages */
   @AutoCodec.VisibleForSerialization
   @AutoCodec
-  public static class Key extends AbstractSkyKey<PackageIdentifier> implements CPUHeavySkyKey {
+  public static class Key extends CPUHeavySkyKey {
     private static final Interner<Key> interner = BlazeInterners.newWeakInterner();
 
+    private final PackageIdentifier arg;
+
     private Key(PackageIdentifier arg) {
-      super(arg);
+      this.arg = arg;
+    }
+
+    @Override
+    public PackageIdentifier argument() {
+      return arg;
     }
 
     @AutoCodec.VisibleForSerialization
@@ -80,6 +87,25 @@ static Key create(PackageIdentifier arg) {
     public SkyFunctionName functionName() {
       return SkyFunctions.PACKAGE;
     }
+
+    @Override
+    public int hashCode() {
+      return 31 * functionName().hashCode() + arg.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (!(obj instanceof Key)) {
+        return false;
+      }
+      var that = (Key) obj;
+      return this.arg.equals(that.arg);
+    }
+
+    @Override
+    public String toString() {
+      return functionName() + ":" + arg;
+    }
   }
 
   public static ImmutableList<SkyKey> keys(Iterable<PackageIdentifier> pkgIdentifiers) {

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/RegisteredExecutionPlatformsValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/RegisteredExecutionPlatformsValue.java
@@ -19,6 +19,7 @@
 import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.skyframe.CPUHeavySkyKey;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
@@ -38,7 +39,7 @@ public static SkyKey key(BuildConfigurationKey configurationKey) {
   /** {@link SkyKey} implementation used for {@link RegisteredExecutionPlatformsFunction}. */
   @AutoCodec
   @AutoCodec.VisibleForSerialization
-  static class Key implements SkyKey {
+  static class Key extends CPUHeavySkyKey {
     private static final Interner<Key> interners = BlazeInterners.newWeakInterner();
 
     private final BuildConfigurationKey configurationKey;

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/RegisteredToolchainsValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/RegisteredToolchainsValue.java
@@ -20,6 +20,7 @@
 import com.google.devtools.build.lib.analysis.platform.DeclaredToolchainInfo;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.skyframe.CPUHeavySkyKey;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
@@ -36,9 +37,13 @@ public static Key key(BuildConfigurationKey configurationKey) {
     return Key.of(configurationKey);
   }
 
-  /** A {@link SkyKey} for {@code RegisteredToolchainsValue}. */
+  /**
+   * A {@link SkyKey} for {@code RegisteredToolchainsValue}.
+   *
+   * <p>This is marked CPU-Heavy because it causes package loading.
+   */
   @AutoCodec
-  static class Key implements SkyKey {
+  static class Key extends CPUHeavySkyKey {
     private static final Interner<Key> interners = BlazeInterners.newWeakInterner();
 
     private final BuildConfigurationKey configurationKey;

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/RepositoryMappingValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/RepositoryMappingValue.java
@@ -23,6 +23,7 @@
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.util.DetailedExitCode;
 import com.google.devtools.build.lib.util.ExitCode;
+import com.google.devtools.build.skyframe.CPUHeavySkyKey;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
@@ -99,9 +100,13 @@ public String toString() {
     return repositoryMapping.toString();
   }
 
-  /** {@link SkyKey} for {@link RepositoryMappingValue}. */
+  /**
+   * {@link SkyKey} for {@link RepositoryMappingValue}.
+   *
+   * <p>Marked CPU heavy because it causes package loading.
+   */
   @AutoValue
-  public abstract static class Key implements SkyKey {
+  public abstract static class Key extends CPUHeavySkyKey {
 
     private static final Interner<Key> interner = BlazeInterners.newWeakInterner();
 

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionValue.java
@@ -21,6 +21,7 @@
 import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.platform.ToolchainTypeInfo;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.skyframe.CPUHeavySkyKey;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
@@ -65,9 +66,13 @@ public static SingleToolchainResolutionKey key(
         debugTarget);
   }
 
-  /** {@link SkyKey} implementation used for {@link SingleToolchainResolutionFunction}. */
+  /**
+   * {@link SkyKey} implementation used for {@link SingleToolchainResolutionFunction}.
+   *
+   * <p>Marked CPU heavy because it transitively causes package loading.
+   */
   @AutoValue
-  public abstract static class SingleToolchainResolutionKey implements SkyKey {
+  public abstract static class SingleToolchainResolutionKey extends CPUHeavySkyKey {
 
     @Override
     public SkyFunctionName functionName() {

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/StarlarkBuiltinsValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/StarlarkBuiltinsValue.java
@@ -17,8 +17,8 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
+import com.google.devtools.build.skyframe.CPUHeavySkyKey;
 import com.google.devtools.build.skyframe.SkyFunctionName;
-import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import net.starlark.java.eval.StarlarkSemantics;
 
@@ -132,8 +132,10 @@ public static Key key() {
    * Skyframe key for retrieving the {@code @_builtins} definitions.
    *
    * <p>This has no fields since there is only one {@code StarlarkBuiltinsValue} at a time.
+   *
+   * <p>It is marked CPU heavy because it causes package loading.
    */
-  static final class Key implements SkyKey {
+  static final class Key extends CPUHeavySkyKey {
 
     private static final Key INSTANCE = new Key();
 

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainContextKey.java b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainContextKey.java
@@ -19,17 +19,19 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.skyframe.CPUHeavySkyKey;
 import com.google.devtools.build.skyframe.SkyFunctionName;
-import com.google.devtools.build.skyframe.SkyKey;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.Optional;
 
 /**
  * {@link SkyKey} implementation used for {@link ToolchainResolutionFunction} to produce {@link
  * UnloadedToolchainContextImpl} instances.
+ *
+ * <p>Marked CPU heavy because it may cause package loading.
  */
 @AutoValue
-public abstract class ToolchainContextKey implements SkyKey {
+public abstract class ToolchainContextKey extends CPUHeavySkyKey {
 
   /** Returns a new {@link Builder}. */
   public static Builder key() {

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceStatusValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceStatusValue.java
@@ -54,7 +54,7 @@ public Artifact getVolatileArtifact() {
   }
 
   /** {@link com.google.devtools.build.skyframe.SkyKey} for {@link WorkspaceStatusValue}. */
-  public static final class BuildInfoKey implements ActionLookupKey {
+  public static final class BuildInfoKey extends ActionLookupKey {
     private BuildInfoKey() {}
 
     @Override

diff --git a/src/main/java/com/google/devtools/build/skyframe/CPUHeavySkyKey.java b/src/main/java/com/google/devtools/build/skyframe/CPUHeavySkyKey.java
@@ -14,9 +14,11 @@
 package com.google.devtools.build.skyframe;
 
 /**
- * An empty interface used to annotate whether the evaluation of a SkyKey contributes significantly
- * to the CPU footprint of Skyframe.
+ * A {@link SkyKey} for a {@link SkyFunction} that causes heavy resource consumption.
+ *
+ * <p>This applies to both {@link SkyKey}s that have a high resource footprint and ancestors of
+ * those {@link SkyKey}s that depend on them, transitively.
  *
  * <p>This is currently only applicable to the loading/analysis phase of Skyframe.
  */
-public interface CPUHeavySkyKey extends SkyKey {}
+public abstract class CPUHeavySkyKey implements SkyKey {}
diff --git a/src/test/java/com/google/devtools/build/lib/actions/util/InjectedActionLookupKey.java b/src/test/java/com/google/devtools/build/lib/actions/util/InjectedActionLookupKey.java
@@ -24,7 +24,7 @@
  * An {@link ActionLookupKey} with a non-hermetic {@link SkyFunctionName} so that its value can be
  * directly injected during tests.
  */
-public final class InjectedActionLookupKey implements ActionLookupKey {
+public final class InjectedActionLookupKey extends ActionLookupKey {
   public static final SkyFunctionName INJECTED_ACTION_LOOKUP =
       SkyFunctionName.createNonHermetic("INJECTED_ACTION_LOOKUP");
 

diff --git a/src/test/java/com/google/devtools/build/lib/analysis/RunfilesTest.java b/src/test/java/com/google/devtools/build/lib/analysis/RunfilesTest.java
@@ -155,7 +155,7 @@ private void checkConflictError() {
     assertThat(Iterables.getOnlyElement(eventCollector).getKind()).isEqualTo(EventKind.ERROR);
   }
 
-  private static final class SimpleActionLookupKey implements ActionLookupKey {
+  private static final class SimpleActionLookupKey extends ActionLookupKey {
     private final String name;
 
     SimpleActionLookupKey(String name) {