Set the task context in writer benchmarks (apache-spark-on-k8s#529)

Author: mccheah

core/src/test/scala/org/apache/spark/shuffle/sort/ShuffleWriterBenchmarkBase.scala

@@ -111,6 +111,7 @@ abstract class ShuffleWriterBenchmarkBase extends BenchmarkBase {
   when(rpcEnv.setupEndpoint(any[String], any[RpcEndpoint])).thenReturn(rpcEndpointRef)
 
   def setup(): Unit = {
+    TaskContext.setTaskContext(taskContext)
     memoryManager = new TestMemoryManager(defaultConf)
     memoryManager.limit(PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES)
     taskMemoryManager = new TaskMemoryManager(memoryManager, 0)

core/src/test/scala/org/apache/spark/shuffle/sort/SortShuffleWriterBenchmark.scala

@@ -19,10 +19,9 @@ package org.apache.spark.shuffle.sort
 
 import org.mockito.Mockito.when
 
-import org.apache.spark.{Aggregator, SparkEnv}
+import org.apache.spark.{Aggregator, SparkEnv, TaskContext}
 import org.apache.spark.benchmark.Benchmark
 import org.apache.spark.shuffle.BaseShuffleHandle
-import org.apache.spark.shuffle.sort.io.DefaultShuffleDataIO
 
 /**
  * Benchmark to measure performance for aggregate primitives.
@@ -76,6 +75,7 @@ object SortShuffleWriterBenchmark extends ShuffleWriterBenchmarkBase {
     }
 
     when(taskContext.taskMemoryManager()).thenReturn(taskMemoryManager)
+    TaskContext.setTaskContext(taskContext)
 
     val shuffleWriter = new SortShuffleWriter[String, String, String](
       blockResolver,

core/src/test/scala/org/apache/spark/shuffle/sort/UnsafeShuffleWriterBenchmark.scala

@@ -16,9 +16,8 @@
  */
 package org.apache.spark.shuffle.sort
 
-import org.apache.spark.SparkConf
+import org.apache.spark.{SparkConf, TaskContext}
 import org.apache.spark.benchmark.Benchmark
-import org.apache.spark.util.Utils
 
 /**
  * Benchmark to measure performance for aggregate primitives.
@@ -44,6 +43,7 @@ object UnsafeShuffleWriterBenchmark extends ShuffleWriterBenchmarkBase {
     val conf = new SparkConf(loadDefaults = false)
     conf.set("spark.file.transferTo", String.valueOf(transferTo))
 
+    TaskContext.setTaskContext(taskContext)
     new UnsafeShuffleWriter[String, String](
       blockManager,
       blockResolver,