Add an IdentityHashMap to the BuildOptions.OptionsDiffForReconstructi…

…on codec.

PiperOrigin-RevId: 196310244

src/main/java/com/google/devtools/build/lib/analysis/config/BuildOptions.java

@@ -27,6 +27,10 @@
 import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.runtime.proto.InvocationPolicyOuterClass.InvocationPolicy;
+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.util.Fingerprint;
 import com.google.devtools.build.lib.util.OrderedSetMultimap;
@@ -36,13 +40,18 @@
 import com.google.devtools.common.options.OptionsClassProvider;
 import com.google.devtools.common.options.OptionsParser;
 import com.google.devtools.common.options.OptionsParsingException;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -424,11 +433,11 @@ public static OptionsDiffForReconstruction diffForReconstruction(
  if (diff.areSame()) {
  return OptionsDiffForReconstruction.getEmpty(first.fingerprint, second.computeChecksum());
  }
- HashMap<Class<? extends FragmentOptions>, Map<String, Object>> differingOptions =
- new HashMap<>(diff.differingOptions.keySet().size());
+ LinkedHashMap<Class<? extends FragmentOptions>, Map<String, Object>> differingOptions =
+ new LinkedHashMap<>(diff.differingOptions.keySet().size());
  for (Class<? extends FragmentOptions> clazz : diff.differingOptions.keySet()) {
  Collection<OptionDefinition> fields = diff.differingOptions.get(clazz);
- HashMap<String, Object> valueMap = new HashMap<>(fields.size());
+ LinkedHashMap<String, Object> valueMap = new LinkedHashMap<>(fields.size());
  for (OptionDefinition optionDefinition : fields) {
  Object secondValue;
  try {
@@ -543,15 +552,13 @@ public List<String> getPrettyPrintList() {
  * another: the full fragments of the second one, the fragment classes of the first that should be
  * omitted, and the values of any fields that should be changed.
  */
- @AutoCodec
  public static class OptionsDiffForReconstruction {
  private final Map<Class<? extends FragmentOptions>, Map<String, Object>> differingOptions;
  private final ImmutableSet<Class<? extends FragmentOptions>> extraFirstFragmentClasses;
  private final ImmutableList<FragmentOptions> extraSecondFragments;
  private final byte[] baseFingerprint;
  private final String checksum;
 
- @AutoCodec.VisibleForSerialization
  OptionsDiffForReconstruction(
  Map<Class<? extends FragmentOptions>, Map<String, Object>> differingOptions,
  ImmutableSet<Class<? extends FragmentOptions>> extraFirstFragmentClasses,
@@ -613,7 +620,9 @@ public boolean equals(Object o) {
  OptionsDiffForReconstruction that = (OptionsDiffForReconstruction) o;
  return differingOptions.equals(that.differingOptions)
  && extraFirstFragmentClasses.equals(that.extraFirstFragmentClasses)
- && this.extraSecondFragments.equals(that.extraSecondFragments);
+ && this.extraSecondFragments.equals(that.extraSecondFragments)
+ && Arrays.equals(this.baseFingerprint, that.baseFingerprint)
+ && this.checksum.equals(that.checksum);
  }
 
  @Override
@@ -626,7 +635,91 @@ public String toString() {
 
  @Override
  public int hashCode() {
- return Objects.hash(differingOptions, extraFirstFragmentClasses, extraSecondFragments);
+ return Objects.hash(
+ differingOptions,
+ extraFirstFragmentClasses,
+ extraSecondFragments,
+ Arrays.hashCode(baseFingerprint),
+ checksum);
+ }
+ }
+
+ /**
+ * Hand-rolled Codec so we can cache the byte representation of a {@link
+ * BuildOptions.OptionsDiffForReconstruction} object because serialization is expensive.
+ */
+ @VisibleForTesting
+ static class OptionsDiffForReconstructionCodec
+ implements ObjectCodec<OptionsDiffForReconstruction> {
+
+ @Override
+ public void serialize(
+ SerializationContext context,
+ BuildOptions.OptionsDiffForReconstruction input,
+ CodedOutputStream codedOut)
+ throws SerializationException, IOException {
+ context = context.getNewNonMemoizingContext();
+ // We get this cache from our context because there can be different ObjectCodecRegistry's for
+ // SkyKeys and SkyValues.
+ @SuppressWarnings("unchecked")
+ IdentityHashMap<OptionsDiffForReconstruction, byte[]> cache =
+ context.getDependency(IdentityHashMap.class);
+ if (cache.containsKey(input)) {
+ byte[] rawBytes = cache.get(input);
+ codedOut.writeInt32NoTag(rawBytes.length);
+ codedOut.writeRawBytes(cache.get(input));
+ } else {
+ ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+ CodedOutputStream codedOutputStream = CodedOutputStream.newInstance(byteArrayOutputStream);
+ context.serialize(input.differingOptions, codedOutputStream);
+ context.serialize(input.extraFirstFragmentClasses, codedOutputStream);
+ context.serialize(input.extraSecondFragments, codedOutputStream);
+ if (input.baseFingerprint != null) {
+ codedOutputStream.writeBoolNoTag(true);
+ codedOutputStream.writeInt32NoTag(input.baseFingerprint.length);
+ codedOutputStream.writeRawBytes(input.baseFingerprint);
+ } else {
+ codedOutputStream.writeBoolNoTag(false);
+ }
+ context.serialize(input.checksum, codedOutputStream);
+ codedOutputStream.flush();
+ byteArrayOutputStream.flush();
+ byte[] serializedBytes = byteArrayOutputStream.toByteArray();
+ cache.put(input, serializedBytes);
+ codedOut.writeInt32NoTag(serializedBytes.length);
+ codedOut.writeRawBytes(serializedBytes);
+ codedOut.flush();
+ }
+ }
+
+ @Override
+ public BuildOptions.OptionsDiffForReconstruction deserialize(
+ DeserializationContext context, CodedInputStream codedIn)
+ throws SerializationException, IOException {
+ byte[] serializedBytes = codedIn.readRawBytes(codedIn.readInt32());
+ CodedInputStream codedInputStream = CodedInputStream.newInstance(serializedBytes);
+ context = context.getNewNonMemoizingContext();
+ Map<Class<? extends FragmentOptions>, Map<String, Object>> differingOptions =
+ context.deserialize(codedInputStream);
+ ImmutableSet<Class<? extends FragmentOptions>> extraFirstFragmentClasses =
+ context.deserialize(codedInputStream);
+ ImmutableList<FragmentOptions> extraSecondFragments = context.deserialize(codedInputStream);
+ byte[] baseFingerprint = null;
+ if (codedInputStream.readBool()) {
+ baseFingerprint = codedInputStream.readRawBytes(codedInputStream.readInt32());
+ }
+ String checksum = context.deserialize(codedInputStream);
+ return new OptionsDiffForReconstruction(
+ differingOptions,
+ extraFirstFragmentClasses,
+ extraSecondFragments,
+ baseFingerprint,
+ checksum);
+ }
+
+ @Override
+ public Class<BuildOptions.OptionsDiffForReconstruction> getEncodedClass() {
+ return BuildOptions.OptionsDiffForReconstruction.class;
  }
  }
 }

src/main/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContext.java

@@ -120,4 +120,8 @@ DeserializationContext getMemoizingContext() {
  public DeserializationContext getNewMemoizingContext() {
  return new DeserializationContext(this.registry, this.dependencies, new Deserializer());
  }
+
+ public DeserializationContext getNewNonMemoizingContext() {
+ return new DeserializationContext(this.registry, this.dependencies, null);
+ }
 }

src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java

@@ -131,6 +131,11 @@ private SerializationContext getNewMemoizingContext(boolean allowFuturesToBlockW
  this.registry, this.dependencies, new Memoizer.Serializer(), allowFuturesToBlockWritingOn);
  }
 
+ public SerializationContext getNewNonMemoizingContext() {
+ return new SerializationContext(
+ this.registry, this.dependencies, null, this.allowFuturesToBlockWritingOn);
+ }
+
  /**
  * Register a {@link ListenableFuture} that must complete successfully before the serialized bytes
  * generated using this context can be written remotely. Failure of the future implies a bug or

src/test/java/com/google/devtools/build/lib/BUILD

@@ -704,6 +704,7 @@ java_test(
  "//src/main/java/com/google/devtools/build/lib:java-compilation",
  "//src/main/java/com/google/devtools/build/lib:java-rules",
  "//src/main/java/com/google/devtools/build/lib:packages",
+ "//src/main/java/com/google/devtools/build/lib:proto-rules",
  "//src/main/java/com/google/devtools/build/lib:python-rules",
  "//src/main/java/com/google/devtools/build/lib:util",
  "//src/main/java/com/google/devtools/build/lib/rules/cpp",

src/test/java/com/google/devtools/build/lib/analysis/config/BuildConfigurationTest.java

@@ -32,6 +32,7 @@
 import com.google.devtools.build.lib.skyframe.serialization.testutils.SerializationTester;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.common.options.Options;
+import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.regex.Pattern;
 import org.junit.Test;
@@ -452,6 +453,9 @@ public void testCodec() throws Exception {
  "--define",
  "#a=pounda"))
  .addDependency(FileSystem.class, getScratch().getFileSystem())
+ .addDependency(
+ IdentityHashMap.class,
+ new IdentityHashMap<BuildOptions.OptionsDiffForReconstruction, byte[]>())
  .setVerificationFunction(BuildConfigurationTest::verifyDeserialized)
  .runTests();
  }

src/test/java/com/google/devtools/build/lib/analysis/config/BuildOptionsTest.java

@@ -20,7 +20,10 @@
 import com.google.devtools.build.lib.analysis.config.BuildOptions.OptionsDiff;
 import com.google.devtools.build.lib.analysis.config.BuildOptions.OptionsDiffForReconstruction;
 import com.google.devtools.build.lib.rules.cpp.CppOptions;
+import com.google.devtools.build.lib.rules.proto.ProtoConfiguration;
+import com.google.devtools.build.lib.skyframe.serialization.testutils.ObjectCodecTester;
 import com.google.devtools.common.options.OptionsParser;
+import java.util.IdentityHashMap;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -31,7 +34,8 @@
 @RunWith(JUnit4.class)
 public class BuildOptionsTest {
  private static final ImmutableList<Class<? extends FragmentOptions>> TEST_OPTIONS =
- ImmutableList.<Class<? extends FragmentOptions>>of(BuildConfiguration.Options.class);
+ ImmutableList.<Class<? extends FragmentOptions>>of(
+ BuildConfiguration.Options.class, ProtoConfiguration.Options.class);
 
  @Test
  public void optionSetCaching() {
@@ -165,4 +169,32 @@ public void applyDiff() throws Exception {
  assertThat(otherFragment.applyDiff(BuildOptions.diffForReconstruction(otherFragment, one)))
  .isEqualTo(one);
  }
+
+ @Test
+ public void testCodec() throws Exception {
+ BuildOptions one =
+ BuildOptions.of(
+ TEST_OPTIONS,
+ "--compilation_mode=opt",
+ "cpu=k8",
+ "--proto_compiler=//net/proto2/compiler/public:protocol_compiler",
+ "--proto_toolchain_for_java=//tools/proto/toolchains:java");
+ BuildOptions two =
+ BuildOptions.of(
+ TEST_OPTIONS,
+ "--compilation_mode=dbg",
+ "cpu=k8",
+ "--proto_compiler=@com_google_protobuf//:protoc");
+ ObjectCodecTester.newBuilder(new BuildOptions.OptionsDiffForReconstructionCodec())
+ .addSubjects(BuildOptions.diffForReconstruction(one, two))
+ .addDependency(
+ IdentityHashMap.class,
+ new IdentityHashMap<BuildOptions.OptionsDiffForReconstruction, byte[]>())
+ .skipBadDataTest() // Bad data doesn't make sense with our caching.
+ .verificationFunction(
+ ((original, deserialized) -> {
+ assertThat(original).isEqualTo(deserialized);
+ }))
+ .buildAndRunTests();
+ }
 }

src/test/java/com/google/devtools/build/lib/exec/MiddlemanActionTest.java

@@ -23,14 +23,11 @@
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.MiddlemanAction;
 import com.google.devtools.build.lib.actions.MiddlemanFactory;
-import com.google.devtools.build.lib.actions.OutputBaseSupplier;
 import com.google.devtools.build.lib.analysis.util.AnalysisTestUtil;
 import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
-import com.google.devtools.build.lib.skyframe.serialization.testutils.SerializationTester;
 import com.google.devtools.build.lib.testutil.Suite;
 import com.google.devtools.build.lib.testutil.TestSpec;
-import com.google.devtools.build.lib.vfs.FileSystem;
 import java.util.ArrayList;
 import java.util.Arrays;
 import org.junit.Before;
@@ -138,20 +135,4 @@ public void testDifferentExecutablesForRunfilesMiddleman() throws Exception {
  assertThat(Sets.newHashSet(middlemanActionForD.getOutputs()))
  .isNotEqualTo(Sets.newHashSet(middlemanActionForC.getOutputs()));
  }
-
- @Test
- public void testCodec() throws Exception {
- new SerializationTester(getGeneratingAction(middle))
- .addDependency(FileSystem.class, scratch.getFileSystem())
- .addDependency(OutputBaseSupplier.class, () -> outputBase)
- .setVerificationFunction(MiddlemanActionTest::verifyEquivalent)
- .runTests();
- }
-
- private static void verifyEquivalent(MiddlemanAction first, MiddlemanAction second) {
- assertThat(first.getActionType()).isEqualTo(second.getActionType());
- assertThat(first.getInputs()).isEqualTo(second.getInputs());
- assertThat(first.getOutputs()).isEqualTo(second.getOutputs());
- assertThat(first.getOwner()).isEqualTo(second.getOwner());
- }
 }