PiperOrigin-RevId: 206181674

1 files changed, 79 insertions, 7 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/BuildConfigurationValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/BuildConfigurationValue.java
index 90479a9a54..1e9bc52d30 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/BuildConfigurationValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/BuildConfigurationValue.java
@@ -21,14 +21,23 @@ import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.config.FragmentClassSet;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
+import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
+import com.google.devtools.build.lib.skyframe.serialization.SerializationContext;
+import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
-import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.VisibleForSerialization;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
 import java.io.Serializable;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.logging.Logger;
 
 /** A Skyframe value representing a {@link BuildConfiguration}. */
@@ -59,7 +68,7 @@ public class BuildConfigurationValue implements SkyValue {
   public static Key key(
       Set<Class<? extends BuildConfiguration.Fragment>> fragments,
       BuildOptions.OptionsDiffForReconstruction optionsDiff) {
-    return Key.create(
+    return key(
         FragmentClassSet.of(
             ImmutableSortedSet.copyOf(BuildConfiguration.lexicalFragmentSorter, fragments)),
         optionsDiff);
@@ -75,18 +84,15 @@ public class BuildConfigurationValue implements SkyValue {
   }
 
   /** {@link SkyKey} for {@link BuildConfigurationValue}. */
-  @AutoCodec
   public static final class Key implements SkyKey, Serializable {
     private static final Interner<Key> keyInterner = BlazeInterners.newWeakInterner();
 
     private final FragmentClassSet fragments;
-    final BuildOptions.OptionsDiffForReconstruction optionsDiff;
+    private final BuildOptions.OptionsDiffForReconstruction optionsDiff;
     // If hashCode really is -1, we'll recompute it from scratch each time. Oh well.
     private volatile int hashCode = -1;
 
-    @AutoCodec.Instantiator
-    @VisibleForSerialization
-    static Key create(
+    private static Key create(
         FragmentClassSet fragments, BuildOptions.OptionsDiffForReconstruction optionsDiff) {
       return keyInterner.intern(new Key(fragments, optionsDiff));
     }
@@ -135,5 +141,71 @@ public class BuildConfigurationValue implements SkyValue {
     public String toString() {
       return "BuildConfigurationValue.Key[" + optionsDiff.getChecksum() + "]";
     }
+
+    private static class Codec implements ObjectCodec<Key> {
+      @Override
+      public Class<Key> getEncodedClass() {
+        return Key.class;
+      }
+
+      @Override
+      public void serialize(SerializationContext context, Key obj, CodedOutputStream codedOut)
+          throws SerializationException, IOException {
+        @SuppressWarnings("unchecked")
+        ConcurrentMap<BuildConfigurationValue.Key, ByteString> cache =
+            context.getDependency(KeyCodecCache.class).map;
+        ByteString bytes = cache.get(obj);
+        if (bytes == null) {
+          context = context.getNewNonMemoizingContext();
+          ByteString.Output byteStringOut = ByteString.newOutput();
+          CodedOutputStream bytesOut = CodedOutputStream.newInstance(byteStringOut);
+          context.serialize(obj.optionsDiff, bytesOut);
+          bytesOut.flush();
+          byteStringOut.flush();
+          int optionsDiffSerializedSize = byteStringOut.size();
+          context.serialize(obj.fragments, bytesOut);
+          bytesOut.flush();
+          byteStringOut.flush();
+          bytes = byteStringOut.toByteString();
+          cache.put(obj, bytes);
+          logger.info(
+              "Serialized "
+                  + obj.optionsDiff
+                  + " and "
+                  + obj.fragments
+                  + " to "
+                  + bytes.size()
+                  + " bytes (optionsDiff took "
+                  + optionsDiffSerializedSize
+                  + " bytes)");
+        }
+        codedOut.writeBytesNoTag(bytes);
+      }
+
+      @Override
+      public Key deserialize(DeserializationContext context, CodedInputStream codedIn)
+          throws SerializationException, IOException {
+        codedIn = codedIn.readBytes().newCodedInput();
+        context = context.getNewNonMemoizingContext();
+        BuildOptions.OptionsDiffForReconstruction optionsDiff = context.deserialize(codedIn);
+        FragmentClassSet fragmentClassSet = context.deserialize(codedIn);
+        return key(fragmentClassSet, optionsDiff);
+      }
+    }
+  }
+
+  /**
+   * Injected cache for {@code Codec}, so that we don't have to repeatedly serialize the same
+   * object. We still incur the over-the-wire cost of the bytes, but we don't use CPU to repeatedly
+   * compute it.
+   *
+   * <p>We provide the cache as an injected dependency so that different serializers' caches are
+   * isolated.
+   */
+  public static class KeyCodecCache {
+    private final ConcurrentMap<BuildConfigurationValue.Key, ByteString> map =
+        new ConcurrentHashMap<>();
+
+    public KeyCodecCache() {}
   }
 }