Open-source some more serialization codecs, and create a PrecomputedValue codec. Since PrecomputedValues can contain any value, give them access to an ObjectCodecs instance so we don't have to have a whitelist inside PrecomputedValueCodec.

PiperOrigin-RevId: 168624137

6 files changed, 462 insertions, 5 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValueCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValueCodec.java
new file mode 100644
index 0000000000..cae35b468a
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValueCodec.java
@@ -0,0 +1,64 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.skyframe;
+
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecs;
+import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
+import com.google.devtools.build.lib.util.Preconditions;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
+import java.util.function.Supplier;
+
+/**
+ * {@link ObjectCodec} for {@link PrecomputedValue} objects. Because {@link PrecomputedValue}
+ * objects can theoretically contain any kind of value object as their value, an {@link
+ * ObjectCodecs} instance must be accessible to this codec during serialization/deserialization, so
+ * that it can handle the arbitrary objects it's given.
+ */
+public class PrecomputedValueCodec implements ObjectCodec<PrecomputedValue> {
+  private final Supplier<ObjectCodecs> objectCodecsSupplier;
+
+  public PrecomputedValueCodec(Supplier<ObjectCodecs> objectCodecsSupplier) {
+    this.objectCodecsSupplier = objectCodecsSupplier;
+  }
+
+  @Override
+  public Class<PrecomputedValue> getEncodedClass() {
+    return PrecomputedValue.class;
+  }
+
+  @Override
+  public void serialize(PrecomputedValue obj, CodedOutputStream codedOut)
+      throws SerializationException, IOException {
+    ObjectCodecs objectCodecs = objectCodecsSupplier.get();
+    Object val = obj.get();
+    Preconditions.checkState(!(val instanceof PrecomputedValue), "recursive precomputed: %s", obj);
+    // TODO(janakr): this assumes the classifier is the class of the object. This should be enforced
+    // by the ObjectCodecs instance.
+    String classifier = val.getClass().getName();
+    codedOut.writeStringNoTag(classifier);
+    objectCodecs.serialize(classifier, val, codedOut);
+  }
+
+  @Override
+  public PrecomputedValue deserialize(CodedInputStream codedIn)
+      throws SerializationException, IOException {
+    ObjectCodecs objectCodecs = objectCodecsSupplier.get();
+    Object val = objectCodecs.deserialize(codedIn.readBytes(), codedIn);
+    return new PrecomputedValue(val);
+  }
+}
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 9f58769888..b747a567af 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
@@ -9,6 +9,7 @@ java_library(
     name = "serialization",
     srcs = glob(["*.java"]),
     deps = [
+        "//src/main/java/com/google/devtools/build/lib:preconditions",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/JavaSerializableCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/JavaSerializableCodec.java
new file mode 100644
index 0000000000..f3ce577ccc
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/JavaSerializableCodec.java
@@ -0,0 +1,65 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.skyframe.serialization;
+
+import com.google.protobuf.ByteString;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.NotSerializableException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.nio.ByteBuffer;
+
+/** Naive ObjectCodec using Java native Serialization. Not performant, but a good fallback */
+class JavaSerializableCodec implements ObjectCodec<Object> {
+
+  @Override
+  public Class<Object> getEncodedClass() {
+    return Object.class;
+  }
+
+  @Override
+  public void serialize(Object obj, CodedOutputStream codedOut)
+      throws SerializationException, IOException {
+    ByteString.Output out = ByteString.newOutput();
+    ObjectOutputStream objOut = new ObjectOutputStream(out);
+    try {
+      objOut.writeObject(obj);
+    } catch (NotSerializableException e) {
+      throw new SerializationException.NoCodecException("Object " + obj + " not serializable", e);
+    } catch (NotSerializableRuntimeException e) {
+      // Values that inherit from Serializable but actually aren't serializable.
+      throw new SerializationException.NoCodecException("Object " + obj + " not serializable", e);
+    }
+    codedOut.writeBytesNoTag(out.toByteString());
+  }
+
+  @Override
+  public Object deserialize(CodedInputStream codedIn) throws SerializationException, IOException {
+    try {
+      // Get the ByteBuffer as it is potentially a view of the underlying bytes (not a copy), which
+      // is more efficient.
+      ByteBuffer buffer = codedIn.readByteBuffer();
+      ObjectInputStream objIn =
+          new ObjectInputStream(
+              new ByteArrayInputStream(buffer.array(), buffer.arrayOffset(), buffer.remaining()));
+      return objIn.readObject();
+    } catch (ClassNotFoundException e) {
+      throw new SerializationException("Java deserialization failed", e);
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
new file mode 100644
index 0000000000..5778cdb50d
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
@@ -0,0 +1,272 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.skyframe.serialization;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.skyframe.SkyFunctionName;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
+import java.util.Map;
+import java.util.Map.Entry;
+
+/**
+ * Wrapper for the minutiae of serializing and deserializing objects using {@link ObjectCodec}s,
+ * serving as a layer between the streaming-oriented {@link ObjectCodec} interface and users.
+ * Handles the mapping and selection of custom serialization implementations, falling back on less
+ * performant java serialization by default when no better option is available and it is allowed by
+ * the configuration.
+ *
+ * <p>To use, create a {@link ObjectCodecs.Builder} and add custom classifier to {@link ObjectCodec}
+ * mappings using {@link ObjectCodecs.Builder#add}. The provided mappings used to determine
+ * serialization/deserialization logic. For example:
+ *
+ * <pre>{@code
+ * // Create an instance for which anything identified as "foo" will use FooCodec.
+ * ObjectCodecs objectCodecs = ObjectCodecs.newBuilder()
+ *     .add("foo", new FooCodec())
+ *     .build();
+ *
+ * // This will use the custom supplied FooCodec to serialize obj:
+ * ByteString serialized = objectCodecs.serialize("foo", obj);
+ * Object deserialized = objectCodecs.deserialize(ByteString.copyFromUtf8("foo"), serialized);
+ *
+ * // This will use default java object serialization to serialize obj:
+ * ByteString serialized = objectCodecs.serialize("bar", obj);
+ * Object deserialized = objectCodecs.deserialize(ByteString.copyFromUtf8("bar"), serialized);
+ * }</pre>
+ *
+ * <p>Classifiers will typically be class names or SkyFunction names.
+ */
+public class ObjectCodecs {
+
+  private static final ObjectCodec<Object> DEFAULT_CODEC = new JavaSerializableCodec();
+
+  /** Create new ObjectCodecs.Builder, the preferred instantiation method. */
+  // TODO(janakr,michajlo): Specialize builders into ones keyed by class (even if the class isn't
+  // the one specified by the codec) and ones keyed by string, and expose a getClassifier() method
+  // for ObjectCodecs keyed by class.
+  public static ObjectCodecs.Builder newBuilder() {
+    return new Builder();
+  }
+
+  private final Map<String, ObjectCodec<?>> stringMappedCodecs;
+  private final Map<ByteString, ObjectCodec<?>> byteStringMappedCodecs;
+  private final boolean allowDefaultCodec;
+
+  private ObjectCodecs(Map<String, ObjectCodec<?>> codecs, boolean allowDefaultCodec) {
+    this.stringMappedCodecs = codecs;
+    this.byteStringMappedCodecs = makeByteStringMappedCodecs(codecs);
+    this.allowDefaultCodec = allowDefaultCodec;
+  }
+
+  /**
+   * Serialize {@code subject}, using the serialization strategy determined by {@code classifier},
+   * returning a {@link ByteString} containing the serialized representation.
+   */
+  public ByteString serialize(String classifier, Object subject) throws SerializationException {
+    ByteString.Output resultOut = ByteString.newOutput();
+    CodedOutputStream codedOut = CodedOutputStream.newInstance(resultOut);
+    ObjectCodec<?> codec = getCodec(classifier);
+    try {
+      doSerialize(classifier, codec, subject, codedOut);
+      codedOut.flush();
+      return resultOut.toByteString();
+    } catch (IOException e) {
+      throw new SerializationException(
+          "Failed to serialize " + subject + " using " + codec + " for " + classifier, e);
+    }
+  }
+
+  /**
+   * Similar to {@link #serialize(String, Object)}, except allows the caller to specify a {@link
+   * CodedOutputStream} to serialize {@code subject} to. Has less object overhead than {@link
+   * #serialize(String, Object)} and as such is preferrable when serializing objects in bulk.
+   *
+   * <p>{@code codedOut} is not flushed by this method.
+   */
+  public void serialize(String classifier, Object subject, CodedOutputStream codedOut)
+      throws SerializationException {
+    ObjectCodec<?> codec = getCodec(classifier);
+    try {
+      doSerialize(classifier, codec, subject, codedOut);
+    } catch (IOException e) {
+      throw new SerializationException(
+          "Failed to serialize " + subject + " using " + codec + " for " + classifier, e);
+    }
+  }
+
+  /**
+   * Deserialize {@code data} using the serialization strategy determined by {@code classifier}.
+   * {@code classifier} should be the utf-8 encoded {@link ByteString} representation of the {@link
+   * String} classifier used to serialize {@code data}. This is preferred since callers typically
+   * have parsed {@code classifier} from a protocol buffer, for which {@link ByteString}s are
+   * cheaper to use.
+   */
+  public Object deserialize(ByteString classifier, ByteString data) throws SerializationException {
+    return deserialize(classifier, data.newCodedInput());
+  }
+
+  /**
+   * Similar to {@link #deserialize(ByteString, ByteString)}, except allows the caller to specify a
+   * {@link CodedInputStream} to deserialize data from. This is useful for decoding objects
+   * serialized in bulk by {@link #serialize(String, Object, CodedOutputStream)}.
+   */
+  public Object deserialize(ByteString classifier, CodedInputStream codedIn)
+      throws SerializationException {
+    ObjectCodec<?> codec = getCodec(classifier);
+    // If safe, this will allow CodedInputStream to return a direct view of the underlying bytes
+    // in some situations, bypassing a copy.
+    codedIn.enableAliasing(true);
+    try {
+      Object result = codec.deserialize(codedIn);
+      if (result == null) {
+        throw new NullPointerException(
+            "ObjectCodec " + codec + " for " + classifier.toStringUtf8() + " returned null");
+      }
+      return result;
+    } catch (IOException e) {
+      throw new SerializationException(
+          "Failed to deserialize data using " + codec + " for " + classifier.toStringUtf8(), e);
+    }
+  }
+
+  private ObjectCodec<?> getCodec(String classifier)
+      throws SerializationException.NoCodecException {
+    ObjectCodec<?> result = stringMappedCodecs.get(classifier);
+    if (result != null) {
+      return result;
+    } else if (allowDefaultCodec) {
+      return DEFAULT_CODEC;
+    } else {
+      throw new SerializationException.NoCodecException(
+          "No codec available for " + classifier + " and default fallback disabled");
+    }
+  }
+
+  private ObjectCodec<?> getCodec(ByteString classifier) throws SerializationException {
+    ObjectCodec<?> result = byteStringMappedCodecs.get(classifier);
+    if (result != null) {
+      return result;
+    } else if (allowDefaultCodec) {
+      return DEFAULT_CODEC;
+    } else {
+      throw new SerializationException(
+          "No codec available for " + classifier.toStringUtf8() + " and default fallback disabled");
+    }
+  }
+
+  private static <T> void doSerialize(
+      String classifier, ObjectCodec<T> codec, Object subject, CodedOutputStream codedOut)
+      throws SerializationException, IOException {
+    try {
+      codec.serialize(codec.getEncodedClass().cast(subject), codedOut);
+    } catch (ClassCastException e) {
+      throw new SerializationException(
+          "Codec "
+              + codec
+              + " for "
+              + classifier
+              + " is incompatible with "
+              + subject
+              + " (of type "
+              + subject.getClass().getName()
+              + ")",
+          e);
+    }
+  }
+
+  /** Builder for {@link ObjectCodecs}. */
+  static class Builder {
+    private final ImmutableMap.Builder<String, ObjectCodec<?>> codecsBuilder =
+        ImmutableMap.builder();
+    private boolean allowDefaultCodec = true;
+
+    private Builder() {}
+
+    /** Add custom serialization strategy ({@code codec}) for {@code classifier}. */
+    public Builder add(String classifier, ObjectCodec<?> codec) {
+      codecsBuilder.put(classifier, codec);
+      return this;
+    }
+
+    /** Set whether or not we allow fallback to the default codec, java serialization. */
+    public Builder setAllowDefaultCodec(boolean allowDefaultCodec) {
+      this.allowDefaultCodec = allowDefaultCodec;
+      return this;
+    }
+
+    /** Wrap this builder with a {@link ClassKeyedBuilder}. */
+    public ClassKeyedBuilder asClassKeyedBuilder() {
+      return new ClassKeyedBuilder(this);
+    }
+
+    /** Wrap this builder with a {@link SkyFunctionNameKeyedBuilder}. */
+    public SkyFunctionNameKeyedBuilder asSkyFunctionNameKeyedBuilder() {
+      return new SkyFunctionNameKeyedBuilder(this);
+    }
+
+    public ObjectCodecs build() {
+      return new ObjectCodecs(codecsBuilder.build(), allowDefaultCodec);
+    }
+  }
+
+  /** Convenience builder for adding codecs classified by class name. */
+  static class ClassKeyedBuilder {
+    private final Builder underlying;
+
+    private ClassKeyedBuilder(Builder underlying) {
+      this.underlying = underlying;
+    }
+
+    public <T> ClassKeyedBuilder add(Class<T> clazz, ObjectCodec<? extends T> codec) {
+      underlying.add(clazz.getName(), codec);
+      return this;
+    }
+
+    public ObjectCodecs build() {
+      return underlying.build();
+    }
+  }
+
+  /** Convenience builder for adding codecs classified by SkyFunctionName. */
+  static class SkyFunctionNameKeyedBuilder {
+    private final Builder underlying;
+
+    private SkyFunctionNameKeyedBuilder(Builder underlying) {
+      this.underlying = underlying;
+    }
+
+    public SkyFunctionNameKeyedBuilder add(SkyFunctionName skyFuncName, ObjectCodec<?> codec) {
+      underlying.add(skyFuncName.getName(), codec);
+      return this;
+    }
+
+    public ObjectCodecs build() {
+      return underlying.build();
+    }
+  }
+
+  private static Map<ByteString, ObjectCodec<?>> makeByteStringMappedCodecs(
+      Map<String, ObjectCodec<?>> stringMappedCodecs) {
+    ImmutableMap.Builder<ByteString, ObjectCodec<?>> result = ImmutableMap.builder();
+    for (Entry<String, ObjectCodec<?>> entry : stringMappedCodecs.entrySet()) {
+      result.put(ByteString.copyFromUtf8(entry.getKey()), entry.getValue());
+    }
+    return result.build();
+  }
+
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PathCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PathCodec.java
new file mode 100644
index 0000000000..4aca0b07ed
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PathCodec.java
@@ -0,0 +1,57 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.skyframe.serialization;
+
+import com.google.devtools.build.lib.util.Preconditions;
+import com.google.devtools.build.lib.vfs.FileSystem;
+import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
+
+/** Custom serialization for {@link Path}s. */
+public class PathCodec implements ObjectCodec<Path> {
+
+  private final FileSystem fileSystem;
+  private final PathFragmentCodec pathFragmentCodec;
+
+  /** Create an instance for serializing and deserializing {@link Path}s on {@code fileSystem}. */
+  PathCodec(FileSystem fileSystem) {
+    this.fileSystem = fileSystem;
+    this.pathFragmentCodec = new PathFragmentCodec();
+  }
+
+  @Override
+  public Class<Path> getEncodedClass() {
+    return Path.class;
+  }
+
+  @Override
+  public void serialize(Path path, CodedOutputStream codedOut) throws IOException {
+    Preconditions.checkState(
+        path.getFileSystem() == fileSystem,
+        "Path's FileSystem (%s) did not match the configured FileSystem (%s)",
+        path.getFileSystem(),
+        fileSystem);
+    pathFragmentCodec.serialize(path.asFragment(), codedOut);
+  }
+
+  @Override
+  public Path deserialize(CodedInputStream codedIn) throws IOException {
+    PathFragment pathFragment = pathFragmentCodec.deserialize(codedIn);
+    return fileSystem.getPath(pathFragment);
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PathFragmentCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PathFragmentCodec.java
index 7438f6de0b..8cd06f2249 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PathFragmentCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/PathFragmentCodec.java
@@ -22,7 +22,7 @@ import java.io.IOException;
 /** Custom serialization for {@link PathFragment}s. */
 class PathFragmentCodec implements ObjectCodec<PathFragment> {
 
-  private final ObjectCodec<String> stringCodec = new FastStringCodec();
+  private final FastStringCodec stringCodec = new FastStringCodec();
 
   @Override
   public Class<PathFragment> getEncodedClass() {
@@ -30,8 +30,7 @@ class PathFragmentCodec implements ObjectCodec<PathFragment> {
   }
 
   @Override
-  public void serialize(PathFragment pathFragment, CodedOutputStream codedOut)
-      throws IOException, SerializationException {
+  public void serialize(PathFragment pathFragment, CodedOutputStream codedOut) throws IOException {
     codedOut.writeInt32NoTag(pathFragment.getDriveLetter());
     codedOut.writeBoolNoTag(pathFragment.isAbsolute());
     codedOut.writeInt32NoTag(pathFragment.segmentCount());
@@ -41,8 +40,7 @@ class PathFragmentCodec implements ObjectCodec<PathFragment> {
   }
 
   @Override
-  public PathFragment deserialize(CodedInputStream codedIn)
-      throws IOException, SerializationException {
+  public PathFragment deserialize(CodedInputStream codedIn) throws IOException {
     char driveLetter = (char) codedIn.readInt32();
     boolean isAbsolute = codedIn.readBool();
     int segmentCount = codedIn.readInt32();