Allow NestedSetCodec to share members across multiple deserializations.

This avoids redundancy in memory when multiple NestedSets share a member

PiperOrigin-RevId: 190085907

7 files changed, 79 insertions, 12 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/BUILD b/src/main/java/com/google/devtools/build/lib/BUILD
index 528d15aca2..b793c30850 100644
--- a/src/main/java/com/google/devtools/build/lib/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/BUILD
@@ -169,6 +169,7 @@ java_library(
             "util/JavaClock.java",
             "util/OS.java",
             "util/ProcessUtils.java",
+            "util/ResourceUsage.java",
             "util/SingleLineFormatter.java",
             "util/StringCanonicalizer.java",
             "util/StringTrie.java",
@@ -181,6 +182,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib:exitcode-external",
         "//src/main/java/com/google/devtools/build/lib:filetype",
         "//src/main/java/com/google/devtools/build/lib:os_util",
+        "//src/main/java/com/google/devtools/build/lib:resource_usage",
         "//src/main/java/com/google/devtools/build/lib:string_util",
         "//src/main/java/com/google/devtools/build/lib/clock",
         "//src/main/java/com/google/devtools/build/lib/collect",
@@ -218,6 +220,16 @@ java_library(
 )
 
 java_library(
+    name = "resource_usage",
+    srcs = [
+        "util/ResourceUsage.java",
+    ],
+    deps = [
+        "//third_party:guava",
+    ],
+)
+
+java_library(
     name = "string_util",
     srcs = [
         "util/StringCanonicalizer.java",
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD b/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
index bd372ff043..09be2fbb75 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
@@ -21,6 +21,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/collect/compacthashset",
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization:constants",
         "//third_party:guava",
         "//third_party:jsr305",
         "//third_party/protobuf:protobuf_java",
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodec.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodec.java
index a4954ab740..b8d8a2207e 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodec.java
@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.collect.nestedset;
 
 import com.google.common.base.Preconditions;
+import com.google.common.cache.CacheBuilder;
 import com.google.common.hash.Hashing;
 import com.google.common.hash.HashingOutputStream;
 import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
@@ -28,10 +29,10 @@ import com.google.protobuf.CodedOutputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.LinkedHashSet;
 import java.util.Map;
+import java.util.concurrent.ConcurrentMap;
 
 /**
  * Codec for {@link NestedSet}.
@@ -44,6 +45,15 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
 
   private static final EnumCodec<Order> orderCodec = new EnumCodec<>(Order.class);
 
+  private static final ConcurrentMap<ByteString, Object> digestToChild =
+      CacheBuilder.newBuilder()
+          .concurrencyLevel(SerializationConstants.DESERIALIZATION_POOL_SIZE)
+          .weakValues()
+          .<ByteString, Object>build()
+          .asMap();
+
+  public NestedSetCodec() {}
+
   @SuppressWarnings("unchecked")
   @Override
   public Class<NestedSet<T>> getEncodedClass() {
@@ -78,7 +88,6 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
       return NestedSetBuilder.emptySet(Order.STABLE_ORDER);
     }
 
-    Map<ByteString, Object> digestToChild = new HashMap<>();
     int nestedSetCount = codedIn.readInt32();
     Preconditions.checkState(
         nestedSetCount >= 1,
@@ -88,7 +97,7 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
     Object children = null;
     for (int i = 0; i < nestedSetCount; ++i) {
       // Update var, the last one in the list is the top level nested set
-      children = deserializeOneNestedSet(context, codedIn, digestToChild);
+      children = deserializeOneNestedSet(context, codedIn);
     }
     return createNestedSet(order, children);
   }
@@ -150,10 +159,7 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
     context.serialize(singleChild, childCodedOut);
   }
 
-  private Object deserializeOneNestedSet(
-      DeserializationContext context,
-      CodedInputStream codedIn,
-      Map<ByteString, Object> digestToChild)
+  private Object deserializeOneNestedSet(DeserializationContext context, CodedInputStream codedIn)
       throws SerializationException, IOException {
     ByteString digest = codedIn.readBytes();
     CodedInputStream childCodedIn = codedIn.readBytes().newCodedInput();
@@ -161,19 +167,22 @@ public class NestedSetCodec<T> implements ObjectCodec<NestedSet<T>> {
     int childCount = childCodedIn.readInt32();
     final Object result;
     if (childCount > 1) {
-      result = deserializeMultipleItemChildArray(context, digestToChild, childCodedIn, childCount);
+      result = deserializeMultipleItemChildArray(context, childCodedIn, childCount);
     } else if (childCount == 1) {
       result = context.deserialize(childCodedIn);
     } else {
       result = NestedSet.EMPTY_CHILDREN;
     }
-    digestToChild.put(digest, result);
-    return result;
+    // If this member has been deserialized already, use it, so that NestedSets that share members
+    // will point at the same object in memory instead of duplicating them.  Unfortunately, we had
+    // to do the work of deserializing the member that we will now throw away, in order to ensure
+    // that the pointer in the codedIn buffer was incremented appropriately.
+    Object oldValue = digestToChild.putIfAbsent(digest, result);
+    return oldValue == null ? result : oldValue;
   }
 
   private Object deserializeMultipleItemChildArray(
       DeserializationContext context,
-      Map<ByteString, Object> digestToChild,
       CodedInputStream childCodedIn,
       int childCount)
       throws IOException, SerializationException {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
index 8bd800b8d1..63fd3de187 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
@@ -11,7 +11,13 @@ filegroup(
 
 java_library(
     name = "serialization",
-    srcs = glob(["**/*.java"]),
+    srcs = glob(
+        ["**/*.java"],
+        exclude = ["SerializationConstants.java"],
+    ),
+    exports = [
+        ":constants",
+    ],
     deps = [
         ":kryo",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:registered-singleton",
@@ -22,6 +28,14 @@ java_library(
     ],
 )
 
+java_library(
+    name = "constants",
+    srcs = ["SerializationConstants.java"],
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib:resource_usage",
+    ],
+)
+
 # A convenience library that bundles together kryo and its runtime dependencies.
 java_library(
     name = "kryo",
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationConstants.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationConstants.java
index e3886dd276..b9873d6ee4 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationConstants.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationConstants.java
@@ -14,11 +14,16 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import com.google.devtools.build.lib.util.ResourceUsage;
+
 /**
  * Some static constants for deciding serialization behavior.
  */
 public class SerializationConstants {
 
+  /** Number of threads in deserialization pools. */
+  public static final int DESERIALIZATION_POOL_SIZE = 2 * ResourceUsage.getAvailableProcessors();
+
   /**
    * If true, we attempt to to serialize ConfiguredTargetValue in testing.
    */
diff --git a/src/test/java/com/google/devtools/build/lib/BUILD b/src/test/java/com/google/devtools/build/lib/BUILD
index e487d55123..205bf8f020 100644
--- a/src/test/java/com/google/devtools/build/lib/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/BUILD
@@ -224,6 +224,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset:fingerprint_cache",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils",
+        "//third_party/protobuf:protobuf_java",
     ],
 )
 
diff --git a/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java b/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
index 978d7a0bb5..74eea2176d 100644
--- a/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
+++ b/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
@@ -15,8 +15,12 @@ package com.google.devtools.build.lib.collect.nestedset;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.skyframe.serialization.AutoRegistry;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecs;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationConstants;
 import com.google.devtools.build.lib.skyframe.serialization.testutils.SerializationTester;
+import com.google.protobuf.ByteString;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -66,6 +70,27 @@ public class NestedSetCodecTest {
         .runTests();
   }
 
+  @SuppressWarnings("unchecked")
+  @Test
+  public void testDeserializedNestedSetsShareChildren() throws Exception {
+    ObjectCodecs objectCodecs =
+        new ObjectCodecs(
+            AutoRegistry.get().getBuilder().setAllowDefaultCodec(true).build(), ImmutableMap.of());
+    NestedSet<String> originalChild = NestedSetBuilder.create(Order.STABLE_ORDER, "a", "b", "c");
+    NestedSet<String> originalA =
+        new NestedSetBuilder<String>(Order.STABLE_ORDER).addTransitive(originalChild).build();
+    NestedSet<String> originalB =
+        new NestedSetBuilder<String>(Order.STABLE_ORDER).addTransitive(originalChild).build();
+
+    ByteString serializedA = objectCodecs.serialize(originalA);
+    ByteString serializedB = objectCodecs.serialize(originalB);
+
+    NestedSet<String> deserializedA = (NestedSet<String>) objectCodecs.deserialize(serializedA);
+    NestedSet<String> deserializedB = (NestedSet<String>) objectCodecs.deserialize(serializedB);
+
+    assertThat(deserializedA.rawChildren()).isSameAs(deserializedB.rawChildren());
+  }
+
   private static void verifyDeserialization(
       NestedSet<String> subject, NestedSet<String> deserialized) {
     assertThat(subject.getOrder()).isEqualTo(deserialized.getOrder());