Serialize lambdas when they are cast to Serializable.

We could conceivably do some monkey-patching at server startup to make all lambdas act Serializable, but one step at a time.

PiperOrigin-RevId: 200509321

6 files changed, 202 insertions, 19 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java
index d78ba32791..de3d34ebcf 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java
@@ -82,6 +82,7 @@ import com.google.devtools.build.lib.rules.java.ProguardHelper.ProguardOutput;
 import com.google.devtools.build.lib.rules.java.ProguardSpecProvider;
 import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
@@ -1373,13 +1374,7 @@ public abstract class AndroidBinary implements RuleConfiguredTargetFactory {
             .setExecutable(ruleContext.getExecutablePrerequisite("$dexmerger", Mode.HOST))
             .setMnemonics("DexShardsToMerge", "DexMerger")
             .setOutputPathMapper(
-                // Use an anonymous inner class for serialization.
-                new OutputPathMapper() {
-                  @Override
-                  public PathFragment parentRelativeOutputPath(TreeFileArtifact input) {
-                    return input.getParentRelativePath();
-                  }
-                });
+                (OutputPathMapper & Serializable) TreeFileArtifact::getParentRelativePath);
     CustomCommandLine.Builder commandLine =
         CustomCommandLine.builder()
             .addPlaceholderTreeArtifactExecPath("--input", inputTree)
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
index 87be12e504..c4e880881e 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
@@ -46,7 +46,10 @@ public class AutoRegistry {
 
   /** Classes outside {@link AutoRegistry#PACKAGE_PREFIX} that need to be serialized. */
   private static final ImmutableList<String> EXTERNAL_CLASS_NAMES_TO_REGISTER =
-      ImmutableList.of("java.io.FileNotFoundException", "java.io.IOException");
+      ImmutableList.of(
+          "java.io.FileNotFoundException",
+          "java.io.IOException",
+          "java.lang.invoke.SerializedLambda");
 
   private static final ImmutableList<Object> REFERENCE_CONSTANTS_TO_REGISTER =
       ImmutableList.of(
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/LambdaCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/LambdaCodec.java
new file mode 100644
index 0000000000..1127308ca9
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/LambdaCodec.java
@@ -0,0 +1,93 @@
+// Copyright 2018 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.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
+import java.io.Serializable;
+import java.lang.invoke.SerializedLambda;
+import java.lang.reflect.Method;
+
+/**
+ * A codec for Java 8 serializable lambdas. Lambdas that are tagged as {@link Serializable} have a
+ * generated method, {@code writeReplace}, that converts them into a {@link SerializedLambda}, which
+ * can then be serialized like a normal object. On deserialization, we call {@link
+ * SerializedLambda#readResolve}, which converts the object back into a lambda.
+ *
+ * <p>Since lambdas do not share a common base class, choosing this codec for serializing them must
+ * be special-cased in {@link ObjectCodecRegistry}. We must also make a somewhat arbitrary choice
+ * around the generic parameter. Since all of our lambdas are {@link Serializable}, we use that.
+ * Because {@link Serializable} is an interface, not a class, this codec will never be chosen for
+ * any object without special-casing.
+ */
+class LambdaCodec implements ObjectCodec<Serializable> {
+  private final Method readResolveMethod;
+
+  LambdaCodec() {
+    try {
+      this.readResolveMethod = SerializedLambda.class.getDeclaredMethod("readResolve");
+    } catch (NoSuchMethodException e) {
+      throw new IllegalStateException(e);
+    }
+    readResolveMethod.setAccessible(true);
+  }
+
+  static boolean isProbablyLambda(Class<?> type) {
+    return type.isSynthetic() && !type.isLocalClass() && !type.isAnonymousClass();
+  }
+
+  @Override
+  public Class<? extends Serializable> getEncodedClass() {
+    return Serializable.class;
+  }
+
+  @Override
+  public void serialize(SerializationContext context, Serializable obj, CodedOutputStream codedOut)
+      throws SerializationException, IOException {
+    Class<?> objClass = obj.getClass();
+    if (!isProbablyLambda(objClass)) {
+      throw new SerializationException(obj + " is not a lambda: " + objClass);
+    }
+    Method writeReplaceMethod;
+    try {
+      // TODO(janakr): We could cache these methods if retrieval shows up as a hotspot.
+      writeReplaceMethod = objClass.getDeclaredMethod("writeReplace");
+    } catch (NoSuchMethodException e) {
+      throw new SerializationException(
+          "No writeReplace method for " + obj + " with " + objClass, e);
+    }
+    writeReplaceMethod.setAccessible(true);
+    SerializedLambda serializedLambda;
+    try {
+      serializedLambda = (SerializedLambda) writeReplaceMethod.invoke(obj);
+    } catch (ReflectiveOperationException e) {
+      throw new SerializationException(
+          "Exception invoking writeReplace for " + obj + " with " + objClass, e);
+    }
+    context.serialize(serializedLambda, codedOut);
+  }
+
+  @Override
+  public Serializable deserialize(DeserializationContext context, CodedInputStream codedIn)
+      throws SerializationException, IOException {
+    SerializedLambda serializedLambda = context.deserialize(codedIn);
+    try {
+      return (Serializable) readResolveMethod.invoke(serializedLambda);
+    } catch (ReflectiveOperationException e) {
+      throw new IllegalStateException("Error read-resolving " + serializedLambda, e);
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodec.java
index ab004def86..7ab9de5959 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodec.java
@@ -30,6 +30,11 @@ public interface ObjectCodec<T> {
    *
    * <p>This is useful for automatically dispatching to the correct codec, e.g. in {@link
    * ObjectCodecs}.
+   *
+   * <p>If {@link T} is an interface, then this codec will never be used by the auto-registration
+   * framework in {@link ObjectCodecRegistry} unless it is explicitly invoked or {@link
+   * #additionalEncodedClasses} is non-empty, since the {@link ObjectCodecRegistry} traverses the
+   * concrete class hierarchy looking for matches, and will never come to an interface.
    */
   Class<? extends T> getEncodedClass();
 
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
index 2fb0ac39e7..7c28f6a581 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
@@ -23,6 +23,7 @@ import com.google.common.collect.ImmutableSortedSet;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.HashMap;
@@ -106,7 +107,7 @@ public class ObjectCodecRegistry {
       // Enums must be serialized using declaring class.
       type = ((Enum) obj).getDeclaringClass();
     }
-    return getDynamicCodecDescriptor(type.getName());
+    return getDynamicCodecDescriptor(type.getName(), type);
   }
 
   /**
@@ -157,7 +158,7 @@ public class ObjectCodecRegistry {
     if (!allowDefaultCodec || tagOffset < 0 || tagOffset >= classNames.size()) {
       throw new SerializationException.NoCodecException("No codec available for tag " + tag);
     }
-    return getDynamicCodecDescriptor(classNames.get(tagOffset));
+    return getDynamicCodecDescriptor(classNames.get(tagOffset), /*type=*/ null);
   }
 
   /**
@@ -382,18 +383,32 @@ public class ObjectCodecRegistry {
     return new TypedCodecDescriptor(tag, new EnumCodec(enumType));
   }
 
-  private CodecDescriptor getDynamicCodecDescriptor(String className)
+  private CodecDescriptor getDynamicCodecDescriptor(String className, @Nullable Class<?> type)
       throws SerializationException.NoCodecException {
     Supplier<CodecDescriptor> supplier = dynamicCodecs.get(className);
-    if (supplier == null) {
-      throw new SerializationException.NoCodecException(
-          "No default codec available for " + className);
+    if (supplier != null) {
+      CodecDescriptor descriptor = supplier.get();
+      if (descriptor == null) {
+        throw new SerializationException.NoCodecException(
+            "There was a problem creating a codec for " + className + ". Check logs for details");
+      }
+      return descriptor;
     }
-    CodecDescriptor descriptor = supplier.get();
-    if (descriptor == null) {
-      throw new SerializationException.NoCodecException(
-          "There was a problem creating a codec for " + className + " check logs for details.");
+    if (type != null && LambdaCodec.isProbablyLambda(type)) {
+      if (Serializable.class.isAssignableFrom(type)) {
+        // LambdaCodec is hidden away as a codec for Serializable. This avoids special-casing it in
+        // all places we look up a codec, and doesn't clash with anything else because Serializable
+        // is an interface, not a class.
+        return classMappedCodecs.get(Serializable.class);
+      } else {
+        throw new SerializationException.NoCodecException(
+            "No default codec available for "
+                + className
+                + ". If this is a lambda, try casting it to (type & Serializable), like "
+                + "(Supplier<String> & Serializable)");
+      }
     }
-    return descriptor;
+    throw new SerializationException.NoCodecException(
+        "No default codec available for " + className);
   }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/LambdaCodecTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/LambdaCodecTest.java
new file mode 100644
index 0000000000..4d987e2567
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/LambdaCodecTest.java
@@ -0,0 +1,72 @@
+// Copyright 2018 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 static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.testutil.MoreAsserts.assertThrows;
+
+import com.google.devtools.build.lib.skyframe.serialization.testutils.SerializationTester;
+import com.google.devtools.build.lib.skyframe.serialization.testutils.TestUtils;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link LambdaCodec}. */
+@RunWith(JUnit4.class)
+public class LambdaCodecTest {
+  private interface MyInterface {
+    boolean func(String arg);
+  }
+
+  @Test
+  public void smoke() throws Exception {
+    List<Boolean> returnValue = new ArrayList<>();
+    new SerializationTester(
+            (Supplier<Void> & Serializable) () -> null,
+            (Function<Object, String> & Serializable) Object::toString,
+            (MyInterface & Serializable) arg -> "hello".equals(arg),
+            (MyInterface & Serializable) "hello"::equals,
+            (MyInterface & Serializable) (arg) -> !returnValue.isEmpty())
+        // We can't compare lambdas for equality, just make sure they get deserialized.
+        .setVerificationFunction((original, deserialized) -> assertThat(deserialized).isNotNull())
+        .runTests();
+  }
+
+  @Test
+  public void lambdaBehaviorPreserved() throws Exception {
+    List<Boolean> returnValue = new ArrayList<>();
+    MyInterface lambda = (MyInterface & Serializable) (arg) -> returnValue.isEmpty();
+    MyInterface deserializedLambda = TestUtils.roundTrip(lambda);
+    assertThat(lambda.func("any")).isTrue();
+    assertThat(deserializedLambda.func("any")).isTrue();
+    returnValue.add(true);
+    assertThat(lambda.func("any")).isFalse();
+    // Deserialized object's list is not the same as original's. Changes to original aren't seen.
+    assertThat(deserializedLambda.func("any")).isTrue();
+  }
+
+  @Test
+  public void onlySerializableWorks() {
+    MyInterface unserializableLambda = (arg) -> true;
+    assertThrows(
+        SerializationException.class,
+        () -> TestUtils.toBytesMemoized(unserializableLambda, AutoRegistry.get()));
+  }
+}