Add unit test and address comments

core/src/main/scala/org/apache/spark/CacheManager.scala

@@ -66,11 +66,13 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging {
         }
         try {
           // If we got here, we have to load the split
-          val elements = new ArrayBuffer[Any]
           logInfo("Computing partition " + split)
-          elements ++= rdd.computeOrReadCheckpoint(split, context)
+          val computedValues = rdd.computeOrReadCheckpoint(split, context)
           // Persist the result, so long as the task is not running locally
-          if (!context.runningLocally) blockManager.put(key, elements, storageLevel, true)
+          if (context.runningLocally) return computedValues
+          val elements = new ArrayBuffer[Any]
+          elements ++= computedValues
+          blockManager.put(key, elements, storageLevel, true)
           return elements.iterator.asInstanceOf[Iterator[T]]
         } finally {
           loading.synchronized {

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

@@ -478,7 +478,8 @@ class DAGScheduler(
       SparkEnv.set(env)
       val rdd = job.finalStage.rdd
       val split = rdd.partitions(job.partitions(0))
-      val taskContext = new TaskContext(job.finalStage.id, job.partitions(0), 0, true)
+      val taskContext =
+        new TaskContext(job.finalStage.id, job.partitions(0), 0, runningLocally = true)
       try {
         val result = job.func(taskContext, rdd.iterator(split, taskContext))
         job.listener.taskSucceeded(0, result)

core/src/main/scala/org/apache/spark/scheduler/ResultTask.scala

@@ -93,7 +93,7 @@ private[spark] class ResultTask[T, U](
   }
 
   override def run(attemptId: Long): U = {
-    val context = new TaskContext(stageId, partition, attemptId)
+    val context = new TaskContext(stageId, partition, attemptId, runningLocally = false)
     metrics = Some(context.taskMetrics)
     try {
       func(context, rdd.iterator(split, context))

core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala

@@ -132,7 +132,7 @@ private[spark] class ShuffleMapTask(
   override def run(attemptId: Long): MapStatus = {
     val numOutputSplits = dep.partitioner.numPartitions
 
-    val taskContext = new TaskContext(stageId, partition, attemptId)
+    val taskContext = new TaskContext(stageId, partition, attemptId, runningLocally = false)
     metrics = Some(taskContext.taskMetrics)
 
     val blockManager = SparkEnv.get.blockManager

core/src/test/scala/org/apache/spark/CacheManagerSuite.scala

+/*
+ * 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
+
+import scala.Array
+import scala.collection.mutable.ArrayBuffer
+
+import org.scalatest.{BeforeAndAfter, FunSuite}
+import org.scalatest.mock.EasyMockSugar
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.storage.{BlockManager, StorageLevel}
+import org.junit.Assert
+
+// TODO: Test the CacheManager's thread-safety aspects
+class CacheManagerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar {
+  val sc = new SparkContext("local", "test")
+
+  val split = new Partition { override def index: Int = 0 }
+
+  // An RDD which returns the values [1, 2, 3, 4].
+  val rdd = new RDD[Int](sc, Nil) {
+    override def getPartitions: Array[Partition] = Array(split)
+    override val getDependencies = List[Dependency[_]]()
+    override def compute(split: Partition, context: TaskContext) = Array(1, 2, 3, 4).iterator
+  }
+
+  var blockManager: BlockManager = _
+  var cacheManager: CacheManager = _
+
+  before {
+    blockManager = mock[BlockManager]
+    cacheManager = new CacheManager(blockManager)
+  }
+
+  test("get uncached rdd") {
+    expecting {
+      blockManager.get("rdd_0_0").andReturn(None)
+      blockManager.put("rdd_0_0", ArrayBuffer[Any](1, 2, 3, 4), StorageLevel.MEMORY_ONLY, true).
+        andReturn(0)
+    }
+
+    whenExecuting(blockManager) {
+      val context = new TaskContext(0, 0, 0, runningLocally = false, null)
+      val value = cacheManager.getOrCompute(rdd, split, context, StorageLevel.MEMORY_ONLY)
+      Assert.assertEquals(value.toList, List(1, 2, 3, 4))
+    }
+  }
+
+  test("get cached rdd") {
+    expecting {
+      blockManager.get("rdd_0_0").andReturn(Some(ArrayBuffer(5, 6, 7).iterator))
+    }
+
+    whenExecuting(blockManager) {
+      val context = new TaskContext(0, 0, 0, runningLocally = false, null)
+      val value = cacheManager.getOrCompute(rdd, split, context, StorageLevel.MEMORY_ONLY)
+      Assert.assertEquals(value.toList, List(5, 6, 7))
+    }
+  }
+
+  test("get uncached local rdd") {
+    expecting {
+      // Local computation should not persist the resulting value, so don't expect a put().
+      blockManager.get("rdd_0_0").andReturn(None)
+    }
+
+    whenExecuting(blockManager) {
+      val context = new TaskContext(0, 0, 0, runningLocally = true, null)
+      val value = cacheManager.getOrCompute(rdd, split, context, StorageLevel.MEMORY_ONLY)
+      Assert.assertEquals(value.toList, List(1, 2, 3, 4))
+    }
+  }
+}