Changed TimeStampedHashMap to use WrappedJavaHashMap.

Author: tdas

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -225,6 +225,7 @@ private[spark] class BlockManager(
    * the slave needs to re-register.
    */
   private def tryToReportBlockStatus(blockId: BlockId, info: BlockInfo, droppedMemorySize: Long = 0L): Boolean = {
+    logInfo("Reporting " + blockId)
     val (curLevel, inMemSize, onDiskSize, tellMaster) = info.synchronized {
       info.level match {
         case null =>
@@ -770,7 +771,7 @@ private[spark] class BlockManager(
     val iterator = blockInfo.internalMap.entrySet().iterator()
     while (iterator.hasNext) {
       val entry = iterator.next()
-      val (id, info, time) = (entry.getKey, entry.getValue._1, entry.getValue._2)
+      val (id, info, time) = (entry.getKey, entry.getValue.value, entry.getValue.timestamp)
       if (time < cleanupTime && shouldDrop(id)) {
         info.synchronized {
           val level = info.level

core/src/main/scala/org/apache/spark/util/BoundedHashMap.scala

@@ -1,3 +1,20 @@
+/*
+ * 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.spark.util
 
 import scala.collection.mutable.{ArrayBuffer, SynchronizedMap}
@@ -7,16 +24,20 @@ import java.util.Map.{Entry => JMapEntry}
 import scala.reflect.ClassTag
 
 /**
- * A map that bounds the number of key-value pairs present in it. It can be configured to
- * drop least recently inserted or used pair. It exposes a scala.collection.mutable.Map interface
- * to allow it to be a drop-in replacement of Scala HashMaps. Internally, a Java LinkedHashMap is
- * used to get insert-order or access-order behavior. Note that the LinkedHashMap is not
- * thread-safe and hence, it is wrapped in a Collections.synchronizedMap.
- * However, getting the Java HashMap's iterator and using it can still lead to
- * ConcurrentModificationExceptions. Hence, the iterator() function is overridden to copy the
- * all pairs into an ArrayBuffer and then return the iterator to the ArrayBuffer. Also,
- * the class apply the trait SynchronizedMap which ensures that all calls to the Scala Map API
- * are synchronized. This together ensures that ConcurrentModificationException is never thrown.
+ * A map that upper bounds the number of key-value pairs present in it. It can be configured to
+ * drop the least recently user pair or the earliest inserted pair. It exposes a
+ * scala.collection.mutable.Map interface to allow it to be a drop-in replacement for Scala
+ * HashMaps.
+ *
+ * Internally, a Java LinkedHashMap is used to get insert-order or access-order behavior.
+ * Note that the LinkedHashMap is not thread-safe and hence, it is wrapped in a
+ * Collections.synchronizedMap. However, getting the Java HashMap's iterator and
+ * using it can still lead to ConcurrentModificationExceptions. Hence, the iterator()
+ * function is overridden to copy the all pairs into an ArrayBuffer and then return the
+ * iterator to the ArrayBuffer. Also, the class apply the trait SynchronizedMap which
+ * ensures that all calls to the Scala Map API are synchronized. This together ensures
+ * that ConcurrentModificationException is never thrown.
+ *
  * @param bound   max number of key-value pairs
  * @param useLRU  true = least recently used/accessed will be dropped when bound is reached,
  *                false = earliest inserted will be dropped
@@ -37,7 +58,8 @@ private[spark] class BoundedHashMap[A, B](bound: Int, useLRU: Boolean)
 
   /**
    * Overriding iterator to make sure that the internal Java HashMap's iterator
-   * is not concurrently modified.
+   * is not concurrently modified. This can be a performance issue and this should be overridden
+   * if it is known that this map will not be used in a multi-threaded environment.
    */
   override def iterator: Iterator[(A, B)] = {
     (new ArrayBuffer[(A, B)] ++= super.iterator).iterator

core/src/main/scala/org/apache/spark/util/TimeStampedHashMap.scala

@@ -18,108 +18,66 @@
 package org.apache.spark.util
 
 import java.util.concurrent.ConcurrentHashMap
-import scala.collection.JavaConversions
-import scala.collection.mutable.Map
-import scala.collection.immutable
-import org.apache.spark.scheduler.MapStatus
+
 import org.apache.spark.Logging
 
+private[util] case class TimeStampedValue[T](timestamp: Long, value: T)
+
 /**
- * This is a custom implementation of scala.collection.mutable.Map which stores the insertion
- * timestamp along with each key-value pair. If specified, the timestamp of each pair can be
- * updated every time it is accessed. Key-value pairs whose timestamp are older than a particular
- * threshold time can then be removed using the clearOldValues method. This is intended to
- * be a drop-in replacement of scala.collection.mutable.HashMap.
+ * A map that stores the timestamp of when a key was inserted along with the value. If specified,
+ * the timestamp of each pair can be updated every time it is accessed.
+ * Key-value pairs whose timestamps are older than a particular
+ * threshold time can then be removed using the clearOldValues method. It exposes a
+ * scala.collection.mutable.Map interface to allow it to be a drop-in replacement for Scala
+ * HashMaps.
+ *
+ * Internally, it uses a Java ConcurrentHashMap, so all operations on this HashMap are thread-safe.
+ *
  * @param updateTimeStampOnGet When enabled, the timestamp of a pair will be
  *                             updated when it is accessed
  */
-class TimeStampedHashMap[A, B](updateTimeStampOnGet: Boolean = false)
-  extends Map[A, B]() with Logging {
-  val internalMap = new ConcurrentHashMap[A, (B, Long)]()
-
-  def get(key: A): Option[B] = {
-    val value = internalMap.get(key)
-    if (value != null && updateTimeStampOnGet) {
-      internalMap.replace(key, value, (value._1, currentTime))
-    }
-    Option(value).map(_._1)
-  }
-
-  def iterator: Iterator[(A, B)] = {
-    val jIterator = internalMap.entrySet().iterator()
-    JavaConversions.asScalaIterator(jIterator).map(kv => (kv.getKey, kv.getValue._1))
-  }
-
-  override def + [B1 >: B](kv: (A, B1)): Map[A, B1] = {
-    val newMap = new TimeStampedHashMap[A, B1]
-    newMap.internalMap.putAll(this.internalMap)
-    newMap.internalMap.put(kv._1, (kv._2, currentTime))
-    newMap
-  }
-
-  override def - (key: A): Map[A, B] = {
-    val newMap = new TimeStampedHashMap[A, B]
-    newMap.internalMap.putAll(this.internalMap)
-    newMap.internalMap.remove(key)
-    newMap
-  }
+private[spark] class TimeStampedHashMap[A, B](updateTimeStampOnGet: Boolean = false)
+  extends WrappedJavaHashMap[A, B, A, TimeStampedValue[B]] with Logging {
 
-  override def += (kv: (A, B)): this.type = {
-    internalMap.put(kv._1, (kv._2, currentTime))
-    this
-  }
+  protected[util] val internalJavaMap = new ConcurrentHashMap[A, TimeStampedValue[B]]()
 
-  // Should we return previous value directly or as Option ?
-  def putIfAbsent(key: A, value: B): Option[B] = {
-    val prev = internalMap.putIfAbsent(key, (value, currentTime))
-    if (prev != null) Some(prev._1) else None
+  protected[util] def newInstance[K1, V1](): WrappedJavaHashMap[K1, V1, _, _] = {
+    new TimeStampedHashMap[K1, V1]()
   }
 
+  def internalMap = internalJavaMap
 
-  override def -= (key: A): this.type = {
-    internalMap.remove(key)
-    this
-  }
-
-  override def update(key: A, value: B) {
-    this += ((key, value))
+  override def get(key: A): Option[B] = {
+    val timeStampedValue = internalMap.get(key)
+    if (updateTimeStampOnGet && timeStampedValue != null) {
+      internalJavaMap.replace(key, timeStampedValue, TimeStampedValue(currentTime, timeStampedValue.value))
+    }
+    Option(timeStampedValue).map(_.value)
   }
-
-  override def apply(key: A): B = {
-    val value = internalMap.get(key)
-    if (value == null) throw new NoSuchElementException()
-    value._1
+  @inline override protected def externalValueToInternalValue(v: B): TimeStampedValue[B] = {
+    new TimeStampedValue(currentTime, v)
   }
 
-  override def filter(p: ((A, B)) => Boolean): Map[A, B] = {
-    JavaConversions.mapAsScalaConcurrentMap(internalMap).map(kv => (kv._1, kv._2._1)).filter(p)
+  @inline override protected def internalValueToExternalValue(iv: TimeStampedValue[B]): B = {
+    iv.value
   }
 
-  override def empty: Map[A, B] = new TimeStampedHashMap[A, B]()
-
-  override def size: Int = internalMap.size
-
-  override def foreach[U](f: ((A, B)) => U) {
-    val iterator = internalMap.entrySet().iterator()
-    while(iterator.hasNext) {
-      val entry = iterator.next()
-      val kv = (entry.getKey, entry.getValue._1)
-      f(kv)
-    }
+  /** Atomically put if a key is absent. This exposes the existing API of ConcurrentHashMap. */
+  def putIfAbsent(key: A, value: B): Option[B] = {
+    val prev = internalJavaMap.putIfAbsent(key, TimeStampedValue(currentTime, value))
+    Option(prev).map(_.value)
   }
 
-  def toMap: immutable.Map[A, B] = iterator.toMap
-
   /**
    * Removes old key-value pairs that have timestamp earlier than `threshTime`,
    * calling the supplied function on each such entry before removing.
    */
   def clearOldValues(threshTime: Long, f: (A, B) => Unit) {
-    val iterator = internalMap.entrySet().iterator()
+    val iterator = internalJavaMap.entrySet().iterator()
     while (iterator.hasNext) {
       val entry = iterator.next()
-      if (entry.getValue._2 < threshTime) {
-        f(entry.getKey, entry.getValue._1)
+      if (entry.getValue.timestamp < threshTime) {
+        f(entry.getKey, entry.getValue.value)
         logDebug("Removing key " + entry.getKey)
         iterator.remove()
       }
@@ -134,5 +92,4 @@ class TimeStampedHashMap[A, B](updateTimeStampOnGet: Boolean = false)
   }
 
   private def currentTime: Long = System.currentTimeMillis()
-
 }

core/src/main/scala/org/apache/spark/util/TimeStampedWeakValueHashMap.scala

@@ -1,3 +1,20 @@
+/*
+ * 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.spark.util
 
 import scala.collection.{JavaConversions, immutable}
@@ -12,8 +29,19 @@ private[util] case class TimeStampedWeakValue[T](timestamp: Long, weakValue: Wea
   def this(timestamp: Long, value: T) = this(timestamp, new WeakReference[T](value))
 }
 
+/**
+ * A map that stores the timestamp of when a key was inserted along with the value,
+ * while ensuring that the values are weakly referenced. If the value is garbage collected and
+ * the weak reference is null, get() operation returns the key be non-existent. However,
+ * the key is actually not remmoved in the current implementation. Key-value pairs whose
+ * timestamps are older than a particular threshold time can then be removed using the
+ * clearOldValues method. It exposes a scala.collection.mutable.Map interface to allow it to be a
+ * drop-in replacement for Scala HashMaps.
+ *
+ * Internally, it uses a Java ConcurrentHashMap, so all operations on this HashMap are thread-safe.
+ */
 
-private[spark] class TimeStampedWeakValueHashMap[A, B]
+private[spark] class TimeStampedWeakValueHashMap[A, B]()
   extends WrappedJavaHashMap[A, B, A, TimeStampedWeakValue[B]] with Logging {
 
   protected[util] val internalJavaMap: util.Map[A, TimeStampedWeakValue[B]] = {

core/src/main/scala/org/apache/spark/util/WrappedJavaHashMap.scala

@@ -1,3 +1,20 @@
+/*
+ * 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.spark.util
 
 import scala.collection.mutable.Map
@@ -8,8 +25,8 @@ import scala.reflect.ClassTag
 
 /**
  * Convenient wrapper class for exposing Java HashMaps as Scala Maps even if the
- * exposed key-value type is different from the internal type. This allows Scala HashMaps to be
- * hot replaceable with these Java HashMaps.
+ * exposed key-value type is different from the internal type. This allows these
+ * implementations of WrappedJavaHashMap to be drop-in replacements for Scala HashMaps.
  *
  * While Java <-> Scala conversion methods exists, its hard to understand the performance
  * implications and thread safety of the Scala wrapper. This class allows you to convert
@@ -62,7 +79,7 @@ private[spark] abstract class WrappedJavaHashMap[K, V, IK, IV] extends Map[K, V]
     (internalKeyToExternalKey(ip.getKey), internalValueToExternalValue(ip.getValue) )
   }
 
-  /* Implicit functions to convert the types. */
+  /* Implicit methods to convert the types. */
 
   @inline implicit private def convExtKeyToIntKey(k: K) = externalKeyToInternalKey(k)
 
@@ -76,6 +93,8 @@ private[spark] abstract class WrappedJavaHashMap[K, V, IK, IV] extends Map[K, V]
     internalPairToExternalPair(ip)
   }
 
+  /* Methods that must be implemented for a scala.collection.mutable.Map */
+
   def get(key: K): Option[V] = {
     Option(internalJavaMap.get(key))
   }
@@ -85,6 +104,8 @@ private[spark] abstract class WrappedJavaHashMap[K, V, IK, IV] extends Map[K, V]
     JavaConversions.asScalaIterator(jIterator).map(kv => convIntPairToExtPair(kv))
   }
 
+  /* Other methods that are implemented to ensure performance. */
+
   def +=(kv: (K, V)): this.type = {
     internalJavaMap.put(kv._1, kv._2)
     this
@@ -109,8 +130,9 @@ private[spark] abstract class WrappedJavaHashMap[K, V, IK, IV] extends Map[K, V]
   }
 
   override def foreach[U](f: ((K, V)) => U) {
-    while(iterator.hasNext) {
-      f(iterator.next())
+    val jIterator = internalJavaMap.entrySet().iterator()
+    while(jIterator.hasNext) {
+      f(jIterator.next())
     }
   }