Merge pull request #1501 from twitter/ianoc/cacheBoxedClasses

Cache boxed classes

scalding-core/src/main/scala/com/twitter/scalding/typed/Grouped.scala

@@ -27,6 +27,7 @@ import com.twitter.scalding.serialization.{
   Boxed,
   BoxedOrderedSerialization,
   CascadingBinaryComparator,
+  EquivSerialization,
   OrderedSerialization,
   WrappedSerialization
 }
@@ -97,7 +98,8 @@ object Grouped {
    * to prevent other serializers from handling the key
    */
   private[scalding] def getBoxFnAndOrder[K](ordser: OrderedSerialization[K], flowDef: FlowDef): (K => Boxed[K], BoxedOrderedSerialization[K]) = {
-    val (boxfn, cls) = Boxed.next[K]
+    // We can only supply a cacheKey if the equals and hashcode are known sane
+    val (boxfn, cls) = Boxed.nextCached[K](if (ordser.isInstanceOf[EquivSerialization[_]]) Some(ordser) else None)
     val boxordSer = BoxedOrderedSerialization(boxfn, ordser)
 
     WrappedSerialization.rawSetBinary(List((cls, boxordSer)),

scalding-core/src/test/scala/com/twitter/scalding/ExecutionTest.scala

@@ -29,6 +29,9 @@ import scala.util.Try
 // this is the scalding ExecutionContext
 import ExecutionContext._
 
+import com.twitter.scalding.serialization.OrderedSerialization
+import com.twitter.scalding.serialization.macros.impl.ordered_serialization.runtime_helpers.MacroEqualityOrderedSerialization
+
 object ExecutionTestJobs {
   def wordCount(in: String, out: String) =
     TypedPipe.from(TextLine(in))
@@ -66,6 +69,8 @@ class WordCountEc(args: Args) extends ExecutionJob[Unit](args) {
   }
 }
 
+case class MyCustomType(s: String)
+
 class ExecutionTest extends WordSpec with Matchers {
   "An Execution" should {
     "run" in {
@@ -169,6 +174,42 @@ class ExecutionTest extends WordSpec with Matchers {
       assert((first, second, third) == (1, 1, 1))
     }
 
+    "Running a large loop won't exhaust boxed instances" in {
+      var timesEvaluated = 0
+      import com.twitter.scalding.serialization.macros.impl.BinaryOrdering._
+      // Attempt to use up 4 boxed classes for every execution
+      def baseExecution(idx: Int): Execution[Unit] = TypedPipe.from(0 until 1000).map(_.toShort).flatMap { i =>
+        timesEvaluated += 1
+        List((i, i), (i, i))
+      }.sumByKey.map {
+        case (k, v) =>
+          (k.toInt, v)
+      }.sumByKey.map {
+        case (k, v) =>
+          (k.toLong, v)
+      }.sumByKey.map {
+        case (k, v) =>
+          (k.toString, v)
+      }.sumByKey.map {
+        case (k, v) =>
+          (MyCustomType(k), v)
+      }.sumByKey.writeExecution(TypedTsv(s"/tmp/asdf_${idx}"))
+
+      implicitly[OrderedSerialization[MyCustomType]] match {
+        case mos: MacroEqualityOrderedSerialization[_] => assert(mos.uniqueId == "com.twitter.scalding.typed.MyCustomType")
+        case _ => sys.error("Ordered serialization should have been the MacroEqualityOrderedSerialization for this test")
+      }
+      def executionLoop(idx: Int): Execution[Unit] = {
+        if (idx > 0)
+          baseExecution(idx).flatMap(_ => executionLoop(idx - 1))
+        else
+          Execution.unit
+      }
+
+      executionLoop(55).waitFor(Config.default, Local(true))
+      assert(timesEvaluated == 55 * 1000, "Should run the 55 execution loops for 1000 elements")
+    }
+
     "evaluate shared portions just once, writeExecution" in {
 
       var timesEvaluated = 0

scalding-serialization/src/main/scala/com/twitter/scalding/serialization/Boxed.scala

@@ -818,10 +818,26 @@ object Boxed {
 
   def allClasses: Seq[Class[_ <: Boxed[_]]] = allBoxes.map(_._2)
 
-  def next[K]: (K => Boxed[K], Class[Boxed[K]]) = boxes.get match {
+  private[this] val boxedCache = new java.util.concurrent.ConcurrentHashMap[AnyRef, (Any => Boxed[Any], Class[Boxed[Any]])]()
+
+  private[scalding] def nextCached[K](cacheKey: Option[AnyRef]): (K => Boxed[K], Class[Boxed[K]]) =
+    cacheKey match {
+      case Some(cls) =>
+        val untypedRes = Option(boxedCache.get(cls)) match {
+          case Some(r) => r
+          case None =>
+            val r = next[Any]()
+            boxedCache.putIfAbsent(cls, r)
+            r
+        }
+        untypedRes.asInstanceOf[(K => Boxed[K], Class[Boxed[K]])]
+      case None => next[K]()
+    }
+
+  def next[K](): (K => Boxed[K], Class[Boxed[K]]) = boxes.get match {
     case list @ (h :: tail) if boxes.compareAndSet(list, tail) =>
       h.asInstanceOf[(K => Boxed[K], Class[Boxed[K]])]
-    case (h :: tail) => next[K] // Try again
+    case (h :: tail) => next[K]() // Try again
     case Nil => sys.error(
       """|Scalding's ordered serialization logic exhausted the finite supply of boxed classes.
          |

scalding-serialization/src/main/scala/com/twitter/scalding/serialization/Serialization.scala

@@ -55,6 +55,12 @@ trait Serialization[T] extends Equiv[T] with Hashing[T] with Serializable {
   def dynamicSize(t: T): Option[Int]
 }
 
+/**
+ * In order to cache Serializations having equality and hashes can be useful.
+ * Extend this trait when those two properties can be satisfied
+ */
+trait EquivSerialization[T] extends Serialization[T]
+
 object Serialization {
   import JavaStreamEnrichments._
   /**

scalding-serialization/src/main/scala/com/twitter/scalding/serialization/macros/impl/ordered_serialization/TreeOrderedBuf.scala

@@ -228,10 +228,12 @@ object TreeOrderedBuf {
     val lenB = freshT("lenB")
 
     t.ctx.Expr[OrderedSerialization[T]](q"""
-      new _root_.com.twitter.scalding.serialization.OrderedSerialization[$T] {
+      new _root_.com.twitter.scalding.serialization.macros.impl.ordered_serialization.runtime_helpers.MacroEqualityOrderedSerialization[$T] {
         // Ensure macro hygene for Option/Some/None
         import _root_.scala.{Option, Some, None}
 
+        override val uniqueId: String = ${T.tpe.toString}
+
         private[this] var lengthCalculationAttempts: Long = 0L
         private[this] var couldNotLenCalc: Long = 0L
         private[this] var skipLenCalc: Boolean = false

scalding-serialization/src/main/scala/com/twitter/scalding/serialization/macros/impl/ordered_serialization/runtime_helpers/MacroEqualityOrderedSerialization.scala

@@ -0,0 +1,31 @@
+/*
+ Copyright 2016 Twitter, Inc.
+
+ 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.twitter.scalding.serialization.macros.impl.ordered_serialization.runtime_helpers
+
+import com.twitter.scalding.serialization.{ EquivSerialization, OrderedSerialization }
+
+object MacroEqualityOrderedSerialization {
+  private val seed = "MacroEqualityOrderedSerialization".hashCode
+}
+
+abstract class MacroEqualityOrderedSerialization[T] extends OrderedSerialization[T] with EquivSerialization[T] {
+  def uniqueId: String
+  override def hashCode = MacroEqualityOrderedSerialization.seed ^ uniqueId.hashCode
+  override def equals(other: Any): Boolean = other match {
+    case o: MacroEqualityOrderedSerialization[_] => o.uniqueId == uniqueId
+    case _ => false
+  }
+}