spark存储管理源码分析系列之BlockInfoManager

BlockId定义了块的唯一标识，BlockInfo是块信息的管理，本节将介绍块信息管理器BlockInfoManager，BlockInfoManager对Block的锁管理，采用了共享锁与排他锁，其中读锁是共享锁，写锁是排他锁。

BlockInfoManager属性

BlockInfoManager的成员属性有以下几个：

1. infos：存储BlockId -> BlockInfo对应关系，包含了MemoryStore以及DiskStore管理的Block。
2. writeLocksByTask：每次TaskAttempt的标识TaskAttemptId与执行获取的Block的写锁之间的映射关系；TaskAttemptId与写锁之间是一对多的关系，即一次TaskAttempt执行会获取零到多个Block的写锁。
3. readLocksByTask： 每次TaskAttempt执行的标识TaskAttemptId与获取的Block的读锁之间的映射关系；TaskAttemptId与读锁之间是一对多的关系，即一次TaskAttempt执行会获取零到多个Block的读锁，并且会记录对于同一个Block的读锁的占用次数。

private[this] val infos = new mutable.HashMap[BlockId, BlockInfo]

private[this] val writeLocksByTask =
    new mutable.HashMap[TaskAttemptId, mutable.Set[BlockId]]
      with mutable.MultiMap[TaskAttemptId, BlockId]

private[this] val readLocksByTask =
    new mutable.HashMap[TaskAttemptId, ConcurrentHashMultiset[BlockId]]

一个任务尝试执行线程可以同时获得零到多个不同 Block 的写锁或零到多个一同 Block 的读锁，但不能同时获得同一个Block的读锁与写锁。读锁是可以重入的，但是写锁不能重入。

注册任务

注册任务，简单的将TaskAttemptId放入写锁的Map中，供后续获取写锁时候，记录TaskId与Block读锁的使用关系。

def registerTask(taskAttemptId: TaskAttemptId): Unit = synchronized {
  require(!readLocksByTask.contains(taskAttemptId),
          s"Task attempt $taskAttemptId is already registered")
  readLocksByTask(taskAttemptId) = ConcurrentHashMultiset.create()
}

获取读锁

为某个TaskAttempt获取读锁需要经过以下步骤：

1. 获取要读的BlockInfo，判断它是否正在被写。
2. 如果没有正在被写，就将BlockInfo的读锁次数加1，然后将维护readLocksByTask字典中的记录，并返回BlockInfo。
3. 如果正在被写，就判断是否指定了阻塞等待:
   1. 如果指定了阻塞等待，则阻塞等待直到写锁释放后被唤醒，然后重新获取读锁；
   2. 如果没有指定阻塞等待，就放弃，返回NONE。

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 { // 从infos中获取BlockId对应的BlockInfo
      // 获取不到返回None
      case None => return None
      case Some(info) =>
      if (info.writerTask == BlockInfo.NO_WRITER) { // 获取要读的BlockInfo，判断它是否正在被写
        info.readerCount += 1
        readLocksByTask(currentTaskAttemptId).add(blockId)
        logTrace(s"Task $currentTaskAttemptId acquired read lock for $blockId")
        return Some(info)
      }
    }
    if (blocking) { // 被占有了写锁，正在被写，则不允许读，如果指定阻塞获取，就wait等待释放写锁时候进行重试
      wait() // 如果设置了阻塞，则等待，阻塞在BlockManager对象上
    }
  } while (blocking)
  None
}

获取写锁

为某个TaskAttempt获取写锁需要经过以下步骤：

1. 获取要读的BlockInfo，判断它是否正在被写或者被读。
2. 如果没有正在被写或被读，就使用BlockInfo的writerTask记录当前TaskAttempt的ID，然后将维护writeLocksByTask字典中的记录，并返回BlockInfo。
3. 如果正在被写或被读，就判断是否指定了阻塞等待；
   1. 如果指定了阻塞等待，则阻塞等待直到读锁和写锁都释放后被唤醒，然后重新获取写锁；
   2. 如果没有指定阻塞等待，就放弃，返回NONE

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) {
        // 没有写锁，且没有读锁重入，则由当前TaskAttempt线程持有写锁并返回BlockInfo
        info.writerTask = currentTaskAttemptId
        writeLocksByTask.addBinding(currentTaskAttemptId, blockId)
        logTrace(s"Task $currentTaskAttemptId acquired write lock for $blockId")
        return Some(info)
      }
    }
    if (blocking) { // 走到这里说明无法获取写锁，有其他TaskAttempt线程正在写或读
      wait()
    }
  } while (blocking)
  None
}

释放锁

释放锁会释放掉当前TaskAttemptId所持有读锁以及写锁:

1. 获取BlockId对应的BlockInfo。
2. 如果当前任务尝试线程已经获得了Block的写锁，则释放当前Block的写锁。
3. 如果当前任务尝试线程没有获得Block的写锁，则释放当前Block的读锁，释放读锁实际是减少了当前任务尝试线程已经获得的Block的读锁次数。

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) // 持有的读锁
    // newPinCountForTask表示当前TaskAttempt持有BlockId对应的Block的读锁次数与1的差值
    // 如果newPinCountForTask次数小于0，表示读锁释放次数大于加锁次数，会抛出异常
    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() // 唤醒等待读或者写的请求
}

锁降级

锁降级是将持有的写锁变为读锁，逻辑比较简单，先释放掉锁，然后获取读锁即可。

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)
}

删除Block

删除Block信息需要经过以下步骤：

1. 获取BlockId对应的BlockInfo。

2. 如果对BlockInfo正在写入的任务尝试线程是当前线程的话，当前线程才有权利去移除BlockInfo。移除BlockInfo操作如下：

   1. 将BlockInfo从infos中移除；
   2. 将BlockInfo的读线程数清零；
   3. 将BlockInfo的writeTask置为BlockInfo.NO_WRITER；
   4. 将任务尝试线程与BlockId的关系清除。

3. 通知所有的BlockId对应的Block的锁上等待的线程。

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()
}

参考

1. https://cloud.tencent.com/developer/article/1491382
2. https://mzl9039.github.io/2018/02/07/spark-core-BlockInfoManager-%E6%BA%90%E7%A0%81%E9%98%85%E8%AF%BB.html