[FLINK-11865] Improve code generation speed in TraversableSerializer

flink-scala/src/main/scala/org/apache/flink/api/scala/typeutils/TraversableSerializer.scala

@@ -17,17 +17,20 @@
  */
 package org.apache.flink.api.scala.typeutils
 
-import java.io.{ObjectInputStream, ObjectOutputStream}
+import java.io.ObjectInputStream
+import java.util.concurrent.Callable
 
 import org.apache.flink.annotation.Internal
 import org.apache.flink.api.common.typeutils._
 import org.apache.flink.core.memory.{DataInputView, DataOutputView}
+import org.apache.flink.shaded.guava18.com.google.common.cache.{Cache, CacheBuilder}
 
 import scala.collection.generic.CanBuildFrom
+import scala.ref.WeakReference
 
 /**
- * Serializer for Scala Collections.
- */
+  * Serializer for Scala Collections.
+  */
 @Internal
 @SerialVersionUID(7522917416391312410L)
 class TraversableSerializer[T <: TraversableOnce[E], E](
@@ -44,15 +47,8 @@ class TraversableSerializer[T <: TraversableOnce[E], E](
   protected def legacyCbfCode: String = null
 
   def compileCbf(code: String): CanBuildFrom[T, E, T] = {
-
-    import scala.reflect.runtime.universe._
-    import scala.tools.reflect.ToolBox
-
-    val tb = runtimeMirror(Thread.currentThread().getContextClassLoader).mkToolBox()
-    val tree = tb.parse(code)
-    val compiled = tb.compile(tree)
-    val cbf = compiled()
-    cbf.asInstanceOf[CanBuildFrom[T, E, T]]
+    val cl = Thread.currentThread().getContextClassLoader
+    TraversableSerializer.compileCbf(cl, code)
   }
 
   override def duplicate = {
@@ -173,3 +169,69 @@ class TraversableSerializer[T <: TraversableOnce[E], E](
     new TraversableSerializerSnapshot[T, E](this)
   }
 }
+
+object TraversableSerializer {
+
+  private val CACHE: Cache[Key, CanBuildFrom[_, _, _]] = CacheBuilder.newBuilder()
+    .weakValues()
+    .maximumSize(128)
+    .build()
+
+  def compileCbf[T, E](classLoader: ClassLoader, cbfCode: String): CanBuildFrom[T, E, T] = {
+    val key = Key(classLoader, cbfCode)
+
+    CACHE
+      .get(key, LazyRuntimeCompiler(classLoader, cbfCode))
+      .asInstanceOf[CanBuildFrom[T, E, T]]
+  }
+
+  object Key {
+
+    def apply(classLoader: ClassLoader, cbfCode: String): Key = {
+      val hashCode = System.identityHashCode(classLoader)
+      val weakReference = WeakReference(classLoader)
+      Key(hashCode, weakReference, cbfCode)
+    }
+  }
+
+  case class Key(classLoaderHash: Int,
+                 classLoaderRef: WeakReference[ClassLoader],
+                 cbfCode: String) {
+
+    override def hashCode(): Int = classLoaderHash * 37 + cbfCode.hashCode
+
+    override def equals(obj: Any): Boolean = {
+      obj match {
+        case Key(thatHashCode, thatClassLoaderRef, thatCbfCode) => 
+          (this.classLoaderHash == thatHashCode) && 
+          (this.classLoaderRef.get == thatClassLoaderRef.get) &&
+          (this.cbfCode == thatCbfCode)
+
+        case _ =>
+          false
+      }
+    }
+  }
+
+  private case class LazyRuntimeCompiler[T, E](classLoader: ClassLoader,
+                                               code: String)
+    extends Callable[CanBuildFrom[T, E, T]] {
+
+    override def call(): CanBuildFrom[T, E, T] = compileCbfInternal(classLoader, code)
+
+    private def compileCbfInternal(classLoader: ClassLoader, code: String):
+    CanBuildFrom[T, E, T] = {
+
+      import scala.reflect.runtime.universe._
+      import scala.tools.reflect.ToolBox
+
+      val tb = runtimeMirror(classLoader).mkToolBox()
+      val tree = tb.parse(code)
+      val compiled = tb.compile(tree)
+      val cbf = compiled()
+      cbf.asInstanceOf[CanBuildFrom[T, E, T]]
+    }
+  }
+
+}
+

flink-scala/src/test/scala/org/apache/flink/api/scala/runtime/TraversableSerializerTest.scala

@@ -24,10 +24,12 @@ import org.apache.flink.api.java.typeutils.runtime.kryo.KryoSerializer
 import org.apache.flink.api.scala._
 import org.apache.flink.api.scala.typeutils.TraversableSerializer
 import org.junit.Assert._
-import org.junit.{Assert, Ignore, Test}
+import org.junit.{Assert, Test}
 
-import scala.collection.immutable.{BitSet, LinearSeq, SortedSet}
-import scala.collection.{SortedMap, mutable}
+import scala.collection.immutable.{BitSet, LinearSeq}
+import scala.collection.mutable
+import scala.ref.WeakReference
+import scala.reflect.internal.util.ScalaClassLoader.URLClassLoader
 
 class TraversableSerializerTest {
 
@@ -97,6 +99,55 @@ class TraversableSerializerTest {
     runTests(testData)
   }
 
+  @Test
+  def sameClassLoaderAndCodeShouldProvideEqualKeys(): Unit = {
+    val classLoaderA = new URLClassLoader(Seq.empty[java.net.URL], null)
+
+    val keyA = TraversableSerializer.Key(classLoaderA, "code")
+    val keyB = TraversableSerializer.Key(classLoaderA, "code")
+
+    assertEquals(keyA, keyB)
+  }
+
+  @Test
+  def differentClassLoadersProvideNonEqualKeys(): Unit = {
+    val classLoaderA = new URLClassLoader(Seq.empty[java.net.URL], null)
+    val classLoaderB = new URLClassLoader(Seq.empty[java.net.URL], null)
+
+    val keyA = TraversableSerializer.Key(classLoaderA, "code")
+    val keyB = TraversableSerializer.Key(classLoaderB, "code")
+
+    assertNotEquals(keyA, keyB)
+  }
+
+  @Test
+  def expiredReferenceShouldProduceNonEqualKeys(): Unit = {
+    val classLoaderA = new URLClassLoader(Seq.empty[java.net.URL], null)
+
+    val keyA = TraversableSerializer.Key(classLoaderA, "code")
+    val keyB = keyA.copy(classLoaderRef = WeakReference(null)) 
+
+    assertNotEquals(keyA, keyB)
+  }
+
+  @Test
+  def bootStrapClassLoaderShouldProduceTheSameKeys(): Unit = {
+    val keyA = TraversableSerializer.Key(null, "a")
+    val keyB = TraversableSerializer.Key(null, "a")
+
+    assertEquals(keyA, keyB)
+  }
+
+  @Test
+  def differentCanBuildFromCodeShouldProduceDifferentKeys(): Unit = {
+    val classLoaderA = new URLClassLoader(Seq.empty[java.net.URL], null)
+
+    val keyA = TraversableSerializer.Key(classLoaderA, "a")
+    val keyB = TraversableSerializer.Key(classLoaderA, "b")
+
+    assertNotEquals(keyA, keyB)
+  }
+
   private final def runTests[T : TypeInformation](instances: Array[T]) {
     try {
       val typeInfo = implicitly[TypeInformation[T]]