Spark Core源码精读计划#22：BlockInfoMan
目录
前言
在上一篇文章中，我们对与块相关的BlockId、BlockData和BlockInfo有了比较全面的理解。前面已经提到过，块在读写时有锁机制，并且委托给BlockInfoManager来管理。虽然BlockInfoManager的字面意思是“块信息管理器”，但管理块信息的意图并不明显，管理块的锁才是真正主要的任务。本文就来研究BlockInfoManager的具体实现。
BlockInfoManager的成员属性及构造方法
代码#22.1 - o.a.s.storage.BlockInfoManager的成员属性及构造方法
private[storage] class BlockInfoManager extends Logging {
private type TaskAttemptId = Long
@GuardedBy("this")
private[this] val infos = new mutable.HashMap[BlockId, BlockInfo]
@GuardedBy("this")
private[this] val writeLocksByTask =
new mutable.HashMap[TaskAttemptId, mutable.Set[BlockId]]
with mutable.MultiMap[TaskAttemptId, BlockId]
@GuardedBy("this")
private[this] val readLocksByTask =
new mutable.HashMap[TaskAttemptId, ConcurrentHashMultiset[BlockId]]
registerTask(BlockInfo.NON_TASK_WRITER)
def registerTask(taskAttemptId: TaskAttemptId): Unit = synchronized {
require(!readLocksByTask.contains(taskAttemptId),
s"Task attempt $taskAttemptId is already registered")
readLocksByTask(taskAttemptId) = ConcurrentHashMultiset.create()
}
- TaskAttemptId：实际上就是对Long型的重命名，用来表示一次Task尝试的ID。
- infos：存储BlockId与BlockInfo的映射关系，这就是为什么BlockInfo结构中并没有包含BlockId对应的字段。
- writeLocksByTask：存储TaskAttemptId与该Task获取写锁的块之间的映射关系。注意BlockId存储在集合中，也就是说一次Task尝试可以获取多个块的写锁。
- readLocksByTask：存储TaskAttemptId与该Task获取读锁的块之间的映射关系。一次Task尝试也可以获取多个块的读锁。
在BlockInfoManager构造时，会调用registerTask()方法注册任务，其实就是将NON_TASK_WRITER这个TaskAttemptId对应的BlockId集合初始化好。NON_TASK_WRITER在BlockInfo伴生对象里定义，是一个特殊的标记（-1024），表示当前持有写锁的并非一个具体的Task，而是其他线程。registerTask()也会被BlockManager调用，这是后话。
下面我们来看看BlockInfoManager提供的与锁相关的操作。
BlockInfoManager提供的锁方法
注意这些方法都是同步方法（被synchronized关键字修饰的）。
获取读锁
lockForReading()方法为一个块加读锁，其代码如下。
代码#21.2 - o.a.s.storage.BlockInfoManager.lockForReading()方法
def lockForReading(
blockId: BlockId,
blocking: Boolean = true): Option[BlockInfo] = synchronized {
logTrace(s"Task $currentTaskAttemptId trying to acquire read lock for $blockId")
do {
infos.get(blockId) match {
case None => return None
case Some(info) =>
if (info.writerTask == BlockInfo.NO_WRITER) {
info.readerCount += 1
readLocksByTask(currentTaskAttemptId).add(blockId)
logTrace(s"Task $currentTaskAttemptId acquired read lock for $blockId")
return Some(info)
}
}
if (blocking) {
wait()
}
} while (blocking)
None
}
注意blocking参数，它表示加读锁的过程是否阻塞（默认阻塞）。如果不阻塞的话，获取读锁失败就会立即返回。
该方法的执行流程是：根据块ID获取它对应的BlockInfo，检查它的writerTask是否为NO_WRITER（值为-1，表示该BlockInfo的写锁没有被占用）。如果是，就自增BlockInfo结构中的readerCount计数，并将块ID加入readLocksByTask映射，视为加锁成功。若blocking为true的话，就会调用Object.wait()方法等待，直到该块的写锁释放后被notify()/notifyAll()方法唤醒。可见，如果该块的写锁一直不释放，那么lockForReading()方法可能会无限等待下去。
获取写锁
与lockForReading()方法相对地，lockForWriting()方法为一个块加写锁，其代码如下。
代码#21.3 - o.a.s.storage.BlockInfoManager.lockForWriting()方法
def lockForWriting(
blockId: BlockId,
blocking: Boolean = true): Option[BlockInfo] = synchronized {
logTrace(s"Task $currentTaskAttemptId trying to acquire write lock for $blockId")
do {
infos.get(blockId) match {
case None => return None
case Some(info) =>
if (info.writerTask == BlockInfo.NO_WRITER && info.readerCount == 0) {
info.writerTask = currentTaskAttemptId
writeLocksByTask.addBinding(currentTaskAttemptId, blockId)
logTrace(s"Task $currentTaskAttemptId acquired write lock for $blockId")
return Some(info)
}
}
if (blocking) {
wait()
}
} while (blocking)
None
}
这个方法的执行流程与lockForReading()方法相似，不过会将BlockInfo中的writerTask字段设为Task尝试ID，将块ID加入writeLocksByTask映射，并且判断条件是没有读锁也没有写锁。也就是说，块的读锁和写锁、写锁和写锁之间是互斥的，而读锁和读锁之间是可以共享的，并且读锁可重入，写锁不可重入。
同样地，如果该块的其他写锁一直不释放，那么lockForWriting()方法也有可能会无限等待下去。
另外，还有一个lockNewBlockForWriting()方法用来获取一个新块的写锁。
代码#21.4 - o.a.s.storage.BlockInfoManager.lockNewBlockForWriting()方法
def lockNewBlockForWriting(
blockId: BlockId,
newBlockInfo: BlockInfo): Boolean = synchronized {
logTrace(s"Task $currentTaskAttemptId trying to put $blockId")
lockForReading(blockId) match {
case Some(info) =>
false
case None =>
infos(blockId) = newBlockInfo
lockForWriting(blockId)
true
}
}
该方法先试图持有blockId对应的块的读锁。如果能获取到，说明该块已经存在了，亦即已经有其他线程赢得竞争并写了这个块，没有必要再写，直接返回false（表示返回读锁）。反之，就将这个新的块放入infos映射，然后获取其对应的写锁，并返回true。
释放锁
释放单个块的锁的逻辑由unlock()方法实现。
代码#21.5 - o.a.s.storage.BlockInfoManager.unlock()方法
def unlock(blockId: BlockId, taskAttemptId: Option[TaskAttemptId] = None): Unit = synchronized {
val taskId = taskAttemptId.getOrElse(currentTaskAttemptId)
logTrace(s"Task $taskId releasing lock for $blockId")
val info = get(blockId).getOrElse {
throw new IllegalStateException(s"Block $blockId not found")
}
if (info.writerTask != BlockInfo.NO_WRITER) {
info.writerTask = BlockInfo.NO_WRITER
writeLocksByTask.removeBinding(taskId, blockId)
} else {
assert(info.readerCount > 0, s"Block $blockId is not locked for reading")
info.readerCount -= 1
val countsForTask = readLocksByTask(taskId)
val newPinCountForTask: Int = countsForTask.remove(blockId, 1) - 1
assert(newPinCountForTask >= 0,
s"Task $taskId release lock on block $blockId more times than it acquired it")
}
notifyAll()
}
该方法首先获取Task尝试ID与对应的块信息（get()方法就负责从infos映射中取得块信息），然后检查当前Task如果已经持有块的写锁，就将writerTask置为NO_WRITER，即释放写锁。如果未持有写锁，就将readerCount自减，即释放读锁。最后，调用notifyAll()方法唤醒所有块上等待的线程。
另外，还有一个releaseAllLocksForTask()方法，它会释放当前TaskAttemptId对应的所有锁，并返回所有块ID的序列。它的实现如下，没有什么特殊的点，看官可以自行参考。
代码#21.6 - o.a.s.storage.BlockInfoManager.releaseAllLocksForTask()方法
def releaseAllLocksForTask(taskAttemptId: TaskAttemptId): Seq[BlockId] = synchronized {
val blocksWithReleasedLocks = mutable.ArrayBuffer[BlockId]()
val readLocks = readLocksByTask.remove(taskAttemptId).getOrElse(ImmutableMultiset.of[BlockId]())
val writeLocks = writeLocksByTask.remove(taskAttemptId).getOrElse(Seq.empty)
for (blockId <- writeLocks) {
infos.get(blockId).foreach { info =>
assert(info.writerTask == taskAttemptId)
info.writerTask = BlockInfo.NO_WRITER
}
blocksWithReleasedLocks += blockId
}
readLocks.entrySet().iterator().asScala.foreach { entry =>
val blockId = entry.getElement
val lockCount = entry.getCount
blocksWithReleasedLocks += blockId
get(blockId).foreach { info =>
info.readerCount -= lockCount
assert(info.readerCount >= 0)
}
}
notifyAll()
blocksWithReleasedLocks
}
锁降级
锁降级的标准定义就是写线程在持有写锁的情况下去获取读锁，然后释放写锁。BlockInfoManager中的块锁降级实现如下。
代码#21.7 - o.a.s.storage.BlockInfoManager.downgradeLock()方法
def downgradeLock(blockId: BlockId): Unit = synchronized {
logTrace(s"Task $currentTaskAttemptId downgrading write lock for $blockId")
val info = get(blockId).get
require(info.writerTask == currentTaskAttemptId,
s"Task $currentTaskAttemptId tried to downgrade a write lock that it does not hold on" +
s" block $blockId")
unlock(blockId)
val lockOutcome = lockForReading(blockId, blocking = false)
assert(lockOutcome.isDefined)
}
可见，这个降级的过程与上面的标准定义有所出入，实际上是先释放了写锁，然后重新获取了读锁，但结果是相同的。
删除BlockInfo
removeBlock()方法从infos映射中删掉对应的BlockInfo，同时释放它对应的所有锁。代码如下。
代码#21.8 - o.a.s.storage.BlockInfoManager.removeBlock()方法
def removeBlock(blockId: BlockId): Unit = synchronized {
logTrace(s"Task $currentTaskAttemptId trying to remove block $blockId")
infos.get(blockId) match {
case Some(blockInfo) =>
if (blockInfo.writerTask != currentTaskAttemptId) {
throw new IllegalStateException(
s"Task $currentTaskAttemptId called remove() on block $blockId without a write lock")
} else {
infos.remove(blockId)
blockInfo.readerCount = 0
blockInfo.writerTask = BlockInfo.NO_WRITER
writeLocksByTask.removeBinding(currentTaskAttemptId, blockId)
}
case None =>
throw new IllegalArgumentException(
s"Task $currentTaskAttemptId called remove() on non-existent block $blockId")
}
notifyAll()
}
可见，只有在持有BlockInfo写锁的Task是当前Task的情况下，才可以真正释放锁，包括将readerCount清零，将writerTask置为NO_WRITER。最后仍然要调用notifyAll()方法唤醒所有块上等待的线程。
总结
本文通过块信息管理器BlockInfoManager的源码，详细解释了Spark块的锁机制，包含获取读锁、获取写锁、释放锁和锁降级的细节。
