Merge pull request #192 from rxin/asJavaIterable

Support serializing asJavaIterable.

Author: ianoc

chill-scala/src/main/scala/com/twitter/chill/JavaIterableWrapperSerializer.scala

@@ -0,0 +1,48 @@
+package com.twitter.chill
+
+import _root_.java.lang.{ Iterable => JIterable }
+
+/**
+ * A Kryo serializer for serializing results returned by asJavaIterable.
+ *
+ * The underlying object is scala.collection.convert.Wrappers$IterableWrapper.
+ * Kryo deserializes this into an AbstractCollection, which unfortunately doesn't work.
+ *
+ * Ported from Apache Spark's KryoSerializer.scala.
+ */
+private class JavaIterableWrapperSerializer extends KSerializer[JIterable[_]] {
+
+  import JavaIterableWrapperSerializer._
+
+  override def write(kryo: Kryo, out: Output, obj: JIterable[_]): Unit = {
+    // If the object is the wrapper, simply serialize the underlying Scala Iterable object.
+    // Otherwise, serialize the object itself.
+    if (obj.getClass == wrapperClass && underlyingMethodOpt.isDefined) {
+      kryo.writeClassAndObject(out, underlyingMethodOpt.get.invoke(obj))
+    } else {
+      kryo.writeClassAndObject(out, obj)
+    }
+  }
+
+  override def read(kryo: Kryo, in: Input, clz: Class[JIterable[_]]): JIterable[_] = {
+    kryo.readClassAndObject(in) match {
+      case scalaIterable: Iterable[_] =>
+        scala.collection.JavaConversions.asJavaIterable(scalaIterable)
+      case javaIterable: JIterable[_] =>
+        javaIterable
+    }
+  }
+}
+
+private object JavaIterableWrapperSerializer {
+  // The class returned by asJavaIterable (scala.collection.convert.Wrappers$IterableWrapper).
+  val wrapperClass = scala.collection.JavaConversions.asJavaIterable(Seq(1)).getClass
+
+  // Get the underlying method so we can use it to get the Scala collection for serialization.
+  private val underlyingMethodOpt = {
+    try Some(wrapperClass.getDeclaredMethod("underlying")) catch {
+      case e: Exception =>
+        None
+    }
+  }
+}

chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala

@@ -163,11 +163,21 @@ class ScalaCollectionsRegistrar extends IKryoRegistrar {
   }
 }
 
+class JavaWrapperCollectionRegistrar extends IKryoRegistrar {
+  def apply(newK: Kryo) {
+    newK.register(JavaIterableWrapperSerializer.wrapperClass, new JavaIterableWrapperSerializer)
+  }
+}
+
 /** Registers all the scala (and java) serializers we have */
 class AllScalaRegistrar extends IKryoRegistrar {
   def apply(k: Kryo) {
     val col = new ScalaCollectionsRegistrar
     col(k)
+
+    val jcol = new JavaWrapperCollectionRegistrar
+    jcol(k)
+
     // Register all 22 tuple serializers and specialized serializers
     ScalaTupleSerialization.register(k)
     k.forClass[Symbol](new KSerializer[Symbol] {

chill-scala/src/test/scala/com/twitter/chill/KryoSpec.scala

@@ -172,6 +172,10 @@ class KryoSpec extends Specification with BaseProperties {
     "handle scala enums" in {
       WeekDay.values.foreach { _ must roundtrip }
     }
+    "handle asJavaIterable" in {
+      val col = scala.collection.JavaConversions.asJavaIterable(Seq(12345))
+      col must roundtrip
+    }
     "use java serialization" in {
       val kinst = { () => getKryo.javaForClass[TestCaseClassForSerialization] }
       rtEquiv(kinst, TestCaseClassForSerialization("hey", 42)) must beTrue