improve OnHeapSpillManager concurrency:

1. use ConcurrentHashMap to avoid creating multiple memmgr in one task.
2. avoid deadlock with a spill spills itself.

Author: zhangli20

spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpill.scala

@@ -28,16 +28,24 @@ case class OnHeapSpill(hsm: OnHeapSpillManager, id: Int) extends Logging {
   def diskIOTime: Long = spillBuf.diskIOTime
 
   def write(buf: ByteBuffer): Unit = {
+    var needSpill = false
     synchronized {
       spillBuf match {
         case _: MemBasedSpillBuf =>
           val acquiredMemory = hsm.acquireMemory(buf.capacity())
           if (acquiredMemory < buf.capacity()) { // cannot allocate memory, will spill buffer
             hsm.freeMemory(acquiredMemory)
-            spillInternal()
+            needSpill = true
           }
-        case _ =>
+        case _ => false
       }
+    }
+
+    if (needSpill) {
+      spillInternal()
+    }
+
+    synchronized {
       spillBuf.write(buf)
     }
   }

spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpillManager.scala

@@ -17,6 +17,7 @@ package org.apache.spark.sql.blaze.memory
 
 import java.nio.ByteBuffer
 
+import scala.collection.concurrent
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
@@ -180,10 +181,10 @@ class OnHeapSpillManager(taskContext: TaskContext)
 }
 
 object OnHeapSpillManager extends Logging {
-  val all: mutable.Map[Long, OnHeapSpillManager] = mutable.Map()
+  val all: mutable.Map[Long, OnHeapSpillManager] = concurrent.TrieMap[Long, OnHeapSpillManager]()
 
   def current: OnHeapSpillManager = {
-    val taskContext = TaskContext.get
-    all.getOrElseUpdate(taskContext.taskAttemptId(), new OnHeapSpillManager(taskContext))
+    val tc = TaskContext.get
+    all.getOrElseUpdate(tc.taskAttemptId(), new OnHeapSpillManager(tc))
   }
 }