1 files changed, 272 insertions, 0 deletions
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();
+  }
+
+}