Cache StarlarkDefinedAspect AspectDefinitons

While replacing `java_lite_proto_library` native implementation with Starlarks, a significant regression occurred in retained heap. Inspecting the heap, there were more instances of AspectDefinition and instances of classes used within it. Further inspection of unrelated AspectDefinition showed, that an instance is created for every target that uses a rule with a Starlark aspect.

The native aspects already used a cache - implemented in https://cs.opensource.google/bazel/bazel/+/master:src/main/java/com/google/devtools/build/lib/packages/Aspect.java;l=47-54;drc=a4063f6928965e3cc55a18a6b3efff032abb7311

Implement similar cache except for using `weakKeys` (because moving to a different part of the depot, should free those aspects).

I considered putting the cache into the same location (`Aspect` class), but that would require refactoring signature of `Aspect.forStarlark` method and there were some preliminary problems with `StarlarkNativeAspect`-s.

PiperOrigin-RevId: 447918213

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

@@ -41,8 +41,6 @@ public final class Aspect implements DependencyFilter.AttributeInfoProvider {
    *
    * <p>The native aspects are loaded with blaze and are not stateful. Reference equality works fine
    * in this case.
-   *
-   * <p>Caching of Starlark aspects is not yet implemented.
    */
   private static final LoadingCache<
           NativeAspectClass, LoadingCache<AspectParameters, AspectDefinition>>

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

@@ -14,6 +14,8 @@
 
 package com.google.devtools.build.lib.packages;
 
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.LoadingCache;
 import com.google.common.base.Ascii;
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
@@ -135,7 +137,31 @@ public void export(EventHandler handler, Label extensionLabel, String name) {
 
   private static final ImmutableList<String> ALL_ATTR_ASPECTS = ImmutableList.of("*");
 
+  /**
+   * The <code>AspectDefinition</code> is a function of the aspect's definition + its parameters, so
+   * we can cache that.
+   *
+   * <p>StarlarkDefinedAspect key: Weak keys use reference comparison, so the values are freed when
+   * aspect's definition changes or when it's evicted from the skyframe.
+   *
+   * <p>AspectParameters key: parameters of Starlark aspects are combinatorially limited (only bool,
+   * int and enum types). Using strong keys possibly results in a small memory leak. Weak keys don't
+   * work because reference equality is used and AspectParameters are created per target.
+   */
+  private static final LoadingCache<
+          StarlarkDefinedAspect, LoadingCache<AspectParameters, AspectDefinition>>
+      definitionCache =
+          Caffeine.newBuilder()
+              .weakKeys()
+              .build(
+                  starlarkDefinedAspect ->
+                      Caffeine.newBuilder().build(starlarkDefinedAspect::buildDefinition));
+
   public AspectDefinition getDefinition(AspectParameters aspectParams) {
+    return definitionCache.get(this).get(aspectParams);
+  }
+
+  private AspectDefinition buildDefinition(AspectParameters aspectParams) {
     AspectDefinition.Builder builder = new AspectDefinition.Builder(aspectClass);
     if (ALL_ATTR_ASPECTS.equals(attributeAspects)) {
       builder.propagateAlongAllAttributes();