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,
+  EquivOrderedSerialization,
   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.next[K](if (ordser.isInstanceOf[EquivOrderedSerialization[_]]) Some(ordser) else None)
     val boxordSer = BoxedOrderedSerialization(boxfn, ordser)
 
     WrappedSerialization.rawSetBinary(List((cls, boxordSer)),

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

@@ -66,6 +66,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 {
@@ -152,6 +154,38 @@ 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}"))
+
+      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

@@ -801,10 +801,27 @@ 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]])]()
+
+  def next[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 = fetchNextBoxed[K]()
+            boxedCache.putIfAbsent(cls, r.asInstanceOf[(Any => Boxed[Any], Class[Boxed[Any]])])
+            r
+        }
+        untypedRes.asInstanceOf[(K => Boxed[K], Class[Boxed[K]])]
+      case None =>
+        fetchNextBoxed[K]()
+    }
+
+  private[scalding] def fetchNextBoxed[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) => fetchNextBoxed[K] // Try again
     case Nil => sys.error("Exhausted the boxed classes")
   }
 }

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

@@ -34,6 +34,12 @@ trait OrderedSerialization[T] extends Ordering[T] with Serialization[T] {
   def compareBinary(a: InputStream, b: InputStream): OrderedSerialization.Result
 }
 
+/**
+ * In order to cache OrderedSerializations having equality and hashes can be useful.
+ * Extend this trait when those two properties can be satisfied
+ */
+trait EquivOrderedSerialization[T] extends OrderedSerialization[T]
+
 object OrderedSerialization {
   /**
    * Represents the result of a comparison that might fail due

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 2014 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.EquivOrderedSerialization
+
+object MacroEqualityOrderedSerialization {
+  private val seed = "MacroEqualityOrderedSerialization".hashCode
+}
+
+abstract class MacroEqualityOrderedSerialization[T] extends EquivOrderedSerialization[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
+  }
+}