Implement full synchronization for SkyKeyInterner to ensure that it…

… fulfills the contract of `Interner` - no duplicate instances possible.

Prior to this change, there was the possibility of a race condition leading to a thread missing the canonical instance in both the weak interner and the global pool because another thread was simultaneously transferring it.

It turns out that weak interning under the synchronization provided by `ConcurrentHashMap#computeIfAbsent` is reasonably efficient. Benchmarks show that `-DBAZEL_USE_POOLED_SKY_KEY_INTERNER=1` with this change performs on-par with the current state (special `SkyKey` interning disabled). My brief explanation for this is:

* In the common case where a node is already present in the graph, interning is performing lock-free lookups on `ConcurrentHashMap`, which is much faster than the archaic `MapMakerInternalMap` that backs Guava's interner.
* The optimization in this change to only call `removeWeak` when a node is newly created helps avoid a lot of unnecessary interaction with the Guava interner.

PiperOrigin-RevId: 512666596
Change-Id: I67c28a7bfab6d49e65b2aba59ec9d19074ea493d

src/main/java/com/google/devtools/build/skyframe/InMemoryGraphImpl.java

@@ -29,7 +29,6 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import javax.annotation.Nullable;
 
 /**
@@ -125,13 +124,6 @@ protected InMemoryNodeEntry newNodeEntry(SkyKey key) {
     return new InMemoryNodeEntry(key);
   }
 
-  /**
-   * This is used to call newNodeEntry() from within computeIfAbsent. Instantiated here to avoid
-   * lambda instantiation overhead.
-   */
-  @SuppressWarnings("UnnecessaryLambda")
-  private final Function<SkyKey, InMemoryNodeEntry> newNodeEntryFunction = this::newNodeEntry;
-
   @Override
   @CanIgnoreReturnValue
   public NodeBatch createIfAbsentBatch(
@@ -144,14 +136,25 @@ public NodeBatch createIfAbsentBatch(
 
   @CanIgnoreReturnValue
   private InMemoryNodeEntry createIfAbsent(SkyKey skyKey) {
-    InMemoryNodeEntry inMemoryNodeEntry = nodeMap.computeIfAbsent(skyKey, newNodeEntryFunction);
-    if (usePooledSkyKeyInterning) {
-      SkyKeyInterner<?> interner = skyKey.getSkyKeyInterner();
-      if (interner != null) {
-        interner.removeWeak(skyKey);
-      }
+    SkyKeyInterner<?> interner = skyKey.getSkyKeyInterner();
+    if (!usePooledSkyKeyInterning || interner == null) {
+      return nodeMap.computeIfAbsent(skyKey, this::newNodeEntry);
+    }
+
+    // The key is typically already present. Record whether this thread newly created a node so that
+    // we can skip calling removeWeak if it was already present.
+    boolean[] newlyCreated = new boolean[1];
+    InMemoryNodeEntry nodeEntry =
+        nodeMap.computeIfAbsent(
+            skyKey,
+            k -> {
+              newlyCreated[0] = true;
+              return newNodeEntry(k);
+            });
+    if (newlyCreated[0]) {
+      interner.removeWeak(skyKey);
     }
-    return inMemoryNodeEntry;
+    return nodeEntry;
   }
 
   @Override
@@ -187,11 +190,21 @@ public void parallelForEach(Consumer<InMemoryNodeEntry> consumer) {
     nodeMap.forEachValue(PARALLELISM_THRESHOLD, consumer);
   }
 
-  @Nullable
   @Override
-  public SkyKey canonicalize(SkyKey key) {
-    InMemoryNodeEntry node = nodeMap.get(key);
-    return node != null ? node.getKey() : null;
+  public SkyKey getOrWeakIntern(SkyKey key) {
+    // Use computeIfAbsent not to mutate the map, but to call weakIntern under synchronization. This
+    // ensures that the canonical instance isn't being transferred to the node map concurrently in
+    // createIfAbsent. In the common case that the key is already present in the node map, this is a
+    // lock-free lookup.
+    SkyKey[] weakInterned = new SkyKey[1];
+    InMemoryNodeEntry nodeEntry =
+        nodeMap.computeIfAbsent(
+            key,
+            k -> {
+              weakInterned[0] = k.getSkyKeyInterner().weakIntern(k);
+              return null; // Don't actually store a mapping.
+            });
+    return nodeEntry != null ? nodeEntry.getKey() : weakInterned[0];
   }
 
   /**

src/main/java/com/google/devtools/build/skyframe/SkyKey.java

@@ -16,6 +16,7 @@
 import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.concurrent.ThreadSafety;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.io.Serializable;
 import java.lang.reflect.Field;
 import java.util.Map;
@@ -87,11 +88,14 @@ static <T extends SkyKey> SkyKeyInterner<T> newInterner() {
   interface SkyKeyPool {
 
     /**
-     * Returns the canonical instance of the {@link SkyKey} if found in the pool, otherwise returns
-     * {@code null}.
+     * Returns the canonical instance for the given key in the pool if it is present, otherwise
+     * interns the key using its {@linkplain SkyKeyInterner#weakIntern weak interner}.
+     *
+     * <p>To ensure a single canonical instance, if the key is not present in the pool, it should be
+     * weakly interned using synchronization so that it is not concurrently {@linkplain
+     * SkyKeyInterner#removeWeak removed from the weak interner}.
      */
-    @Nullable
-    SkyKey canonicalize(SkyKey key);
+    SkyKey getOrWeakIntern(SkyKey key);
 
     /**
      * Cleans up {@link SkyKey}s stored in the {@link SkyKeyPool} if necessary and uninstalls the in
@@ -161,22 +165,17 @@ public static void setGlobalPool(@Nullable SkyKeyPool pool) {
     @SuppressWarnings("unchecked")
     public T intern(T sample) {
       SkyKeyPool pool = globalPool;
-      if (pool != null) {
-        SkyKey result = pool.canonicalize(sample);
-        if (result != null) {
-          return (T) result;
-        }
-      }
-      return weakInterner.intern(sample);
+      return pool != null ? (T) pool.getOrWeakIntern(sample) : weakInterner.intern(sample);
     }
 
     /**
      * Interns {@code sample} directly into {@link #weakInterner} without checking {@link
      * #globalPool} and returns the canonical instance of {@code sample}.
      */
     @SuppressWarnings("unchecked")
-    public void weakIntern(SkyKey sample) {
-      var unused = weakInterner.intern((T) sample);
+    @CanIgnoreReturnValue
+    public T weakIntern(SkyKey sample) {
+      return weakInterner.intern((T) sample);
     }
 
     /**

src/test/java/com/google/devtools/build/skyframe/DeterministicInMemoryGraph.java

@@ -80,8 +80,8 @@ public void parallelForEach(Consumer<InMemoryNodeEntry> consumer) {
   }
 
   @Override
-  public SkyKey canonicalize(SkyKey key) {
-    return ((InMemoryGraph) delegate).canonicalize(key);
+  public SkyKey getOrWeakIntern(SkyKey key) {
+    return ((InMemoryGraph) delegate).getOrWeakIntern(key);
   }
 
   @Override

src/test/java/com/google/devtools/build/skyframe/NotifyingInMemoryGraph.java

@@ -87,8 +87,8 @@ public void parallelForEach(Consumer<InMemoryNodeEntry> consumer) {
   }
 
   @Override
-  public SkyKey canonicalize(SkyKey key) {
-    return ((InMemoryGraph) delegate).canonicalize(key);
+  public SkyKey getOrWeakIntern(SkyKey key) {
+    return ((InMemoryGraph) delegate).getOrWeakIntern(key);
   }
 
   @Override

src/test/java/com/google/devtools/build/skyframe/SkyKeyTest.java

@@ -22,7 +22,6 @@
 import com.google.devtools.build.lib.skyframe.serialization.testutils.TestUtils;
 import com.google.devtools.build.skyframe.SkyKey.SkyKeyInterner;
 import com.google.devtools.build.skyframe.SkyKey.SkyKeyPool;
-import javax.annotation.Nullable;
 import org.junit.After;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -103,46 +102,39 @@ public SkyFunctionName functionName() {
 
   @Test
   public void skyKeyInterner_noGlobalPoolTestIntern() {
-    SkyKeyInterner<SkyKey> interner = SkyKey.newInterner();
-    SkyKey keyToIntern1 = new SkyKeyForInternerTests("HelloWorld");
-
-    assertThat(interner.intern(keyToIntern1)).isSameInstanceAs(keyToIntern1);
+    SkyKey keyToIntern1 = SkyKeyForInternerTests.createInterned("HelloWorld");
 
     // Interning a duplicate instance will result the same instance to be returned.
-    assertThat(interner.intern(new SkyKeyForInternerTests("HelloWorld")))
-        .isSameInstanceAs(keyToIntern1);
+    assertThat(SkyKeyForInternerTests.createInterned("HelloWorld")).isSameInstanceAs(keyToIntern1);
   }
 
   @Test
   public void skyKeyInterner_noGlobalPoolTestRemoval() {
-    SkyKeyInterner<SkyKey> interner = SkyKey.newInterner();
-    SkyKey keyToIntern1 = new SkyKeyForInternerTests("HelloWorld");
+    SkyKey keyToIntern1 = SkyKeyForInternerTests.createInterned("HelloWorld");
 
-    assertThat(interner.intern(keyToIntern1)).isSameInstanceAs(keyToIntern1);
+    assertThat(SkyKeyForInternerTests.createInterned("HelloWorld")).isSameInstanceAs(keyToIntern1);
 
     // Remove one instance from the interner and re-intern a duplicate one. The newly interned
     // instance is different from the previous one, which confirms that the previous interned
     // instance has already been successfully removed from the interner.
-    interner.removeWeak(new SkyKeyForInternerTests("HelloWorld"));
-    assertThat(interner.intern(new SkyKeyForInternerTests("HelloWorld")))
+    keyToIntern1.getSkyKeyInterner().removeWeak(keyToIntern1);
+    assertThat(SkyKeyForInternerTests.createInterned("HelloWorld"))
         .isNotSameInstanceAs(keyToIntern1);
   }
 
   @Test
   public void skyKeyInterner_withGlobalPool() {
-    SkyKey keyToIntern1 = new SkyKeyForInternerTests("HelloWorld");
-    SkyKey keyToIntern2 = new SkyKeyForInternerTests("FooBar");
+    SkyKey keyToIntern1 = SkyKeyForInternerTests.createInterned("HelloWorld");
     SkyKey keyInPool = new SkyKeyForInternerTests("FooBar");
 
     SkyKeyPool globalPool =
         new SkyKeyPool() {
-          @Nullable
           @Override
-          public SkyKey canonicalize(SkyKey key) {
+          public SkyKey getOrWeakIntern(SkyKey key) {
             if (key.argument() == "FooBar") {
               return keyInPool;
             } else {
-              return null;
+              return key.getSkyKeyInterner().weakIntern(key);
             }
           }
 
@@ -153,10 +145,9 @@ public void cleanupPool() {
         };
 
     SkyKeyInterner.setGlobalPool(globalPool);
-    SkyKeyInterner<SkyKey> interner = SkyKey.newInterner();
 
-    assertThat(interner.intern(keyToIntern1)).isSameInstanceAs(keyToIntern1);
-    assertThat(interner.intern(keyToIntern2)).isSameInstanceAs(keyInPool);
+    assertThat(SkyKeyForInternerTests.createInterned("HelloWorld")).isSameInstanceAs(keyToIntern1);
+    assertThat(SkyKeyForInternerTests.createInterned("FooBar")).isSameInstanceAs(keyInPool);
 
     globalPool.cleanupPool();
   }
@@ -166,16 +157,26 @@ public void cleanUpGlobalPool() {
     SkyKeyInterner.setGlobalPool(null);
   }
 
-  @AutoCodec
   static final class SkyKeyForInternerTests extends AbstractSkyKey<String> {
 
-    SkyKeyForInternerTests(String arg) {
+    private static final SkyKeyInterner<SkyKeyForInternerTests> interner = SkyKey.newInterner();
+
+    static SkyKeyForInternerTests createInterned(String arg) {
+      return interner.intern(new SkyKeyForInternerTests(arg));
+    }
+
+    private SkyKeyForInternerTests(String arg) {
       super(arg);
     }
 
     @Override
     public SkyFunctionName functionName() {
       return SkyFunctionName.FOR_TESTING;
     }
+
+    @Override
+    public SkyKeyInterner<?> getSkyKeyInterner() {
+      return interner;
+    }
   }
 }