Reduce interning memory overhead Part III -- Applies SkyKeyInterner to all SkyKeys that are using weak interner

836c608 implements how `InMemoryGraph` interacts with `SkyKeyInterner`, and applies this feature only to `FileValue#Key`.

In this commit, we want to apply `SkyKeyInterner` to all `SkyKey`s using weak interner. Currently there are ~30 types of `SkyKey`s that use bazel weak interner.

PiperOrigin-RevId: 511521005
Change-Id: Ia6fb0fc2d494ac1b470268f7e263f7564d64a56b

Author: yuyue730

src/main/java/com/google/devtools/build/lib/concurrent/BlazeInterners.java

@@ -35,6 +35,12 @@ private static InternerBuilder setConcurrencyLevel(InternerBuilder builder) {
     return builder.concurrencyLevel(CONCURRENCY_LEVEL);
   }
 
+  /**
+   * Creates an interner which retains a weak reference to each instance it has interned.
+   *
+   * <p>It is preferred to use {@code SkyKey#SkyKeyInterner} instead for interning {@code SkyKey}
+   * types.
+   */
   public static <T> Interner<T> newWeakInterner() {
     return setConcurrencyLevel(Interners.newBuilder().weak()).build();
   }

src/main/java/com/google/devtools/build/lib/packages/WorkspaceFileValue.java

@@ -16,9 +16,7 @@
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.vfs.RootedPath;
@@ -43,7 +41,7 @@ public class WorkspaceFileValue implements SkyValue {
   @Immutable
   @AutoCodec
   public static class WorkspaceFileKey implements SkyKey {
-    private static final Interner<WorkspaceFileKey> interner = BlazeInterners.newWeakInterner();
+    private static final SkyKeyInterner<WorkspaceFileKey> interner = SkyKey.newInterner();
 
     private final RootedPath path;
     private final int idx;
@@ -93,6 +91,11 @@ public int hashCode() {
     public String toString() {
       return path + ", " + idx;
     }
+
+    @Override
+    public SkyKeyInterner<WorkspaceFileKey> getSkyKeyInterner() {
+      return interner;
+    }
   }
 
   private final Package pkg;

src/main/java/com/google/devtools/build/lib/skyframe/AspectKeyCreator.java

@@ -16,16 +16,15 @@
 import com.google.common.base.Objects;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.actions.ActionLookupKey;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.packages.AspectClass;
 import com.google.devtools.build.lib.packages.AspectDescriptor;
 import com.google.devtools.build.lib.packages.AspectParameters;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.skyframe.SkyFunctionName;
+import com.google.devtools.build.skyframe.SkyKey;
 import java.util.Comparator;
 import javax.annotation.Nullable;
 
@@ -85,7 +84,7 @@ public final int hashCode() {
   /** Represents an aspect applied to a particular target. */
   @AutoCodec
   public static final class AspectKey extends AspectBaseKey {
-    private static final Interner<AspectKey> interner = BlazeInterners.newWeakInterner();
+    private static final SkyKeyInterner<AspectKey> interner = SkyKey.newInterner();
 
     private final ImmutableList<AspectKey> baseKeys;
     private final AspectDescriptor aspectDescriptor;
@@ -234,6 +233,11 @@ AspectKey withLabel(Label label) {
           newBaseKeys.build(),
           aspectDescriptor);
     }
+
+    @Override
+    public SkyKeyInterner<AspectKey> getSkyKeyInterner() {
+      return interner;
+    }
   }
 
   /**
@@ -242,7 +246,7 @@ AspectKey withLabel(Label label) {
    */
   @AutoCodec
   public static final class TopLevelAspectsKey extends AspectBaseKey {
-    private static final Interner<TopLevelAspectsKey> interner = BlazeInterners.newWeakInterner();
+    private static final SkyKeyInterner<TopLevelAspectsKey> interner = SkyKey.newInterner();
 
     private final ImmutableList<AspectClass> topLevelAspectsClasses;
     private final Label targetLabel;
@@ -324,5 +328,10 @@ public boolean equals(Object o) {
           && Objects.equal(topLevelAspectsClasses, that.topLevelAspectsClasses)
           && Objects.equal(topLevelAspectsParameters, that.topLevelAspectsParameters);
     }
+
+    @Override
+    public SkyKeyInterner<TopLevelAspectsKey> getSkyKeyInterner() {
+      return interner;
+    }
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/BuildConfigurationKey.java

@@ -15,9 +15,7 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.base.Preconditions;
-import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
@@ -57,7 +55,7 @@ public static BuildConfigurationKey withoutPlatformMapping(BuildOptions options)
     return interner.intern(new BuildConfigurationKey(options));
   }
 
-  private static final Interner<BuildConfigurationKey> interner = BlazeInterners.newWeakInterner();
+  private static final SkyKeyInterner<BuildConfigurationKey> interner = SkyKey.newInterner();
 
   private final BuildOptions options;
 
@@ -96,4 +94,9 @@ public String toString() {
     // This format is depended on by integration tests.
     return "BuildConfigurationKey[" + options.checksum() + "]";
   }
+
+  @Override
+  public SkyKeyInterner<BuildConfigurationKey> getSkyKeyInterner() {
+    return interner;
+  }
 }

src/main/java/com/google/devtools/build/lib/skyframe/BuildInfoCollectionValue.java

@@ -14,17 +14,16 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.base.Preconditions;
-import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.actions.ActionLookupKey;
 import com.google.devtools.build.lib.actions.Actions.GeneratingActions;
 import com.google.devtools.build.lib.actions.BasicActionLookupValue;
 import com.google.devtools.build.lib.analysis.buildinfo.BuildInfoCollection;
 import com.google.devtools.build.lib.analysis.buildinfo.BuildInfoKey;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.skyframe.SkyFunctionName;
+import com.google.devtools.build.skyframe.SkyKey;
 import java.util.Objects;
 import javax.annotation.Nullable;
 
@@ -58,8 +57,7 @@ public static BuildInfoKeyAndConfig key(
   /** Key for BuildInfoCollectionValues. */
   @AutoCodec
   public static final class BuildInfoKeyAndConfig extends ActionLookupKey {
-    private static final Interner<BuildInfoKeyAndConfig> keyInterner =
-        BlazeInterners.newWeakInterner();
+    private static final SkyKeyInterner<BuildInfoKeyAndConfig> keyInterner = SkyKey.newInterner();
 
     private final BuildInfoKey infoKey;
     private final BuildConfigurationKey configKey;
@@ -114,5 +112,10 @@ public boolean equals(Object other) {
       return Objects.equals(this.infoKey, that.infoKey)
           && Objects.equals(this.configKey, that.configKey);
     }
+
+    @Override
+    public SkyKeyInterner<BuildInfoKeyAndConfig> getSkyKeyInterner() {
+      return keyInterner;
+    }
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/BuildTopLevelAspectsDetailsFunction.java

@@ -19,10 +19,8 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.analysis.AspectCollection;
 import com.google.devtools.build.lib.analysis.AspectCollection.AspectCycleOnPathException;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.packages.Aspect;
 import com.google.devtools.build.lib.packages.AspectClass;
@@ -244,8 +242,8 @@ public ImmutableList<AspectDetails> getUsedAspects() {
   /** {@link SkyKey} for building top-level aspects details. */
   @AutoCodec
   static final class BuildTopLevelAspectsDetailsKey implements SkyKey {
-    private static final Interner<BuildTopLevelAspectsDetailsKey> interner =
-        BlazeInterners.newWeakInterner();
+    private static final SkyKeyInterner<BuildTopLevelAspectsDetailsKey> interner =
+        SkyKey.newInterner();
 
     private final ImmutableList<AspectClass> topLevelAspectsClasses;
     private final ImmutableMap<String, String> topLevelAspectsParameters;
@@ -311,6 +309,11 @@ public String toString() {
           .add("topLevelAspectsParameters", topLevelAspectsParameters)
           .toString();
     }
+
+    @Override
+    public SkyKeyInterner<BuildTopLevelAspectsDetailsKey> getSkyKeyInterner() {
+      return interner;
+    }
   }
 
   /**

src/main/java/com/google/devtools/build/lib/skyframe/BzlCompileValue.java

@@ -15,9 +15,7 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.base.Preconditions;
-import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
 import com.google.devtools.build.lib.vfs.Root;
@@ -132,8 +130,6 @@ public static BzlCompileValue withProgram(Program prog, byte[] digest) {
     return new Success(prog, digest);
   }
 
-  private static final Interner<Key> keyInterner = BlazeInterners.newWeakInterner();
-
   /** Types of bzl files we may encounter. */
   enum Kind {
     /** A regular .bzl file loaded on behalf of a BUILD or WORKSPACE file. */
@@ -159,6 +155,8 @@ enum Kind {
   /** SkyKey for retrieving a compiled .bzl program. */
   @AutoCodec
   public static class Key implements SkyKey {
+    private static final SkyKeyInterner<Key> interner = SkyKey.newInterner();
+
     /** The root in which the .bzl file is to be found. Null for EMPTY_PRELUDE. */
     @Nullable final Root root;
 
@@ -180,7 +178,7 @@ private Key(Root root, Label label, Kind kind) {
     @AutoCodec.VisibleForSerialization
     @AutoCodec.Instantiator
     static Key create(Root root, Label label, Kind kind) {
-      return keyInterner.intern(new Key(root, label, kind));
+      return interner.intern(new Key(root, label, kind));
     }
 
     /** Returns whether this key is for a {@code @_builtins} .bzl file. */
@@ -225,6 +223,11 @@ public SkyFunctionName functionName() {
     public String toString() {
       return String.format("%s:[%s]%s", functionName(), root, label);
     }
+
+    @Override
+    public SkyKeyInterner<Key> getSkyKeyInterner() {
+      return interner;
+    }
   }
 
   /** Constructs a key for loading a regular (non-prelude) .bzl. */

src/main/java/com/google/devtools/build/lib/skyframe/BzlLoadValue.java

@@ -17,16 +17,15 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
-import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.packages.BzlVisibility;
 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.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyKey.SkyKeyInterner;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.util.Objects;
 import net.starlark.java.eval.Module;
@@ -70,7 +69,7 @@ public BzlVisibility getBzlVisibility() {
     return bzlVisibility;
   }
 
-  private static final Interner<Key> keyInterner = BlazeInterners.newWeakInterner();
+  private static final SkyKeyInterner<Key> keyInterner = SkyKey.newInterner();
 
   /** SkyKey for a Starlark load. */
   public abstract static class Key implements SkyKey {
@@ -157,6 +156,11 @@ protected final MoreObjects.ToStringHelper toStringHelper() {
     public String toString() {
       return toStringHelper().toString();
     }
+
+    @Override
+    public SkyKeyInterner<Key> getSkyKeyInterner() {
+      return keyInterner;
+    }
   }
 
   /** A key for loading a .bzl during package loading (BUILD evaluation). */

src/main/java/com/google/devtools/build/lib/skyframe/ClientEnvironmentFunction.java

@@ -14,8 +14,6 @@
 
 package com.google.devtools.build.lib.skyframe;
 
-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.AbstractSkyKey;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -35,7 +33,7 @@ public static Key key(String keyString) {
   @AutoCodec.VisibleForSerialization
   @AutoCodec
   static class Key extends AbstractSkyKey<String> {
-    private static final Interner<Key> interner = BlazeInterners.newWeakInterner();
+    private static final SkyKeyInterner<Key> interner = SkyKey.newInterner();
 
     private Key(String arg) {
       super(arg);
@@ -51,6 +49,11 @@ static Key create(String arg) {
     public SkyFunctionName functionName() {
       return SkyFunctions.CLIENT_ENVIRONMENT_VARIABLE;
     }
+
+    @Override
+    public SkyKeyInterner<Key> getSkyKeyInterner() {
+      return interner;
+    }
   }
 
   private final AtomicReference<Map<String, String>> clientEnv;

src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetKey.java

@@ -17,14 +17,13 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.base.MoreObjects;
-import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.actions.ActionLookupKey;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.skyframe.SkyFunctionName;
+import com.google.devtools.build.skyframe.SkyKey;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.Objects;
 import javax.annotation.Nullable;
@@ -63,7 +62,7 @@ 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)}.
    */
-  private static final Interner<ConfiguredTargetKey> interner = BlazeInterners.newWeakInterner();
+  private static final SkyKeyInterner<ConfiguredTargetKey> interner = SkyKey.newInterner();
 
   private final Label label;
   @Nullable private final BuildConfigurationKey configurationKey;
@@ -161,11 +160,16 @@ public final String toString() {
     return helper.toString();
   }
 
+  @Override
+  public SkyKeyInterner<? extends ConfiguredTargetKey> getSkyKeyInterner() {
+    return interner;
+  }
+
   @AutoCodec.VisibleForSerialization
   @AutoCodec
   static class ToolchainDependencyConfiguredTargetKey extends ConfiguredTargetKey {
-    private static final Interner<ToolchainDependencyConfiguredTargetKey>
-        toolchainDependencyConfiguredTargetKeyInterner = BlazeInterners.newWeakInterner();
+    private static final SkyKeyInterner<ToolchainDependencyConfiguredTargetKey>
+        toolchainDependencyConfiguredTargetKeyInterner = SkyKey.newInterner();
 
     private final Label executionPlatformLabel;
 
@@ -194,6 +198,11 @@ static ToolchainDependencyConfiguredTargetKey create(
     public final Label getExecutionPlatformLabel() {
       return executionPlatformLabel;
     }
+
+    @Override
+    public final SkyKeyInterner<? extends ConfiguredTargetKey> getSkyKeyInterner() {
+      return toolchainDependencyConfiguredTargetKeyInterner;
+    }
   }
 
   /** Returns a new {@link Builder} to create instances of {@link ConfiguredTargetKey}. */