Mark key as not loading in CacheTracker even when compute() fails

core/src/main/scala/spark/CacheTracker.scala

@@ -202,26 +202,26 @@ private[spark] class CacheTracker(actorSystem: ActorSystem, isMaster: Boolean, b
             loading.add(key)
           }
         }
-        // If we got here, we have to load the split
-        // Tell the master that we're doing so
-        //val host = System.getProperty("spark.hostname", Utils.localHostName)
-        //val future = trackerActor !! AddedToCache(rdd.id, split.index, host)
-        // TODO: fetch any remote copy of the split that may be available
-        // TODO: also register a listener for when it unloads
-        logInfo("Computing partition " + split)
-        val elements = new ArrayBuffer[Any]
-        elements ++= rdd.compute(split, context)
         try {
+          // If we got here, we have to load the split
+          // Tell the master that we're doing so
+          //val host = System.getProperty("spark.hostname", Utils.localHostName)
+          //val future = trackerActor !! AddedToCache(rdd.id, split.index, host)
+          // TODO: fetch any remote copy of the split that may be available
+          // TODO: also register a listener for when it unloads
+          val elements = new ArrayBuffer[Any]
+          logInfo("Computing partition " + split)
+          elements ++= rdd.compute(split, context)
           // Try to put this block in the blockManager
           blockManager.put(key, elements, storageLevel, true)
           //future.apply() // Wait for the reply from the cache tracker
+          return elements.iterator.asInstanceOf[Iterator[T]]
         } finally {
           loading.synchronized {
             loading.remove(key)
             loading.notifyAll()
           }
         }
-        return elements.iterator.asInstanceOf[Iterator[T]]
     }
   }