2011-02-27 22:15:52 -05:00
|
|
|
package spark
|
|
|
|
|
|
|
|
import scala.actors._
|
|
|
|
import scala.actors.Actor._
|
|
|
|
import scala.actors.remote._
|
2011-03-06 15:16:38 -05:00
|
|
|
import scala.collection.mutable.HashMap
|
|
|
|
import scala.collection.mutable.HashSet
|
2011-02-27 22:15:52 -05:00
|
|
|
|
2011-05-17 15:41:13 -04:00
|
|
|
sealed trait CacheTrackerMessage
|
|
|
|
case class AddedToCache(rddId: Int, partition: Int, host: String) extends CacheTrackerMessage
|
|
|
|
case class DroppedFromCache(rddId: Int, partition: Int, host: String) extends CacheTrackerMessage
|
|
|
|
case class MemoryCacheLost(host: String) extends CacheTrackerMessage
|
|
|
|
case class RegisterRDD(rddId: Int, numPartitions: Int) extends CacheTrackerMessage
|
|
|
|
case object GetCacheLocations extends CacheTrackerMessage
|
|
|
|
case object StopCacheTracker extends CacheTrackerMessage
|
2011-02-27 22:15:52 -05:00
|
|
|
|
2011-05-17 15:41:13 -04:00
|
|
|
class CacheTrackerActor extends DaemonActor with Logging {
|
2011-03-06 15:16:38 -05:00
|
|
|
val locs = new HashMap[Int, Array[List[String]]]
|
|
|
|
// TODO: Should probably store (String, CacheType) tuples
|
|
|
|
|
2011-02-27 22:15:52 -05:00
|
|
|
def act() {
|
2011-05-17 15:41:13 -04:00
|
|
|
val port = System.getProperty("spark.master.port").toInt
|
2011-02-27 22:15:52 -05:00
|
|
|
RemoteActor.alive(port)
|
2011-05-17 15:41:13 -04:00
|
|
|
RemoteActor.register('CacheTracker, self)
|
2011-03-06 15:16:38 -05:00
|
|
|
logInfo("Registered actor on port " + port)
|
2011-02-27 22:15:52 -05:00
|
|
|
|
|
|
|
loop {
|
|
|
|
react {
|
2011-03-06 15:16:38 -05:00
|
|
|
case RegisterRDD(rddId: Int, numPartitions: Int) =>
|
|
|
|
logInfo("Registering RDD " + rddId + " with " + numPartitions + " partitions")
|
|
|
|
locs(rddId) = Array.fill[List[String]](numPartitions)(Nil)
|
2011-05-13 14:43:52 -04:00
|
|
|
reply('OK)
|
2011-03-06 15:16:38 -05:00
|
|
|
|
|
|
|
case AddedToCache(rddId, partition, host) =>
|
|
|
|
logInfo("Cache entry added: (%s, %s) on %s".format(rddId, partition, host))
|
|
|
|
locs(rddId)(partition) = host :: locs(rddId)(partition)
|
2011-05-13 14:43:52 -04:00
|
|
|
reply('OK)
|
2011-02-27 22:15:52 -05:00
|
|
|
|
2011-03-06 15:16:38 -05:00
|
|
|
case DroppedFromCache(rddId, partition, host) =>
|
|
|
|
logInfo("Cache entry removed: (%s, %s) on %s".format(rddId, partition, host))
|
|
|
|
locs(rddId)(partition) -= host
|
|
|
|
|
|
|
|
case MemoryCacheLost(host) =>
|
|
|
|
logInfo("Memory cache lost on " + host)
|
|
|
|
// TODO: Drop host from the memory locations list of all RDDs
|
|
|
|
|
|
|
|
case GetCacheLocations =>
|
|
|
|
logInfo("Asked for current cache locations")
|
|
|
|
val locsCopy = new HashMap[Int, Array[List[String]]]
|
|
|
|
for ((rddId, array) <- locs) {
|
|
|
|
locsCopy(rddId) = array.clone()
|
|
|
|
}
|
|
|
|
reply(locsCopy)
|
2011-05-13 15:03:58 -04:00
|
|
|
|
|
|
|
case StopCacheTracker =>
|
|
|
|
reply('OK)
|
|
|
|
exit()
|
2011-02-27 22:15:52 -05:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2011-05-17 15:41:13 -04:00
|
|
|
class CacheTracker(isMaster: Boolean, theCache: Cache) extends Logging {
|
2011-02-27 22:15:52 -05:00
|
|
|
// Tracker actor on the master, or remote reference to it on workers
|
|
|
|
var trackerActor: AbstractActor = null
|
|
|
|
|
2011-05-17 15:41:13 -04:00
|
|
|
if (isMaster) {
|
|
|
|
val tracker = new CacheTrackerActor
|
2011-05-19 14:19:25 -04:00
|
|
|
tracker.start()
|
2011-05-17 15:41:13 -04:00
|
|
|
trackerActor = tracker
|
|
|
|
} else {
|
|
|
|
val host = System.getProperty("spark.master.host")
|
|
|
|
val port = System.getProperty("spark.master.port").toInt
|
|
|
|
trackerActor = RemoteActor.select(Node(host, port), 'CacheTracker)
|
2011-02-27 22:15:52 -05:00
|
|
|
}
|
2011-05-17 15:41:13 -04:00
|
|
|
|
|
|
|
val registeredRddIds = new HashSet[Int]
|
|
|
|
|
|
|
|
// Stores map results for various splits locally
|
|
|
|
val cache = theCache.newKeySpace()
|
|
|
|
|
|
|
|
// Remembers which splits are currently being loaded (on worker nodes)
|
|
|
|
val loading = new HashSet[(Int, Int)]
|
2011-02-27 22:15:52 -05:00
|
|
|
|
2011-03-06 15:16:38 -05:00
|
|
|
// Registers an RDD (on master only)
|
|
|
|
def registerRDD(rddId: Int, numPartitions: Int) {
|
|
|
|
registeredRddIds.synchronized {
|
|
|
|
if (!registeredRddIds.contains(rddId)) {
|
2011-03-06 19:16:38 -05:00
|
|
|
logInfo("Registering RDD ID " + rddId + " with cache")
|
2011-03-06 15:16:38 -05:00
|
|
|
registeredRddIds += rddId
|
|
|
|
trackerActor !? RegisterRDD(rddId, numPartitions)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// Get a snapshot of the currently known locations
|
|
|
|
def getLocationsSnapshot(): HashMap[Int, Array[List[String]]] = {
|
|
|
|
(trackerActor !? GetCacheLocations) match {
|
|
|
|
case h: HashMap[Int, Array[List[String]]] => h
|
|
|
|
case _ => throw new SparkException(
|
2011-05-17 15:41:13 -04:00
|
|
|
"Internal error: CacheTrackerActor did not reply with a HashMap")
|
2011-03-06 15:16:38 -05:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2011-02-27 22:15:52 -05:00
|
|
|
// Gets or computes an RDD split
|
|
|
|
def getOrCompute[T](rdd: RDD[T], split: Split)(implicit m: ClassManifest[T])
|
|
|
|
: Iterator[T] = {
|
|
|
|
val key = (rdd.id, split.index)
|
2011-03-06 15:16:38 -05:00
|
|
|
logInfo("CachedRDD partition key is " + key)
|
2011-02-27 22:15:52 -05:00
|
|
|
val cachedVal = cache.get(key)
|
|
|
|
if (cachedVal != null) {
|
|
|
|
// Split is in cache, so just return its values
|
2011-03-06 15:16:38 -05:00
|
|
|
logInfo("Found partition in cache!")
|
2011-02-27 22:15:52 -05:00
|
|
|
return Iterator.fromArray(cachedVal.asInstanceOf[Array[T]])
|
|
|
|
} else {
|
|
|
|
// Mark the split as loading (unless someone else marks it first)
|
|
|
|
loading.synchronized {
|
|
|
|
if (loading.contains(key)) {
|
|
|
|
while (loading.contains(key)) {
|
|
|
|
try {loading.wait()} catch {case _ =>}
|
|
|
|
}
|
|
|
|
return Iterator.fromArray(cache.get(key).asInstanceOf[Array[T]])
|
|
|
|
} else {
|
|
|
|
loading.add(key)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
// If we got here, we have to load the split
|
2011-03-06 15:16:38 -05:00
|
|
|
// Tell the master that we're doing so
|
|
|
|
val host = System.getProperty("spark.hostname", Utils.localHostName)
|
2011-05-13 14:43:52 -04:00
|
|
|
val future = trackerActor !! AddedToCache(rdd.id, split.index, host)
|
2011-02-27 22:15:52 -05:00
|
|
|
// TODO: fetch any remote copy of the split that may be available
|
|
|
|
// TODO: also register a listener for when it unloads
|
2011-03-06 15:16:38 -05:00
|
|
|
logInfo("Computing partition " + split)
|
2011-02-27 22:15:52 -05:00
|
|
|
val array = rdd.compute(split).toArray(m)
|
|
|
|
cache.put(key, array)
|
|
|
|
loading.synchronized {
|
|
|
|
loading.remove(key)
|
|
|
|
loading.notifyAll()
|
|
|
|
}
|
2011-05-13 14:43:52 -04:00
|
|
|
future.apply() // Wait for the reply from the cache tracker
|
2011-02-27 22:15:52 -05:00
|
|
|
return Iterator.fromArray(array)
|
|
|
|
}
|
|
|
|
}
|
2011-05-13 15:03:58 -04:00
|
|
|
|
|
|
|
def stop() {
|
|
|
|
trackerActor !? StopCacheTracker
|
|
|
|
registeredRddIds.clear()
|
|
|
|
trackerActor = null
|
|
|
|
}
|
2011-05-13 14:43:52 -04:00
|
|
|
}
|