[BEAM-1144] Spark runner fails to deserialize MicrobatchSource in cluster mode

runners/spark/pom.xml

@@ -134,14 +134,11 @@
       <version>${hadoop.version}</version>
       <scope>provided</scope>
     </dependency>
-    <!-- Kryo bugfix version needed due to a state re-use issue in Kryo version 2.21 used in Spark 1.x
-    See: https://issues.apache.org/jira/browse/SPARK-7708
-    See: https://github.com/EsotericSoftware/kryo/issues/312
-    -->
     <dependency>
       <groupId>com.esotericsoftware.kryo</groupId>
       <artifactId>kryo</artifactId>
-      <version>2.21.1</version>
+      <version>2.21</version>
+      <scope>provided</scope>
     </dependency>
     <dependency>
       <groupId>com.google.code.findbugs</groupId>
@@ -396,10 +393,6 @@
                     <pattern>com.google.thirdparty</pattern>
                     <shadedPattern>org.apache.beam.spark.relocated.com.google.thirdparty</shadedPattern>
                   </relocation>
-                  <relocation>
-                    <pattern>com.esotericsoftware.kryo</pattern>
-                    <shadedPattern>org.apache.beam.spark.relocated.com.esotericsoftware.kryo</shadedPattern>
-                  </relocation>
                 </relocations>
                 <shadedArtifactAttached>true</shadedArtifactAttached>
                 <shadedClassifierName>spark-app</shadedClassifierName>

runners/spark/src/main/java/org/apache/beam/runners/spark/coders/BeamSparkRunnerRegistrator.java

@@ -40,7 +40,7 @@ public class BeamSparkRunnerRegistrator implements KryoRegistrator {
   @Override
   public void registerClasses(Kryo kryo) {
     for (Class<?> clazz : ClassesForJavaSerialization.getClasses()) {
-      kryo.register(clazz, new JavaSerializer());
+      kryo.register(clazz, new StatelessJavaSerializer());
     }
   }
 

runners/spark/src/main/java/org/apache/beam/runners/spark/coders/StatelessJavaSerializer.java

@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.beam.runners.spark.coders;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.KryoException;
+import com.esotericsoftware.kryo.Serializer;
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.Output;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.ObjectStreamClass;
+
+
+/**
+ * Stateless Java Serializer.
+ *
+ * <p>
+ * Solves state re-use issue in Kryo version 2.21 used in Spark 1.x
+ * See:
+ * https://issues.apache.org/jira/browse/SPARK-7708
+ * https://github.com/EsotericSoftware/kryo/issues/312
+ * </p>
+ *
+ * <p>
+ * Also, solves class loading issue in cluster caused by ${@link ObjectInputStream}
+ * by using ${@link ObjectInputStreamWithClassLoader}
+ * ${@link ObjectInputStream} uses the last user-defined class loader in the stack which can be the
+ * wrong class loader.
+ * This is a known Java issue and a similar solution is often used.
+ * See:
+ * https://github.com/apache/spark/blob/v1.6.3/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala#L154
+ * https://issues.apache.org/jira/browse/GROOVY-1627
+ * https://github.com/spring-projects/spring-loaded/issues/107
+ * </p>
+ */
+class StatelessJavaSerializer extends Serializer {
+  @SuppressWarnings("unchecked")
+  public void write(Kryo kryo, Output output, Object object) {
+    try {
+      ObjectOutputStream objectStream = new ObjectOutputStream(output);
+      objectStream.writeObject(object);
+      objectStream.flush();
+    } catch (Exception e) {
+      throw new KryoException("Error during Java serialization.", e);
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  public Object read (Kryo kryo, Input input, Class type) {
+    try {
+      return new ObjectInputStreamWithClassLoader(input, kryo.getClassLoader()).readObject();
+    } catch (Exception e) {
+      throw new KryoException("Error during Java deserialization.", e);
+    }
+  }
+
+  /**
+   * ObjectInputStream with specific ClassLoader.
+   */
+  private static class ObjectInputStreamWithClassLoader extends ObjectInputStream {
+    private final ClassLoader classLoader;
+
+    ObjectInputStreamWithClassLoader(InputStream in, ClassLoader classLoader) throws IOException {
+      super(in);
+      this.classLoader = classLoader;
+    }
+
+    @Override
+    protected Class<?> resolveClass(ObjectStreamClass desc) {
+      try {
+        return Class.forName(desc.getName(), false, classLoader);
+      } catch (ClassNotFoundException e) {
+        throw new RuntimeException("Could not find class: " + desc.getName(), e);
+      }
+    }
+  }
+}

runners/spark/src/test/java/org/apache/beam/runners/spark/coders/BeamSparkRunnerRegistratorTest.java

@@ -19,7 +19,7 @@
 package org.apache.beam.runners.spark.coders;
 
 import com.esotericsoftware.kryo.Kryo;
-import com.esotericsoftware.kryo.serializers.JavaSerializer;
+
 import com.google.common.collect.Iterables;
 import java.io.Serializable;
 import org.apache.beam.sdk.coders.Coder;
@@ -49,9 +49,9 @@ public void testCodersAndSourcesRegistration() {
 
     for (Class<?> clazz : classesForJavaSerialization) {
       Assert.assertThat("Registered serializer for class " + clazz.getName()
-              + " was not an instance of " + JavaSerializer.class.getName(),
+          + " was not an instance of " + StatelessJavaSerializer.class.getName(),
           kryo.getSerializer(clazz),
-          Matchers.instanceOf(JavaSerializer.class));
+          Matchers.instanceOf(StatelessJavaSerializer.class));
     }
   }
 }