akka cluster sharding source code 学习 (2/5) handle off

一旦 shard coordinator（相当于分布式系统的 zookeeper） 启动，它就会启动一个定时器，每隔一定的时间尝试平衡一下集群中各个节点的负载，平衡的办法是把那些负载较重的 actor 移动到负载较轻的节点上。在这一点上，我以前的理解有误，我以为 shardRegion 是移动的最小单位。

val rebalanceTask = context.system.scheduler.schedule(rebalanceInterval, rebalanceInterval, self, RebalanceTick)

当 coordinator 收到 ReblanceTick 后，就开始尝试平衡系统负载

case RebalanceTick ⇒
if (persistentState.regions.nonEmpty) {
val shardsFuture = allocationStrategy.rebalance(persistentState.regions, rebalanceInProgress)
shardsFuture.value match {
case Some(Success(shards)) ⇒
continueRebalance(shards)
case _ ⇒
// continue when future is completed
shardsFuture.map { shards ⇒ RebalanceResult(shards)
}.recover {
case _ ⇒ RebalanceResult(Set.empty)
}.pipeTo(self)
}
}

上面的逻辑我看懂了，但是 Future 的用法没看明白。按照一般的写法，当 shardsFuture 返回 Failure 以后，应该直接执行 RebalanceResut(Set.empty).pipeTo(self)，不知道为什么失败以后还要尝试等待 Future

allocationStrategy 提供了默认的实现，也可以自定义负载均衡策略。rebalance 函数返回的是 Set(ShardId)，即那些要被移动的 shards

当 coordinator 收到 RebalanceResult 后，开始 启动 balance 逻辑

def continueRebalance(shards: Set[ShardId]): Unit =
shards.foreach { shard ⇒
if (!rebalanceInProgress(shard)) {
persistentState.shards.get(shard) match {
case Some(rebalanceFromRegion) ⇒
rebalanceInProgress += shard
log.debug("Rebalance shard [{}] from [{}]", shard, rebalanceFromRegion)
context.actorOf(rebalanceWorkerProps(shard, rebalanceFromRegion, handOffTimeout,
persistentState.regions.keySet ++ persistentState.regionProxies)
.withDispatcher(context.props.dispatcher))
case None ⇒
log.debug("Rebalance of non-existing shard [{}] is ignored", shard)
}

}
}

rebalanceInProcess 是一个 Set，记录正在被移动的 shard，我想，在新一轮 balance 开始时， rebalanceInProcess 为空的情况只会发生在上次 balance 还没有做完。不知道这个时候，是应该报错还是继续 balance 更好，因为 balanceStrategy 应该不会考虑吧到 上一轮 balance 还没做完这种可能性。

然后， coordinator 启动 rebalanceWorker，也就是上篇提到的替身 actor。

private[akka] class RebalanceWorker(shard: String, from: ActorRef, handOffTimeout: FiniteDuration,
regions: Set[ActorRef]) extends Actor {
import Internal._
regions.foreach(_ ! BeginHandOff(shard))
var remaining = regions

import context.dispatcher
context.system.scheduler.scheduleOnce(handOffTimeout, self, ReceiveTimeout)

def receive = {
case BeginHandOffAck(`shard`) ⇒
remaining -= sender()
if (remaining.isEmpty) {
from ! HandOff(shard)
context.become(stoppingShard, discardOld = true)
}
case ReceiveTimeout ⇒ done(ok = false)
}

def stoppingShard: Receive = {
case ShardStopped(shard) ⇒ done(ok = true)
case ReceiveTimeout      ⇒ done(ok = false)
}

def done(ok: Boolean): Unit = {
context.parent ! RebalanceDone(shard, ok)
context.stop(self)
}
}

akka 的逻辑是基于消息传递的，这种代码其实是很难去读的。在 rebalanceWorker 运行时，牵扯到很多个 actor。首先是，coordinator，其次是 shardRegion，也就是 host 待迁移 shard actor 的那个 region，然后是 shard actor 本身，最后是系统里所有的 shardRegion，他们也要参与进来。写到这里，我不禁把电脑屏幕竖了起来。

1. RebalanceWorker 首先给所有的 ShardRegion BeginHandOff 消息，告诉大家，hand off 开始，然后等待大家的回复

2. ShardRegion 收到 BeginHandOff 后，开始更新自己的知识库，将 HostShardRegion 和 shardActor 的记忆从自己的知识库中抹去

case BeginHandOff(shard) ⇒
log.debug("BeginHandOff shard [{}]", shard)
if (regionByShard.contains(shard)) {
val regionRef = regionByShard(shard)
val updatedShards = regions(regionRef) - shard
if (updatedShards.isEmpty) regions -= regionRef
else regions = regions.updated(regionRef, updatedShards)
regionByShard -= shard
}
sender() ! BeginHandOffAck(shard)

最后，发送 BeginHandOffAck 消息，告诉 rebalanceWorker 自己准备完毕（这些 shardRegion 以后也没事干了）

3. 继续回到 rebalanceWorker，它发送 HandOff 告诉 Host shard actor 的 ShardRegion，你可以做自己的清理工作了。然后将自己的状态设置成 stoppingShard，等待 ShardStopped 消息，这个消息的来源有两个，一个是 HostShardRegion，另外一个是 shard actor

4. HostShardRegion 收到 HandOff 消息后

case msg @ HandOff(shard) ⇒
log.debug("HandOff shard [{}]", shard)

// must drop requests that came in between the BeginHandOff and now,
// because they might be forwarded from other regions and there
// is a risk or message re-ordering otherwise
if (shardBuffers.contains(shard)) {
shardBuffers -= shard
loggedFullBufferWarning = false
}

if (shards.contains(shard)) {
handingOff += shards(shard)
shards(shard) forward msg
} else
sender() ! ShardStopped(shard)

如果 HostShardRegion 已经不再含有 shard actor，那么直接返回 ShardStopped，否则 HandOff 这个 Set 加入 shard actor，并将 HandOff 传给 shard actor

5. 又看了一遍代码，发现 shard actor 和 entity actor 又是两种东西，shard actor 存在于 entity actor 和 shard region 之间

目前还不知道 entity actor 和 shard region 之间的关系

def getEntity(id: EntityId): ActorRef = {
val name = URLEncoder.encode(id, "utf-8")
context.child(name).getOrElse {
log.debug("Starting entity [{}] in shard [{}]", id, shardId)

val a = context.watch(context.actorOf(entityProps, name))
idByRef = idByRef.updated(a, id)
refById = refById.updated(id, a)
state = state.copy(state.entities + id)
a
}
}

从这段代码来看， shard actor 与 entity actor 是一对多的关系。

def receiveCoordinatorMessage(msg: CoordinatorMessage): Unit = msg match {
case HandOff(`shardId`) ⇒ handOff(sender())
case HandOff(shard)     ⇒ log.warning("Shard [{}] can not hand off for another Shard [{}]", shardId, shard)
case _                  ⇒ unhandled(msg)
}

def handOff(replyTo: ActorRef): Unit = handOffStopper match {
case Some(_) ⇒ log.warning("HandOff shard [{}] received during existing handOff", shardId)
case None ⇒
log.debug("HandOff shard [{}]", shardId)

if (state.entities.nonEmpty) {
handOffStopper = Some(context.watch(context.actorOf(
handOffStopperProps(shardId, replyTo, idByRef.keySet, handOffStopMessage))))

//During hand off we only care about watching for termination of the hand off stopper
context become {
case Terminated(ref) ⇒ receiveTerminated(ref)
}
} else {
replyTo ! ShardStopped(shardId)
context stop self
}
}

def receiveTerminated(ref: ActorRef): Unit = {
if (handOffStopper.exists(_ == ref))
context stop self
else if (idByRef.contains(ref) && handOffStopper.isEmpty)
entityTerminated(ref)
}

从这段代码看， shard actor 与 entity actor 的关系是一对一，因为当 entity stop self 了以后， shard actor 也会 stop self。这让我想到 coursera reactive programming 的最后一道作业题，为什么也是类似于 一个 entity 有一个 shard actor 对应。