[转载] Spark源码分析之-Storage模块

def registerOrLookup(name: String, newActor: => Actor): ActorRef = {
  if (isDriver) {
    logInfo("Registering " + name)
    actorSystem.actorOf(Props(newActor), name = name)
  } else {
    val driverHost: String = System.getProperty("spark.driver.host", "localhost")
    val driverPort: Int = System.getProperty("spark.driver.port", "7077").toInt
    Utils.checkHost(driverHost, "Expected hostname")
    val url = "akka://spark@%s:%s/user/%s".format(driverHost, driverPort, name)
    logInfo("Connecting to " + name + ": " + url)
    actorSystem.actorFor(url)
  }
}

val blockManagerMaster = new BlockManagerMaster(registerOrLookup(
  "BlockManagerMaster",
  new BlockManagerMasterActor(isLocal)))
val blockManager = new BlockManager(executorId, actorSystem, blockManagerMaster, serializer)

private def initialize() {
  master.registerBlockManager(blockManagerId, maxMemory, slaveActor)
  ...
  if (!BlockManager.getDisableHeartBeatsForTesting) {
    heartBeatTask = actorSystem.scheduler.schedule(0.seconds, heartBeatFrequency.milliseconds) {
      heartBeat()
    }
  }
}

private def register(id: BlockManagerId, maxMemSize: Long, slaveActor: ActorRef) {
  if (id.executorId == "<driver>" && !isLocal) {
    // Got a register message from the master node; don't register it
  } else if (!blockManagerInfo.contains(id)) {
    blockManagerIdByExecutor.get(id.executorId) match {
      case Some(manager) =>
        // A block manager of the same executor already exists.
        // This should never happen. Let's just quit.
        logError("Got two different block manager registrations on " + id.executorId)
        System.exit(1)
      case None =>
        blockManagerIdByExecutor(id.executorId) = id
    }
    blockManagerInfo(id) = new BlockManagerMasterActor.BlockManagerInfo(
      id, System.currentTimeMillis(), maxMemSize, slaveActor)
  }
}

private[storage] val memoryStore: BlockStore = new MemoryStore(this, maxMemory)
private[storage] val diskStore: DiskStore =
  new DiskStore(this, System.getProperty("spark.local.dir", System.getProperty("java.io.tmpdir")))

private def initialize() {
  ...
  BlockManagerWorker.startBlockManagerWorker(this)
  ...
}

private def createLocalDirs(): Array[File] = {
  logDebug("Creating local directories at root dirs '" + rootDirs + "'")
  val dateFormat = new SimpleDateFormat("yyyyMMddHHmmss")
  rootDirs.split(",").map { rootDir =>
    var foundLocalDir = false
    var localDir: File = null
    var localDirId: String = null
    var tries = 0
    val rand = new Random()
    while (!foundLocalDir && tries < MAX_DIR_CREATION_ATTEMPTS) {
      tries += 1
      try {
        localDirId = "%s-%04x".format(dateFormat.format(new Date), rand.nextInt(65536))
        localDir = new File(rootDir, "spark-local-" + localDirId)
        if (!localDir.exists) {
          foundLocalDir = localDir.mkdirs()
        }
      } catch {
        case e: Exception =>
          logWarning("Attempt " + tries + " to create local dir " + localDir + " failed", e)
      }
    }
    if (!foundLocalDir) {
      logError("Failed " + MAX_DIR_CREATION_ATTEMPTS +
        " attempts to create local dir in " + rootDir)
      System.exit(ExecutorExitCode.DISK_STORE_FAILED_TO_CREATE_DIR)
    }
    logInfo("Created local directory at " + localDir)
    localDir
  }
}

private def getFile(blockId: String): File = {
  logDebug("Getting file for block " + blockId)

  // Figure out which local directory it hashes to, and which subdirectory in that
  val hash = Utils.nonNegativeHash(blockId)
  val dirId = hash % localDirs.length
  val subDirId = (hash / localDirs.length) % subDirsPerLocalDir

  // Create the subdirectory if it doesn't already exist
  var subDir = subDirs(dirId)(subDirId)
  if (subDir == null) {
    subDir = subDirs(dirId).synchronized {
      val old = subDirs(dirId)(subDirId)
      if (old != null) {
        old
      } else {
        val newDir = new File(localDirs(dirId), "%02x".format(subDirId))
        newDir.mkdir()
        subDirs(dirId)(subDirId) = newDir
        newDir
      }
    }
  }

  new File(subDir, blockId)
}

override def putBytes(blockId: String, _bytes: ByteBuffer, level: StorageLevel) {
  // So that we do not modify the input offsets !
  // duplicate does not copy buffer, so inexpensive
  val bytes = _bytes.duplicate()
  logDebug("Attempting to put block " + blockId)
  val startTime = System.currentTimeMillis
  val file = createFile(blockId)
  val channel = new RandomAccessFile(file, "rw").getChannel()
  while (bytes.remaining > 0) {
    channel.write(bytes)
  }
  channel.close()
  val finishTime = System.currentTimeMillis
  logDebug("Block %s stored as %s file on disk in %d ms".format(
    blockId, Utils.bytesToString(bytes.limit), (finishTime - startTime)))
}

override def getBytes(blockId: String): Option[ByteBuffer] = {
  val file = getFile(blockId)
  val bytes = getFileBytes(file)
  Some(bytes)
}

case class Entry(value: Any, size: Long, deserialized: Boolean)

private val entries = new LinkedHashMap[String, Entry](32, 0.75f, true)

override def putBytes(blockId: String, _bytes: ByteBuffer, level: StorageLevel) {
  // Work on a duplicate - since the original input might be used elsewhere.
  val bytes = _bytes.duplicate()
  bytes.rewind()
  if (level.deserialized) {
    val values = blockManager.dataDeserialize(blockId, bytes)
    val elements = new ArrayBuffer[Any]
    elements ++= values
    val sizeEstimate = SizeEstimator.estimate(elements.asInstanceOf[AnyRef])
    tryToPut(blockId, elements, sizeEstimate, true)
  } else {
    tryToPut(blockId, bytes, bytes.limit, false)
  }
}

private def tryToPut(blockId: String, value: Any, size: Long, deserialized: Boolean): Boolean = {
  // TODO: Its possible to optimize the locking by locking entries only when selecting blocks
  // to be dropped. Once the to-be-dropped blocks have been selected, and lock on entries has been
  // released, it must be ensured that those to-be-dropped blocks are not double counted for
  // freeing up more space for another block that needs to be put. Only then the actually dropping
  // of blocks (and writing to disk if necessary) can proceed in parallel.
  putLock.synchronized {
    if (ensureFreeSpace(blockId, size)) {
      val entry = new Entry(value, size, deserialized)
      entries.synchronized {
        entries.put(blockId, entry)
        currentMemory += size
      }
      if (deserialized) {
        logInfo("Block %s stored as values to memory (estimated size %s, free %s)".format(
          blockId, Utils.bytesToString(size), Utils.bytesToString(freeMemory)))
      } else {
        logInfo("Block %s stored as bytes to memory (size %s, free %s)".format(
          blockId, Utils.bytesToString(size), Utils.bytesToString(freeMemory)))
      }
      true
    } else {
      // Tell the block manager that we couldn't put it in memory so that it can drop it to
      // disk if the block allows disk storage.
      val data = if (deserialized) {
        Left(value.asInstanceOf[ArrayBuffer[Any]])
      } else {
        Right(value.asInstanceOf[ByteBuffer].duplicate())
      }
      blockManager.dropFromMemory(blockId, data)
      false
    }
  }
}

override def getValues(blockId: String): Option[Iterator[Any]] = {
  val entry = entries.synchronized {
    entries.get(blockId)
  }
  if (entry == null) {
    None
  } else if (entry.deserialized) {
    Some(entry.value.asInstanceOf[ArrayBuffer[Any]].iterator)
  } else {
    val buffer = entry.value.asInstanceOf[ByteBuffer].duplicate() // Doesn't actually copy data
    Some(blockManager.dataDeserialize(blockId, buffer))
  }
}

def put(blockId: String, values: ArrayBuffer[Any], level: StorageLevel,
  tellMaster: Boolean = true) : Long = {

  ...

  // Remember the block's storage level so that we can correctly drop it to disk if it needs
  // to be dropped right after it got put into memory. Note, however, that other threads will
  // not be able to get() this block until we call markReady on its BlockInfo.
  val myInfo = {
    val tinfo = new BlockInfo(level, tellMaster)
    // Do atomically !
    val oldBlockOpt = blockInfo.putIfAbsent(blockId, tinfo)

    if (oldBlockOpt.isDefined) {
      if (oldBlockOpt.get.waitForReady()) {
        logWarning("Block " + blockId + " already exists on this machine; not re-adding it")
        return oldBlockOpt.get.size
      }

      // TODO: So the block info exists - but previous attempt to load it (?) failed. What do we do now ? Retry on it ?
      oldBlockOpt.get
    } else {
      tinfo
    }
  }

  val startTimeMs = System.currentTimeMillis

  // If we need to replicate the data, we'll want access to the values, but because our
  // put will read the whole iterator, there will be no values left. For the case where
  // the put serializes data, we'll remember the bytes, above; but for the case where it
  // doesn't, such as deserialized storage, let's rely on the put returning an Iterator.
  var valuesAfterPut: Iterator[Any] = null

  // Ditto for the bytes after the put
  var bytesAfterPut: ByteBuffer = null

  // Size of the block in bytes (to return to caller)
  var size = 0L

  myInfo.synchronized {
    logTrace("Put for block " + blockId + " took " + Utils.getUsedTimeMs(startTimeMs)
      + " to get into synchronized block")

    var marked = false
    try {
      if (level.useMemory) {
        // Save it just to memory first, even if it also has useDisk set to true; we will later
        // drop it to disk if the memory store can't hold it.
        val res = memoryStore.putValues(blockId, values, level, true)
        size = res.size
        res.data match {
          case Right(newBytes) => bytesAfterPut = newBytes
          case Left(newIterator) => valuesAfterPut = newIterator
        }
      } else {
        // Save directly to disk.
        // Don't get back the bytes unless we replicate them.
        val askForBytes = level.replication > 1
        val res = diskStore.putValues(blockId, values, level, askForBytes)
        size = res.size
        res.data match {
          case Right(newBytes) => bytesAfterPut = newBytes
          case _ =>
        }
      }

      // Now that the block is in either the memory or disk store, let other threads read it,
      // and tell the master about it.
      marked = true
      myInfo.markReady(size)
      if (tellMaster) {
        reportBlockStatus(blockId, myInfo)
      }
    } finally {
      // If we failed at putting the block to memory/disk, notify other possible readers
      // that it has failed, and then remove it from the block info map.
      if (! marked) {
        // Note that the remove must happen before markFailure otherwise another thread
        // could've inserted a new BlockInfo before we remove it.
        blockInfo.remove(blockId)
        myInfo.markFailure()
        logWarning("Putting block " + blockId + " failed")
      }
    }
  }
  logDebug("Put block " + blockId + " locally took " + Utils.getUsedTimeMs(startTimeMs))

  // Replicate block if required
  if (level.replication > 1) {
    val remoteStartTime = System.currentTimeMillis
    // Serialize the block if not already done
    if (bytesAfterPut == null) {
      if (valuesAfterPut == null) {
        throw new SparkException(
          "Underlying put returned neither an Iterator nor bytes! This shouldn't happen.")
      }
      bytesAfterPut = dataSerialize(blockId, valuesAfterPut)
    }
    replicate(blockId, bytesAfterPut, level)
    logDebug("Put block " + blockId + " remotely took " + Utils.getUsedTimeMs(remoteStartTime))
  }
  BlockManager.dispose(bytesAfterPut)

  return size
}

def get(blockId: String): Option[Iterator[Any]] = {
  val local = getLocal(blockId)
  if (local.isDefined) {
    logInfo("Found block %s locally".format(blockId))
    return local
  }
  val remote = getRemote(blockId)
  if (remote.isDefined) {
    logInfo("Found block %s remotely".format(blockId))
    return remote
  }
  None
}

def getLocal(blockId: String): Option[Iterator[Any]] = {
  logDebug("Getting local block " + blockId)
  val info = blockInfo.get(blockId).orNull
  if (info != null) {
    info.synchronized {

      // In the another thread is writing the block, wait for it to become ready.
      if (!info.waitForReady()) {
        // If we get here, the block write failed.
        logWarning("Block " + blockId + " was marked as failure.")
        return None
      }

      val level = info.level
      logDebug("Level for block " + blockId + " is " + level)

      // Look for the block in memory
      if (level.useMemory) {
        logDebug("Getting block " + blockId + " from memory")
        memoryStore.getValues(blockId) match {
          case Some(iterator) =>
            return Some(iterator)
          case None =>
            logDebug("Block " + blockId + " not found in memory")
        }
      }

      // Look for block on disk, potentially loading it back into memory if required
      if (level.useDisk) {
        logDebug("Getting block " + blockId + " from disk")
        if (level.useMemory && level.deserialized) {
          diskStore.getValues(blockId) match {
            case Some(iterator) =>
              // Put the block back in memory before returning it
              // TODO: Consider creating a putValues that also takes in a iterator ?
              val elements = new ArrayBuffer[Any]
              elements ++= iterator
              memoryStore.putValues(blockId, elements, level, true).data match {
                case Left(iterator2) =>
                  return Some(iterator2)
                case _ =>
                  throw new Exception("Memory store did not return back an iterator")
              }
            case None =>
              throw new Exception("Block " + blockId + " not found on disk, though it should be")
          }
        } else if (level.useMemory && !level.deserialized) {
          // Read it as a byte buffer into memory first, then return it
          diskStore.getBytes(blockId) match {
            case Some(bytes) =>
              // Put a copy of the block back in memory before returning it. Note that we can't
              // put the ByteBuffer returned by the disk store as that's a memory-mapped file.
              // The use of rewind assumes this.
              assert (0 == bytes.position())
              val copyForMemory = ByteBuffer.allocate(bytes.limit)
              copyForMemory.put(bytes)
              memoryStore.putBytes(blockId, copyForMemory, level)
              bytes.rewind()
              return Some(dataDeserialize(blockId, bytes))
            case None =>
              throw new Exception("Block " + blockId + " not found on disk, though it should be")
          }
        } else {
          diskStore.getValues(blockId) match {
            case Some(iterator) =>
              return Some(iterator)
            case None =>
              throw new Exception("Block " + blockId + " not found on disk, though it should be")
          }
        }
      }
    }
  } else {
    logDebug("Block " + blockId + " not registered locally")
  }
  return None
}

def getRemote(blockId: String): Option[Iterator[Any]] = {
  if (blockId == null) {
    throw new IllegalArgumentException("Block Id is null")
  }
  logDebug("Getting remote block " + blockId)
  // Get locations of block
  val locations = master.getLocations(blockId)

  // Get block from remote locations
  for (loc <- locations) {
    logDebug("Getting remote block " + blockId + " from " + loc)
    val data = BlockManagerWorker.syncGetBlock(
      GetBlock(blockId), ConnectionManagerId(loc.host, loc.port))
    if (data != null) {
      return Some(dataDeserialize(blockId, data))
    }
    logDebug("The value of block " + blockId + " is null")
  }
  logDebug("Block " + blockId + " not found")
  return None
}

final def iterator(split: Partition, context: TaskContext): Iterator[T] = {
  if (storageLevel != StorageLevel.NONE) {
    SparkEnv.get.cacheManager.getOrCompute(this, split, context, storageLevel)
  } else {
    computeOrReadCheckpoint(split, context)
  }
}

def getOrCompute[T](rdd: RDD[T], split: Partition, context: TaskContext, storageLevel: StorageLevel)
    : Iterator[T] = {
  val key = "rdd_%d_%d".format(rdd.id, split.index)
  logDebug("Looking for partition " + key)
  blockManager.get(key) match {
    case Some(values) =>
      // Partition is already materialized, so just return its values
      return values.asInstanceOf[Iterator[T]]

    case None =>
      // Mark the split as loading (unless someone else marks it first)
      loading.synchronized {
        if (loading.contains(key)) {
          logInfo("Another thread is loading %s, waiting for it to finish...".format (key))
          while (loading.contains(key)) {
            try {loading.wait()} catch {case _ : Throwable =>}
          }
          logInfo("Finished waiting for %s".format(key))
          // See whether someone else has successfully loaded it. The main way this would fail
          // is for the RDD-level cache eviction policy if someone else has loaded the same RDD
          // partition but we didn't want to make space for it. However, that case is unlikely
          // because it's unlikely that two threads would work on the same RDD partition. One
          // downside of the current code is that threads wait serially if this does happen.
          blockManager.get(key) match {
            case Some(values) =>
              return values.asInstanceOf[Iterator[T]]
            case None =>
              logInfo("Whoever was loading %s failed; we'll try it ourselves".format  (key))
              loading.add(key)
          }
        } else {
          loading.add(key)
        }
      }
      try {
        // If we got here, we have to load the split
        logInfo("Partition %s not found, computing it".format(key))
        val computedValues = rdd.computeOrReadCheckpoint(split, context)
        // Persist the result, so long as the task is not running locally
        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 {
          loading.remove(key)
          loading.notifyAll()
        }
      }
  }
}