Allow BzlLoadFunction to fully reset the cache.

Some tests (specifically those that use `usesInliningBzlLoadFunction`) mutate the builtins, and so need to be able to reset the cache, making the comment at [`BzlLoadFunction$InlineCacheManager.reset`](https://cs.opensource.google/bazel/bazel/+/master:src/main/java/com/google/devtools/build/lib/skyframe/BzlLoadFunction.java;drc=2c0faf93deae61bfc0bb1c42becbca4c49c988e6;l=1551) invalid.

Work towards platform-based flags: #19409.

PiperOrigin-RevId: 640607616
Change-Id: I3c30b2888c99b46084ea2519870dc6d61c5a2a79

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

@@ -15,6 +15,7 @@
 
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableList;
@@ -171,6 +172,11 @@ public static BzlLoadFunction create(
         /* inlineCacheManager= */ null);
   }
 
+  /**
+   * Constructs a new instance that uses bzl inlining.
+   *
+   * <p>Must call {@link #resetInliningCache} on the returned instance before use.
+   */
   public static BzlLoadFunction createForInlining(
       RuleClassProvider ruleClassProvider,
       BlazeDirectories directories,
@@ -362,8 +368,15 @@ private CachedBzlLoadData computeInlineForCacheMiss(
     return childState.buildCachedData(key, value);
   }
 
+  /** Re-initializes the bzl inlining cache, if this instance uses one. No-op otherwise. */
   public void resetInliningCache() {
-    inlineCacheManager.reset();
+    inlineCacheManager.reset(/* resetBuiltins= */ false);
+  }
+
+  /** Re-initializes the bzl inlining cache, if this instance uses one. No-op otherwise. */
+  @VisibleForTesting
+  public void resetInliningCacheAndBuiltinsForTesting() {
+    inlineCacheManager.reset(/* resetBuiltins= */ true);
   }
 
   /**
@@ -1516,11 +1529,13 @@ public void doneWithBzlCompileValue(BzlCompileValue.Key key) {
    */
   static class InlineCacheManager {
     private final int bzlLoadCacheSize;
+
+    // Data which will be cleared on #reset().
     private Cache<BzlLoadValue.Key, CachedBzlLoadData> bzlLoadCache;
     private CachedBzlLoadDataBuilderFactory cachedBzlLoadDataBuilderFactory =
         new CachedBzlLoadDataBuilderFactory();
-
-    final AtomicReference<StarlarkBuiltinsValue> builtinsRef = new AtomicReference<>();
+    // Not private so that StarlarkBuiltinsFunction can directly access this.
+    AtomicReference<StarlarkBuiltinsValue> builtinsRef = new AtomicReference<>();
 
     private InlineCacheManager(int bzlLoadCacheSize) {
       this.bzlLoadCacheSize = bzlLoadCacheSize;
@@ -1530,7 +1545,7 @@ private CachedBzlLoadData.Builder cachedDataBuilder() {
       return cachedBzlLoadDataBuilderFactory.newCachedBzlLoadDataBuilder();
     }
 
-    private void reset() {
+    private void reset(boolean resetBuiltins) {
       if (bzlLoadCache != null) {
         logger.atInfo().log(
             "Starlark inlining cache stats from earlier build: %s", bzlLoadCache.stats());
@@ -1546,13 +1561,16 @@ private void reset() {
               .maximumSize(bzlLoadCacheSize)
               .recordStats()
               .build();
-      // Don't reset `builtinsRef`.
-      //
-      // All usages of BzlLoadFunction inlining assume builtins can never change (i.e. no usage of
-      // --experimental_builtins_bzl_path, no inter-invocation flipping of Starlark semantics
-      // options that'd cause evaluation of builtins to differ). If this assumption ever doesn't
-      // hold, we'd want to reset `builtinsRef` and we'd also want to inline deps on the logical
-      // Skyframe subgraph when we get a `builtins` cache hit.
+
+      // All actual usages of BzlLoadFunction inlining assume builtins can never change (i.e. no
+      // usage of --experimental_builtins_bzl_path, no inter-invocation flipping of Starlark
+      // semantics options that'd cause evaluation of builtins to differ). This assumption is only
+      // violated in some tests, which rewrite the builtins to validate the state. If non-test usage
+      // can ever change builtins, we'd also want to inline deps on the logical Skyframe subgraph
+      // when we get a `builtins` cache hit.
+      if (resetBuiltins) {
+        builtinsRef = new AtomicReference<>();
+      }
     }
   }
 

src/test/java/com/google/devtools/build/lib/analysis/util/BuildViewTestCase.java

@@ -247,6 +247,8 @@ public abstract class BuildViewTestCase extends FoundationTestCase {
 
   private ActionLogBufferPathGenerator actionLogBufferPathGenerator;
 
+  @Nullable private BzlLoadFunction inliningBzlLoadFunction;
+
   @After
   public final void cleanupInterningPools() {
     skyframeExecutor.getEvaluator().cleanupInterningPools();
@@ -377,24 +379,25 @@ protected SkyframeExecutorRepositoryHelpersHolder getRepositoryHelpersHolder() {
     return null;
   }
 
-  private static void injectInliningBzlLoadFunction(
+  private void injectInliningBzlLoadFunction(
       SkyframeExecutor skyframeExecutor,
       RuleClassProvider ruleClassProvider,
       BlazeDirectories directories) {
     ImmutableMap<SkyFunctionName, SkyFunction> skyFunctions =
         ((InMemoryMemoizingEvaluator) skyframeExecutor.getEvaluator()).getSkyFunctionsForTesting();
-    BzlLoadFunction bzlLoadFunction =
+    inliningBzlLoadFunction =
         BzlLoadFunction.createForInlining(
             ruleClassProvider,
             directories,
             // Use a cache size of 2 for testing to balance coverage for where loads are present and
             // aren't present in the cache.
             /* bzlLoadValueCacheSize= */ 2);
-    bzlLoadFunction.resetInliningCache();
+    // The builtins should be empty since this was just created but reset it anyway to be sure.
+    inliningBzlLoadFunction.resetInliningCacheAndBuiltinsForTesting();
     // This doesn't override the BZL_LOAD -> BzlLoadFunction mapping, but nothing besides
     // PackageFunction should be requesting that key while using the inlining code path.
     ((PackageFunction) skyFunctions.get(SkyFunctions.PACKAGE))
-        .setBzlLoadFunctionForInliningForTesting(bzlLoadFunction);
+        .setBzlLoadFunctionForInliningForTesting(inliningBzlLoadFunction);
   }
 
   /**
@@ -562,19 +565,27 @@ protected PackageManager getPackageManager() {
     return skyframeExecutor.getPackageManager();
   }
 
+  /**
+   * Invalidates all existing packages, clears the cache for inlined bzl loads (including builtins),
+   * and invalidates configurations.
+   */
   protected void invalidatePackages() throws InterruptedException, AbruptExitException {
     invalidatePackages(true);
   }
 
   /**
-   * Invalidates all existing packages. Optionally invalidates configurations too.
+   * Invalidates all existing packages and clears the cache for inlined bzl loads (including
+   * builtins). Optionally also invalidates configurations.
    *
    * <p>Tests should invalidate both unless they have specific reason not to.
    */
   protected void invalidatePackages(boolean alsoConfigs)
       throws InterruptedException, AbruptExitException {
     skyframeExecutor.invalidateFilesUnderPathForTesting(
         reporter, ModifiedFileSet.EVERYTHING_MODIFIED, Root.fromPath(rootDirectory));
+    if (inliningBzlLoadFunction != null) {
+      inliningBzlLoadFunction.resetInliningCacheAndBuiltinsForTesting();
+    }
     if (alsoConfigs) {
       try {
         // Also invalidate all configurations. This is important: by invalidating all files we

src/test/java/com/google/devtools/build/lib/skyframe/BuiltinsInjectionTest.java

@@ -264,6 +264,8 @@ private void setBuildLanguageOptionsWithBuiltinsStaging(String... options) throw
    */
   private void writeExportsBzl(String... lines) throws Exception {
     scratch.overwriteFile("tools/builtins_staging/exports.bzl", lines);
+    // Since builtins have changed, we need to be sure the cache is reset to re-load them.
+    invalidatePackages(/* alsoConfigs= */ false);
   }
 
   /**

src/test/java/com/google/devtools/build/lib/starlark/StarlarkIntegrationTest.java

@@ -3473,6 +3473,7 @@ public void testLicenseType() throws Exception {
     // (See --incompatible_no_attr_license). However, this verifies that until the attribute
     // is removed, values of the attribute are a valid Starlark type.
     setBuildLanguageOptions("--incompatible_no_attr_license=false");
+    invalidatePackages();
     scratch.file(
         "test/rule.bzl",
         """