Make the builtins registry thread-safe

It was previously assumed that safety wasn't needed because

  1) all builtins should be registered in static initializer blocks, and
  2) all retrievals should occur during Skylark evaluation, after static initialization completes.

It turns out these assumptions aren't actually true (Who would've thunk it!). SkylarkActionFactory has been observed to be initialized as late as analysis time, and retrievals occur as early as constructing a PackageFactory (when scanning the native module). The failure mode is particularly ugly: Random Skylark method lookups will fail non-deterministically.

This change guards against this by making the builtins registry implement a form of freezing. Before freezing, reads and writes are allowed and are synchronized. After freezing, only reads are allowed and they are unsynchronized for performance. BlazeRuntime is responsible for flipping the bit, and for ensuring classes like SkylarkActionFactory run their initialization by that point. Unit tests don't need to worry, since they just stay unfrozen and synchronized throughout.

RELNOTES: None
PiperOrigin-RevId: 188080136

5 files changed, 158 insertions, 16 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkActionFactory.java b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkActionFactory.java
index e14bfc0046..914be9b43f 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkActionFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkActionFactory.java
@@ -1142,6 +1142,10 @@ public class SkylarkActionFactory implements SkylarkValue {
     static {
       SkylarkSignatureProcessor.configureSkylarkFunctions(Args.class);
     }
+
+    /** No-op method that can be called to ensure the above static initializer runs. */
+    public static void forceStaticInitialization() {
+    }
   }
 
   @SkylarkSignature(
@@ -1185,4 +1189,12 @@ public class SkylarkActionFactory implements SkylarkValue {
   static {
     SkylarkSignatureProcessor.configureSkylarkFunctions(SkylarkActionFactory.class);
   }
+
+  /**
+   * No-op method that can be called to ensure the above static initializer runs, as well as the
+   * initializer for nested classes.
+   */
+  public static void forceStaticInitialization() {
+    Args.forceStaticInitialization();
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
index 6932d9ea18..dab37d82cb 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
@@ -30,6 +30,7 @@ import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider;
 import com.google.devtools.build.lib.analysis.ServerDirectories;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.config.ConfigurationFragmentFactory;
+import com.google.devtools.build.lib.analysis.skylark.SkylarkActionFactory;
 import com.google.devtools.build.lib.analysis.test.CoverageReportActionFactory;
 import com.google.devtools.build.lib.buildeventstream.PathConverter;
 import com.google.devtools.build.lib.buildtool.BuildRequestOptions;
@@ -1073,11 +1074,46 @@ public final class BlazeRuntime {
     runtime.initWorkspace(directories, binTools);
     CustomExitCodePublisher.setAbruptExitStatusFileDir(serverDirectories.getOutputBase());
 
+    // Most static initializers for @SkylarkSignature-containing classes have already run by this
+    // point, but this will pick up the stragglers.
+    initSkylarkBuiltinsRegistry();
+
     AutoProfiler.setClock(runtime.getClock());
     BugReport.setRuntime(runtime);
     return runtime;
   }
 
+  /**
+   * Configures the Skylark builtins registry.
+   *
+   * <p>Any class containing {@link SkylarkSignature}-annotated fields should call
+   * {@link SkylarkSignatureProcessor#configureSkylarkFunctions} on itself. This serves two
+   * purposes: 1) it initializes those fields for use, and 2) it registers them with the Skylark
+   * builtins registry object
+   * ({@link com.google.devtools.build.lib.syntax.Runtime#getBuiltinRegistry}). Unfortunately
+   * there's some technical debt here: The registry object is static and the registration occurs
+   * inside static initializer blocks.
+   *
+   * <p>The registry supports concurrent read/write access, but read access is not actually
+   * efficient (lockless) until write access is disallowed by calling its
+   * {@link com.google.devtools.build.lib.syntax.Runtime.BuiltinRegistry#freeze freeze} method.
+   * We want to freeze before the build begins, but not before all classes have had a chance to run
+   * their static initializers.
+   *
+   * <p>Therefore, this method first ensures that the initializers have run, and then explicitly
+   * freezes the registry. It ensures initialization by calling a no-op static method on the class.
+   * Only classes whose initializers have been observed to cause {@code BuiltinRegistry} to throw an
+   * exception need to be included here, since that indicates that their initialization did not
+   * happen by this point in time.
+   *
+   * <p>Unit tests don't need to worry about registry freeze exceptions, since the registry isn't
+   * frozen at all for them. They just pay the cost of extra synchronization on every access.
+   */
+  private static void initSkylarkBuiltinsRegistry() {
+    SkylarkActionFactory.forceStaticInitialization();
+    com.google.devtools.build.lib.syntax.Runtime.getBuiltinRegistry().freeze();
+  }
+
   private static String maybeGetPidString() {
     Integer pid = maybeForceJNIByGettingPid(null);
     return pid == null ? "" : "pid " + pid + " and ";
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Runtime.java b/src/main/java/com/google/devtools/build/lib/syntax/Runtime.java
index 3decbebff5..2c7b0ee978 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Runtime.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Runtime.java
@@ -160,11 +160,24 @@ public final class Runtime {
    * A registry of builtins, including both global builtins and builtins that are under some
    * namespace.
    *
-   * <p>This object is unsynchronized, but concurrent reads are fine.
+   * <p>Concurrency model: This object is thread-safe. Read accesses are always allowed, while write
+   * accesses are only allowed before this object has been frozen ({@link #freeze}). Prior to
+   * freezing, all operations are synchronized, while after freezing they are lockless.
    */
   public static class BuiltinRegistry {
 
     /**
+     * Whether the registry's construction has completed.
+     *
+     * <p>Mutating methods may only be called while this is still false. Accessor methods may be
+     * called at any time.
+     *
+     * <p>We use {@code volatile} rather than {@link AtomicBoolean} because the bit flip only
+     * happens once, and doesn't require correlated reads and writes.
+     */
+    private volatile boolean frozen = false;
+
+    /**
      * All registered builtins, keyed and sorted by an identifying (but otherwise unimportant)
      * string.
      *
@@ -177,13 +190,36 @@ public final class Runtime {
     /** All non-global builtin functions, keyed by their namespace class and their name. */
     private final Map<Class<?>, Map<String, BaseFunction>> functions = new HashMap<>();
 
+    /**
+     * Marks the registry as initialized, if it wasn't already.
+     *
+     * <p>It is guaranteed that after this method returns, all accessor methods are safe without
+     * synchronization; i.e. no mutation operation can touch the data structures.
+     */
+    public void freeze() {
+      // Similar to double-checked locking, but no need to check again on the inside since we don't
+      // care if two threads set the bit at once. The synchronized block is only to provide
+      // exclusion with mutations.
+      if (!this.frozen) {
+        synchronized (this) {
+          this.frozen = true;
+        }
+      }
+    }
+
     /** Registers a builtin with the given simple name, that was defined in the given Java class. */
-    public void registerBuiltin(Class<?> definingClass, String name, Object builtin) {
+    public synchronized void registerBuiltin(Class<?> definingClass, String name, Object builtin) {
       String key = String.format("%s.%s", definingClass.getName(), name);
       Preconditions.checkArgument(
           !allBuiltins.containsKey(key),
           "Builtin '%s' registered multiple times",
           key);
+
+      Preconditions.checkState(
+          !frozen,
+          "Attempted to register builtin '%s' after registry has already been frozen",
+          key);
+
       allBuiltins.put(key, builtin);
     }
 
@@ -192,7 +228,7 @@ public final class Runtime {
      *
      * <p>This is independent of {@link #registerBuiltin}.
      */
-    public void registerFunction(Class<?> namespace, BaseFunction function) {
+    public synchronized void registerFunction(Class<?> namespace, BaseFunction function) {
       Preconditions.checkNotNull(namespace);
       Preconditions.checkNotNull(function.getObjectType());
       Class<?> skylarkNamespace = getSkylarkNamespace(namespace);
@@ -200,13 +236,26 @@ public final class Runtime {
       Class<?> objType = getSkylarkNamespace(function.getObjectType());
       Preconditions.checkArgument(objType.equals(skylarkNamespace));
 
+      Preconditions.checkState(
+          !frozen,
+          "Attempted to register function '%s' in namespace '%s' after registry has already been "
+              + "frozen",
+          function,
+          namespace);
+
       functions.computeIfAbsent(namespace, k -> new HashMap<>());
       functions.get(namespace).put(function.getName(), function);
     }
 
     /** Returns a list of all registered builtins, in a deterministic order. */
     public ImmutableList<Object> getBuiltins() {
-      return ImmutableList.copyOf(allBuiltins.values());
+      if (frozen) {
+        return ImmutableList.copyOf(allBuiltins.values());
+      } else {
+        synchronized (this) {
+          return ImmutableList.copyOf(allBuiltins.values());
+        }
+      }
     }
 
     @Nullable
@@ -220,8 +269,15 @@ public final class Runtime {
      * <p>If the namespace does not exist or has no function with that name, returns null.
      */
     public BaseFunction getFunction(Class<?> namespace, String name) {
-      Map<String, BaseFunction> namespaceFunctions = getFunctionsInNamespace(namespace);
-      return namespaceFunctions != null ? namespaceFunctions.get(name) : null;
+      if (frozen) {
+        Map<String, BaseFunction> namespaceFunctions = getFunctionsInNamespace(namespace);
+        return namespaceFunctions != null ? namespaceFunctions.get(name) : null;
+      } else {
+        synchronized (this) {
+          Map<String, BaseFunction> namespaceFunctions = getFunctionsInNamespace(namespace);
+          return namespaceFunctions != null ? namespaceFunctions.get(name) : null;
+        }
+      }
     }
 
     /**
@@ -230,11 +286,21 @@ public final class Runtime {
      * <p>If the namespace does not exist, returns an empty set.
      */
     public ImmutableSet<String> getFunctionNames(Class<?> namespace) {
-      Map<String, BaseFunction> namespaceFunctions = getFunctionsInNamespace(namespace);
-      if (namespaceFunctions == null) {
-        return ImmutableSet.of();
+      if (frozen) {
+        Map<String, BaseFunction> namespaceFunctions = getFunctionsInNamespace(namespace);
+        if (namespaceFunctions == null) {
+          return ImmutableSet.of();
+        }
+        return ImmutableSet.copyOf(namespaceFunctions.keySet());
+      } else {
+        synchronized (this) {
+          Map<String, BaseFunction> namespaceFunctions = getFunctionsInNamespace(namespace);
+          if (namespaceFunctions == null) {
+            return ImmutableSet.of();
+          }
+          return ImmutableSet.copyOf(namespaceFunctions.keySet());
+        }
       }
-      return ImmutableSet.copyOf(namespaceFunctions.keySet());
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/SkylarkSignatureProcessor.java b/src/main/java/com/google/devtools/build/lib/syntax/SkylarkSignatureProcessor.java
index 346a913820..6424a47e44 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/SkylarkSignatureProcessor.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/SkylarkSignatureProcessor.java
@@ -215,17 +215,20 @@ public class SkylarkSignatureProcessor {
   }
 
   /**
-   * Processes all {@link SkylarkSignature}-annotated fields in a class. This function should only
-   * be called once per class, and not concurrently.
+   * Processes all {@link SkylarkSignature}-annotated fields in a class.
    *
    * <p>This includes registering these fields as builtins using {@link Runtime}, and for {@link
    * BaseFunction} instances, calling {@link BaseFunction#configure(SkylarkSignature)}. The fields
    * will be picked up by reflection even if they are not public.
    *
-   * <p>A class typically calls this function from within a static initializer block, passing itself
-   * as the {@code type}. E.g., A class {@code Foo} containing {@code @SkylarkSignature} annotations
-   * would end with {@code static { SkylarkSignatureProcessor.configureSkylarkFunctions(Foo.class);
-   * }}
+   * <p>This function should be called once per class, before the builtins registry is frozen. In
+   * practice, this is usually called from the class's own static initializer block. E.g., a class
+   * {@code Foo} containing {@code @SkylarkSignature} annotations would end with
+   * {@code static { SkylarkSignatureProcessor.configureSkylarkFunctions(Foo.class); }}.
+   *
+   * <p><b>If you see exceptions from {@link Runtime.BuiltinRegistry} here:</b> Be sure the class's
+   * static initializer has in fact been called before the registry was frozen. In Bazel, see
+   * {@link com.google.devtools.build.lib.runtime.BlazeRuntime#initSkylarkBuiltinsRegistry}.
    */
   public static void configureSkylarkFunctions(Class<?> type) {
     Runtime.BuiltinRegistry builtins = Runtime.getBuiltinRegistry();
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/RuntimeTest.java b/src/test/java/com/google/devtools/build/lib/syntax/RuntimeTest.java
index 60df36ba70..9da686707a 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/RuntimeTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/RuntimeTest.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.syntax;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.testutil.MoreAsserts.assertThrows;
 
 import com.google.devtools.build.lib.skylarkinterface.SkylarkPrinter;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkValue;
@@ -83,6 +84,30 @@ public final class RuntimeTest {
   }
 
   @Test
+  public void checkRegistry_WriteAfterFreezeFails_Builtin() {
+    Runtime.BuiltinRegistry reg = new Runtime.BuiltinRegistry();
+    reg.freeze();
+    IllegalStateException expected = assertThrows(
+        IllegalStateException.class,
+        () -> reg.registerBuiltin(DummyType.class, "dummy", "foo"));
+    assertThat(expected).hasMessageThat()
+        .matches("Attempted to register builtin '(.*)DummyType.dummy' after registry has already "
+            + "been frozen");
+  }
+
+  @Test
+  public void checkRegistry_WriteAfterFreezeFails_Function() {
+    Runtime.BuiltinRegistry reg = new Runtime.BuiltinRegistry();
+    reg.freeze();
+    IllegalStateException expected = assertThrows(
+        IllegalStateException.class,
+        () -> reg.registerFunction(DummyType.class, DUMMY_FUNC));
+    assertThat(expected).hasMessageThat()
+        .matches("Attempted to register function 'dummyFunc' in namespace '(.*)DummyType' after "
+            + "registry has already been frozen");
+  }
+
+  @Test
   public void checkStaticallyRegistered_Method() throws Exception {
     Field splitField = MethodLibrary.class.getDeclaredField("split");
     splitField.setAccessible(true);