use WeakHashMap

apache · Oct 20, 2016 · a2b0358 · a2b0358
1 parent 635424f
commit a2b0358
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/core/src/main/scala/org/apache/spark/api/r/RBackendHandler.scala b/core/src/main/scala/org/apache/spark/api/r/RBackendHandler.scala
@@ -20,9 +20,8 @@ package org.apache.spark.api.r
 import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream}
 import java.util.concurrent.atomic.AtomicLong
 
-import scala.collection.mutable.HashMap
+import scala.collection.mutable.WeakHashMap
 import scala.language.existentials
-import scala.ref.WeakReference
 
 import io.netty.channel.{ChannelHandlerContext, SimpleChannelInboundHandler}
 import io.netty.channel.ChannelHandler.Sharable
@@ -263,7 +262,7 @@ private[r] class RBackendHandler(server: RBackend)
  */
 private[r] object JVMObjectTracker {
 
-  private[this] val objMap = new HashMap[String, WeakReference[Object]]
+  private[this] val objMap = new WeakHashMap[String, Object]
 
   private[this] val objCounter: AtomicLong = new AtomicLong(0L)
 
@@ -278,7 +277,7 @@ private[r] object JVMObjectTracker {
   def put(obj: Object): String = {
     val objId = objCounter.toString
     objCounter.incrementAndGet()
-    objMap.put(objId, new WeakReference(obj))
+    objMap.put(objId, obj)
     objId
   }